From 51b6c1985685192c881d15c6caba6c8e72e6314b Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 09:28:32 -0700
Subject: [PATCH 01/11] Add CUDA kernels for IQ packing

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../torch/kernels/quantization/ggml/iq1_s.cpp |  34 ++
 .../torch/kernels/quantization/ggml/iq1_s.cu  | 275 ++++++++++++++++
 .../kernels/quantization/ggml/iq2_xs.cpp      |  33 ++
 .../torch/kernels/quantization/ggml/iq2_xs.cu | 302 ++++++++++++++++++
 modelopt/torch/quantization/extensions.py     |  48 ++-
 5 files changed, 691 insertions(+), 1 deletion(-)
 create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
 create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_s.cu
 create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
 create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_xs.cu

diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
new file mode 100644
index 00000000000..71b6c2e2fbc
--- /dev/null
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
@@ -0,0 +1,34 @@
+/*
+ * 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 <ATen/ATen.h>
+#include <torch/extension.h>
+
+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");
+  return iq1_s_pack_cuda(input.contiguous(), grid.contiguous());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("pack", &iq1_s_pack,
+             "Pack a float32, float64, float16, or bfloat16 CUDA tensor whose numel is a positive "
+             "multiple of 256. The grid must be float32 [2048, 8]. Returns uint8 [numel / 256, "
+             "50] on the input device.");
+}
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
new file mode 100644
index 00000000000..1b88086a361
--- /dev/null
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -0,0 +1,275 @@
+/*
+ * 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 <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+
+#include <cuda_fp16.h>
+
+#include <cfloat>
+#include <cstdint>
+#include <limits>
+
+namespace {
+
+constexpr int kBlockSize = 256;
+constexpr int kVectorSize = 8;
+constexpr int kEntries = 2048;
+constexpr int kGroups = 8;
+constexpr int kLocalScales = 8;
+constexpr int kChoices = 16;
+constexpr int kPayloadBytes = 50;
+constexpr float kDelta = 0.125f;
+constexpr float kNativeMax = 16.875f;
+
+template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
+  return static_cast<float>(*input);
+}
+
+__device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x,
+                                             const float *q, float scale, float delta) {
+  float dot = 0.0f;
+  float qnorm = 0.0f;
+  float qsum = 0.0f;
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j) {
+    dot = fmaf(x[j], q[j], dot);
+    qnorm = fmaf(q[j], q[j], qnorm);
+    qsum += q[j];
+  }
+  const float shifted_dot = dot + delta * xsum;
+  const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta;
+  return fmaxf(fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f);
+}
+
+template <typename scalar_t>
+__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) {
+  const int64_t block = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (block >= num_blocks)
+    return;
+
+  float amax = 0.0f;
+  const scalar_t *values = input + block * kBlockSize;
+#pragma unroll 1
+  for (int i = 0; i < kBlockSize; ++i)
+    amax = fmaxf(amax, fabsf(load_float(values + i)));
+  const __half scale = __float2half_rn(fminf((amax / kNativeMax) * 0.61f, 65504.0f));
+  scale_bits[block] = static_cast<int64_t>(__half_as_ushort(scale));
+}
+
+template <typename scalar_t>
+__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid,
+                       const int64_t *scale_bits, uint8_t *output) {
+  __shared__ float warp_best[8 * kChoices];
+  __shared__ float group_error[kChoices];
+  __shared__ unsigned long long warp_keys[8];
+  __shared__ int selected_choice;
+  __shared__ uint16_t selected_entries[4];
+
+  const int tid = threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+  const int64_t block = blockIdx.x;
+  if (block >= num_blocks)
+    return;
+
+  const scalar_t *source = input + block * kBlockSize;
+  uint8_t *payload = output + block * kPayloadBytes;
+  const uint16_t d_bits = static_cast<uint16_t>(scale_bits[block]);
+  const float d = __half2float(__ushort_as_half(d_bits));
+  if (d_bits == 0) {
+    if (tid < kPayloadBytes)
+      payload[tid] = 0;
+    return;
+  }
+  if (tid == 0) {
+    payload[0] = static_cast<uint8_t>(d_bits);
+    payload[1] = static_cast<uint8_t>(d_bits >> 8);
+  }
+
+#pragma unroll 1
+  for (int group = 0; group < kGroups; ++group) {
+    if (tid < kChoices)
+      group_error[tid] = 0.0f;
+    __syncthreads();
+
+#pragma unroll
+    for (int vector = 0; vector < 4; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      float xsum = 0.0f;
+      const int offset = group * 32 + vector * 8;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+        xsum += x[j];
+      }
+      float local_best[kChoices];
+#pragma unroll
+      for (int choice = 0; choice < kChoices; ++choice)
+        local_best[choice] = FLT_MAX;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float *q = grid + entry * kVectorSize;
+        float dot = 0.0f;
+        float qnorm = 0.0f;
+        float qsum = 0.0f;
+#pragma unroll
+        for (int j = 0; j < kVectorSize; ++j) {
+          dot = fmaf(x[j], q[j], dot);
+          qnorm = fmaf(q[j], q[j], qnorm);
+          qsum += q[j];
+        }
+#pragma unroll
+        for (int choice = 0; choice < kChoices; ++choice) {
+          const int local = choice & 7;
+          const float delta = choice < 8 ? kDelta : -kDelta;
+          const float scale = d * (2 * local + 1);
+          const float shifted_dot = dot + delta * xsum;
+          const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta;
+          const float error = fmaxf(
+              fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f);
+          local_best[choice] = fminf(local_best[choice], error);
+        }
+      }
+#pragma unroll
+      for (int choice = 0; choice < kChoices; ++choice) {
+        float value = local_best[choice];
+#pragma unroll
+        for (int delta = 16; delta > 0; delta >>= 1)
+          value = fminf(value, __shfl_down_sync(0xffffffff, value, delta));
+        if (lane == 0)
+          warp_best[warp * kChoices + choice] = value;
+      }
+      __syncthreads();
+      if (tid < kChoices) {
+        float value = warp_best[tid];
+#pragma unroll
+        for (int w = 1; w < 8; ++w)
+          value = fminf(value, warp_best[w * kChoices + tid]);
+        group_error[tid] += value;
+      }
+      __syncthreads();
+    }
+
+    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;
+        }
+      }
+    }
+    __syncthreads();
+    const int selected_local = selected_choice & 7;
+    const float selected_delta = selected_choice < 8 ? kDelta : -kDelta;
+    const float selected_scale = d * (2 * selected_local + 1);
+
+#pragma unroll
+    for (int vector = 0; vector < 4; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      float xsum = 0.0f;
+      const int offset = group * 32 + vector * 8;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+        xsum += x[j];
+      }
+      unsigned long long key = ~0ULL;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float error =
+            quant_error(xnorm, xsum, x, grid + entry * kVectorSize, selected_scale, selected_delta);
+        const unsigned long long candidate =
+            (static_cast<unsigned long long>(__float_as_uint(error)) << 32) |
+            static_cast<unsigned long long>(entry);
+        key = candidate < key ? candidate : key;
+      }
+#pragma unroll
+      for (int delta = 16; delta > 0; delta >>= 1) {
+        const auto other = __shfl_down_sync(0xffffffff, key, delta);
+        key = other < key ? other : key;
+      }
+      if (lane == 0)
+        warp_keys[warp] = key;
+      __syncthreads();
+      if (tid == 0) {
+        key = warp_keys[0];
+#pragma unroll
+        for (int w = 1; w < 8; ++w)
+          key = warp_keys[w] < key ? warp_keys[w] : key;
+        const uint16_t entry = static_cast<uint16_t>(key & 0x7ff);
+        selected_entries[vector] = entry;
+        payload[2 + group * 4 + vector] = static_cast<uint8_t>(entry);
+      }
+      __syncthreads();
+    }
+
+    if (tid == 0) {
+      const uint16_t qh = static_cast<uint16_t>(
+          ((selected_entries[0] >> 8) & 7) | (((selected_entries[1] >> 8) & 7) << 3) |
+          (((selected_entries[2] >> 8) & 7) << 6) | (((selected_entries[3] >> 8) & 7) << 9) |
+          (selected_local << 12) | ((selected_choice >> 3) << 15));
+      payload[34 + 2 * group] = static_cast<uint8_t>(qh);
+      payload[35 + 2 * group] = static_cast<uint8_t>(qh >> 8);
+    }
+    __syncthreads();
+  }
+}
+
+} // namespace
+
+at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) {
+  TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous");
+  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,
+              "IQ1_S packing supports float32, float64, float16, and bfloat16 inputs");
+  TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0,
+              "input size must be a positive multiple of 256");
+  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
+                  grid.size(1) == kVectorSize,
+              "grid must be float32 [2048, 8]");
+  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
+  c10::cuda::CUDAGuard guard(input.device());
+  const int64_t num_blocks = input.numel() / kBlockSize;
+  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ1_S CUDA grid is too large");
+  auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong));
+  auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte));
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+  const int scale_grid = static_cast<int>((num_blocks + 255) / 256);
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_s_pack", [&] {
+        find_scale<scalar_t><<<scale_grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(), num_blocks,
+                                                             scales.data_ptr<int64_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+        encode<scalar_t><<<static_cast<int>(num_blocks), 256, 0, stream>>>(
+            input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
+            scales.data_ptr<int64_t>(), output.data_ptr<uint8_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
new file mode 100644
index 00000000000..447be7dcae5
--- /dev/null
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
@@ -0,0 +1,33 @@
+/*
+ * 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 <torch/extension.h>
+
+at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid);
+
+at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid) {
+  TORCH_CHECK(input.is_cuda(), "IQ2_XS packing requires a CUDA input");
+  TORCH_CHECK(grid.is_cuda(), "IQ2_XS packing requires a CUDA grid");
+  return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("pack", &iq2_xs_pack,
+             "Pack a float32, float64, float16, or bfloat16 CUDA tensor whose numel is a positive "
+             "multiple of 256. The grid must be float32 [512, 8]. Returns uint8 [numel / 256, "
+             "74] on the input device.");
+}
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
new file mode 100644
index 00000000000..5351eec5c8a
--- /dev/null
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -0,0 +1,302 @@
+/*
+ * 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 <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+
+#include <cuda_fp16.h>
+
+#include <cfloat>
+#include <cstdint>
+#include <limits>
+
+namespace {
+
+constexpr int kBlockSize = 256;
+constexpr int kVectorSize = 8;
+constexpr int kEntries = 512;
+constexpr int kGroups = 16;
+constexpr int kLocalScales = 16;
+constexpr int kPayloadBytes = 74;
+constexpr float kNativeMax = 166.625f;
+
+template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
+  return static_cast<float>(*input);
+}
+
+__device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm, float scale) {
+  return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f);
+}
+
+__device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) {
+  float dot = 0.0f;
+  float weakest = FLT_MAX;
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j) {
+    const float term = fabsf(x[j]) * q[j];
+    dot += term;
+    weakest = fminf(weakest, term);
+  }
+  return odd_parity ? dot - 2.0f * weakest : dot;
+}
+
+template <typename scalar_t>
+__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) {
+  const int64_t block = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (block >= num_blocks)
+    return;
+
+  float amax = 0.0f;
+  float sumsq = 0.0f;
+  const scalar_t *values = input + block * kBlockSize;
+#pragma unroll 1
+  for (int i = 0; i < kBlockSize; ++i) {
+    const float value = load_float(values + i);
+    amax = fmaxf(amax, fabsf(value));
+    sumsq = fmaf(value, value, sumsq);
+  }
+  if (amax == 0.0f) {
+    scale_bits[block] = 0;
+    return;
+  }
+  const float rms = sqrtf(sumsq / kBlockSize);
+  const float peak_to_rms = rms > 0.0f ? amax / rms : 0.0f;
+  const float anchor = fminf(0.92f, fmaxf(0.65f, 1.0f - 0.035f * peak_to_rms));
+  const __half scale = __float2half_rn(fminf((amax / kNativeMax) * anchor, 65504.0f));
+  scale_bits[block] = static_cast<int64_t>(__half_as_ushort(scale));
+}
+
+template <typename scalar_t>
+__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid,
+                       const int64_t *scale_bits, uint8_t *output) {
+  __shared__ float shared_grid[kEntries * kVectorSize];
+  __shared__ float grid_norm[kEntries];
+  __shared__ float warp_best[8 * kLocalScales];
+  __shared__ float group_error[kLocalScales];
+  __shared__ unsigned long long warp_keys[8];
+  __shared__ int selected_local;
+  __shared__ uint8_t locals[kGroups];
+
+  const int tid = threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+  for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x)
+    shared_grid[i] = grid[i];
+  __syncthreads();
+  for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+    float norm = 0.0f;
+#pragma unroll
+    for (int j = 0; j < kVectorSize; ++j) {
+      const float q = shared_grid[entry * kVectorSize + j];
+      norm = fmaf(q, q, norm);
+    }
+    grid_norm[entry] = norm;
+  }
+  __syncthreads();
+
+  const int64_t block = blockIdx.x;
+  if (block >= num_blocks)
+    return;
+  const scalar_t *source = input + block * kBlockSize;
+  uint8_t *payload = output + block * kPayloadBytes;
+  const uint16_t d_bits = static_cast<uint16_t>(scale_bits[block]);
+  const float d = __half2float(__ushort_as_half(d_bits));
+  if (d_bits == 0) {
+    if (tid < kPayloadBytes)
+      payload[tid] = 0;
+    return;
+  }
+  if (tid == 0) {
+    payload[0] = static_cast<uint8_t>(d_bits);
+    payload[1] = static_cast<uint8_t>(d_bits >> 8);
+  }
+
+#pragma unroll 1
+  for (int group = 0; group < kGroups; ++group) {
+    if (tid < kLocalScales)
+      group_error[tid] = 0.0f;
+    __syncthreads();
+
+#pragma unroll
+    for (int vector = 0; vector < 2; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      int negative_count = 0;
+      const int offset = group * 16 + vector * 8;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+        negative_count += x[j] < 0.0f;
+      }
+      const bool odd_parity = (negative_count & 1) != 0;
+      float local_best[kLocalScales];
+#pragma unroll
+      for (int local = 0; local < kLocalScales; ++local)
+        local_best[local] = FLT_MAX;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float *q = shared_grid + entry * kVectorSize;
+        const float dot = even_parity_dot(x, q, odd_parity);
+#pragma unroll
+        for (int local = 0; local < kLocalScales; ++local) {
+          const float scale = d * (2 * local + 1) * 0.125f;
+          local_best[local] =
+              fminf(local_best[local], quant_error(xnorm, dot, grid_norm[entry], scale));
+        }
+      }
+#pragma unroll
+      for (int local = 0; local < kLocalScales; ++local) {
+        float value = local_best[local];
+#pragma unroll
+        for (int delta = 16; delta > 0; delta >>= 1)
+          value = fminf(value, __shfl_down_sync(0xffffffff, value, delta));
+        if (lane == 0)
+          warp_best[warp * kLocalScales + local] = value;
+      }
+      __syncthreads();
+      if (tid < kLocalScales) {
+        float value = warp_best[tid];
+#pragma unroll
+        for (int w = 1; w < 8; ++w)
+          value = fminf(value, warp_best[w * kLocalScales + tid]);
+        group_error[tid] += value;
+      }
+      __syncthreads();
+    }
+
+    if (tid == 0) {
+      selected_local = 0;
+      float best = group_error[0];
+#pragma unroll
+      for (int local = 1; local < kLocalScales; ++local) {
+        if (group_error[local] < best) {
+          best = group_error[local];
+          selected_local = local;
+        }
+      }
+      locals[group] = static_cast<uint8_t>(selected_local);
+    }
+    __syncthreads();
+    const float selected_scale = d * (2 * selected_local + 1) * 0.125f;
+
+#pragma unroll
+    for (int vector = 0; vector < 2; ++vector) {
+      float x[kVectorSize];
+      float xnorm = 0.0f;
+      int negative_count = 0;
+      const int offset = group * 16 + vector * 8;
+#pragma unroll
+      for (int j = 0; j < kVectorSize; ++j) {
+        x[j] = load_float(source + offset + j);
+        xnorm = fmaf(x[j], x[j], xnorm);
+        negative_count += x[j] < 0.0f;
+      }
+      const bool odd_parity = (negative_count & 1) != 0;
+      unsigned long long key = ~0ULL;
+      for (int entry = tid; entry < kEntries; entry += blockDim.x) {
+        const float error =
+            quant_error(xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity),
+                        grid_norm[entry], selected_scale);
+        const unsigned long long candidate =
+            (static_cast<unsigned long long>(__float_as_uint(error)) << 32) |
+            static_cast<unsigned long long>(entry);
+        key = candidate < key ? candidate : key;
+      }
+#pragma unroll
+      for (int delta = 16; delta > 0; delta >>= 1) {
+        const auto other = __shfl_down_sync(0xffffffff, key, delta);
+        key = other < key ? other : key;
+      }
+      if (lane == 0)
+        warp_keys[warp] = key;
+      __syncthreads();
+      if (tid == 0) {
+        key = warp_keys[0];
+#pragma unroll
+        for (int w = 1; w < 8; ++w)
+          key = warp_keys[w] < key ? warp_keys[w] : key;
+        const int entry = static_cast<int>(key & 0x1ff);
+        const float *q = shared_grid + entry * kVectorSize;
+        int flip_index = 0;
+        float weakest = fabsf(x[0]) * q[0];
+#pragma unroll
+        for (int j = 1; j < kVectorSize; ++j) {
+          const float term = fabsf(x[j]) * q[j];
+          if (term < weakest) {
+            weakest = term;
+            flip_index = j;
+          }
+        }
+        int sign_mask = 0;
+#pragma unroll
+        for (int j = 0; j < kVectorSize; ++j) {
+          bool is_negative = x[j] < 0.0f;
+          if (odd_parity && j == flip_index)
+            is_negative = !is_negative;
+          sign_mask |= static_cast<int>(is_negative) << j;
+        }
+        const uint16_t code = static_cast<uint16_t>(entry | ((sign_mask & 0x7f) << 9));
+        const int code_offset = 2 + 2 * (group * 2 + vector);
+        payload[code_offset] = static_cast<uint8_t>(code);
+        payload[code_offset + 1] = static_cast<uint8_t>(code >> 8);
+      }
+      __syncthreads();
+    }
+  }
+
+  if (tid < 8)
+    payload[66 + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4);
+}
+
+} // namespace
+
+at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) {
+  TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous");
+  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,
+              "IQ2_XS packing supports float32, float64, float16, and bfloat16 inputs");
+  TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0,
+              "input size must be a positive multiple of 256");
+  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
+                  grid.size(1) == kVectorSize,
+              "grid must be float32 [512, 8]");
+  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
+  c10::cuda::CUDAGuard guard(input.device());
+  const int64_t num_blocks = input.numel() / kBlockSize;
+  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ2_XS CUDA grid is too large");
+  auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong));
+  auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte));
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+  const int scale_grid = static_cast<int>((num_blocks + 255) / 256);
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] {
+        find_scale<scalar_t><<<scale_grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(), num_blocks,
+                                                             scales.data_ptr<int64_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+        encode<scalar_t><<<static_cast<int>(num_blocks), 256, 0, stream>>>(
+            input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
+            scales.data_ptr<int64_t>(), output.data_ptr<uint8_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}
diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py
index a65396d64ff..1c802d7b1d0 100644
--- a/modelopt/torch/quantization/extensions.py
+++ b/modelopt/torch/quantization/extensions.py
@@ -19,10 +19,18 @@
 
 from modelopt.torch.utils import load_cpp_extension
 
-__all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"]
+__all__ = [
+    "get_cuda_ext",
+    "get_cuda_ext_fp8",
+    "get_cuda_ext_iq1_s",
+    "get_cuda_ext_iq2_xs",
+    "get_cuda_ext_mx",
+    "precompile",
+]
 
 path = Path(__file__).parent
 kernels_gemm = path.parent / "kernels" / "quantization" / "gemm"
+kernels_ggml = path.parent / "kernels" / "quantization" / "ggml"
 
 
 def get_cuda_ext(raise_if_failed: bool = False):
@@ -72,6 +80,38 @@ 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 is unavailable; using the PyTorch reference encoder.",
+            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_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
+    ):
+        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"],
+            cuda_version_specifiers=">=11.8",
+            fail_msg="IQ2_XS CUDA packing is unavailable; using the PyTorch reference encoder.",
+            extra_cuda_cflags=["-O3"],
+            raise_if_failed=raise_if_failed,
+        )
+    return get_cuda_ext_iq2_xs.extension  # type:ignore[attr-defined]
+
+
 def __getattr__(name):
     if name == "cuda_ext":
         return get_cuda_ext()
@@ -79,6 +119,10 @@ 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()
     else:
         raise AttributeError(f"module {__name__} has no attribute {name}")
 
@@ -88,3 +132,5 @@ 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())

From b39dca34ca18eded4f1c32627287b2e75e632978 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 10:51:42 -0700
Subject: [PATCH 02/11] Document IQ CUDA format constants

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../torch/kernels/quantization/ggml/iq1_s.cu  | 30 ++++++++++++-------
 .../torch/kernels/quantization/ggml/iq2_xs.cu | 30 ++++++++++++++-----
 2 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index 1b88086a361..772e8743ba1 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -30,15 +30,23 @@
 
 namespace {
 
+// Packed layout and format constants follow the GGML definition at:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kBlockSize = 256;
 constexpr int kVectorSize = 8;
 constexpr int kEntries = 2048;
 constexpr int kGroups = 8;
-constexpr int kLocalScales = 8;
 constexpr int kChoices = 16;
-constexpr int kPayloadBytes = 50;
-constexpr float kDelta = 0.125f;
-constexpr float kNativeMax = 16.875f;
+constexpr int kScaleOffset = 0;
+constexpr int kIndexOffset = 2;
+constexpr int kIndexBytes = kBlockSize / kVectorSize;
+constexpr int kMetadataOffset = kIndexOffset + kIndexBytes;
+constexpr int kPayloadBytes = kMetadataOffset + 2 * kGroups;
+constexpr float kDelta = 0.125f;        // The metadata shift bit selects +1/8 or -1/8.
+constexpr float kMaxLocalScale = 15.0f; // Largest multiplier: 2 * 7 + 1.
+constexpr float kMaxShiftedMagnitude = 1.0f + kDelta;
+constexpr float kNativeMax = kMaxLocalScale * kMaxShiftedMagnitude; // 16.875.
+constexpr float kScaleAnchor = 0.61f;
 
 template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
   return static_cast<float>(*input);
@@ -71,7 +79,9 @@ __global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *s
 #pragma unroll 1
   for (int i = 0; i < kBlockSize; ++i)
     amax = fmaxf(amax, fabsf(load_float(values + i)));
-  const __half scale = __float2half_rn(fminf((amax / kNativeMax) * 0.61f, 65504.0f));
+  // Match the reference encoder's empirical predictor. The 0.61 anchor favors most values
+  // instead of forcing the block's largest value to be exactly representable.
+  const __half scale = __float2half_rn(fminf((amax / kNativeMax) * kScaleAnchor, 65504.0f));
   scale_bits[block] = static_cast<int64_t>(__half_as_ushort(scale));
 }
 
@@ -101,8 +111,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     return;
   }
   if (tid == 0) {
-    payload[0] = static_cast<uint8_t>(d_bits);
-    payload[1] = static_cast<uint8_t>(d_bits >> 8);
+    payload[kScaleOffset] = static_cast<uint8_t>(d_bits);
+    payload[kScaleOffset + 1] = static_cast<uint8_t>(d_bits >> 8);
   }
 
 #pragma unroll 1
@@ -222,7 +232,7 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
           key = warp_keys[w] < key ? warp_keys[w] : key;
         const uint16_t entry = static_cast<uint16_t>(key & 0x7ff);
         selected_entries[vector] = entry;
-        payload[2 + group * 4 + vector] = static_cast<uint8_t>(entry);
+        payload[kIndexOffset + group * 4 + vector] = static_cast<uint8_t>(entry);
       }
       __syncthreads();
     }
@@ -232,8 +242,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
           ((selected_entries[0] >> 8) & 7) | (((selected_entries[1] >> 8) & 7) << 3) |
           (((selected_entries[2] >> 8) & 7) << 6) | (((selected_entries[3] >> 8) & 7) << 9) |
           (selected_local << 12) | ((selected_choice >> 3) << 15));
-      payload[34 + 2 * group] = static_cast<uint8_t>(qh);
-      payload[35 + 2 * group] = static_cast<uint8_t>(qh >> 8);
+      payload[kMetadataOffset + 2 * group] = static_cast<uint8_t>(qh);
+      payload[kMetadataOffset + 2 * group + 1] = static_cast<uint8_t>(qh >> 8);
     }
     __syncthreads();
   }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index 5351eec5c8a..b5f0ce9c255 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -30,13 +30,24 @@
 
 namespace {
 
+// Packed layout and format constants follow the GGML definition at:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kBlockSize = 256;
 constexpr int kVectorSize = 8;
 constexpr int kEntries = 512;
 constexpr int kGroups = 16;
 constexpr int kLocalScales = 16;
-constexpr int kPayloadBytes = 74;
-constexpr float kNativeMax = 166.625f;
+constexpr int kScaleOffset = 0;
+constexpr int kCodeOffset = 2;
+constexpr int kCodeBytes = 2 * (kBlockSize / kVectorSize);
+constexpr int kLocalScaleOffset = kCodeOffset + kCodeBytes;
+constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2;
+constexpr float kMaxMagnitude = 43.0f;
+constexpr float kMaxLocalScale = 31.0f / 8.0f;
+constexpr float kNativeMax = kMaxMagnitude * kMaxLocalScale; // 166.625.
+constexpr float kPeakToRmsSlope = 0.035f;
+constexpr float kMinScaleAnchor = 0.65f;
+constexpr float kMaxScaleAnchor = 0.92f;
 
 template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
   return static_cast<float>(*input);
@@ -47,6 +58,8 @@ __device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm
 }
 
 __device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) {
+  // The format stores seven sign bits. For odd parity, flip the coordinate with the smallest
+  // |x| * q penalty; the eighth sign is recovered from even parity during decoding.
   float dot = 0.0f;
   float weakest = FLT_MAX;
 #pragma unroll
@@ -79,7 +92,10 @@ __global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *s
   }
   const float rms = sqrtf(sumsq / kBlockSize);
   const float peak_to_rms = rms > 0.0f ? amax / rms : 0.0f;
-  const float anchor = fminf(0.92f, fmaxf(0.65f, 1.0f - 0.035f * peak_to_rms));
+  // Match the reference encoder's empirical predictor. Peaky blocks get a smaller anchor so
+  // outliers do not set the entire scale, while the clamp bounds the adjustment.
+  const float anchor =
+      fminf(kMaxScaleAnchor, fmaxf(kMinScaleAnchor, 1.0f - kPeakToRmsSlope * peak_to_rms));
   const __half scale = __float2half_rn(fminf((amax / kNativeMax) * anchor, 65504.0f));
   scale_bits[block] = static_cast<int64_t>(__half_as_ushort(scale));
 }
@@ -125,8 +141,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     return;
   }
   if (tid == 0) {
-    payload[0] = static_cast<uint8_t>(d_bits);
-    payload[1] = static_cast<uint8_t>(d_bits >> 8);
+    payload[kScaleOffset] = static_cast<uint8_t>(d_bits);
+    payload[kScaleOffset + 1] = static_cast<uint8_t>(d_bits >> 8);
   }
 
 #pragma unroll 1
@@ -254,7 +270,7 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
           sign_mask |= static_cast<int>(is_negative) << j;
         }
         const uint16_t code = static_cast<uint16_t>(entry | ((sign_mask & 0x7f) << 9));
-        const int code_offset = 2 + 2 * (group * 2 + vector);
+        const int code_offset = kCodeOffset + 2 * (group * 2 + vector);
         payload[code_offset] = static_cast<uint8_t>(code);
         payload[code_offset + 1] = static_cast<uint8_t>(code >> 8);
       }
@@ -263,7 +279,7 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
   }
 
   if (tid < 8)
-    payload[66 + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4);
+    payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4);
 }
 
 } // namespace

From d554c36879de502ccac59d7dc0eddd9fab922a53 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 11:32:50 -0700
Subject: [PATCH 03/11] Clarify IQ CUDA extension failure messages

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/quantization/extensions.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py
index 1c802d7b1d0..900bf666588 100644
--- a/modelopt/torch/quantization/extensions.py
+++ b/modelopt/torch/quantization/extensions.py
@@ -89,7 +89,7 @@ def get_cuda_ext_iq1_s(raise_if_failed: bool = False):
             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 is unavailable; using the PyTorch reference encoder.",
+            fail_msg="IQ1_S CUDA packing extension is unavailable.",
             extra_cuda_cflags=["-O3"],
             raise_if_failed=raise_if_failed,
         )
@@ -105,7 +105,7 @@ def get_cuda_ext_iq2_xs(raise_if_failed: bool = False):
             name="modelopt_cuda_ext_iq2_xs",
             sources=[kernels_ggml / "iq2_xs.cpp", kernels_ggml / "iq2_xs.cu"],
             cuda_version_specifiers=">=11.8",
-            fail_msg="IQ2_XS CUDA packing is unavailable; using the PyTorch reference encoder.",
+            fail_msg="IQ2_XS CUDA packing extension is unavailable.",
             extra_cuda_cflags=["-O3"],
             raise_if_failed=raise_if_failed,
         )

From 333ef0e611d0ac1d6452b0da40d5d4bb5ef1cfb5 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 12:24:02 -0700
Subject: [PATCH 04/11] Enforce per-row IQ CUDA block alignment

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/kernels/quantization/ggml/iq1_s.cpp  | 6 +++---
 modelopt/torch/kernels/quantization/ggml/iq1_s.cu   | 6 ++++--
 modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp | 6 +++---
 modelopt/torch/kernels/quantization/ggml/iq2_xs.cu  | 6 ++++--
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
index 71b6c2e2fbc..76a1085aec7 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
@@ -28,7 +28,7 @@ at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) {
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("pack", &iq1_s_pack,
-             "Pack a float32, float64, float16, or bfloat16 CUDA tensor whose numel is a positive "
-             "multiple of 256. The grid must be float32 [2048, 8]. Returns uint8 [numel / 256, "
-             "50] on the input device.");
+             "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.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index 772e8743ba1..e338449f372 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -257,8 +257,10 @@ at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) {
   TORCH_CHECK(input_type == at::kFloat || input_type == at::kDouble || input_type == at::kHalf ||
                   input_type == at::kBFloat16,
               "IQ1_S packing supports float32, float64, float16, and bfloat16 inputs");
-  TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0,
-              "input size must be a positive multiple of 256");
+  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 256 so blocks do not straddle "
+              "rows");
   TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
                   grid.size(1) == kVectorSize,
               "grid must be float32 [2048, 8]");
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
index 447be7dcae5..52b811b7509 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
@@ -27,7 +27,7 @@ at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid) {
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("pack", &iq2_xs_pack,
-             "Pack a float32, float64, float16, or bfloat16 CUDA tensor whose numel is a positive "
-             "multiple of 256. The grid must be float32 [512, 8]. Returns uint8 [numel / 256, "
-             "74] on the input device.");
+             "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]. Returns uint8 "
+             "[numel / 256, 74] on the input device.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index b5f0ce9c255..2a5ec94e62f 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -290,8 +290,10 @@ at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) {
   TORCH_CHECK(input_type == at::kFloat || input_type == at::kDouble || input_type == at::kHalf ||
                   input_type == at::kBFloat16,
               "IQ2_XS packing supports float32, float64, float16, and bfloat16 inputs");
-  TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0,
-              "input size must be a positive multiple of 256");
+  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 256 so blocks do not straddle "
+              "rows");
   TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
                   grid.size(1) == kVectorSize,
               "grid must be float32 [512, 8]");

From b8b3047b67bd9ce1fcf33ad322c95f344cb1323d Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 14:50:24 -0700
Subject: [PATCH 05/11] Test IQ CUDA extension boundaries

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../gpu/_extensions/test_torch_extensions.py  | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index 4c104952897..6071595f0dd 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -15,6 +15,7 @@
 
 
 import pytest
+import torch
 
 import modelopt.torch.quantization.extensions as ext
 
@@ -33,3 +34,53 @@ def test_cuda_ext_fp8():
 
 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
+
+
+_IQ_EXTENSIONS = (
+    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), 50, id="iq1_s"),
+    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), 74, id="iq2_xs"),
+)
+_IQ_EXTENSION_GRIDS = (
+    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), id="iq1_s"),
+    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), id="iq2_xs"),
+)
+
+
+@pytest.mark.parametrize(("get_extension", "grid_shape", "payload_bytes"), _IQ_EXTENSIONS)
+def test_cuda_ext_iq_zero_block_layout(get_extension, grid_shape, payload_bytes):
+    extension = get_extension(raise_if_failed=True)
+    weight = torch.zeros((2, 256), device="cuda", dtype=torch.bfloat16)
+    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
+
+    packed = extension.pack(weight, grid)
+
+    assert packed.shape == (2, payload_bytes)
+    assert not packed.any()
+
+
+@pytest.mark.parametrize(("get_extension", "grid_shape"), _IQ_EXTENSION_GRIDS)
+def test_cuda_ext_iq_rejects_unsupported_dtype(get_extension, grid_shape):
+    extension = get_extension(raise_if_failed=True)
+    weight = torch.ones((1, 256), device="cuda").to(torch.float8_e4m3fn)
+    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
+
+    with pytest.raises(RuntimeError, match="supports float32, float64, float16, and bfloat16"):
+        extension.pack(weight, grid)
+
+
+@pytest.mark.parametrize(("get_extension", "grid_shape"), _IQ_EXTENSION_GRIDS)
+def test_cuda_ext_iq_rejects_row_straddling_input(get_extension, grid_shape):
+    extension = get_extension(raise_if_failed=True)
+    weight = torch.ones((512, 384), device="cuda", dtype=torch.bfloat16)
+    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
+
+    with pytest.raises(RuntimeError, match="innermost dimension must be a multiple of 256"):
+        extension.pack(weight, grid)

From 46afe88fc0de68427593ce55b6d90db9a46e3da3 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 15:57:40 -0700
Subject: [PATCH 06/11] Deduplicate IQ extension test cases

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 tests/gpu/_extensions/test_torch_extensions.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index 6071595f0dd..4641f2e202b 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -48,10 +48,6 @@ def test_cuda_ext_iq2_xs():
     pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), 50, id="iq1_s"),
     pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), 74, id="iq2_xs"),
 )
-_IQ_EXTENSION_GRIDS = (
-    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), id="iq1_s"),
-    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), id="iq2_xs"),
-)
 
 
 @pytest.mark.parametrize(("get_extension", "grid_shape", "payload_bytes"), _IQ_EXTENSIONS)
@@ -66,8 +62,8 @@ def test_cuda_ext_iq_zero_block_layout(get_extension, grid_shape, payload_bytes)
     assert not packed.any()
 
 
-@pytest.mark.parametrize(("get_extension", "grid_shape"), _IQ_EXTENSION_GRIDS)
-def test_cuda_ext_iq_rejects_unsupported_dtype(get_extension, grid_shape):
+@pytest.mark.parametrize(("get_extension", "grid_shape", "_payload_bytes"), _IQ_EXTENSIONS)
+def test_cuda_ext_iq_rejects_unsupported_dtype(get_extension, grid_shape, _payload_bytes):
     extension = get_extension(raise_if_failed=True)
     weight = torch.ones((1, 256), device="cuda").to(torch.float8_e4m3fn)
     grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
@@ -76,8 +72,8 @@ def test_cuda_ext_iq_rejects_unsupported_dtype(get_extension, grid_shape):
         extension.pack(weight, grid)
 
 
-@pytest.mark.parametrize(("get_extension", "grid_shape"), _IQ_EXTENSION_GRIDS)
-def test_cuda_ext_iq_rejects_row_straddling_input(get_extension, grid_shape):
+@pytest.mark.parametrize(("get_extension", "grid_shape", "_payload_bytes"), _IQ_EXTENSIONS)
+def test_cuda_ext_iq_rejects_row_straddling_input(get_extension, grid_shape, _payload_bytes):
     extension = get_extension(raise_if_failed=True)
     weight = torch.ones((512, 384), device="cuda", dtype=torch.bfloat16)
     grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)

From 7643d7353bc88bc1fce773626329fdb3fab45b09 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 18:50:34 -0700
Subject: [PATCH 07/11] Harden IQ CUDA packing contract

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../torch/kernels/quantization/ggml/iq1_s.cpp | 19 +++++-
 .../torch/kernels/quantization/ggml/iq1_s.cu  |  3 +-
 .../kernels/quantization/ggml/iq2_xs.cpp      | 32 +++++++--
 .../torch/kernels/quantization/ggml/iq2_xs.cu | 66 +++++--------------
 .../gpu/_extensions/test_torch_extensions.py  | 39 +++++++----
 5 files changed, 92 insertions(+), 67 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
index 76a1085aec7..adecc464aa3 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
@@ -18,11 +18,27 @@
 #include <ATen/ATen.h>
 #include <torch/extension.h>
 
+#include <limits>
+
 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");
+  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,
+              "IQ1_S 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) % 256 == 0,
+              "input's innermost dimension must be a multiple of 256 so blocks do not straddle "
+              "rows");
+  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == 2048 &&
+                  grid.size(1) == 8,
+              "grid must be float32 [2048, 8]");
+  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
+  TORCH_CHECK(input.numel() / 256 <= std::numeric_limits<int>::max(),
+              "IQ1_S CUDA grid is too large");
   return iq1_s_pack_cuda(input.contiguous(), grid.contiguous());
 }
 
@@ -30,5 +46,6 @@ 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.");
+             "[numel / 256, 50] on the input device. Non-finite input elements are treated as "
+             "zero during packing.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index e338449f372..883b1435885 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -49,7 +49,8 @@ constexpr float kNativeMax = kMaxLocalScale * kMaxShiftedMagnitude; // 16.875.
 constexpr float kScaleAnchor = 0.61f;
 
 template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
-  return static_cast<float>(*input);
+  const float value = static_cast<float>(*input);
+  return isfinite(value) ? value : 0.0f;
 }
 
 __device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x,
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
index 52b811b7509..007b4f96fb7 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
@@ -17,17 +17,39 @@
 
 #include <torch/extension.h>
 
-at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid);
+#include <limits>
 
-at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid) {
+at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales);
+
+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");
-  return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous());
+  TORCH_CHECK(scales.is_cuda(), "IQ2_XS packing requires CUDA scales");
+  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,
+              "IQ2_XS 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) % 256 == 0,
+              "input's innermost dimension must be a multiple of 256 so blocks do not straddle "
+              "rows");
+  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == 512 &&
+                  grid.size(1) == 8,
+              "grid must be float32 [512, 8]");
+  const auto num_blocks = input.numel() / 256;
+  TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
+                  scales.numel() == num_blocks,
+              "scales must be float16 [numel / 256]");
+  TORCH_CHECK(input.get_device() == grid.get_device() && input.get_device() == scales.get_device(),
+              "input, grid, and scales must share a device");
+  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ2_XS CUDA grid is too large");
+  return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("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]. Returns uint8 "
-             "[numel / 256, 74] on the input device.");
+             "dimension is a multiple of 256. The grid must be float32 [512, 8], and scales must "
+             "be float16 [numel / 256]. Returns uint8 [numel / 256, 74] on the input device. "
+             "Non-finite input elements are treated as zero during packing.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index 2a5ec94e62f..e60c849a67b 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -42,15 +42,10 @@ constexpr int kCodeOffset = 2;
 constexpr int kCodeBytes = 2 * (kBlockSize / kVectorSize);
 constexpr int kLocalScaleOffset = kCodeOffset + kCodeBytes;
 constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2;
-constexpr float kMaxMagnitude = 43.0f;
-constexpr float kMaxLocalScale = 31.0f / 8.0f;
-constexpr float kNativeMax = kMaxMagnitude * kMaxLocalScale; // 166.625.
-constexpr float kPeakToRmsSlope = 0.035f;
-constexpr float kMinScaleAnchor = 0.65f;
-constexpr float kMaxScaleAnchor = 0.92f;
 
 template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
-  return static_cast<float>(*input);
+  const float value = static_cast<float>(*input);
+  return isfinite(value) ? value : 0.0f;
 }
 
 __device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm, float scale) {
@@ -71,38 +66,9 @@ __device__ __forceinline__ float even_parity_dot(const float *x, const float *q,
   return odd_parity ? dot - 2.0f * weakest : dot;
 }
 
-template <typename scalar_t>
-__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) {
-  const int64_t block = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-  if (block >= num_blocks)
-    return;
-
-  float amax = 0.0f;
-  float sumsq = 0.0f;
-  const scalar_t *values = input + block * kBlockSize;
-#pragma unroll 1
-  for (int i = 0; i < kBlockSize; ++i) {
-    const float value = load_float(values + i);
-    amax = fmaxf(amax, fabsf(value));
-    sumsq = fmaf(value, value, sumsq);
-  }
-  if (amax == 0.0f) {
-    scale_bits[block] = 0;
-    return;
-  }
-  const float rms = sqrtf(sumsq / kBlockSize);
-  const float peak_to_rms = rms > 0.0f ? amax / rms : 0.0f;
-  // Match the reference encoder's empirical predictor. Peaky blocks get a smaller anchor so
-  // outliers do not set the entire scale, while the clamp bounds the adjustment.
-  const float anchor =
-      fminf(kMaxScaleAnchor, fmaxf(kMinScaleAnchor, 1.0f - kPeakToRmsSlope * peak_to_rms));
-  const __half scale = __float2half_rn(fminf((amax / kNativeMax) * anchor, 65504.0f));
-  scale_bits[block] = static_cast<int64_t>(__half_as_ushort(scale));
-}
-
 template <typename scalar_t>
 __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid,
-                       const int64_t *scale_bits, uint8_t *output) {
+                       const __half *scales, uint8_t *output) {
   __shared__ float shared_grid[kEntries * kVectorSize];
   __shared__ float grid_norm[kEntries];
   __shared__ float warp_best[8 * kLocalScales];
@@ -133,8 +99,9 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     return;
   const scalar_t *source = input + block * kBlockSize;
   uint8_t *payload = output + block * kPayloadBytes;
-  const uint16_t d_bits = static_cast<uint16_t>(scale_bits[block]);
-  const float d = __half2float(__ushort_as_half(d_bits));
+  const __half d_half = scales[block];
+  const uint16_t d_bits = __half_as_ushort(d_half);
+  const float d = __half2float(d_half);
   if (d_bits == 0) {
     if (tid < kPayloadBytes)
       payload[tid] = 0;
@@ -284,8 +251,9 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
 
 } // namespace
 
-at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) {
-  TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous");
+at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) {
+  TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(),
+              "inputs must be contiguous");
   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,
@@ -297,23 +265,23 @@ at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) {
   TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
                   grid.size(1) == kVectorSize,
               "grid must be float32 [512, 8]");
-  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
-  c10::cuda::CUDAGuard guard(input.device());
   const int64_t num_blocks = input.numel() / kBlockSize;
+  TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
+                  scales.numel() == num_blocks,
+              "scales must be float16 [numel / 256]");
+  TORCH_CHECK(input.get_device() == grid.get_device() && input.get_device() == scales.get_device(),
+              "input, grid, and scales must share a device");
+  c10::cuda::CUDAGuard guard(input.device());
   TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ2_XS CUDA grid is too large");
-  auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong));
   auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte));
   const auto stream = c10::cuda::getCurrentCUDAStream();
-  const int scale_grid = static_cast<int>((num_blocks + 255) / 256);
 
   AT_DISPATCH_FLOATING_TYPES_AND2(
       at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] {
-        find_scale<scalar_t><<<scale_grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(), num_blocks,
-                                                             scales.data_ptr<int64_t>());
-        C10_CUDA_KERNEL_LAUNCH_CHECK();
         encode<scalar_t><<<static_cast<int>(num_blocks), 256, 0, stream>>>(
             input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
-            scales.data_ptr<int64_t>(), output.data_ptr<uint8_t>());
+            reinterpret_cast<const __half *>(scales.data_ptr<at::Half>()),
+            output.data_ptr<uint8_t>());
         C10_CUDA_KERNEL_LAUNCH_CHECK();
       });
   return output;
diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index 4641f2e202b..e1d1c47b463 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -45,38 +45,55 @@ def test_cuda_ext_iq2_xs():
 
 
 _IQ_EXTENSIONS = (
-    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), 50, id="iq1_s"),
-    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), 74, id="iq2_xs"),
+    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), 50, False, id="iq1_s"),
+    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), 74, True, id="iq2_xs"),
 )
 
 
-@pytest.mark.parametrize(("get_extension", "grid_shape", "payload_bytes"), _IQ_EXTENSIONS)
-def test_cuda_ext_iq_zero_block_layout(get_extension, grid_shape, payload_bytes):
+def _pack(extension, weight, grid, needs_scales):
+    if needs_scales:
+        scales = torch.zeros(weight.numel() // 256, device=weight.device, dtype=torch.float16)
+        return extension.pack(weight, grid, scales)
+    return extension.pack(weight, grid)
+
+
+@pytest.mark.parametrize(
+    ("get_extension", "grid_shape", "payload_bytes", "needs_scales"), _IQ_EXTENSIONS
+)
+def test_cuda_ext_iq_zero_block_layout(get_extension, grid_shape, payload_bytes, needs_scales):
     extension = get_extension(raise_if_failed=True)
     weight = torch.zeros((2, 256), device="cuda", dtype=torch.bfloat16)
     grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
-    packed = extension.pack(weight, grid)
+    packed = _pack(extension, weight, grid, needs_scales)
 
     assert packed.shape == (2, payload_bytes)
     assert not packed.any()
 
 
-@pytest.mark.parametrize(("get_extension", "grid_shape", "_payload_bytes"), _IQ_EXTENSIONS)
-def test_cuda_ext_iq_rejects_unsupported_dtype(get_extension, grid_shape, _payload_bytes):
+@pytest.mark.parametrize(
+    ("get_extension", "grid_shape", "_payload_bytes", "needs_scales"), _IQ_EXTENSIONS
+)
+def test_cuda_ext_iq_rejects_unsupported_dtype(
+    get_extension, grid_shape, _payload_bytes, needs_scales
+):
     extension = get_extension(raise_if_failed=True)
     weight = torch.ones((1, 256), device="cuda").to(torch.float8_e4m3fn)
     grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
     with pytest.raises(RuntimeError, match="supports float32, float64, float16, and bfloat16"):
-        extension.pack(weight, grid)
+        _pack(extension, weight, grid, needs_scales)
 
 
-@pytest.mark.parametrize(("get_extension", "grid_shape", "_payload_bytes"), _IQ_EXTENSIONS)
-def test_cuda_ext_iq_rejects_row_straddling_input(get_extension, grid_shape, _payload_bytes):
+@pytest.mark.parametrize(
+    ("get_extension", "grid_shape", "_payload_bytes", "needs_scales"), _IQ_EXTENSIONS
+)
+def test_cuda_ext_iq_rejects_row_straddling_input(
+    get_extension, grid_shape, _payload_bytes, needs_scales
+):
     extension = get_extension(raise_if_failed=True)
     weight = torch.ones((512, 384), device="cuda", dtype=torch.bfloat16)
     grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
     with pytest.raises(RuntimeError, match="innermost dimension must be a multiple of 256"):
-        extension.pack(weight, grid)
+        _pack(extension, weight, grid, needs_scales)

From 74240e17454159b39f9297c5d321243ba6fb56e8 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 17 Sep 2026 06:20:10 +0000
Subject: [PATCH 08/11] Share IQ CUDA packing helpers and harden the input
 contract

Both IQ packers carried near-identical device code that had already started
to drift, so extract it into a shared ggml/common.cuh: the include block, the
block geometry and codebook sizes, load_float, the clamped squared-error form,
the (error << 32) | entry search key, the warp-then-block reductions, and the
fp16 scale / zero-payload write. The reductions now fold over a kWarps derived
from one kThreads constant instead of a literal 8 repeated in four places, and
each kernel static_asserts that its codebook divides evenly among the threads
and is a power of two, so the index masks follow kEntries instead of repeating
it as 0x7ff / 0x1ff.

The validation duplicated between each pybind wrapper and its CUDA entry point
becomes one check_pack_inputs, so the enforced rule and the message it reports
are written once instead of as literals in the .cpp and named constants in the
.cu.

Fixes found while consolidating:

- load_float tested finiteness after narrowing to float32, so a finite float64
  such as 1e100 became inf and was dropped to zero. Finiteness is now tested at
  the source precision and finite out-of-range values saturate.
- IQ2_XS accepted non-finite caller scales and copied their bits straight into
  the GGML block scale field; they are now rejected before dispatch.
- IQ2_XS returned on its out-of-range block guard after two __syncthreads()
  used to stage the codebook; the guard moves above the staging.

Tests gain a non-zero case per format that actually runs the encode loop --
every existing passing case packed zeros and took the early zero-payload
return -- asserting the fp16 block scale, that the codebook indices and signs
past it are written, and that two different blocks encode differently.

Both .cu files and the shared header carry the GGML MIT notice below the
NVIDIA Apache-2.0 header, matching how gemm/fp8_kernel.py attributes DeepSeek,
since the packed block layouts and format constants are GGML's.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../kernels/quantization/ggml/common.cuh      | 207 ++++++++++++++++++
 .../torch/kernels/quantization/ggml/iq1_s.cpp |  22 +-
 .../torch/kernels/quantization/ggml/iq1_s.cu  | 151 +++++--------
 .../kernels/quantization/ggml/iq2_xs.cpp      |  34 ++-
 .../torch/kernels/quantization/ggml/iq2_xs.cu | 170 ++++++--------
 pyproject.toml                                |   2 +-
 .../gpu/_extensions/test_torch_extensions.py  | 122 ++++++++---
 7 files changed, 436 insertions(+), 272 deletions(-)
 create mode 100644 modelopt/torch/kernels/quantization/ggml/common.cuh

diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh
new file mode 100644
index 00000000000..ed330cef456
--- /dev/null
+++ b/modelopt/torch/kernels/quantization/ggml/common.cuh
@@ -0,0 +1,207 @@
+/*
+ * 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 packed block layouts and format constants used by this header and by the IQ1_S / IQ2_XS
+ * packers beside it are defined by GGML, pinned at
+ * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
+ *
+ * MIT License
+ *
+ * Copyright (c) 2023-2026 The ggml authors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#pragma once
+
+#include <ATen/ATen.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <limits>
+
+#ifdef __CUDACC__
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_fp16.h>
+
+#include <cfloat>
+#endif
+
+namespace modelopt::ggml {
+
+// Block geometry shared by every IQ format: 256 values are encoded as 8-element codebook vectors
+// behind one fp16 block scale that occupies the first two payload bytes.
+constexpr int kBlockSize = 256;
+constexpr int kVectorSize = 8;
+constexpr int kScaleOffset = 0;
+constexpr int kScaleBytes = 2;
+
+// Codebook sizes. Defined here so the pybind wrappers that validate them and the kernels that
+// index with them cannot drift apart.
+constexpr int kIq1sEntries = 2048;
+constexpr int kIq2xsEntries = 512;
+
+// One CUDA block encodes one GGML block. The reductions below fold over exactly this many warps,
+// and each kernel static_asserts that its codebook divides evenly among the threads.
+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) {
+  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,
+              " so blocks do not straddle rows");
+  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__
+
+// Reads one input element as float32. Non-finite elements are treated as zero, and finiteness is
+// tested at the source precision so that a finite float64 such as 1e100 saturates at the float32
+// maximum instead of overflowing to infinity and being dropped to zero.
+template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
+  if constexpr (sizeof(scalar_t) > sizeof(float)) {
+    constexpr double kFloatMax = static_cast<double>(FLT_MAX);
+    const double value = static_cast<double>(*input);
+    if (!isfinite(value))
+      return 0.0f;
+    return static_cast<float>(fmin(fmax(value, -kFloatMax), kFloatMax));
+  } else {
+    const float value = static_cast<float>(*input);
+    return isfinite(value) ? value : 0.0f;
+  }
+}
+
+// Squared error of approximating x by scale * q, given |x|^2, x . q and |q|^2. The clamp keeps the
+// result non-negative so that its bit pattern orders the same way the value does inside error_key.
+__device__ __forceinline__ float clamped_quant_error(float xnorm, float dot, float qnorm,
+                                                     float scale) {
+  return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f);
+}
+
+// Orders candidates by error first and codebook index second, so the lowest index wins a tie --
+// the rule the PyTorch reference encoder applies.
+__device__ __forceinline__ unsigned long long error_key(float error, int entry) {
+  return (static_cast<unsigned long long>(__float_as_uint(error)) << 32) |
+         static_cast<unsigned long long>(entry);
+}
+
+// Adds the block-wide minimum of local[slot] to accum[slot] for every slot. scratch must hold
+// kWarps * kSlots floats and accum kSlots floats. Barriers are internal, so every thread of the
+// block must call this.
+template <int kSlots>
+__device__ __forceinline__ void block_min_accumulate(const float (&local)[kSlots], float *scratch,
+                                                     float *accum) {
+  const int tid = threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+#pragma unroll
+  for (int slot = 0; slot < kSlots; ++slot) {
+    float value = local[slot];
+#pragma unroll
+    for (int delta = 16; delta > 0; delta >>= 1)
+      value = fminf(value, __shfl_down_sync(0xffffffff, value, delta));
+    if (lane == 0)
+      scratch[warp * kSlots + slot] = value;
+  }
+  __syncthreads();
+  if (tid < kSlots) {
+    float value = scratch[tid];
+#pragma unroll
+    for (int w = 1; w < kWarps; ++w)
+      value = fminf(value, scratch[w * kSlots + tid]);
+    accum[tid] += value;
+  }
+  __syncthreads();
+}
+
+// Block-wide minimum of key, valid on thread 0 only. scratch must hold kWarps entries. Barriers
+// are internal, so every thread of the block must call this.
+__device__ __forceinline__ unsigned long long block_min_key(unsigned long long key,
+                                                            unsigned long long *scratch) {
+  const int tid = threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+#pragma unroll
+  for (int delta = 16; delta > 0; delta >>= 1) {
+    const unsigned long long other = __shfl_down_sync(0xffffffff, key, delta);
+    key = other < key ? other : key;
+  }
+  if (lane == 0)
+    scratch[warp] = key;
+  __syncthreads();
+  if (tid == 0) {
+#pragma unroll
+    for (int w = 1; w < kWarps; ++w)
+      key = scratch[w] < key ? scratch[w] : key;
+  }
+  return key;
+}
+
+// Writes the fp16 block scale into the payload, or zeroes the whole payload when the block scale
+// rounded to zero. Returns false once the payload is final and the caller should stop. The branch
+// is uniform across the block, so returning on false is barrier-safe.
+template <int kPayloadBytes>
+__device__ __forceinline__ bool store_block_scale(uint8_t *payload, uint16_t d_bits) {
+  if (d_bits == 0) {
+    if (threadIdx.x < kPayloadBytes)
+      payload[threadIdx.x] = 0;
+    return false;
+  }
+  if (threadIdx.x == 0) {
+    payload[kScaleOffset] = static_cast<uint8_t>(d_bits);
+    payload[kScaleOffset + 1] = static_cast<uint8_t>(d_bits >> 8);
+  }
+  return true;
+}
+
+#endif // __CUDACC__
+
+} // namespace modelopt::ggml
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
index adecc464aa3..dbc187d295d 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp
@@ -15,30 +15,14 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <torch/extension.h>
-
-#include <limits>
+#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");
-  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,
-              "IQ1_S 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) % 256 == 0,
-              "input's innermost dimension must be a multiple of 256 so blocks do not straddle "
-              "rows");
-  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == 2048 &&
-                  grid.size(1) == 8,
-              "grid must be float32 [2048, 8]");
-  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
-  TORCH_CHECK(input.numel() / 256 <= std::numeric_limits<int>::max(),
-              "IQ1_S CUDA grid is too large");
+  modelopt::ggml::check_pack_inputs("IQ1_S", input, grid, modelopt::ggml::kIq1sEntries);
   return iq1_s_pack_cuda(input.contiguous(), grid.contiguous());
 }
 
@@ -47,5 +31,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
              "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.");
+             "zero during packing, and finite elements outside the float32 range saturate.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index 883b1435885..165aeebee37 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -15,30 +15,45 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/extension.h>
-
-#include <cuda_fp16.h>
+/*
+ * The IQ1_S packed block layout and format constants implemented here are defined by GGML, pinned
+ * at
+ * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
+ *
+ * MIT License
+ *
+ * Copyright (c) 2023-2026 The ggml authors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
 
-#include <cfloat>
-#include <cstdint>
-#include <limits>
+#include "common.cuh"
 
 namespace {
 
-// Packed layout and format constants follow the GGML definition at:
-// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
-constexpr int kBlockSize = 256;
-constexpr int kVectorSize = 8;
-constexpr int kEntries = 2048;
+using namespace modelopt::ggml;
+
+constexpr int kEntries = kIq1sEntries;
 constexpr int kGroups = 8;
+constexpr int kVectorsPerGroup = 4;
 constexpr int kChoices = 16;
-constexpr int kScaleOffset = 0;
-constexpr int kIndexOffset = 2;
+constexpr int kIndexOffset = kScaleBytes;
 constexpr int kIndexBytes = kBlockSize / kVectorSize;
 constexpr int kMetadataOffset = kIndexOffset + kIndexBytes;
 constexpr int kPayloadBytes = kMetadataOffset + 2 * kGroups;
@@ -48,10 +63,8 @@ constexpr float kMaxShiftedMagnitude = 1.0f + kDelta;
 constexpr float kNativeMax = kMaxLocalScale * kMaxShiftedMagnitude; // 16.875.
 constexpr float kScaleAnchor = 0.61f;
 
-template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
-  const float value = static_cast<float>(*input);
-  return isfinite(value) ? value : 0.0f;
-}
+static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries");
+static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two");
 
 __device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x,
                                              const float *q, float scale, float delta) {
@@ -66,7 +79,7 @@ __device__ __forceinline__ float quant_error(float xnorm, float xsum, const floa
   }
   const float shifted_dot = dot + delta * xsum;
   const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta;
-  return fmaxf(fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f);
+  return clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale);
 }
 
 template <typename scalar_t>
@@ -89,15 +102,13 @@ __global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *s
 template <typename scalar_t>
 __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid,
                        const int64_t *scale_bits, uint8_t *output) {
-  __shared__ float warp_best[8 * kChoices];
+  __shared__ float warp_best[kWarps * kChoices];
   __shared__ float group_error[kChoices];
-  __shared__ unsigned long long warp_keys[8];
+  __shared__ unsigned long long warp_keys[kWarps];
   __shared__ int selected_choice;
-  __shared__ uint16_t selected_entries[4];
+  __shared__ uint16_t selected_entries[kVectorsPerGroup];
 
   const int tid = threadIdx.x;
-  const int lane = tid & 31;
-  const int warp = tid >> 5;
   const int64_t block = blockIdx.x;
   if (block >= num_blocks)
     return;
@@ -106,15 +117,8 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
   uint8_t *payload = output + block * kPayloadBytes;
   const uint16_t d_bits = static_cast<uint16_t>(scale_bits[block]);
   const float d = __half2float(__ushort_as_half(d_bits));
-  if (d_bits == 0) {
-    if (tid < kPayloadBytes)
-      payload[tid] = 0;
+  if (!store_block_scale<kPayloadBytes>(payload, d_bits))
     return;
-  }
-  if (tid == 0) {
-    payload[kScaleOffset] = static_cast<uint8_t>(d_bits);
-    payload[kScaleOffset + 1] = static_cast<uint8_t>(d_bits >> 8);
-  }
 
 #pragma unroll 1
   for (int group = 0; group < kGroups; ++group) {
@@ -123,11 +127,11 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     __syncthreads();
 
 #pragma unroll
-    for (int vector = 0; vector < 4; ++vector) {
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
       float x[kVectorSize];
       float xnorm = 0.0f;
       float xsum = 0.0f;
-      const int offset = group * 32 + vector * 8;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
 #pragma unroll
       for (int j = 0; j < kVectorSize; ++j) {
         x[j] = load_float(source + offset + j);
@@ -156,29 +160,11 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
           const float scale = d * (2 * local + 1);
           const float shifted_dot = dot + delta * xsum;
           const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta;
-          const float error = fmaxf(
-              fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f);
-          local_best[choice] = fminf(local_best[choice], error);
+          local_best[choice] = fminf(local_best[choice],
+                                     clamped_quant_error(xnorm, shifted_dot, shifted_norm, scale));
         }
       }
-#pragma unroll
-      for (int choice = 0; choice < kChoices; ++choice) {
-        float value = local_best[choice];
-#pragma unroll
-        for (int delta = 16; delta > 0; delta >>= 1)
-          value = fminf(value, __shfl_down_sync(0xffffffff, value, delta));
-        if (lane == 0)
-          warp_best[warp * kChoices + choice] = value;
-      }
-      __syncthreads();
-      if (tid < kChoices) {
-        float value = warp_best[tid];
-#pragma unroll
-        for (int w = 1; w < 8; ++w)
-          value = fminf(value, warp_best[w * kChoices + tid]);
-        group_error[tid] += value;
-      }
-      __syncthreads();
+      block_min_accumulate<kChoices>(local_best, warp_best, group_error);
     }
 
     if (tid == 0) {
@@ -198,11 +184,11 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     const float selected_scale = d * (2 * selected_local + 1);
 
 #pragma unroll
-    for (int vector = 0; vector < 4; ++vector) {
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
       float x[kVectorSize];
       float xnorm = 0.0f;
       float xsum = 0.0f;
-      const int offset = group * 32 + vector * 8;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
 #pragma unroll
       for (int j = 0; j < kVectorSize; ++j) {
         x[j] = load_float(source + offset + j);
@@ -213,27 +199,14 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
       for (int entry = tid; entry < kEntries; entry += blockDim.x) {
         const float error =
             quant_error(xnorm, xsum, x, grid + entry * kVectorSize, selected_scale, selected_delta);
-        const unsigned long long candidate =
-            (static_cast<unsigned long long>(__float_as_uint(error)) << 32) |
-            static_cast<unsigned long long>(entry);
+        const unsigned long long candidate = error_key(error, entry);
         key = candidate < key ? candidate : key;
       }
-#pragma unroll
-      for (int delta = 16; delta > 0; delta >>= 1) {
-        const auto other = __shfl_down_sync(0xffffffff, key, delta);
-        key = other < key ? other : key;
-      }
-      if (lane == 0)
-        warp_keys[warp] = key;
-      __syncthreads();
+      key = block_min_key(key, warp_keys);
       if (tid == 0) {
-        key = warp_keys[0];
-#pragma unroll
-        for (int w = 1; w < 8; ++w)
-          key = warp_keys[w] < key ? warp_keys[w] : key;
-        const uint16_t entry = static_cast<uint16_t>(key & 0x7ff);
+        const uint16_t entry = static_cast<uint16_t>(key & (kEntries - 1));
         selected_entries[vector] = entry;
-        payload[kIndexOffset + group * 4 + vector] = static_cast<uint8_t>(entry);
+        payload[kIndexOffset + group * kVectorsPerGroup + vector] = static_cast<uint8_t>(entry);
       }
       __syncthreads();
     }
@@ -254,32 +227,20 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
 
 at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) {
   TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous");
-  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,
-              "IQ1_S 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 256 so blocks do not straddle "
-              "rows");
-  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
-                  grid.size(1) == kVectorSize,
-              "grid must be float32 [2048, 8]");
-  TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
+  check_pack_inputs("IQ1_S", input, grid, kEntries);
   c10::cuda::CUDAGuard guard(input.device());
   const int64_t num_blocks = input.numel() / kBlockSize;
-  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ1_S CUDA grid is too large");
   auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong));
   auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte));
   const auto stream = c10::cuda::getCurrentCUDAStream();
-  const int scale_grid = static_cast<int>((num_blocks + 255) / 256);
+  const int scale_grid = static_cast<int>((num_blocks + kThreads - 1) / kThreads);
 
   AT_DISPATCH_FLOATING_TYPES_AND2(
       at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_s_pack", [&] {
-        find_scale<scalar_t><<<scale_grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(), num_blocks,
-                                                             scales.data_ptr<int64_t>());
+        find_scale<scalar_t><<<scale_grid, kThreads, 0, stream>>>(
+            input.data_ptr<scalar_t>(), num_blocks, scales.data_ptr<int64_t>());
         C10_CUDA_KERNEL_LAUNCH_CHECK();
-        encode<scalar_t><<<static_cast<int>(num_blocks), 256, 0, stream>>>(
+        encode<scalar_t><<<static_cast<int>(num_blocks), kThreads, 0, stream>>>(
             input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
             scales.data_ptr<int64_t>(), output.data_ptr<uint8_t>());
         C10_CUDA_KERNEL_LAUNCH_CHECK();
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
index 007b4f96fb7..755226373fa 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
@@ -15,9 +15,7 @@
  * limitations under the License.
  */
 
