From 6956a28b66b4529a232056baf986ccdebc40520e Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 1 Oct 2026 06:21:48 +0000 Subject: [PATCH 1/2] Share one CUDA encoder per IQ family The five IQ encoders were five copies of the same search. IQ2_XS and IQ2_XXS shared 139 of their roughly 150 lines of encoder and launcher, and IQ2_S 113 of them; IQ1_S and IQ1_M had the same structure with a different choice space. Every scaled packer also validated its scales twice, in the pybind wrapper and again in the CUDA entry point. iq2_family.cuh now holds the one IQ2 encoder: the grid in shared memory, the 16 local scales scored per group, then the best entry per vector under the chosen scale. A format supplies its group shape, whether it stores seven sign bits and recovers the eighth from parity, and a store() that writes the chosen entries, sign masks and local scales into its layout. iq1_family.cuh holds the one IQ1 encoder over the shared ternary grid. Each group picks one of kChoices options; with kSharedShift the option also fixes the +/-1/8 delta (IQ1_S: shift * 8 + local), otherwise each vector picks its own (IQ1_M). IQ1_S's scale kernel now writes FP16 scales, so both IQ1 formats take the same input. Each format file is now one Format struct -- layout constants and store() -- plus its entry point. Validation lives once in common.cuh, as check_pack_inputs and check_scaled_pack_inputs, and ggml.cpp binds the CUDA entry points directly instead of through five wrappers. Nothing changes in the output. Every encoder hashes the same as before on float32, bfloat16, float16 and float64 inputs with zero, tiny, oversized and non-finite blocks, and encode speed is unchanged within 1% on an idle RTX PRO 6000. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../kernels/quantization/ggml/common.cuh | 25 +- .../torch/kernels/quantization/ggml/ggml.cpp | 89 +----- .../kernels/quantization/ggml/iq1_family.cuh | 191 +++++++++++++ .../torch/kernels/quantization/ggml/iq1_m.cu | 215 ++++----------- .../torch/kernels/quantization/ggml/iq1_s.cu | 193 ++++--------- .../kernels/quantization/ggml/iq2_family.cuh | 256 ++++++++++++++++++ .../torch/kernels/quantization/ggml/iq2_s.cu | 207 +++----------- .../torch/kernels/quantization/ggml/iq2_xs.cu | 216 ++------------- .../kernels/quantization/ggml/iq2_xxs.cu | 233 ++-------------- 9 files changed, 665 insertions(+), 960 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_family.cuh create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_family.cuh diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 5eccec6cd03..4ac702bfc96 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -41,6 +41,7 @@ namespace modelopt::ggml { // https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h constexpr int kBlockSize = 256; constexpr int kVectorSize = 8; +constexpr int kVectorsPerBlock = kBlockSize / kVectorSize; constexpr int kScaleOffset = 0; constexpr int kScaleBytes = 2; @@ -72,11 +73,11 @@ inline void check_scalar_pack_input(const char *format, const at::Tensor &input, " CUDA grid is too large"); } -// Validates the codebook contract every IQ format shares. Called from the pybind wrapper on the -// caller's tensors and again from the CUDA entry point on the materialized contiguous tensors, so -// the enforced rule and the message it reports are written once. +// Validates the codebook contract every IQ format shares, once, at each format's CUDA entry point. inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, int64_t entries) { + TORCH_CHECK(input.is_cuda(), format, " packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), format, " packing requires a CUDA grid"); check_scalar_pack_input(format, input, kBlockSize); TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == entries && grid.size(1) == kVectorSize, @@ -84,6 +85,24 @@ inline void check_pack_inputs(const char *format, const at::Tensor &input, const TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); } +// Validates a format that takes one precomputed FP16 block scale per 256 values. The kernels copy +// these bits straight into the block scale field: a non-finite entry produces a payload that +// decodes to garbage, and a negative one inverts the sign of every decoded element while still +// packing cleanly -- GGML's own encoders assert a non-negative block scale. One fused reduction, so +// the synchronization is paid once per packed tensor, on an export path. +inline void check_scaled_pack_inputs(const char *format, const at::Tensor &input, + const at::Tensor &grid, int64_t entries, + const at::Tensor &scales) { + TORCH_CHECK(scales.is_cuda(), format, " packing requires CUDA scales"); + check_pack_inputs(format, input, grid, entries); + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == input.numel() / kBlockSize, + "scales must be float16 [numel / 256]"); + TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), + "scales must be finite and non-negative"); + TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); +} + #ifdef __CUDACC__ // Reads one input element as float32. Non-finite elements are treated as zero, and finiteness is diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index e3d16b2c462..37bd7e8c3be 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -16,8 +16,9 @@ */ // Every GGML IQ format shares common.cuh, the same CUDA version gate, and the same build flags, -// so they compile into one extension and bind here. Each format keeps its kernels in its own -// translation unit and exposes a single host entry point. +// so they compile into one extension and bind here. Each IQ format keeps its block layout in its +// own translation unit, on the encoder its family shares (iq1_family.cuh or iq2_family.cuh), and +// exposes its packer. #include "common.cuh" @@ -30,80 +31,6 @@ at::Tensor q8_0_pack_cuda(at::Tensor input); namespace { -at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) { - TORCH_CHECK(input.is_cuda(), "IQ1_S packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ1_S packing requires a CUDA grid"); - modelopt::ggml::check_pack_inputs("IQ1_S", input, grid, modelopt::ggml::kIq1sEntries); - return iq1_s_pack_cuda(input.contiguous(), grid.contiguous()); -} - -at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_cuda(), "IQ2_XS packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ2_XS packing requires a CUDA grid"); - TORCH_CHECK(scales.is_cuda(), "IQ2_XS packing requires CUDA scales"); - modelopt::ggml::check_pack_inputs("IQ2_XS", input, grid, modelopt::ggml::kIq2xsEntries); - const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - // The kernel copies these bits straight into the GGML block scale field. A non-finite entry - // would produce a payload that decodes to garbage, and a negative one inverts the sign of every - // decoded element while still packing cleanly -- GGML's own encoders assert a non-negative block - // scale. One fused reduction, so the synchronization is paid once per packed tensor, on an - // export path. - TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), - "scales must be finite and non-negative"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); -} - -at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_cuda(), "IQ2_XXS packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ2_XXS packing requires a CUDA grid"); - TORCH_CHECK(scales.is_cuda(), "IQ2_XXS packing requires CUDA scales"); - modelopt::ggml::check_pack_inputs("IQ2_XXS", input, grid, modelopt::ggml::kIq2xxsEntries); - const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - // Same rule as IQ2_XS: these bits become the block scale verbatim, and a negative or non-finite - // one packs cleanly while decoding to garbage. - TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), - "scales must be finite and non-negative"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - return iq2_xxs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); -} - -at::Tensor iq2_s_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_cuda(), "IQ2_S packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ2_S packing requires a CUDA grid"); - TORCH_CHECK(scales.is_cuda(), "IQ2_S packing requires CUDA scales"); - modelopt::ggml::check_pack_inputs("IQ2_S", input, grid, modelopt::ggml::kIq2sEntries); - const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), - "scales must be finite and non-negative"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - return iq2_s_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); -} - -at::Tensor iq1_m_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_cuda(), "IQ1_M packing requires a CUDA input"); - TORCH_CHECK(grid.is_cuda(), "IQ1_M packing requires a CUDA grid"); - TORCH_CHECK(scales.is_cuda(), "IQ1_M packing requires CUDA scales"); - modelopt::ggml::check_pack_inputs("IQ1_M", input, grid, modelopt::ggml::kIq1sEntries); - const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), - "scales must be finite and non-negative"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - return iq1_m_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); -} - at::Tensor q8_0_pack(at::Tensor input) { TORCH_CHECK(input.is_cuda(), "Q8_0 packing requires a CUDA input"); modelopt::ggml::check_scalar_pack_input("Q8_0", input, 32); @@ -113,12 +40,12 @@ at::Tensor q8_0_pack(at::Tensor input) { } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { - module.def("iq1_s_pack", &iq1_s_pack, + module.def("iq1_s_pack", &iq1_s_pack_cuda, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 256. The grid must be float32 [2048, 8]. Returns uint8 " "[numel / 256, 50] on the input device. Non-finite input elements are treated as " "zero during packing, and finite elements outside the float32 range saturate."); - module.def("iq2_xs_pack", &iq2_xs_pack, + module.def("iq2_xs_pack", &iq2_xs_pack_cuda, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 256. The grid must be float32 [512, 8] holding " "non-negative codebook magnitudes, and scales must be finite non-negative float16 " @@ -126,7 +53,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 74] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); - module.def("iq2_xxs_pack", &iq2_xxs_pack, + module.def("iq2_xxs_pack", &iq2_xxs_pack_cuda, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 256. The grid must be float32 [256, 8] holding " "non-negative codebook magnitudes, and scales must be finite non-negative float16 " @@ -134,7 +61,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 66] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); - module.def("iq2_s_pack", &iq2_s_pack, + module.def("iq2_s_pack", &iq2_s_pack_cuda, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 256. The grid must be float32 [1024, 8] holding " "non-negative codebook magnitudes, and scales must be finite non-negative float16 " @@ -142,7 +69,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 82] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); - module.def("iq1_m_pack", &iq1_m_pack, + module.def("iq1_m_pack", &iq1_m_pack_cuda, "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " "dimension is a multiple of 256. The grid must be float32 [2048, 8], and scales must " "be finite non-negative float16 [numel / 256]; IQ1_M carries that scale in the four " diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_family.cuh b/modelopt/torch/kernels/quantization/ggml/iq1_family.cuh new file mode 100644 index 00000000000..f15f317e95e --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_family.cuh @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The IQ1 formats share one encoder. Every 8-value vector is approximated by scale * (q + delta), +// with q from the 2048-entry ternary grid and delta = +/- 1/8; the grid is 64 KiB, past the 48 KiB +// static shared-memory limit, so it is read from global memory and left to the cache. Every group +// of vectors picks one of Format::kChoices options, with local scale choice % 8. A format +// describes the rest with a Format type: +// +// kGroups, kVectorsPerGroup, kChoices, kPayloadBytes block layout +// kSharedShift the choice also fixes the delta for the whole group, as choice / 8 (IQ1_S); +// otherwise each vector picks its own (IQ1_M) +// begin(payload, d_bits) false when the block is all zero and needs nothing more written +// store(payload, picks, choices, d_bits) writes every vector's pick -- its grid entry, with +// its own shift in bit kIq1EntryBits when the shift is per vector -- and every +// group's choice; called by every thread once the search is done + +#pragma once + +#include "common.cuh" + +namespace modelopt::ggml { + +#ifdef __CUDACC__ + +constexpr float kIq1Delta = 0.125f; +constexpr int kIq1EntryBits = 11; // 2048 entries; a per-vector shift sits just above them + +__device__ __forceinline__ float iq1_delta(int shift) { return shift ? -kIq1Delta : kIq1Delta; } + +template +__global__ void iq1_encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + constexpr int kGroups = Format::kGroups; + constexpr int kVectorsPerGroup = Format::kVectorsPerGroup; + constexpr int kChoices = Format::kChoices; + constexpr bool kSharedShift = Format::kSharedShift; + static_assert(kIq1sEntries % kThreads == 0, "every thread must visit the same number of entries"); + static_assert(kIq1sEntries == 1 << kIq1EntryBits, "the shift bit sits just above the entry"); + static_assert(kGroups * kVectorsPerGroup == kVectorsPerBlock, "groups must cover the block"); + + __shared__ float warp_best[kWarps * kChoices]; + __shared__ float group_error[kChoices]; + __shared__ unsigned long long warp_keys[kWarps]; + __shared__ int selected_choice; + __shared__ uint8_t choices[kGroups]; + __shared__ uint16_t picks[kVectorsPerBlock]; + + const int tid = threadIdx.x; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * Format::kPayloadBytes; + const __half d_half = scales[block]; + const uint16_t d_bits = __half_as_ushort(d_half); + const float d = __half2float(d_half); + if (!Format::begin(payload, d_bits)) + return; + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kChoices) + group_error[tid] = 0.0f; + __syncthreads(); + + // Score every choice: each vector's best error under it, summed over the group. With a + // per-vector shift, a vector takes the better of the two before the group chooses. +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm, xsum; + load_vector(source + (group * kVectorsPerGroup + vector) * kVectorSize, x, xnorm, xsum); + float local_best[kChoices]; +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) + local_best[choice] = FLT_MAX; + for (int entry = tid; entry < kIq1sEntries; entry += blockDim.x) { + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) { + const float scale = d * (2 * (choice & 7) + 1); + if constexpr (kSharedShift) { + local_best[choice] = + fminf(local_best[choice], + shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, iq1_delta(choice >> 3))); + } else { +#pragma unroll + for (int shift = 0; shift < 2; ++shift) + local_best[choice] = + fminf(local_best[choice], + shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, iq1_delta(shift))); + } + } + } + block_min_accumulate(local_best, warp_best, group_error); + } + + if (tid == 0) { + selected_choice = 0; + float best = group_error[0]; +#pragma unroll + for (int choice = 1; choice < kChoices; ++choice) { + if (group_error[choice] < best) { + best = group_error[choice]; + selected_choice = choice; + } + } + choices[group] = static_cast(selected_choice); + } + __syncthreads(); + const float selected_scale = d * (2 * (selected_choice & 7) + 1); + + // Under the chosen option, each vector takes its best entry. A per-vector shift sits above + // the entry index in the key, so a tie prefers the lower shift and then the lower entry. +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + const int slot = group * kVectorsPerGroup + vector; + float x[kVectorSize]; + float xnorm, xsum; + load_vector(source + slot * kVectorSize, x, xnorm, xsum); + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kIq1sEntries; entry += blockDim.x) { + float dot, qnorm, qsum; + grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); + if constexpr (kSharedShift) { + const float error = shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, + iq1_delta(selected_choice >> 3)); + const unsigned long long candidate = error_key(error, entry); + key = candidate < key ? candidate : key; + } else { +#pragma unroll + for (int shift = 0; shift < 2; ++shift) { + const float error = + shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, iq1_delta(shift)); + const unsigned long long candidate = error_key(error, (shift << kIq1EntryBits) | entry); + key = candidate < key ? candidate : key; + } + } + } + key = block_min_key(key, warp_keys); + if (tid == 0) + picks[slot] = static_cast(key & ((1u << (kIq1EntryBits + 1)) - 1)); + } + } + __syncthreads(); + Format::store(payload, picks, choices, d_bits); +} + +// Packs one IQ1 format from per-block FP16 scales, once its inputs have been validated. +template +at::Tensor iq1_encode_blocks(const at::Tensor &input, const at::Tensor &grid, + const at::Tensor &scales) { + const auto values = input.contiguous(); + const auto table = grid.contiguous(); + const auto block_scales = scales.contiguous(); + c10::cuda::CUDAGuard guard(values.device()); + const int64_t num_blocks = values.numel() / kBlockSize; + at::Tensor output = + at::empty({num_blocks, Format::kPayloadBytes}, values.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, values.scalar_type(), "iq1_pack", [&] { + iq1_encode<<(num_blocks), kThreads, 0, stream>>>( + values.data_ptr(), num_blocks, table.data_ptr(), + reinterpret_cast(block_scales.data_ptr()), + output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} + +#endif // __CUDACC__ + +} // namespace modelopt::ggml diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_m.cu b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu index b5eadc85677..b98c1a2a33b 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq1_m.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "common.cuh" +#include "iq1_family.cuh" namespace { @@ -24,182 +24,67 @@ using namespace modelopt::ggml; // The IQ1_M 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 -// IQ1_M shares the IQ1_S ternary grid. At 64 KiB it is past the 48 KiB static shared-memory -// limit, so both kernels read it from global memory and rely on the cache. -constexpr int kEntries = kIq1sEntries; -constexpr int kGroups = 16; // one 3-bit local scale per 16 values -constexpr int kVectorsPerGroup = 2; -constexpr int kVectors = kGroups * kVectorsPerGroup; -constexpr int kLocalScales = 8; -constexpr int kSubBlocks = 8; -constexpr int kScaleWords = 4; -constexpr int kLowOffset = 0; // no leading block scale in this format -constexpr int kHighOffset = kLowOffset + kVectors; // 16 bytes, two vectors per byte -constexpr int kScaleWordOffset = kHighOffset + 2 * kSubBlocks; -constexpr int kPayloadBytes = kScaleWordOffset + 2 * kScaleWords; -constexpr float kDelta = 0.125f; -constexpr int kEntryBits = 11; // 2048 entries; the shift bit sits above them in the sort key - -static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); -static_assert(kPayloadBytes == 56, "IQ1_M blocks are 56 bytes"); - -template -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - - const int tid = threadIdx.x; - const int64_t block = blockIdx.x; - if (block >= num_blocks) - return; - - const scalar_t *source = input + block * kBlockSize; - uint8_t *payload = output + block * kPayloadBytes; - const __half d_half = scales[block]; - const uint16_t d_bits = __half_as_ushort(d_half); - const float d = __half2float(d_half); - // IQ1_M has no leading scale field, so the shared helper does not apply: a zero block is all - // zero bytes, which decodes to a zero scale and therefore zero values. - if ((d_bits & 0x7FFF) == 0) { - if (tid < kPayloadBytes) - payload[tid] = 0; - return; +struct Format { + static constexpr int kEntries = kIq1sEntries; // IQ1_M shares the IQ1_S ternary grid + static constexpr int kGroups = 16; // one 3-bit local scale per 16 values + static constexpr int kVectorsPerGroup = 2; + static constexpr int kChoices = 8; // the local scale; each vector picks its own shift + static constexpr bool kSharedShift = false; + static constexpr int kSubBlocks = 8; + static constexpr int kScaleWords = 4; + static constexpr int kLowOffset = 0; // no leading block scale + static constexpr int kHighOffset = kLowOffset + kVectorsPerBlock; // two vectors per byte + static constexpr int kScaleWordOffset = kHighOffset + 2 * kSubBlocks; + static constexpr int kPayloadBytes = kScaleWordOffset + 2 * kScaleWords; + + // IQ1_M has no leading scale field: a zero block is all zero bytes, which decodes to a zero + // scale and therefore zero values. + __device__ static bool begin(uint8_t *payload, uint16_t d_bits) { + if ((d_bits & 0x7FFF) != 0) + return true; + if (threadIdx.x < kPayloadBytes) + payload[threadIdx.x] = 0; + return false; } - if (tid < kHighOffset + 2 * kSubBlocks) - payload[tid] = 0; // index-high and shift nibbles are OR-ed into below - __syncthreads(); - -#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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm, xsum; - load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, - xnorm, xsum); - 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) { - float dot, qnorm, qsum; - grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); - // The shift is free per vector here, unlike IQ1_S where it is shared, so take the - // better of the two before the local scale is chosen. -#pragma unroll - for (int shift = 0; shift < 2; ++shift) { - const float delta = shift ? -kDelta : kDelta; -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1); - local_best[local] = fminf(local_best[local], - shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, delta)); - } - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; + // A low index byte per vector; a nibble per vector, low nibble first, holding three index-high + // bits and the shift; then four scale words, each with two sub-blocks' 3-bit local scales in + // bits 0..11 and one nibble of the FP16 block scale in bits 12..15. + __device__ static void store(uint8_t *payload, const uint16_t *picks, const uint8_t *choices, + uint16_t d_bits) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) + payload[kLowOffset + tid] = static_cast(picks[tid]); + if (tid < kVectorsPerBlock / 2) { + uint32_t byte = 0; #pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } + for (int half = 0; half < 2; ++half) { + const uint32_t pick = picks[2 * tid + half]; + byte |= (((pick >> 8) & 0x7u) | (((pick >> kIq1EntryBits) & 0x1u) << 3)) << (4 * half); } - locals[group] = static_cast(selected_local); + payload[kHighOffset + tid] = static_cast(byte); } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1); - + if (tid < kScaleWords) { + uint32_t word = 0; #pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm, xsum; - load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, - xnorm, xsum); - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - float dot, qnorm, qsum; - grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); -#pragma unroll - for (int shift = 0; shift < 2; ++shift) { - const float delta = shift ? -kDelta : kDelta; - const float error = shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, delta); - // Shift occupies the bits above the entry index so that a tie prefers the lower - // shift first and then the lower entry, matching the reference encoder. - const unsigned long long candidate = error_key(error, (shift << kEntryBits) | entry); - key = candidate < key ? candidate : key; - } - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int packed_index = static_cast(key & 0xFFFFFFFFULL); - const int entry = packed_index & ((1 << kEntryBits) - 1); - const int shift = (packed_index >> kEntryBits) & 1; - const int slot = group * kVectorsPerGroup + vector; // 0..31 - payload[kLowOffset + slot] = static_cast(entry & 0xFF); - // Two vectors share a qh byte: low nibble first, each holding three index-high bits - // and one shift bit. - const int nibble = (slot % 4) % 2; - const int qh_index = 2 * (slot / 4) + (slot % 4) / 2; - payload[kHighOffset + qh_index] |= - static_cast((((entry >> 8) & 0x7) | (shift << 3)) << (4 * nibble)); + for (int parity = 0; parity < 2; ++parity) { + const int sub = 2 * tid + parity; + const int base = 6 * parity; + word |= static_cast(choices[2 * sub]) << base; + word |= static_cast(choices[2 * sub + 1]) << (base + 3); } - __syncthreads(); + word |= static_cast((d_bits >> (4 * tid)) & 0xF) << 12; + payload[kScaleWordOffset + 2 * tid] = static_cast(word); + payload[kScaleWordOffset + 2 * tid + 1] = static_cast(word >> 8); } } +}; - // Four scale words: each carries two sub-blocks' 3-bit local scales in bits 0..11 and one - // nibble of the FP16 block scale in bits 12..15. - if (tid < kScaleWords) { - uint32_t word = 0; -#pragma unroll - for (int parity = 0; parity < 2; ++parity) { - const int sub = 2 * tid + parity; - const int base = 6 * parity; - word |= static_cast(locals[2 * sub]) << base; - word |= static_cast(locals[2 * sub + 1]) << (base + 3); - } - word |= static_cast((d_bits >> (4 * tid)) & 0xF) << 12; - payload[kScaleWordOffset + 2 * tid] = static_cast(word); - payload[kScaleWordOffset + 2 * tid + 1] = static_cast(word >> 8); - } -} +static_assert(Format::kPayloadBytes == 56, "IQ1_M blocks are 56 bytes"); } // namespace at::Tensor iq1_m_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"); - check_pack_inputs("IQ1_M", 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() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_m_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ1_M", input, grid, Format::kEntries, scales); + return iq1_encode_blocks(input, grid, scales); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu index cfe0aba39f3..3621ba4a994 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "common.cuh" +#include "iq1_family.cuh" namespace { @@ -24,25 +24,49 @@ 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; -constexpr int kChoices = 16; -constexpr int kIndexOffset = kScaleBytes; -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; +struct Format { + static constexpr int kEntries = kIq1sEntries; + static constexpr int kGroups = 8; + static constexpr int kVectorsPerGroup = 4; + static constexpr int kChoices = 16; // shift * 8 + local, shared by the group + static constexpr bool kSharedShift = true; + static constexpr int kIndexOffset = kScaleBytes; + static constexpr int kMetadataOffset = kIndexOffset + kVectorsPerBlock; + static constexpr int kPayloadBytes = kMetadataOffset + 2 * kGroups; + + __device__ static bool begin(uint8_t *payload, uint16_t d_bits) { + return store_block_scale(payload, d_bits); + } -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"); + // One index byte per vector, then one metadata word per group: three index-high bits per + // vector, the 3-bit local scale in bits 12..14 and the shift in bit 15. + __device__ static void store(uint8_t *payload, const uint16_t *picks, const uint8_t *choices, + uint16_t) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) + payload[kIndexOffset + tid] = static_cast(picks[tid]); + if (tid < kGroups) { + uint32_t qh = ((choices[tid] & 0x7u) << 12) | ((choices[tid] >> 3) << 15); +#pragma unroll + for (int k = 0; k < kVectorsPerGroup; ++k) + qh |= ((picks[tid * kVectorsPerGroup + k] >> 8) & 0x7u) << (3 * k); + payload[kMetadataOffset + 2 * tid] = static_cast(qh); + payload[kMetadataOffset + 2 * tid + 1] = static_cast(qh >> 8); + } + } +}; + +static_assert(Format::kPayloadBytes == 50, "IQ1_S blocks are 50 bytes"); + +// IQ1_S derives its own block scale. The 0.61 anchor matches the reference encoder's empirical +// predictor: it favors most values instead of forcing the block's largest to be exactly +// representable. +constexpr float kMaxLocalScale = 15.0f; // 2 * 7 + 1 +constexpr float kNativeMax = kMaxLocalScale * (1.0f + kIq1Delta); // 16.875 +constexpr float kScaleAnchor = 0.61f; template -__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { +__global__ void find_scale(const scalar_t *input, int64_t num_blocks, __half *scales) { const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (block >= num_blocks) return; @@ -52,136 +76,25 @@ __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))); - // 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(__half_as_ushort(scale)); -} - -template -__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[kWarps * kChoices]; - __shared__ float group_error[kChoices]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_choice; - __shared__ uint16_t selected_entries[kVectorsPerGroup]; - - const int tid = threadIdx.x; - const int64_t block = blockIdx.x; - if (block >= num_blocks) - return; - - const scalar_t *source = input + block * kBlockSize; - uint8_t *payload = output + block * kPayloadBytes; - const uint16_t d_bits = static_cast(scale_bits[block]); - const float d = __half2float(__ushort_as_half(d_bits)); - if (!store_block_scale(payload, d_bits)) - return; - -#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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm, xsum; - load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, - xnorm, xsum); - float local_best[kChoices]; -#pragma unroll - for (int choice = 0; choice < kChoices; ++choice) - local_best[choice] = FLT_MAX; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - float dot, qnorm, qsum; - grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); -#pragma unroll - for (int choice = 0; choice < kChoices; ++choice) { - const int local = choice & 7; - const float delta = choice < 8 ? kDelta : -kDelta; - const float scale = d * (2 * local + 1); - local_best[choice] = - fminf(local_best[choice], shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, delta)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - if (tid == 0) { - selected_choice = 0; - float best = group_error[0]; -#pragma unroll - for (int choice = 1; choice < kChoices; ++choice) { - if (group_error[choice] < best) { - best = group_error[choice]; - selected_choice = choice; - } - } - } - __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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm, xsum; - load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, - xnorm, xsum); - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - float dot, qnorm, qsum; - grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); - const float error = - shifted_error(xnorm, xsum, dot, qnorm, qsum, selected_scale, selected_delta); - const unsigned long long candidate = error_key(error, entry); - key = candidate < key ? candidate : key; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const uint16_t entry = static_cast(key & (kEntries - 1)); - selected_entries[vector] = entry; - payload[kIndexOffset + group * kVectorsPerGroup + vector] = static_cast(entry); - } - } - - if (tid == 0) { - const uint16_t qh = static_cast( - ((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[kMetadataOffset + 2 * group] = static_cast(qh); - payload[kMetadataOffset + 2 * group + 1] = static_cast(qh >> 8); - } - __syncthreads(); - } + scales[block] = __float2half_rn(fminf((amax / kNativeMax) * kScaleAnchor, 65504.0f)); } } // 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"); - check_pack_inputs("IQ1_S", input, grid, kEntries); - c10::cuda::CUDAGuard guard(input.device()); - const int64_t num_blocks = input.numel() / kBlockSize; - auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong)); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + check_pack_inputs("IQ1_S", input, grid, Format::kEntries); + const auto values = input.contiguous(); + c10::cuda::CUDAGuard guard(values.device()); + const int64_t num_blocks = values.numel() / kBlockSize; + auto scales = at::empty({num_blocks}, values.options().dtype(at::kHalf)); const auto stream = c10::cuda::getCurrentCUDAStream(); - const int scale_grid = static_cast((num_blocks + kThreads - 1) / kThreads); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_s_pack", [&] { - find_scale<<>>( - input.data_ptr(), num_blocks, scales.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - scales.data_ptr(), output.data_ptr()); + at::ScalarType::Half, at::ScalarType::BFloat16, values.scalar_type(), "iq1_s_scale", [&] { + find_scale + <<((num_blocks + kThreads - 1) / kThreads), kThreads, 0, stream>>>( + values.data_ptr(), num_blocks, + reinterpret_cast<__half *>(scales.data_ptr())); C10_CUDA_KERNEL_LAUNCH_CHECK(); }); - return output; + return iq1_encode_blocks(values, grid, scales); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_family.cuh b/modelopt/torch/kernels/quantization/ggml/iq2_family.cuh new file mode 100644 index 00000000000..d79cc897f21 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_family.cuh @@ -0,0 +1,256 @@ +/* + * 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 IQ2 formats share one encoder. Every 8-value vector is matched against a grid of +// non-negative magnitudes held in shared memory, every group of vectors picks one of 16 local +// scales, and the signs come from the input. A format describes the rest with a Format type: +// +// kEntries, kGroups, kVectorsPerGroup, kPayloadBytes grid size and block layout +// kParitySigns stores seven sign bits and lets the eighth make the count of negatives even +// (IQ2_XS, IQ2_XXS), rather than storing all eight (IQ2_S) +// store(payload, entries, signs, locals) writes every vector's grid entry and 8-bit sign mask +// and every group's local scale; called by every thread once the search is done + +#pragma once + +#include "common.cuh" + +namespace modelopt::ggml { + +#ifdef __CUDACC__ + +constexpr int kIq2LocalScales = 16; +constexpr float kIq2LocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. + +// Dot product of |x| against one codebook vector under the even-parity sign rule: with an odd +// number of negatives, the coordinate with the smallest |x| * q penalty gets flipped. +__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; +} + +// Dot product of |x| against one codebook vector when all eight signs are stored, so the best +// signs are simply the input's and the search compares magnitudes directly. +__device__ __forceinline__ float magnitude_dot(const float *x, const float *q) { + float dot = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + dot = fmaf(fabsf(x[j]), q[j], dot); + return dot; +} + +template +__device__ __forceinline__ float iq2_dot(const float *x, const float *q, bool odd_parity) { + if constexpr (kParitySigns) + return even_parity_dot(x, q, odd_parity); + else + return magnitude_dot(x, q); +} + +// The input's sign mask, with the weakest coordinate flipped when the format stores parity and +// the count of negatives is odd. +template +__device__ __forceinline__ uint8_t iq2_sign_mask(const float *x, const float *q, bool odd_parity) { + int flip_index = -1; + if constexpr (kParitySigns) { + if (odd_parity) { + 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) + sign_mask |= static_cast((x[j] < 0.0f) != (j == flip_index)) << j; + return static_cast(sign_mask); +} + +// Loads one vector with its squared norm; returns whether it has an odd number of negatives. +template +__device__ __forceinline__ bool load_signed_vector(const scalar_t *source, float (&x)[kVectorSize], + float &xnorm) { + xnorm = 0.0f; + int negative_count = 0; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + j); + xnorm = fmaf(x[j], x[j], xnorm); + negative_count += x[j] < 0.0f; + } + return (negative_count & 1) != 0; +} + +template +__global__ void iq2_encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + constexpr int kEntries = Format::kEntries; + constexpr int kGroups = Format::kGroups; + constexpr int kVectorsPerGroup = Format::kVectorsPerGroup; + constexpr bool kParitySigns = Format::kParitySigns; + 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"); + static_assert(kGroups * kVectorsPerGroup == kVectorsPerBlock, "groups must cover the block"); + + __shared__ float shared_grid[kEntries * kVectorSize]; + __shared__ float grid_norm[kEntries]; + __shared__ float warp_best[kWarps * kIq2LocalScales]; + __shared__ float group_error[kIq2LocalScales]; + __shared__ unsigned long long warp_keys[kWarps]; + __shared__ int selected_local; + __shared__ uint8_t locals[kGroups]; + __shared__ uint16_t entries[kVectorsPerBlock]; + __shared__ uint8_t signs[kVectorsPerBlock]; + + const int tid = threadIdx.x; + 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(); + 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 scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * Format::kPayloadBytes; + const __half d_half = scales[block]; + const uint16_t d_bits = __half_as_ushort(d_half); + const float d = __half2float(d_half); + if (!store_block_scale(payload, d_bits)) + return; + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kIq2LocalScales) + group_error[tid] = 0.0f; + __syncthreads(); + + // Score every local scale: each vector's best error under it, summed over the group. +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm; + const bool odd_parity = + load_signed_vector(source + (group * kVectorsPerGroup + vector) * kVectorSize, x, xnorm); + float local_best[kIq2LocalScales]; +#pragma unroll + for (int local = 0; local < kIq2LocalScales; ++local) + local_best[local] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float dot = iq2_dot(x, shared_grid + entry * kVectorSize, odd_parity); +#pragma unroll + for (int local = 0; local < kIq2LocalScales; ++local) { + const float scale = d * (2 * local + 1) * kIq2LocalScaleStep; + local_best[local] = + fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } + block_min_accumulate(local_best, warp_best, group_error); + } + + if (tid == 0) { + selected_local = 0; + float best = group_error[0]; +#pragma unroll + for (int local = 1; local < kIq2LocalScales; ++local) { + if (group_error[local] < best) { + best = group_error[local]; + selected_local = local; + } + } + locals[group] = static_cast(selected_local); + } + __syncthreads(); + const float selected_scale = d * (2 * selected_local + 1) * kIq2LocalScaleStep; + + // Under the chosen scale, each vector takes its best entry; ties go to the lowest index. +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + const int slot = group * kVectorsPerGroup + vector; + float x[kVectorSize]; + float xnorm; + const bool odd_parity = load_signed_vector(source + slot * kVectorSize, x, xnorm); + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = clamped_quant_error( + xnorm, iq2_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; + } + key = block_min_key(key, warp_keys); + if (tid == 0) { + const int entry = static_cast(key & (kEntries - 1)); + entries[slot] = static_cast(entry); + signs[slot] = iq2_sign_mask(x, shared_grid + entry * kVectorSize, odd_parity); + } + } + } + __syncthreads(); + Format::store(payload, entries, signs, locals); +} + +// Packs one IQ2 format, once check_scaled_pack_inputs has accepted the arguments. +template +at::Tensor iq2_encode_blocks(const at::Tensor &input, const at::Tensor &grid, + const at::Tensor &scales) { + const auto values = input.contiguous(); + const auto table = grid.contiguous(); + const auto block_scales = scales.contiguous(); + c10::cuda::CUDAGuard guard(values.device()); + const int64_t num_blocks = values.numel() / kBlockSize; + at::Tensor output = + at::empty({num_blocks, Format::kPayloadBytes}, values.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, values.scalar_type(), "iq2_pack", [&] { + iq2_encode<<(num_blocks), kThreads, 0, stream>>>( + values.data_ptr(), num_blocks, table.data_ptr(), + reinterpret_cast(block_scales.data_ptr()), + output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} + +#endif // __CUDACC__ + +} // namespace modelopt::ggml diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_s.cu b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu index 36f9d546e78..c4e9075e6fd 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_s.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "common.cuh" +#include "iq2_family.cuh" namespace { @@ -24,187 +24,44 @@ using namespace modelopt::ggml; // The IQ2_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 = kIq2sEntries; -constexpr int kGroups = 16; // one 4-bit local scale per 16 values -constexpr int kVectorsPerGroup = 2; -constexpr int kVectors = kGroups * kVectorsPerGroup; // 32 codebook vectors per block -constexpr int kLocalScales = 16; -constexpr int kSubBlocks = 8; -constexpr int kLowOffset = kScaleBytes; // 32 low index bytes -constexpr int kSignOffset = kLowOffset + kVectors; // 32 sign masks -constexpr int kHighOffset = kSignOffset + kVectors; // 8 bytes, two high bits per vector -constexpr int kLocalScaleOffset = kHighOffset + kSubBlocks; -constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; -constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. - -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"); -static_assert(kPayloadBytes == 82, "IQ2_S blocks are 82 bytes"); - -// IQ2_S stores all eight sign bits per vector rather than seven plus a parity bit, so unlike -// IQ2_XS and IQ2_XXS the encoder never has to flip a coordinate: the best signs are simply the -// input's, and the search compares magnitudes directly. -__device__ __forceinline__ float magnitude_dot(const float *x, const float *q) { - float dot = 0.0f; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) - dot = fmaf(fabsf(x[j]), q[j], dot); - return dot; -} - -template -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - // The codebook and its norms take 36 KiB, the most in the family but inside the 48 KiB - // static shared-memory limit. - __shared__ float shared_grid[kEntries * kVectorSize]; - __shared__ float grid_norm[kEntries]; - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - - const int tid = threadIdx.x; - 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(); - 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); +struct Format { + static constexpr int kEntries = kIq2sEntries; + static constexpr int kGroups = 16; // one 4-bit local scale per 16 values + static constexpr int kVectorsPerGroup = 2; + static constexpr bool kParitySigns = false; // all eight sign bits are stored + static constexpr int kSubBlocks = 8; + static constexpr int kLowOffset = kScaleBytes; // 32 low index bytes + static constexpr int kSignOffset = kLowOffset + kVectorsPerBlock; // 32 sign masks + static constexpr int kHighOffset = kSignOffset + kVectorsPerBlock; // two high bits per vector + static constexpr int kLocalScaleOffset = kHighOffset + kSubBlocks; + static constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; + + // Low index byte and sign mask per vector, the two high index bits packed four vectors to a + // byte, then the local scales as nibbles, two per byte. + __device__ static void store(uint8_t *payload, const uint16_t *entries, const uint8_t *signs, + const uint8_t *locals) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) { + payload[kLowOffset + tid] = static_cast(entries[tid]); + payload[kSignOffset + tid] = signs[tid]; } - grid_norm[entry] = norm; - } - __syncthreads(); - - 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 (!store_block_scale(payload, d_bits)) - return; - // The high-bit bytes are OR-ed into below, so they must start clear. - if (tid < kSubBlocks) - payload[kHighOffset + tid] = 0; - __syncthreads(); - -#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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; + if (tid < kSubBlocks) { + uint32_t high = 0; #pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - } - 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 dot = magnitude_dot(x, shared_grid + entry * kVectorSize); -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - 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(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - } - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float error = - clamped_quant_error(xnorm, magnitude_dot(x, shared_grid + entry * kVectorSize), - grid_norm[entry], selected_scale); - const unsigned long long candidate = error_key(error, entry); - key = candidate < key ? candidate : key; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - int sign_mask = 0; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) - sign_mask |= static_cast(x[j] < 0.0f) << j; - const int slot = group * kVectorsPerGroup + vector; // 0..31, GGML's qs index - payload[kLowOffset + slot] = static_cast(entry & 0xFF); - payload[kSignOffset + slot] = static_cast(sign_mask); - // Two high index bits per vector, packed four to a byte in vector order. - payload[kHighOffset + slot / 4] |= - static_cast(((entry >> 8) & 0x3) << (2 * (slot % 4))); - } - __syncthreads(); + for (int k = 0; k < 4; ++k) + high |= ((entries[4 * tid + k] >> 8) & 0x3u) << (2 * k); + payload[kHighOffset + tid] = static_cast(high); } + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } +}; - if (tid < kGroups / 2) - payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); -} +static_assert(Format::kPayloadBytes == 82, "IQ2_S blocks are 82 bytes"); } // namespace at::Tensor iq2_s_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"); - check_pack_inputs("IQ2_S", 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() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_s_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_S", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu index c731fe25706..ecc242485ea 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "common.cuh" +#include "iq2_family.cuh" namespace { @@ -24,201 +24,35 @@ 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; -constexpr int kLocalScales = 16; -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. - -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) { - 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 -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - __shared__ float shared_grid[kEntries * kVectorSize]; - __shared__ float grid_norm[kEntries]; - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - - const int tid = threadIdx.x; - 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(); - 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 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 (!store_block_scale(payload, d_bits)) - return; - -#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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#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) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - 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(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#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 = 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; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - 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(is_negative) << j; - } - const uint16_t code = static_cast(entry | ((sign_mask & 0x7f) << 9)); - const int code_offset = kCodeOffset + 2 * (group * kVectorsPerGroup + vector); - payload[code_offset] = static_cast(code); - payload[code_offset + 1] = static_cast(code >> 8); - } +struct Format { + static constexpr int kEntries = kIq2xsEntries; + static constexpr int kGroups = 16; // one 4-bit local scale per 16 values + static constexpr int kVectorsPerGroup = 2; + static constexpr bool kParitySigns = true; + static constexpr int kCodeOffset = kScaleBytes; + static constexpr int kLocalScaleOffset = kCodeOffset + 2 * kVectorsPerBlock; + static constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; + + // One uint16 code per vector, a 9-bit entry under seven sign bits, then the local scales as + // nibbles, two per byte. + __device__ static void store(uint8_t *payload, const uint16_t *entries, const uint8_t *signs, + const uint8_t *locals) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) { + const uint32_t code = entries[tid] | ((signs[tid] & 0x7Fu) << 9); + payload[kCodeOffset + 2 * tid] = static_cast(code); + payload[kCodeOffset + 2 * tid + 1] = static_cast(code >> 8); } + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } +}; - if (tid < kGroups / 2) - payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); -} +static_assert(Format::kPayloadBytes == 74, "IQ2_XS blocks are 74 bytes"); } // namespace 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"); - 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() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_XS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu index c501273a669..31570a11a77 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "common.cuh" +#include "iq2_family.cuh" namespace { @@ -24,217 +24,40 @@ using namespace modelopt::ggml; // The IQ2_XXS 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 = kIq2xxsEntries; -constexpr int kGroups = 8; // one 4-bit local scale per 32 values -constexpr int kVectorsPerGroup = 4; // four 8-value codebook vectors per group -constexpr int kLocalScales = 16; -constexpr int kRecordBytes = 8; // four index bytes then one little-endian uint32 -constexpr int kCodeOffset = kScaleBytes; -constexpr int kPayloadBytes = kCodeOffset + kGroups * kRecordBytes; -constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. - -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"); -static_assert(kPayloadBytes == 66, "IQ2_XXS blocks are 66 bytes"); - -// Dot product of |x| against one codebook vector, under the format's even-parity sign rule. -// IQ2_XXS stores seven sign bits per vector and recovers the eighth from their parity, exactly as -// IQ2_XS does, so an odd sign pattern must flip the coordinate with the smallest |x| * q penalty. -__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 -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - __shared__ float shared_grid[kEntries * kVectorSize]; - __shared__ float grid_norm[kEntries]; - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - __shared__ uint8_t entry_bytes[kGroups * kVectorsPerGroup]; - __shared__ uint8_t sign_bits[kGroups * kVectorsPerGroup]; - - const int tid = threadIdx.x; - 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(); - 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); +struct Format { + static constexpr int kEntries = kIq2xxsEntries; + static constexpr int kGroups = 8; // one 4-bit local scale per 32 values + static constexpr int kVectorsPerGroup = 4; // four 8-value codebook vectors per group + static constexpr bool kParitySigns = true; + static constexpr int kRecordBytes = 8; // four index bytes then one little-endian uint32 + static constexpr int kCodeOffset = kScaleBytes; + static constexpr int kPayloadBytes = kCodeOffset + kGroups * kRecordBytes; + + // One 8-byte record per group: four index bytes, then a uint32 holding four 7-bit sign indices + // in bits 0..27 and the 4-bit local scale in bits 28..31. + __device__ static void store(uint8_t *payload, const uint16_t *entries, const uint8_t *signs, + const uint8_t *locals) { + const int group = threadIdx.x; + if (group >= kGroups) + return; + uint8_t *record = payload + kCodeOffset + group * kRecordBytes; + uint32_t aux = static_cast(locals[group]) << 28; +#pragma unroll + for (int j = 0; j < kVectorsPerGroup; ++j) { + record[j] = static_cast(entries[group * kVectorsPerGroup + j]); + aux |= (signs[group * kVectorsPerGroup + j] & 0x7Fu) << (7 * j); } - grid_norm[entry] = norm; - } - __syncthreads(); - - 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 (!store_block_scale(payload, d_bits)) - return; - -#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 < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#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) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - 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(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#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 = 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; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - 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(is_negative) << j; - } - const int slot = group * kVectorsPerGroup + vector; - entry_bytes[slot] = static_cast(entry); - sign_bits[slot] = static_cast(sign_mask & 0x7f); - } - __syncthreads(); - } - } - - // One 8-byte record per group: four index bytes, then a uint32 holding four 7-bit sign - // indices in bits 0..27 and the 4-bit local scale in bits 28..31. - if (tid < kGroups) { - uint8_t *record = payload + kCodeOffset + tid * kRecordBytes; - const int base = tid * kVectorsPerGroup; -#pragma unroll - for (int j = 0; j < kVectorsPerGroup; ++j) - record[j] = entry_bytes[base + j]; - const uint32_t aux = static_cast(sign_bits[base]) | - (static_cast(sign_bits[base + 1]) << 7) | - (static_cast(sign_bits[base + 2]) << 14) | - (static_cast(sign_bits[base + 3]) << 21) | - (static_cast(locals[tid]) << 28); #pragma unroll for (int j = 0; j < 4; ++j) record[kVectorsPerGroup + j] = static_cast(aux >> (8 * j)); } -} +}; + +static_assert(Format::kPayloadBytes == 66, "IQ2_XXS blocks are 66 bytes"); } // namespace at::Tensor iq2_xxs_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"); - check_pack_inputs("IQ2_XXS", 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() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xxs_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_XXS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } From 3947a86045a07eba05297bcd0e38e97bd2686b62 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 1 Oct 2026 06:36:28 +0000 Subject: [PATCH 2/2] Report packer input errors in the original order check_scaled_pack_inputs checked that the block scales were on CUDA before check_pack_inputs looked at the input and grid, so a call with several bad arguments named the scales first. The pybind wrappers it replaced checked the input, grid and scales devices in that order, then everything else. check_pack_inputs now takes the scales optionally and checks their device right after the grid's, which restores that order. Over 5 formats and 8 combinations of bad arguments, every first error now matches main's. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/kernels/quantization/ggml/common.cuh | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 4ac702bfc96..71e3392b678 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -74,10 +74,13 @@ inline void check_scalar_pack_input(const char *format, const at::Tensor &input, } // Validates the codebook contract every IQ format shares, once, at each format's CUDA entry point. +// A format that takes block scales passes them so their device is checked with the others', before +// any shape or dtype rule. inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, - int64_t entries) { + int64_t entries, const at::Tensor *scales = nullptr) { TORCH_CHECK(input.is_cuda(), format, " packing requires a CUDA input"); TORCH_CHECK(grid.is_cuda(), format, " packing requires a CUDA grid"); + TORCH_CHECK(scales == nullptr || scales->is_cuda(), format, " packing requires CUDA scales"); check_scalar_pack_input(format, input, kBlockSize); TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == entries && grid.size(1) == kVectorSize, @@ -93,8 +96,7 @@ inline void check_pack_inputs(const char *format, const at::Tensor &input, const inline void check_scaled_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, int64_t entries, const at::Tensor &scales) { - TORCH_CHECK(scales.is_cuda(), format, " packing requires CUDA scales"); - check_pack_inputs(format, input, grid, entries); + check_pack_inputs(format, input, grid, entries, &scales); TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && scales.numel() == input.numel() / kBlockSize, "scales must be float16 [numel / 256]");