diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 5eccec6cd03..71e3392b678 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,14 @@ inline void check_scalar_pack_input(const char *format, const at::Tensor &input, " CUDA grid is too large"); } -// Validates the codebook contract every IQ format shares. Called from the pybind wrapper on the -// caller's tensors and again from the CUDA entry point on the materialized contiguous tensors, so -// the enforced rule and the message it reports are written once. +// Validates the codebook contract every IQ format shares, once, at each format's CUDA entry point. +// A format that takes block scales passes them so their device is checked with the others', before +// any shape or dtype rule. inline void check_pack_inputs(const char *format, const at::Tensor &input, const at::Tensor &grid, - int64_t entries) { + int64_t entries, const at::Tensor *scales = nullptr) { + TORCH_CHECK(input.is_cuda(), format, " packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), format, " packing requires a CUDA grid"); + TORCH_CHECK(scales == nullptr || scales->is_cuda(), format, " packing requires CUDA scales"); check_scalar_pack_input(format, input, kBlockSize); TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == entries && grid.size(1) == kVectorSize, @@ -84,6 +88,23 @@ 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) { + check_pack_inputs(format, input, grid, entries, &scales); + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == input.numel() / kBlockSize, + "scales must be float16 [numel / 256]"); + TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), + "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); }