-#include <torch/extension.h>
-
-#include <limits>
+#include "common.cuh"
 
 at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales);
 
@@ -25,31 +23,25 @@ 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");
-  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,
-              "IQ2_XS 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) % 256 == 0,
-              "input's innermost dimension must be a multiple of 256 so blocks do not straddle "
-              "rows");
-  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == 512 &&
-                  grid.size(1) == 8,
-              "grid must be float32 [512, 8]");
-  const auto num_blocks = input.numel() / 256;
+  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]");
-  TORCH_CHECK(input.get_device() == grid.get_device() && input.get_device() == scales.get_device(),
-              "input, grid, and scales must share a device");
-  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ2_XS CUDA grid is too large");
+  // The kernel copies these bits straight into the GGML block scale field, so a non-finite entry
+  // would produce a payload that decodes to garbage. The synchronization this costs is paid once
+  // per packed tensor, on an export path.
+  TORCH_CHECK(scales.isfinite().all().item<bool>(), "scales must be finite");
+  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());
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("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], and scales must "
-             "be float16 [numel / 256]. Returns uint8 [numel / 256, 74] on the input device. "
-             "Non-finite input elements are treated as zero during packing.");
+             "dimension is a multiple of 256. The grid must be float32 [512, 8] holding "
+             "non-negative codebook magnitudes, and scales must be finite 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.");
 }
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index e60c849a67b..c66c118bc2a 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -15,46 +15,58 @@
  * limitations under the License.
  */
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/extension.h>
-
-#include <cuda_fp16.h>
+/*
+ * The IQ2_XS packed block layout and format constants implemented here are defined by GGML, pinned
+ * at
+ * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
+ *
+ * MIT License
+ *
+ * Copyright (c) 2023-2026 The ggml authors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
 
-#include <cfloat>
-#include <cstdint>
-#include <limits>
+#include "common.cuh"
 
 namespace {
 
-// Packed layout and format constants follow the GGML definition at:
-// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
-constexpr int kBlockSize = 256;
-constexpr int kVectorSize = 8;
-constexpr int kEntries = 512;
+using namespace modelopt::ggml;
+
+constexpr int kEntries = kIq2xsEntries;
 constexpr int kGroups = 16;
+constexpr int kVectorsPerGroup = 2;
 constexpr int kLocalScales = 16;
-constexpr int kScaleOffset = 0;
-constexpr int kCodeOffset = 2;
+constexpr int kCodeOffset = kScaleBytes;
 constexpr int kCodeBytes = 2 * (kBlockSize / kVectorSize);
 constexpr int kLocalScaleOffset = kCodeOffset + kCodeBytes;
 constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2;
+constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8.
 
-template <typename scalar_t> __device__ __forceinline__ float load_float(const scalar_t *input) {
-  const float value = static_cast<float>(*input);
-  return isfinite(value) ? value : 0.0f;
-}
-
-__device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm, float scale) {
-  return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f);
-}
+static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries");
+static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two");
 
+// Dot product of |x| against one codebook vector, under the format's even-parity sign rule. The
+// grid must hold non-negative magnitudes: the signs live in the packed 7-bit field, and the
+// eighth sign is recovered from the parity of the other seven during decoding. For odd parity,
+// flip the coordinate with the smallest |x| * q penalty.
 __device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) {
-  // The format stores seven sign bits. For odd parity, flip the coordinate with the smallest
-  // |x| * q penalty; the eighth sign is recovered from even parity during decoding.
   float dot = 0.0f;
   float weakest = FLT_MAX;
 #pragma unroll
@@ -71,15 +83,17 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
                        const __half *scales, uint8_t *output) {
   __shared__ float shared_grid[kEntries * kVectorSize];
   __shared__ float grid_norm[kEntries];
-  __shared__ float warp_best[8 * kLocalScales];
+  __shared__ float warp_best[kWarps * kLocalScales];
   __shared__ float group_error[kLocalScales];
-  __shared__ unsigned long long warp_keys[8];
+  __shared__ unsigned long long warp_keys[kWarps];
   __shared__ int selected_local;
   __shared__ uint8_t locals[kGroups];
 
   const int tid = threadIdx.x;
-  const int lane = tid & 31;
-  const int warp = tid >> 5;
+  const int64_t block = blockIdx.x;
+  if (block >= num_blocks)
+    return;
+
   for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x)
     shared_grid[i] = grid[i];
   __syncthreads();
@@ -94,23 +108,13 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
   }
   __syncthreads();
 
-  const int64_t block = blockIdx.x;
-  if (block >= num_blocks)
-    return;
   const scalar_t *source = input + block * kBlockSize;
   uint8_t *payload = output + block * kPayloadBytes;
   const __half d_half = scales[block];
   const uint16_t d_bits = __half_as_ushort(d_half);
   const float d = __half2float(d_half);
-  if (d_bits == 0) {
-    if (tid < kPayloadBytes)
-      payload[tid] = 0;
+  if (!store_block_scale<kPayloadBytes>(payload, d_bits))
     return;
-  }
-  if (tid == 0) {
-    payload[kScaleOffset] = static_cast<uint8_t>(d_bits);
-    payload[kScaleOffset + 1] = static_cast<uint8_t>(d_bits >> 8);
-  }
 
 #pragma unroll 1
   for (int group = 0; group < kGroups; ++group) {
@@ -119,11 +123,11 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     __syncthreads();
 
 #pragma unroll
-    for (int vector = 0; vector < 2; ++vector) {
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
       float x[kVectorSize];
       float xnorm = 0.0f;
       int negative_count = 0;
-      const int offset = group * 16 + vector * 8;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
 #pragma unroll
       for (int j = 0; j < kVectorSize; ++j) {
         x[j] = load_float(source + offset + j);
@@ -140,29 +144,12 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
         const float dot = even_parity_dot(x, q, odd_parity);
 #pragma unroll
         for (int local = 0; local < kLocalScales; ++local) {
-          const float scale = d * (2 * local + 1) * 0.125f;
+          const float scale = d * (2 * local + 1) * kLocalScaleStep;
           local_best[local] =
-              fminf(local_best[local], quant_error(xnorm, dot, grid_norm[entry], scale));
+              fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale));
         }
       }
-#pragma unroll
-      for (int local = 0; local < kLocalScales; ++local) {
-        float value = local_best[local];
-#pragma unroll
-        for (int delta = 16; delta > 0; delta >>= 1)
-          value = fminf(value, __shfl_down_sync(0xffffffff, value, delta));
-        if (lane == 0)
-          warp_best[warp * kLocalScales + local] = value;
-      }
-      __syncthreads();
-      if (tid < kLocalScales) {
-        float value = warp_best[tid];
-#pragma unroll
-        for (int w = 1; w < 8; ++w)
-          value = fminf(value, warp_best[w * kLocalScales + tid]);
-        group_error[tid] += value;
-      }
-      __syncthreads();
+      block_min_accumulate<kLocalScales>(local_best, warp_best, group_error);
     }
 
     if (tid == 0) {
@@ -178,14 +165,14 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
       locals[group] = static_cast<uint8_t>(selected_local);
     }
     __syncthreads();
-    const float selected_scale = d * (2 * selected_local + 1) * 0.125f;
+    const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep;
 
 #pragma unroll
-    for (int vector = 0; vector < 2; ++vector) {
+    for (int vector = 0; vector < kVectorsPerGroup; ++vector) {
       float x[kVectorSize];
       float xnorm = 0.0f;
       int negative_count = 0;
-      const int offset = group * 16 + vector * 8;
+      const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize;
 #pragma unroll
       for (int j = 0; j < kVectorSize; ++j) {
         x[j] = load_float(source + offset + j);
@@ -195,28 +182,15 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
       const bool odd_parity = (negative_count & 1) != 0;
       unsigned long long key = ~0ULL;
       for (int entry = tid; entry < kEntries; entry += blockDim.x) {
-        const float error =
-            quant_error(xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity),
-                        grid_norm[entry], selected_scale);
-        const unsigned long long candidate =
-            (static_cast<unsigned long long>(__float_as_uint(error)) << 32) |
-            static_cast<unsigned long long>(entry);
+        const float error = clamped_quant_error(
+            xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity),
+            grid_norm[entry], selected_scale);
+        const unsigned long long candidate = error_key(error, entry);
         key = candidate < key ? candidate : key;
       }
-#pragma unroll
-      for (int delta = 16; delta > 0; delta >>= 1) {
-        const auto other = __shfl_down_sync(0xffffffff, key, delta);
-        key = other < key ? other : key;
-      }
-      if (lane == 0)
-        warp_keys[warp] = key;
-      __syncthreads();
+      key = block_min_key(key, warp_keys);
       if (tid == 0) {
-        key = warp_keys[0];
-#pragma unroll
-        for (int w = 1; w < 8; ++w)
-          key = warp_keys[w] < key ? warp_keys[w] : key;
-        const int entry = static_cast<int>(key & 0x1ff);
+        const int entry = static_cast<int>(key & (kEntries - 1));
         const float *q = shared_grid + entry * kVectorSize;
         int flip_index = 0;
         float weakest = fabsf(x[0]) * q[0];
@@ -237,7 +211,7 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
           sign_mask |= static_cast<int>(is_negative) << j;
         }
         const uint16_t code = static_cast<uint16_t>(entry | ((sign_mask & 0x7f) << 9));
-        const int code_offset = kCodeOffset + 2 * (group * 2 + vector);
+        const int code_offset = kCodeOffset + 2 * (group * kVectorsPerGroup + vector);
         payload[code_offset] = static_cast<uint8_t>(code);
         payload[code_offset + 1] = static_cast<uint8_t>(code >> 8);
       }
@@ -245,7 +219,7 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
     }
   }
 
-  if (tid < 8)
+  if (tid < kGroups / 2)
     payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4);
 }
 
@@ -254,31 +228,19 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
 at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) {
   TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(),
               "inputs must be contiguous");
-  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,
-              "IQ2_XS 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 256 so blocks do not straddle "
-              "rows");
-  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
-                  grid.size(1) == kVectorSize,
-              "grid must be float32 [512, 8]");
+  check_pack_inputs("IQ2_XS", input, grid, kEntries);
   const int64_t num_blocks = input.numel() / kBlockSize;
   TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
                   scales.numel() == num_blocks,
               "scales must be float16 [numel / 256]");
-  TORCH_CHECK(input.get_device() == grid.get_device() && input.get_device() == scales.get_device(),
-              "input, grid, and scales must share a device");
+  TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
   c10::cuda::CUDAGuard guard(input.device());
-  TORCH_CHECK(num_blocks <= std::numeric_limits<int>::max(), "IQ2_XS CUDA grid is too large");
   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(), "iq2_xs_pack", [&] {
-        encode<scalar_t><<<static_cast<int>(num_blocks), 256, 0, stream>>>(
+        encode<scalar_t><<<static_cast<int>(num_blocks), kThreads, 0, stream>>>(
             input.data_ptr<scalar_t>(), num_blocks, grid.data_ptr<float>(),
             reinterpret_cast<const __half *>(scales.data_ptr<at::Half>()),
             output.data_ptr<uint8_t>());
diff --git a/pyproject.toml b/pyproject.toml
index e27377033d6..0886e0380b3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -149,7 +149,7 @@ Homepage = "https://github.com/NVIDIA/Model-Optimizer"
 include = ["modelopt*"]
 
 [tool.setuptools.package-data]
-modelopt = ["**/*.h", "**/*.cpp", "**/*.cu"]
+modelopt = ["**/*.cpp", "**/*.cu", "**/*.cuh", "**/*.h"]
 modelopt_recipes = ["**/*.yml", "**/*.yaml"]
 
 [tool.setuptools.exclude-package-data]
diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index e1d1c47b463..38f919763b2 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -14,6 +14,9 @@
 # limitations under the License.
 
 
+from collections.abc import Callable
+from typing import NamedTuple
+
 import pytest
 import torch
 
@@ -44,56 +47,111 @@ def test_cuda_ext_iq2_xs():
     assert ext.get_cuda_ext_iq2_xs() is not None
 
 
+def _generator():
+    """Seeded generator so a failure reproduces exactly."""
+    return torch.Generator(device="cuda").manual_seed(0)
+
+
+class _IqFormat(NamedTuple):
+    """One GGML IQ packing extension and the format constants its contract is defined by."""
+
+    get_extension: Callable
+    entries: int
+    payload_bytes: int
+    needs_scales: bool
+    # Value alphabet the codebook is built from: signed ternary for IQ1_S, and the non-negative
+    # magnitudes IQ2_XS stores (its signs live in the packed code).
+    grid_values: tuple[float, ...]
+    # Largest magnitude the format can represent at a block scale of 1.
+    native_max: float
+
+
 _IQ_EXTENSIONS = (
-    pytest.param(ext.get_cuda_ext_iq1_s, (2048, 8), 50, False, id="iq1_s"),
-    pytest.param(ext.get_cuda_ext_iq2_xs, (512, 8), 74, True, id="iq2_xs"),
+    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, (1.0, 8.0, 25.0, 43.0), 166.625),
+        id="iq2_xs",
+    ),
 )
 
 
