diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 64e158c05bb..6aa1e69f468 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -635,8 +635,9 @@ def _export_quantized_weight( "IQ unified export currently supports modules with a standard 'weight' " f"attribute, got {weight_name!r} on {type(sub_module).__name__}" ) - quantize_iq = IQ_FORMAT_REGISTRY[quantization_format].quantize - packed_weight, _ = quantize_iq(weight.to(dtype)) + packed_weight = IQ_FORMAT_REGISTRY[quantization_format].pack( + weight.to(dtype), getattr(sub_module, quantizer_attrs.weight_quantizer, None) + ) setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) maybe_clear_cuda_cache() return diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 5eccec6cd03..0a98002e699 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -72,11 +72,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 +84,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 @@ -229,6 +247,94 @@ __device__ __forceinline__ bool store_block_scale(uint8_t *payload, uint16_t d_b return true; } +// Decoders run one thread per 8-value vector and follow the PyTorch decoders operation for +// operation. Every float operation is explicitly rounded so the compiler cannot fuse a multiply +// into an add, which keeps them bit-identical to that reference. A format's layout type supplies +// kPayloadBytes, kEntries and decode(block, vector, grid, values), and binds decode_blocks. +constexpr int kVectorsPerBlock = kBlockSize / kVectorSize; + +__device__ __forceinline__ uint32_t load_u16(const uint8_t *bytes) { + return static_cast(bytes[0]) | (static_cast(bytes[1]) << 8); +} + +__device__ __forceinline__ uint32_t load_u32(const uint8_t *bytes) { + return load_u16(bytes) | (load_u16(bytes + 2) << 16); +} + +__device__ __forceinline__ float half_bits_to_float(uint32_t bits) { + return __half2float(__ushort_as_half(static_cast(bits))); +} + +// IQ2_XS and IQ2_XXS store seven sign bits; the eighth makes the count of negatives even. +__device__ __forceinline__ uint32_t with_parity_bit(uint32_t sign_index) { + return sign_index | ((__popc(sign_index) & 1u) << 7); +} + +// x * scale, with coordinate j negated where bit j of sign_mask is set. +__device__ __forceinline__ void signed_scaled(const float *q, uint32_t sign_mask, float scale, + float (&values)[kVectorSize]) { +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + values[j] = __fmul_rn((sign_mask >> j) & 1 ? -q[j] : q[j], scale); +} + +// (q + delta) * scale, the IQ1 formats' shifted ternary grid. +__device__ __forceinline__ void shifted_scaled(const float *q, float delta, float scale, + float (&values)[kVectorSize]) { +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + values[j] = __fmul_rn(__fadd_rn(q[j], delta), scale); +} + +template +__global__ void decode_vectors(const uint8_t *packed, int64_t num_vectors, const float *grid, + out_t *output) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= num_vectors) + return; + float values[kVectorSize]; + Format::decode(packed + (index / kVectorsPerBlock) * Format::kPayloadBytes, + static_cast(index % kVectorsPerBlock), grid, values); + out_t *out = output + index * kVectorSize; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) + out[j] = static_cast(values[j]); +} + +// Decodes uint8 [blocks, kPayloadBytes] into [blocks, 256] of dtype on the payload's device. +template +at::Tensor decode_blocks(const char *format, const at::Tensor &packed, const at::Tensor &grid, + at::ScalarType dtype) { + constexpr int kPayloadBytes = Format::kPayloadBytes; + constexpr int kEntries = Format::kEntries; + TORCH_CHECK(packed.is_cuda() && grid.is_cuda(), format, " decoding requires CUDA tensors"); + TORCH_CHECK(packed.get_device() == grid.get_device(), "payload and grid must share a device"); + TORCH_CHECK(packed.scalar_type() == at::kByte && packed.dim() == 2 && + packed.size(1) == kPayloadBytes, + format, " payloads must be uint8 [blocks, ", kPayloadBytes, "]"); + TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries && + grid.size(1) == kVectorSize, + format, " grid must be float32 [", kEntries, ", ", kVectorSize, "]"); + TORCH_CHECK(at::isFloatingType(dtype), format, " decodes to a floating-point dtype"); + c10::cuda::CUDAGuard guard(packed.device()); + const auto payload = packed.contiguous(); + const auto table = grid.contiguous(); + auto output = at::empty({packed.size(0), kBlockSize}, packed.options().dtype(dtype)); + const int64_t num_vectors = packed.size(0) * kVectorsPerBlock; + if (num_vectors == 0) + return output; + const auto launch_blocks = static_cast((num_vectors + kThreads - 1) / kThreads); + const auto stream = c10::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, dtype, "ggml_decode", [&] { + decode_vectors<<>>( + payload.data_ptr(), num_vectors, table.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/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index e3d16b2c462..e262bf44cc4 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 a packer and an unpacker. #include "common.cuh" @@ -27,83 +28,14 @@ at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scale at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); at::Tensor iq1_m_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); at::Tensor q8_0_pack_cuda(at::Tensor input); +at::Tensor iq1_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype); +at::Tensor iq1_m_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype); +at::Tensor iq2_xxs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype); +at::Tensor iq2_xs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype); +at::Tensor iq2_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype); 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 +45,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 +58,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 +66,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 +74,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 " @@ -155,4 +87,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "dimension is a multiple of 32. Returns uint8 [numel / 32, 34] on the input device. " "Non-finite input elements are treated as zero during packing, and finite elements " "outside the float32 range saturate."); + module.def("iq1_s_unpack", &iq1_s_unpack_cuda, + "Decode uint8 [blocks, 50] IQ1_S payloads on CUDA into [blocks, 256] of the given " + "floating-point dtype, bit-identical to the PyTorch decoder. The grid must be " + "float32 [2048, 8]."); + module.def("iq1_m_unpack", &iq1_m_unpack_cuda, + "Decode uint8 [blocks, 56] IQ1_M payloads on CUDA into [blocks, 256] of the given " + "floating-point dtype, bit-identical to the PyTorch decoder. The grid must be " + "float32 [2048, 8]."); + module.def("iq2_xxs_unpack", &iq2_xxs_unpack_cuda, + "Decode uint8 [blocks, 66] IQ2_XXS payloads on CUDA into [blocks, 256] of the given " + "floating-point dtype, bit-identical to the PyTorch decoder. The grid must be " + "float32 [256, 8]."); + module.def("iq2_xs_unpack", &iq2_xs_unpack_cuda, + "Decode uint8 [blocks, 74] IQ2_XS payloads on CUDA into [blocks, 256] of the given " + "floating-point dtype, bit-identical to the PyTorch decoder. The grid must be " + "float32 [512, 8]."); + module.def("iq2_s_unpack", &iq2_s_unpack_cuda, + "Decode uint8 [blocks, 82] IQ2_S payloads on CUDA into [blocks, 256] of the given " + "floating-point dtype, bit-identical to the PyTorch decoder. The grid must be " + "float32 [1024, 8]."); } 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..0ea242616f1 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,91 @@ 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(); + // 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 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)); - } - } + 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); } - block_min_accumulate(local_best, warp_best, group_error); + payload[kHighOffset + tid] = static_cast(byte); } - - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; + if (tid < kScaleWords) { + uint32_t word = 0; #pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } + 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); } - locals[group] = static_cast(selected_local); - } - __syncthreads(); - 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); -#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)); - } - __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; + // Vector v is low byte v plus nibble v % 2 of qh byte v / 2: three high index bits and the delta + // sign. Its 3-bit local scale is slot 2 * (sub-block % 2) + half of scale word sub-block / 2, and + // d is reassembled from the four words' top nibbles. + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + uint32_t words[kScaleWords]; #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); + for (int word = 0; word < kScaleWords; ++word) + words[word] = load_u16(block + kScaleWordOffset + 2 * word); + const uint32_t d_bits = (words[0] >> 12) | ((words[1] >> 8) & 0x00F0) | + ((words[2] >> 4) & 0x0F00) | (words[3] & 0xF000); + const uint32_t nibble = (block[kHighOffset + vector / 2] >> (4 * (vector % 2))) & 0xF; + const uint32_t entry = block[kLowOffset + vector] | ((nibble & 0x7) << 8); + const int sub = vector / 4; + const uint32_t local = (words[sub / 2] >> (3 * (2 * (sub % 2) + (vector % 4) / 2))) & 0x7; + const float scale = __fmul_rn(half_bits_to_float(d_bits), static_cast(2 * local + 1)); + shifted_scaled(grid + entry * kVectorSize, (nibble & 0x8) ? -kIq1Delta : kIq1Delta, scale, + values); } -} +}; + +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(); + check_scaled_pack_inputs("IQ1_M", input, grid, Format::kEntries, scales); + return iq1_encode_blocks(input, grid, scales); +} - 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; +at::Tensor iq1_m_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { + return decode_blocks("IQ1_M", packed, grid, dtype); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu index cfe0aba39f3..3d7f2e455f7 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,62 @@ 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); + } + + // 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); + } + } + + // Vector v is index byte v plus three high bits from its group's metadata word, which also + // holds the group's 3-bit local scale (bits 12..14) and delta sign (bit 15). + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + const uint32_t qh = load_u16(block + kMetadataOffset + 2 * (vector / kVectorsPerGroup)); + const uint32_t entry = + block[kIndexOffset + vector] | (((qh >> (3 * (vector % kVectorsPerGroup))) & 0x7) << 8); + const float d = half_bits_to_float(load_u16(block + kScaleOffset)); + const float scale = __fmul_rn(d, static_cast(2 * ((qh >> 12) & 0x7) + 1)); + shifted_scaled(grid + entry * kVectorSize, (qh & 0x8000) ? -kIq1Delta : kIq1Delta, scale, + values); + } +}; -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(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 +89,29 @@ __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); +} + +at::Tensor iq1_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { + return decode_blocks("IQ1_S", packed, grid, dtype); } 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..8244006c192 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,61 @@ 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); + for (int k = 0; k < 4; ++k) + high |= ((entries[4 * tid + k] >> 8) & 0x3u) << (2 * k); + payload[kHighOffset + tid] = static_cast(high); } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); + } -#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(); - } + // Vector v is low byte v plus two high bits from byte v / 4 of the high array, with its own full + // sign mask. Its local scale is nibble (v / 2) % 2 of byte v / 4 in the trailing scale array. + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + const uint32_t entry = block[kLowOffset + vector] | + (((block[kHighOffset + vector / 4] >> (2 * (vector % 4))) & 0x3) << 8); + const uint32_t local = + (block[kLocalScaleOffset + vector / 4] >> (4 * ((vector / 2) % 2))) & 0xF; + const float d = half_bits_to_float(load_u16(block + kScaleOffset)); + const float scale = __fmul_rn(__fmul_rn(d, __fadd_rn(0.5f, static_cast(local))), 0.25f); + signed_scaled(grid + entry * kVectorSize, block[kSignOffset + vector], scale, values); } +}; - 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(); + check_scaled_pack_inputs("IQ2_S", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); +} - 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; +at::Tensor iq2_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { + return decode_blocks("IQ2_S", packed, grid, dtype); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu index c731fe25706..8b4725d7ad9 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,51 @@ 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); +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); } - grid_norm[entry] = norm; + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } - __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); - } - } + // Vector v is uint16 code v, a 9-bit entry under a 7-bit sign index. Its local scale is nibble + // (v / 2) % 2 of byte v / 4 in the trailing scale array. + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + const uint32_t code = load_u16(block + kCodeOffset + 2 * vector); + const uint32_t local = + (block[kLocalScaleOffset + vector / 4] >> (4 * ((vector / 2) % 2))) & 0xF; + const float d = half_bits_to_float(load_u16(block + kScaleOffset)); + const float scale = __fdiv_rn(__fmul_rn(d, static_cast(2 * local + 1)), 8.0f); + signed_scaled(grid + (code & 0x1FF) * kVectorSize, with_parity_bit(code >> 9), scale, values); } +}; - 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(); + check_scaled_pack_inputs("IQ2_XS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); +} - 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; +at::Tensor iq2_xs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { + return decode_blocks("IQ2_XS", packed, grid, dtype); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu index c501273a669..c2a11315ee5 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,58 @@ 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); - } - 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); +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 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); + for (int j = 0; j < kVectorsPerGroup; ++j) { + record[j] = static_cast(entries[group * kVectorsPerGroup + j]); + aux |= (signs[group * kVectorsPerGroup + j] & 0x7Fu) << (7 * j); } - - 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)); } -} + + // Vector v is entry byte v % 4 of record v / 4, whose uint32 holds the four vectors' 7-bit sign + // indices and, in its top nibble, the record's local scale. + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + const uint8_t *record = block + kCodeOffset + kRecordBytes * (vector / kVectorsPerGroup); + const uint32_t aux = load_u32(record + kVectorsPerGroup); + const int slot = vector % kVectorsPerGroup; + const float d = half_bits_to_float(load_u16(block + kScaleOffset)); + const float scale = + __fmul_rn(__fmul_rn(d, __fadd_rn(0.5f, static_cast(aux >> 28))), 0.25f); + signed_scaled(grid + record[slot] * kVectorSize, with_parity_bit((aux >> (7 * slot)) & 0x7F), + scale, values); + } +}; + +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(); + check_scaled_pack_inputs("IQ2_XXS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); +} - 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; +at::Tensor iq2_xxs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { + return decode_blocks("IQ2_XXS", packed, grid, dtype); } diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 590569a6f63..7f4f2fc8448 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -86,7 +86,7 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False): build as a single extension: ``iq1_s_pack(input, grid)``, ``iq2_xs_pack(input, grid, scales)``, ``iq2_xxs_pack(input, grid, scales)``, ``iq2_s_pack(input, grid, scales)``, ``iq1_m_pack(input, grid, scales)``, and - ``q8_0_pack(input)``. + ``q8_0_pack(input)``. Each IQ format also has ``_unpack(payload, grid, dtype)``. """ if not hasattr(get_cuda_ext_ggml, "extension") or ( raise_if_failed and get_cuda_ext_ggml.extension is None diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 2b39d3e73b7..6c14a29d0cf 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -19,12 +19,24 @@ import weakref from collections.abc import Callable from dataclasses import dataclass +from typing import NamedTuple import torch GGML_BLOCK_SIZE = 256 +class _CacheKey(NamedTuple): + """What a cached payload was packed from; any change means it must be packed again.""" + + data_ptr: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + device: torch.device + version: int + + @dataclass class _PackedWeightCache: """One weight's packed payload, reused across forwards. @@ -40,7 +52,7 @@ class _PackedWeightCache: """ base_ref: weakref.ReferenceType - input_key: tuple[object, ...] + input_key: _CacheKey format_name: str block_chunk_size: int packed_weights: torch.Tensor @@ -57,13 +69,13 @@ def _cache_base(inputs: torch.Tensor) -> torch.Tensor: return inputs if base is None else base -def _input_cache_key(inputs: torch.Tensor) -> tuple[object, ...] | None: +def _input_cache_key(inputs: torch.Tensor) -> _CacheKey | None: try: version = inputs._version except RuntimeError: # Inference tensors can omit version counters, so changes cannot be detected safely. return None - return ( + return _CacheKey( inputs.data_ptr(), tuple(inputs.shape), tuple(inputs.stride()), @@ -73,6 +85,53 @@ def _input_cache_key(inputs: torch.Tensor) -> tuple[object, ...] | None: ) +def _is_contiguous(shape: tuple[int, ...], stride: tuple[int, ...]) -> bool: + expected = 1 + for size, step in zip(reversed(shape), reversed(stride)): + if size != 1 and step != expected: + return False + expected *= size + return True + + +def _same_elements(key: _CacheKey, other: _CacheKey) -> bool: + """Whether two cache keys are contiguous views of the same values in the same order.""" + return ( + (key.data_ptr, key.dtype, key.device, key.version) + == (other.data_ptr, other.dtype, other.device, other.version) + and math.prod(key.shape) == math.prod(other.shape) + and _is_contiguous(key.shape, key.stride) + and _is_contiguous(other.shape, other.stride) + ) + + +def _matching_cache( + quantizer, inputs: torch.Tensor, format_name: str, *, any_contiguous_view: bool = False +) -> _PackedWeightCache | None: + """The quantizer's cached payload if it was packed from ``inputs``, else None. + + It matches only while the weight is unchanged since it was packed: the key carries the + tensor's version counter, so an in-place update misses. With ``any_contiguous_view`` it also + matches another contiguous view of the same values: TensorQuantizer hands fake quant the + weight reshaped into 256-value blocks, so export, which holds the weight itself, sees a + different shape over the same elements in the same order, and so the same GGML blocks. + """ + cache = getattr(quantizer, "_quantizer_cache", None) + input_key = _input_cache_key(inputs) + if not ( + isinstance(cache, _PackedWeightCache) + and input_key is not None + and cache.base_ref() is _cache_base(inputs) + and cache.format_name == format_name + ): + return None + if cache.input_key == input_key or ( + any_contiguous_view and _same_elements(cache.input_key, input_key) + ): + return cache + return None + + def fake_quantize_with_cache( inputs: torch.Tensor, quantizer, @@ -84,23 +143,15 @@ def fake_quantize_with_cache( dequantize: Callable[..., torch.Tensor], ) -> torch.Tensor: """Fake-quantize a weight while caching its compact packed representation.""" - input_key = _input_cache_key(inputs) - cache_base = _cache_base(inputs) - cache = getattr(quantizer, "_quantizer_cache", None) - if ( - isinstance(cache, _PackedWeightCache) - and input_key is not None - and cache.base_ref() is cache_base - and cache.input_key == input_key - and cache.format_name == format_name - and cache.block_chunk_size == block_chunk_size - ): + cache = _matching_cache(quantizer, inputs, format_name) + if cache is not None and cache.block_chunk_size == block_chunk_size: packed_weights, weight_shape = cache.packed_weights, cache.weight_shape else: packed_weights, weight_shape = quantize(inputs, block_chunk_size=block_chunk_size) + input_key = _input_cache_key(inputs) if input_key is not None: quantizer._quantizer_cache = _PackedWeightCache( - base_ref=weakref.ref(cache_base), + base_ref=weakref.ref(_cache_base(inputs)), input_key=input_key, format_name=format_name, block_chunk_size=block_chunk_size, @@ -173,6 +224,24 @@ def fake_quant( dequantize=self.dequantize, ) + def pack(self, weight: torch.Tensor, quantizer=None) -> torch.Tensor: + """The packed payload of ``weight``, reusing the one fake quant cached if it matches. + + Export calls this. A model that ran a forward since quantization already holds each + weight's payload, and reusing it both skips a second search and exports exactly the + bytes the evaluated model decoded. Any other tensor is packed afresh. + """ + cache = ( + _matching_cache(quantizer, weight, self.name, any_contiguous_view=True) + if quantizer is not None + else None + ) + if cache is None: + return self.quantize(weight)[0] + return cache.packed_weights.reshape( + *weight.shape[:-1], weight.shape[-1] // self.block_size, self.block_bytes + ) + def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does. diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index 31fde77c4c2..0799a7aaeba 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -192,6 +192,11 @@ def dequantize_iq1_m( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ1_M_BLOCK_BYTES) + if blocks.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + grid = iq1_s_grid(blocks.device) + return extension.iq1_m_unpack(blocks, grid, dtype).reshape(shape) grid = iq1_s_grid(blocks.device) scale_shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) decoded = torch.empty((blocks.shape[0], IQ1_M_BLOCK_SIZE), dtype=dtype, device=blocks.device) diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 8ecf801bdf1..90422b7a6ca 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -235,6 +235,11 @@ def dequantize_iq1_s( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES) + if blocks.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + grid = iq1_s_grid(blocks.device) + return extension.iq1_s_unpack(blocks, grid, dtype).reshape(shape) shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) grid = iq1_s_grid(blocks.device) decoded = torch.empty((blocks.shape[0], IQ1_S_BLOCK_SIZE), dtype=dtype, device=blocks.device) diff --git a/modelopt/torch/quantization/ggml/iq2_s.py b/modelopt/torch/quantization/ggml/iq2_s.py index 76067ebedda..9667301acf5 100644 --- a/modelopt/torch/quantization/ggml/iq2_s.py +++ b/modelopt/torch/quantization/ggml/iq2_s.py @@ -205,6 +205,11 @@ def dequantize_iq2_s( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ2_S_BLOCK_BYTES) + if blocks.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + grid = iq2_s_grid(blocks.device) + return extension.iq2_s_unpack(blocks, grid, dtype).reshape(shape) bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) high_shifts = torch.tensor([0, 2, 4, 6], dtype=torch.int64, device=blocks.device) grid = iq2_s_grid(blocks.device) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index a4b7d505bc0..a3cc22d2b54 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -222,6 +222,11 @@ def dequantize_iq2_xs( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES) + if blocks.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + grid = iq2_xs_grid(blocks.device) + return extension.iq2_xs_unpack(blocks, grid, dtype).reshape(shape) bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) grid = iq2_xs_grid(blocks.device) decoded = torch.empty((blocks.shape[0], IQ2_XS_BLOCK_SIZE), dtype=dtype, device=blocks.device) diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py index a27c2679632..ee35d44e3d4 100644 --- a/modelopt/torch/quantization/ggml/iq2_xxs.py +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -237,6 +237,11 @@ def dequantize_iq2_xxs( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, IQ2_XXS_BLOCK_BYTES) + if blocks.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + grid = iq2_xxs_grid(blocks.device) + return extension.iq2_xxs_unpack(blocks, grid, dtype).reshape(shape) bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) sign_shifts = 7 * torch.arange( _IQ2_XXS_GROUPS_PER_SUBBLOCK, dtype=torch.int64, device=blocks.device diff --git a/tests/gpu/torch/quantization/test_iq_formats_cuda.py b/tests/gpu/torch/quantization/test_iq_formats_cuda.py index d0756cd9b6b..37076df607c 100644 --- a/tests/gpu/torch/quantization/test_iq_formats_cuda.py +++ b/tests/gpu/torch/quantization/test_iq_formats_cuda.py @@ -22,6 +22,7 @@ import pytest import torch +from _test_utils.torch.quantization.iq_llama_cpp_vectors import expected_values, packed_blocks import modelopt.torch.quantization.ggml.iq1_m as iq1_m_module import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module @@ -186,6 +187,42 @@ def test_cuda_float64_matches_pytorch_encoder(monkeypatch, name): assert torch.equal(reference, packed) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16, torch.float64]) +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_decoder_matches_pytorch_decoder(monkeypatch, name, dtype): + """Fake quant decodes on every forward; the CUDA decoder must be the PyTorch one, bit for bit.""" + module, _, block_bytes, _ = FORMATS[name] + assert hasattr(_extension(), f"{name}_unpack") + dequantize = getattr(module, f"dequantize_{name}") + generator = torch.Generator().manual_seed(5) + # Every byte pattern, including block scales that decode to inf or NaN. + random_payload = torch.randint(0, 256, (3000, 1, block_bytes), generator=generator).to( + device="cuda", dtype=torch.uint8 + ) + weight = torch.randn((64, 1024), generator=generator).cuda() + encoded, weight_shape = getattr(module, f"quantize_{name}")(weight) + cases = [(random_payload, torch.tensor([3000, 256])), (encoded, weight_shape)] + + decoded = [dequantize(payload, shape, dtype=dtype) for payload, shape in cases] + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + for (payload, shape), got in zip(cases, decoded): + expected = dequantize(payload, shape, dtype=dtype) + assert got.dtype == dtype and got.is_cuda + torch.testing.assert_close(got, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_decoder_matches_llama_cpp_on_captured_blocks(name): + """The CUDA decoder reproduces llama.cpp's own values for bytes we did not produce.""" + module, _, block_bytes, _ = FORMATS[name] + blocks = torch.from_numpy(packed_blocks(name)).cuda() + count = blocks.shape[0] + decoded = getattr(module, f"dequantize_{name}")( + blocks.reshape(count, 1, block_bytes), torch.tensor([count, 256]), dtype=torch.float32 + ) + assert torch.equal(decoded.reshape(count, 256).cpu(), torch.from_numpy(expected_values(name))) + + def test_every_registered_format_is_covered(): """A format registered for dispatch must also be listed here, or it escapes this contract.""" assert sorted(IQ_FORMAT_REGISTRY) == sorted(FORMATS) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 94790846330..d6ae98483c7 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -14,6 +14,8 @@ # limitations under the License. +import dataclasses + import pytest import torch import torch.nn as nn @@ -26,6 +28,7 @@ _process_quantized_modules, ) from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import IQ_FORMAT_REGISTRY from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -105,9 +108,8 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) -@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)]) -def test_export_iq_payload_as_weight(num_bits, payload_bytes): - linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) +def _iq_linear(num_bits, in_features=256): + linear = nn.Linear(in_features, 4, bias=False, dtype=torch.bfloat16) linear.weight_quantizer = TensorQuantizer( QuantizerAttributeConfig( num_bits=num_bits, @@ -115,6 +117,15 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes): backend="ggml", ) ) + return linear + + +@pytest.mark.parametrize( + ("num_bits", "payload_bytes"), + [("iq1_s", 50), ("iq1_m", 56), ("iq2_xxs", 66), ("iq2_xs", 74), ("iq2_s", 82)], +) +def test_export_iq_payload_as_weight(num_bits, payload_bytes): + linear = _iq_linear(num_bits) _export_quantized_weight(linear, torch.bfloat16) state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) @@ -126,6 +137,49 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes): assert "weight_shape" not in state_dict +def _without_search(monkeypatch, num_bits): + """Make the format's encoder fail, so any call proves export ran the search again.""" + + def search(*args, **kwargs): + raise AssertionError(f"export re-ran the {num_bits} search") + + record = dataclasses.replace(IQ_FORMAT_REGISTRY[num_bits], quantize=search) + monkeypatch.setitem(IQ_FORMAT_REGISTRY, num_bits, record) + + +@pytest.mark.parametrize("num_bits", sorted(IQ_FORMAT_REGISTRY)) +def test_export_reuses_the_payload_fake_quant_packed(monkeypatch, num_bits): + """A weight fake quant already packed is exported from those bytes, not packed again. + + Fake quant sees the weight reshaped into 256-value blocks, so the cached payload is keyed on a + different shape than the weight export holds; a wider weight keeps that difference real. + """ + linear = _iq_linear(num_bits, in_features=512) + linear.weight_quantizer(linear.weight) + cache = linear.weight_quantizer._quantizer_cache + assert cache.input_key.shape != tuple(linear.weight.shape) + _without_search(monkeypatch, num_bits) + + _export_quantized_weight(linear, torch.bfloat16) + + assert torch.equal(linear.weight, cache.packed_weights.reshape(linear.weight.shape)) + assert linear.weight.shape[:2] == (4, 2) + + +@pytest.mark.parametrize("num_bits", sorted(IQ_FORMAT_REGISTRY)) +def test_export_repacks_a_weight_changed_since_fake_quant(num_bits): + """An in-place update after the forward invalidates the cached payload.""" + linear = _iq_linear(num_bits, in_features=512) + linear.weight_quantizer(linear.weight) + with torch.no_grad(): + linear.weight.mul_(-1) + expected, _ = IQ_FORMAT_REGISTRY[num_bits].quantize(linear.weight.detach().clone()) + + _export_quantized_weight(linear, torch.bfloat16) + + assert torch.equal(linear.weight, expected) + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__()