-def _pack(extension, weight, grid, needs_scales):
-    if needs_scales:
+def _grid(fmt: _IqFormat, zero: bool = False) -> torch.Tensor:
+    """Codebook of ``fmt.entries`` distinct vectors drawn from the format's value alphabet."""
+    if zero:
+        return torch.zeros((fmt.entries, 8), device="cuda", dtype=torch.float32)
+    values = torch.tensor(fmt.grid_values, device="cuda", dtype=torch.float32)
+    digits = torch.arange(fmt.entries, device="cuda").unsqueeze(1) // len(fmt.grid_values) ** (
+        torch.arange(8, device="cuda")
+    )
+    return values[digits % len(fmt.grid_values)]
+
+
+def _pack(fmt: _IqFormat, extension, weight, grid, scales=None):
+    if not fmt.needs_scales:
+        return extension.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 extension.pack(weight, grid)
+    return extension.pack(weight, grid, scales)
 
 
-@pytest.mark.parametrize(
-    ("get_extension", "grid_shape", "payload_bytes", "needs_scales"), _IQ_EXTENSIONS
-)
-def test_cuda_ext_iq_zero_block_layout(get_extension, grid_shape, payload_bytes, needs_scales):
-    extension = get_extension(raise_if_failed=True)
+@pytest.mark.parametrize("fmt", _IQ_EXTENSIONS)
+def test_cuda_ext_iq_zero_block_layout(fmt):
+    extension = fmt.get_extension(raise_if_failed=True)
     weight = torch.zeros((2, 256), device="cuda", dtype=torch.bfloat16)
-    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
-    packed = _pack(extension, weight, grid, needs_scales)
+    packed = _pack(fmt, extension, weight, _grid(fmt, zero=True))
 
-    assert packed.shape == (2, payload_bytes)
+    assert packed.shape == (2, fmt.payload_bytes)
     assert not packed.any()
 
 
-@pytest.mark.parametrize(
-    ("get_extension", "grid_shape", "_payload_bytes", "needs_scales"), _IQ_EXTENSIONS
-)
-def test_cuda_ext_iq_rejects_unsupported_dtype(
-    get_extension, grid_shape, _payload_bytes, needs_scales
-):
-    extension = get_extension(raise_if_failed=True)
+@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)
+    weight = torch.randn((2, 256), device="cuda", dtype=torch.bfloat16, generator=_generator())
+    scales = (weight.float().abs().amax(dim=-1) / fmt.native_max).half()
+
+    packed = _pack(fmt, extension, weight, _grid(fmt), scales=scales)
+
+    assert packed.shape == (2, fmt.payload_bytes)
+    # The fp16 block scale lands in the first two payload bytes, and the caller supplies it
+    # verbatim for IQ2_XS.
+    block_scale = packed[:, :2].contiguous().view(torch.float16).flatten()
+    assert (block_scale > 0).all()
+    if fmt.needs_scales:
+        assert torch.equal(block_scale, scales)
+    # Codebook indices, signs, and local scales are written past the block scale, and two
+    # different blocks must not encode identically.
+    assert packed[:, 2:].any()
+    assert not torch.equal(packed[0], packed[1])
+
+
+@pytest.mark.parametrize("fmt", _IQ_EXTENSIONS)
+def test_cuda_ext_iq_rejects_unsupported_dtype(fmt):
+    extension = fmt.get_extension(raise_if_failed=True)
     weight = torch.ones((1, 256), device="cuda").to(torch.float8_e4m3fn)
-    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
     with pytest.raises(RuntimeError, match="supports float32, float64, float16, and bfloat16"):
-        _pack(extension, weight, grid, needs_scales)
+        _pack(fmt, extension, weight, _grid(fmt, zero=True))
 
 
-@pytest.mark.parametrize(
-    ("get_extension", "grid_shape", "_payload_bytes", "needs_scales"), _IQ_EXTENSIONS
-)
-def test_cuda_ext_iq_rejects_row_straddling_input(
-    get_extension, grid_shape, _payload_bytes, needs_scales
-):
-    extension = get_extension(raise_if_failed=True)
+@pytest.mark.parametrize("fmt", _IQ_EXTENSIONS)
+def test_cuda_ext_iq_rejects_row_straddling_input(fmt):
+    extension = fmt.get_extension(raise_if_failed=True)
     weight = torch.ones((512, 384), device="cuda", dtype=torch.bfloat16)
-    grid = torch.zeros(grid_shape, device="cuda", dtype=torch.float32)
 
     with pytest.raises(RuntimeError, match="innermost dimension must be a multiple of 256"):
-        _pack(extension, weight, grid, needs_scales)
+        _pack(fmt, extension, weight, _grid(fmt, zero=True))
+
+
+def test_cuda_ext_iq2_xs_rejects_non_finite_scales():
+    extension = ext.get_cuda_ext_iq2_xs(raise_if_failed=True)
+    fmt = _IQ_EXTENSIONS[1].values[0]
+    weight = torch.ones((1, 256), device="cuda", dtype=torch.bfloat16)
+    scales = torch.full((1,), float("nan"), device="cuda", dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="scales must be finite"):
+        extension.pack(weight, _grid(fmt), scales)

From 74e94db9601870e3569c7c8e73506a2f08c29da8 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 17 Sep 2026 06:40:48 +0000
Subject: [PATCH 09/11] Scope the GGML attribution to the constants it covers

Replace the file-level MIT notice added in the previous commit with a pointer
to the pinned ggml-common.h placed on the constants it actually covers: block
geometry and codebook sizes in the shared header, and the payload layout in
each kernel.

A comparison against ggml-quants.c and ggml-common.h at the pinned revision
found no shared source text -- no normalized source line over 25 characters
appears upstream, and none of upstream's encoder identifiers (kmap_q2xs,
kneighbors_q2xs, sumqx/sumq2, nearest_int, x_p/x_m, is_on_grid) appear here.
The encoders are different algorithms: upstream weights each value by
qw*sqrt(sigma2+x^2), sorts the block and exhaustively searches split
boundaries (IQ1_S) or sweeps scale perturbations with nearest_int (IQ2_XS),
both falling back to neighbour tables and refitting the scale by least
squares, while these kernels scan the whole codebook unweighted and reduce a
packed (error, entry) key. llama.cpp also has no GPU IQ encoder to derive
from: its CUDA quantize.cu only encodes activations to Q8_1/NVFP4/MXFP4, the
mmq-instance-iq*.cu files are matmul instantiations that consume packed data,
and the Vulkan dequant_iq*.comp shaders are decoders.

What is shared is the wire format -- payload sizes, field offsets, the qh bit
layout, IQ1S_DELTA, the 7-bit sign field, the scale nibbles -- which any GGUF
writer has to match. A copyright notice would be claiming ggml's copyright
over files that contain none of its expression, so the pinned pointer is the
accurate statement. Two convergences worth a human look remain: 1.125 appears
upstream as an unexplained fudge factor and here as kMaxLocalScale*(1+kDelta),
and both flip the cheapest coordinate to satisfy the even-parity sign rule,
though by different criteria (weight*x^2 upstream, |x|*q here).

Also correct the IQ2_XS codebook alphabet in the new test: iq2xs_grid bytes
are {0x08, 0x19, 0x2b}, with no 1 among them.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../kernels/quantization/ggml/common.cuh      | 37 +++----------------
 .../torch/kernels/quantization/ggml/iq1_s.cu  | 31 ++--------------
 .../torch/kernels/quantization/ggml/iq2_xs.cu | 31 ++--------------
 .../gpu/_extensions/test_torch_extensions.py  |  7 ++--
 4 files changed, 16 insertions(+), 90 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh
index ed330cef456..66314335010 100644
--- a/modelopt/torch/kernels/quantization/ggml/common.cuh
+++ b/modelopt/torch/kernels/quantization/ggml/common.cuh
@@ -15,34 +15,6 @@
  * limitations under the License.
  */
 
-/*
- * The packed block layouts and format constants used by this header and by the IQ1_S / IQ2_XS
- * packers beside it are defined by GGML, pinned at
- * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
- *
- * MIT License
- *
- * Copyright (c) 2023-2026 The ggml authors
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-
 #pragma once
 
 #include <ATen/ATen.h>
@@ -64,14 +36,17 @@
 namespace modelopt::ggml {
 
 // Block geometry shared by every IQ format: 256 values are encoded as 8-element codebook vectors
-// behind one fp16 block scale that occupies the first two payload bytes.
+// behind one fp16 block scale that occupies the first two payload bytes. These follow GGML's
+// QK_K, its uint64 grid entry width, and the leading ggml_half of each block struct:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kBlockSize = 256;
 constexpr int kVectorSize = 8;
 constexpr int kScaleOffset = 0;
 constexpr int kScaleBytes = 2;
 
-// Codebook sizes. Defined here so the pybind wrappers that validate them and the kernels that
-// index with them cannot drift apart.
+// Codebook sizes, from GGML's NGRID_IQ1S and the length of its iq2xs_grid table (see the link
+// above). Defined here so the pybind wrappers that validate them and the kernels that index with
+// them cannot drift apart.
 constexpr int kIq1sEntries = 2048;
 constexpr int kIq2xsEntries = 512;
 
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index 165aeebee37..37f4d1f95ab 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -15,40 +15,15 @@
  * limitations under the License.
  */
 
-/*
- * The IQ1_S packed block layout and format constants implemented here are defined by GGML, pinned
- * at
- * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
- *
- * MIT License
- *
- * Copyright (c) 2023-2026 The ggml authors
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-
 #include "common.cuh"
 
 namespace {
 
 using namespace modelopt::ggml;
 
+// The IQ1_S packed payload layout and format constants below follow the GGML
+// definition at:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kEntries = kIq1sEntries;
 constexpr int kGroups = 8;
 constexpr int kVectorsPerGroup = 4;
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index c66c118bc2a..c38537e4d98 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -15,40 +15,15 @@
  * limitations under the License.
  */
 
-/*
- * The IQ2_XS packed block layout and format constants implemented here are defined by GGML, pinned
- * at
- * https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
- *
- * MIT License
- *
- * Copyright (c) 2023-2026 The ggml authors
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-
 #include "common.cuh"
 
 namespace {
 
 using namespace modelopt::ggml;
 
+// The IQ2_XS packed payload layout and format constants below follow the GGML
+// definition at:
+// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h
 constexpr int kEntries = kIq2xsEntries;
 constexpr int kGroups = 16;
 constexpr int kVectorsPerGroup = 2;
diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index 38f919763b2..c033600c1e2 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -59,8 +59,9 @@ class _IqFormat(NamedTuple):
     entries: int
     payload_bytes: int
     needs_scales: bool
-    # Value alphabet the codebook is built from: signed ternary for IQ1_S, and the non-negative
-    # magnitudes IQ2_XS stores (its signs live in the packed code).
+    # Value alphabet the codebook is built from, as GGML defines it: signed ternary bytes
+    # {0x00, 0x01, 0xff} for IQ1_S, and the non-negative magnitudes {0x08, 0x19, 0x2b} for
+    # IQ2_XS, whose signs live in the packed code instead.
     grid_values: tuple[float, ...]
     # Largest magnitude the format can represent at a block scale of 1.
     native_max: float
@@ -71,7 +72,7 @@ class _IqFormat(NamedTuple):
         _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, (1.0, 8.0, 25.0, 43.0), 166.625),
+        _IqFormat(ext.get_cuda_ext_iq2_xs, 512, 74, True, (8.0, 25.0, 43.0), 166.625),
         id="iq2_xs",
     ),
 )

From bb35d00b786a488a68fe74bd336b64ec9cc11a16 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 17 Sep 2026 06:56:48 +0000
Subject: [PATCH 10/11] Validate the IQ kernels' encoding, not just their
 boundaries

The existing cases only proved the kernels run: any codebook entry, any bit
offset, and any sign convention produce non-zero input-dependent bytes that
differ per block, so a wrong index, a shifted qh field, or an inverted parity
rule all passed.

Add three cases per format that check the encoding itself:

- test_cuda_ext_iq_encoding_is_optimal packs a random tensor, decodes the
  payload from the GGML field positions rather than from the kernel's own
  layout code, and asserts the reconstruction error equals the brute-force
  minimum over every local scale, delta sign, and codebook entry at the block
  scale the payload carries. Comparing achieved error rather than raw indices
  keeps it robust to the kernels' fused-multiply-add ordering reordering
  near-ties. This pins the payload layout, the search, and the block-wide
  reductions in one assertion.
- test_cuda_ext_iq_input_dtype_equivalence packs values exact in every
  accepted dtype and requires byte-identical payloads.
- test_cuda_ext_iq_non_finite_inputs_are_zeroed checks NaN and both
  infinities pack as zeros, and that a finite float64 outside the float32
  range saturates instead -- a regression test for the narrowing fix.

The codebooks are synthetic random grids rather than the GGML tables. The
kernels treat the grid as an opaque argument, so this exercises the search
identically while keeping the tests free of any dependency on the reference
encoder or its tables, and a random grid has no ties to break.

The decoder and the brute-force oracle were checked against an independent
CPU packer written from the same specification, and against four injected
layout bugs -- local scale off by one bit, dropped delta sign, sign mask at
bit 8, swapped scale nibbles -- each of which the new assertion rejects.

Parity against the PyTorch reference encoder still belongs with that encoder
in #2446/#2450, and should compare dequantized error rather than bytes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../gpu/_extensions/test_torch_extensions.py  | 153 ++++++++++++++++++
 1 file changed, 153 insertions(+)

diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index c033600c1e2..7a7a7c54221 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -156,3 +156,156 @@ def test_cuda_ext_iq2_xs_rejects_non_finite_scales():
 
     with pytest.raises(RuntimeError, match="scales must be finite"):
         extension.pack(weight, _grid(fmt), scales)
+
+
+def _random_grid(fmt: _IqFormat) -> torch.Tensor:
+    """Random codebook, so the optimality check below has no ties to break.
+
+    The kernels treat the grid as opaque data, so a synthetic codebook exercises the search
+    exactly as a real one does -- while keeping these tests independent of the GGML tables.
+    IQ2_XS additionally requires non-negative magnitudes, since it carries signs separately.
+    """
+    shape = (fmt.entries, 8)
+    if fmt.needs_scales:
+        return torch.rand(shape, device="cuda", generator=_generator()) * fmt.grid_values[-1]
+    return torch.randn(shape, device="cuda", generator=_generator())
+
+
+def _decode(fmt: _IqFormat, packed: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
+    """Decode a packed payload the way GGML does, from the format definition rather than from
+    the kernel's own layout code, so a misplaced field shows up as a decode mismatch.
+    """
+    payload = packed.cpu()
+    blocks = payload.shape[0]
+    grid = grid.cpu()
+    # The fp16 block scale occupies the first two bytes of every IQ payload.
+    d = payload[:, :2].contiguous().view(torch.float16).float()
+
+    if not fmt.needs_scales:  # IQ1_S: 32 index bytes, then 8 uint16 of per-group metadata.
+        qs = payload[:, 2:34].int().view(blocks, 8, 4)
+        qh = payload[:, 34:50:2].int() | (payload[:, 35:50:2].int() << 8)
+        local = (qh >> 12) & 7
+        delta = torch.where((qh & 0x8000) != 0, -0.125, 0.125)
+        index = qs | (((qh.unsqueeze(-1) >> (3 * torch.arange(4))) & 7) << 8)
+        scale = (d * (2 * local + 1)).unsqueeze(-1).unsqueeze(-1)
+        return (scale * (grid[index] + delta[..., None, None])).reshape(blocks, 256)
+
+    # IQ2_XS: 32 uint16 codes, then 16 four-bit local scales packed two per byte.
+    codes = payload[:, 2:66:2].int() | (payload[:, 3:66:2].int() << 8)
+    index = codes & (fmt.entries - 1)
+    stored = ((codes >> 9).unsqueeze(-1) >> torch.arange(7)) & 1
+    # Only seven sign bits are stored; the eighth restores even parity over all eight.
+    signs = torch.cat([stored, (stored.sum(-1) & 1).unsqueeze(-1)], dim=-1)
+    nibbles = payload[:, 66:74].int()
+    local = torch.stack([nibbles & 0xF, (nibbles >> 4) & 0xF], dim=-1).reshape(blocks, 16)
+    scale = (d * (2 * local + 1) * 0.125).repeat_interleave(2, dim=1).unsqueeze(-1)
+    return (scale * grid[index] * (1.0 - 2.0 * signs.float())).reshape(blocks, 256)
+
+
+def _oracle_group_error(
+    fmt: _IqFormat, values: torch.Tensor, grid: torch.Tensor, d
+) -> torch.Tensor:
+    """Smallest squared error each group can reach at the block scale the payload carries.
+
+    Reproduces the kernels' objective by brute force: every local scale (and, for IQ1_S, every
+    delta sign) against every codebook entry, minimised per vector and summed over the group.
+    """
+    vectors, choices = (4, 16) if not fmt.needs_scales else (2, 16)
+    groups = 32 // vectors
+    x = values.cpu().float().reshape(-1, groups, vectors, 8)
+    grid, d = grid.cpu(), d.cpu()
+    xnorm = x.square().sum(-1)
+
+    errors = []
+    for choice in range(choices):
+        local = choice & 7 if not fmt.needs_scales else choice
+        if not fmt.needs_scales:
+            delta = 0.125 if choice < 8 else -0.125
+            shifted = grid + delta
+            scale = (d * (2 * local + 1)).reshape(-1, 1, 1, 1)
+            dot = x @ shifted.T
+        else:
+            shifted = grid
+            scale = (d * (2 * local + 1) * 0.125).reshape(-1, 1, 1, 1)
+            terms = x.abs().unsqueeze(-2) * grid  # [..., entries, 8]
+            dot = terms.sum(-1)
+            odd = (x < 0).sum(-1, keepdim=True) % 2 != 0
+            dot = torch.where(odd, dot - 2 * terms.min(-1).values, dot)
+        dot = dot.reshape(*x.shape[:3], fmt.entries)
+        error = xnorm.unsqueeze(-1) - 2 * scale * dot + scale.square() * shifted.square().sum(-1)
+        errors.append(error.clamp_min(0).min(-1).values.sum(-1))
+    return torch.stack(errors).min(0).values
+
+
+@pytest.mark.parametrize("fmt", _IQ_EXTENSIONS)
+def test_cuda_ext_iq_encoding_is_optimal(fmt):
+    """Round-trip the payload and check the search actually found the best codes.
+
+    This is the test that pins the bit layout: decoding follows the GGML field positions, so a
+    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)
+    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
+
+    packed = _pack(fmt, extension, weight, grid, scales=scales)
+    decoded = _decode(fmt, packed, grid)
+
+    # The block scale is a fixed heuristic, so compare the search at the scale actually stored.
+    d = packed[:, :2].contiguous().view(torch.float16).float()
+    group_size = 8 * (2 if fmt.needs_scales else 4)
+    achieved = (weight.cpu() - decoded).square().reshape(4, -1, group_size).sum(-1)
+    optimal = _oracle_group_error(fmt, weight, grid, d)
+
+    assert torch.allclose(achieved, optimal, rtol=1e-3, atol=1e-6), (
+        f"max excess {(achieved - optimal).abs().max():.3e}"
+    )
+    # Sanity: the quantizer must be doing better than emitting zeros.
+    assert achieved.sum() < 0.5 * weight.cpu().square().sum()
+
+
+@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)
+    # 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)
+    scales = (weight.abs().amax(dim=-1) / fmt.native_max).half() if fmt.needs_scales else None
+
+    payloads = [
+        _pack(fmt, extension, weight.to(dtype), grid, scales=scales)
+        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"
+
+
+@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)
+    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
+
+    spoiled = clean.clone()
+    spoiled[0, 5], spoiled[0, 200], spoiled[1, 17] = float("nan"), float("inf"), float("-inf")
+    zeroed = clean.clone()
+    zeroed[0, 5], zeroed[0, 200], zeroed[1, 17] = 0.0, 0.0, 0.0
+
+    assert torch.equal(
+        _pack(fmt, extension, spoiled, grid, scales=scales),
+        _pack(fmt, extension, zeroed, grid, scales=scales),
+    )
+
+    # A finite float64 outside the float32 range must saturate, not collapse to zero.
+    huge = zeroed.double()
+    huge[1, 7] = 1e100
+    assert not torch.equal(
+        _pack(fmt, extension, huge, grid, scales=scales),
+        _pack(fmt, extension, zeroed.double(), grid, scales=scales),
+    )

From 7cc1220acb6b3703dc3a2410fa3a945fa06f72c7 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 17 Sep 2026 16:18:24 +0000
Subject: [PATCH 11/11] Align the reduction helpers' barriers and tighten the
 IQ2_XS scale contract

Three review follow-ups, none of which fixed a live bug.

block_min_key now ends with its own __syncthreads(), so it matches the
contract block_min_accumulate beside it already had and the one both comments
advertise. Neither caller was racing -- each had its own trailing barrier --
but two helpers in the same header documenting identical semantics while
having different ones is exactly the drift the shared header exists to
prevent, and a future caller trusting the comment would have raced the next
iteration's scratch write against the previous iteration's thread-0 read. The
callers' barriers are dropped in the same move, so the PTX bar.sync count is
unchanged at 76 for IQ1_S and 48 for IQ2_XS.

IQ2_XS now rejects negative block scales alongside non-finite ones, in one
fused reduction so the boundary still costs a single synchronization. A
negative scale packs cleanly and inverts the sign of every decoded element,
and the search does not even degrade loudly: with a non-negative dot product
the -2*scale*dot term turns positive, so it selects the smallest-magnitude
entries and the smallest local scale. GGML asserts a non-negative block scale
in its own encoders.

store_block_scale treats negative zero as a zero scale. It is reachable only
through caller-supplied IQ2_XS scales, and the output was already correct --
every candidate scores identically, so the block encodes as index 0 with a
-0.0 scale and decodes to zeros either way -- but it ran a full codebook
search to get there. Note the new non-negative check cannot catch this, since
-0.0 >= 0 holds.

Verified on an RTX PRO 6000 Blackwell (sm_120): 25 passed, including both
test_cuda_ext_iq_encoding_is_optimal cases, which are what would catch a
mis-placed barrier.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../kernels/quantization/ggml/common.cuh      | 12 +++++++----
 .../torch/kernels/quantization/ggml/iq1_s.cu  |  1 -
 .../kernels/quantization/ggml/iq2_xs.cpp      | 14 ++++++++-----
 .../torch/kernels/quantization/ggml/iq2_xs.cu |  1 -
 .../gpu/_extensions/test_torch_extensions.py  | 21 ++++++++++++++++---
 5 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh
index 66314335010..30818f0da3b 100644
--- a/modelopt/torch/kernels/quantization/ggml/common.cuh
+++ b/modelopt/torch/kernels/quantization/ggml/common.cuh
@@ -138,7 +138,8 @@ __device__ __forceinline__ void block_min_accumulate(const float (&local)[kSlots
 }
 
 // Block-wide minimum of key, valid on thread 0 only. scratch must hold kWarps entries. Barriers
-// are internal, so every thread of the block must call this.
+// are internal -- including a trailing one, so scratch is free to reuse on return, matching
+// block_min_accumulate above -- and every thread of the block must call this.
 __device__ __forceinline__ unsigned long long block_min_key(unsigned long long key,
                                                             unsigned long long *scratch) {
   const int tid = threadIdx.x;
@@ -157,15 +158,18 @@ __device__ __forceinline__ unsigned long long block_min_key(unsigned long long k
     for (int w = 1; w < kWarps; ++w)
       key = scratch[w] < key ? scratch[w] : key;
   }
+  __syncthreads();
   return key;
 }
 
 // Writes the fp16 block scale into the payload, or zeroes the whole payload when the block scale
-// rounded to zero. Returns false once the payload is final and the caller should stop. The branch
-// is uniform across the block, so returning on false is barrier-safe.
+// rounded to zero. Negative zero counts: it reconstructs every element as zero, so it takes the
+// same branch instead of running a search whose candidates all score identically. Returns false
+// once the payload is final and the caller should stop. The branch is uniform across the block, so
+// returning on false is barrier-safe.
 template <int kPayloadBytes>
 __device__ __forceinline__ bool store_block_scale(uint8_t *payload, uint16_t d_bits) {
-  if (d_bits == 0) {
+  if ((d_bits & 0x7FFF) == 0) {
     if (threadIdx.x < kPayloadBytes)
       payload[threadIdx.x] = 0;
     return false;
diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
index 37f4d1f95ab..666176e951f 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu
@@ -183,7 +183,6 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
         selected_entries[vector] = entry;
         payload[kIndexOffset + group * kVectorsPerGroup + vector] = static_cast<uint8_t>(entry);
       }
-      __syncthreads();
     }
 
     if (tid == 0) {
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
index 755226373fa..0929a961935 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp
@@ -28,10 +28,13 @@ at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
   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, so a non-finite entry
-  // would produce a payload that decodes to garbage. The synchronization this costs is paid once
-  // per packed tensor, on an export path.
-  TORCH_CHECK(scales.isfinite().all().item<bool>(), "scales must be finite");
+  // 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());
 }
@@ -40,7 +43,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
   module.def("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 float16 [numel / 256]. "
+             "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.");
diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
index c38537e4d98..c731fe25706 100644
--- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
+++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu
@@ -190,7 +190,6 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g
         payload[code_offset] = static_cast<uint8_t>(code);
         payload[code_offset + 1] = static_cast<uint8_t>(code >> 8);
       }
-      __syncthreads();
     }
   }
 
diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py
index 7a7a7c54221..ea1e41a5b94 100644
--- a/tests/gpu/_extensions/test_torch_extensions.py
+++ b/tests/gpu/_extensions/test_torch_extensions.py
@@ -148,16 +148,31 @@ def test_cuda_ext_iq_rejects_row_straddling_input(fmt):
         _pack(fmt, extension, weight, _grid(fmt, zero=True))
 
 
-def test_cuda_ext_iq2_xs_rejects_non_finite_scales():
+@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)
     fmt = _IQ_EXTENSIONS[1].values[0]
     weight = torch.ones((1, 256), device="cuda", dtype=torch.bfloat16)
-    scales = torch.full((1,), float("nan"), device="cuda", dtype=torch.float16)
+    scales = torch.full((1,), bad, device="cuda", dtype=torch.float16)
 
-    with pytest.raises(RuntimeError, match="scales must be finite"):
+    with pytest.raises(RuntimeError, match="scales must be finite and non-negative"):
         extension.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)
+    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)
+
+    assert not packed.any()
+
+
 def _random_grid(fmt: _IqFormat) -> torch.Tensor:
     """Random codebook, so the optimality check below has no ties to break.