Skip to content

Share one CUDA encoder per IQ family #2614

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
cjluo-nv wants to merge 3 commits into from
Closed

Share one CUDA encoder per IQ family #2614

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
49f2663
Export IQ weights from the payload fake quant already packed
cjluo-nv Sep 30, 2026
c898758
Decode the IQ formats on CUDA
cjluo-nv Sep 30, 2026
7a32325
Share one CUDA encoder per IQ family
cjluo-nv Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions modelopt/torch/export/unified_export_hf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
112 changes: 109 additions & 3 deletions modelopt/torch/kernels/quantization/ggml/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,18 +72,36 @@ 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,
"grid must be float32 [", entries, ", ", kVectorSize, "]");
TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device");
}

// Validates a format that takes one precomputed FP16 block scale per 256 values. The kernels copy
// these bits straight into the block scale field: a non-finite entry produces a payload that
// decodes to garbage, and a negative one inverts the sign of every decoded element while still
// packing cleanly -- GGML's own encoders assert a non-negative block scale. One fused reduction, so
// the synchronization is paid once per packed tensor, on an export path.
inline void check_scaled_pack_inputs(const char *format, const at::Tensor &input,
const at::Tensor &grid, int64_t entries,
const at::Tensor &scales) {
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<bool>(),
"scales must be finite and non-negative");
TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
}

#ifdef __CUDACC__

// Reads one input element as float32. Non-finite elements are treated as zero, and finiteness is
Expand Down Expand Up @@ -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<uint32_t>(bytes[0]) | (static_cast<uint32_t>(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<unsigned short>(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 <typename Format, typename out_t>
__global__ void decode_vectors(const uint8_t *packed, int64_t num_vectors, const float *grid,
out_t *output) {
const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (index >= num_vectors)
return;
float values[kVectorSize];
Format::decode(packed + (index / kVectorsPerBlock) * Format::kPayloadBytes,
static_cast<int>(index % kVectorsPerBlock), grid, values);
out_t *out = output + index * kVectorSize;
#pragma unroll
for (int j = 0; j < kVectorSize; ++j)
out[j] = static_cast<out_t>(values[j]);
}

// Decodes uint8 [blocks, kPayloadBytes] into [blocks, 256] of dtype on the payload's device.
template <typename Format>
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<unsigned>((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<Format, scalar_t><<<launch_blocks, kThreads, 0, stream>>>(
payload.data_ptr<uint8_t>(), num_vectors, table.data_ptr<float>(),
output.data_ptr<scalar_t>());
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
return output;
}

#endif // __CUDACC__

} // namespace modelopt::ggml
114 changes: 33 additions & 81 deletions modelopt/torch/kernels/quantization/ggml/ggml.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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<bool>(),
"scales must be finite and non-negative");
TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
}

at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
TORCH_CHECK(input.is_cuda(), "IQ2_XXS packing requires a CUDA input");
TORCH_CHECK(grid.is_cuda(), "IQ2_XXS packing requires a CUDA grid");
TORCH_CHECK(scales.is_cuda(), "IQ2_XXS packing requires CUDA scales");
modelopt::ggml::check_pack_inputs("IQ2_XXS", input, grid, modelopt::ggml::kIq2xxsEntries);
const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize;
TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
scales.numel() == num_blocks,
"scales must be float16 [numel / 256]");
// Same rule as IQ2_XS: these bits become the block scale verbatim, and a negative or non-finite
// one packs cleanly while decoding to garbage.
TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item<bool>(),
"scales must be finite and non-negative");
TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
return iq2_xxs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
}

at::Tensor iq2_s_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
TORCH_CHECK(input.is_cuda(), "IQ2_S packing requires a CUDA input");
TORCH_CHECK(grid.is_cuda(), "IQ2_S packing requires a CUDA grid");
TORCH_CHECK(scales.is_cuda(), "IQ2_S packing requires CUDA scales");
modelopt::ggml::check_pack_inputs("IQ2_S", input, grid, modelopt::ggml::kIq2sEntries);
const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize;
TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
scales.numel() == num_blocks,
"scales must be float16 [numel / 256]");
TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item<bool>(),
"scales must be finite and non-negative");
TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
return iq2_s_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
}

at::Tensor iq1_m_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) {
TORCH_CHECK(input.is_cuda(), "IQ1_M packing requires a CUDA input");
TORCH_CHECK(grid.is_cuda(), "IQ1_M packing requires a CUDA grid");
TORCH_CHECK(scales.is_cuda(), "IQ1_M packing requires CUDA scales");
modelopt::ggml::check_pack_inputs("IQ1_M", input, grid, modelopt::ggml::kIq1sEntries);
const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize;
TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 &&
scales.numel() == num_blocks,
"scales must be float16 [numel / 256]");
TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item<bool>(),
"scales must be finite and non-negative");
TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device");
return iq1_m_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous());
}

at::Tensor q8_0_pack(at::Tensor input) {
TORCH_CHECK(input.is_cuda(), "Q8_0 packing requires a CUDA input");
modelopt::ggml::check_scalar_pack_input("Q8_0", input, 32);
Expand All @@ -113,36 +45,36 @@ at::Tensor q8_0_pack(at::Tensor input) {
} // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
module.def("iq1_s_pack", &iq1_s_pack,
module.def("iq1_s_pack", &iq1_s_pack_cuda,
"Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
"dimension is a multiple of 256. The grid must be float32 [2048, 8]. Returns uint8 "
"[numel / 256, 50] on the input device. Non-finite input elements are treated as "
"zero during packing, and finite elements outside the float32 range saturate.");
module.def("iq2_xs_pack", &iq2_xs_pack,
module.def("iq2_xs_pack", &iq2_xs_pack_cuda,
"Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
"dimension is a multiple of 256. The grid must be float32 [512, 8] holding "
"non-negative codebook magnitudes, and scales must be finite non-negative float16 "
"[numel / 256]. "
"Returns uint8 [numel / 256, 74] on the input device. Non-finite input elements are "
"treated as zero during packing, and finite elements outside the float32 range "
"saturate.");
module.def("iq2_xxs_pack", &iq2_xxs_pack,
module.def("iq2_xxs_pack", &iq2_xxs_pack_cuda,
"Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
"dimension is a multiple of 256. The grid must be float32 [256, 8] holding "
"non-negative codebook magnitudes, and scales must be finite non-negative float16 "
"[numel / 256]. "
"Returns uint8 [numel / 256, 66] on the input device. Non-finite input elements are "
"treated as zero during packing, and finite elements outside the float32 range "
"saturate.");
module.def("iq2_s_pack", &iq2_s_pack,
module.def("iq2_s_pack", &iq2_s_pack_cuda,
"Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
"dimension is a multiple of 256. The grid must be float32 [1024, 8] holding "
"non-negative codebook magnitudes, and scales must be finite non-negative float16 "
"[numel / 256]. "
"Returns uint8 [numel / 256, 82] on the input device. Non-finite input elements are "
"treated as zero during packing, and finite elements outside the float32 range "
"saturate.");
module.def("iq1_m_pack", &iq1_m_pack,
module.def("iq1_m_pack", &iq1_m_pack_cuda,
"Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost "
"dimension is a multiple of 256. The grid must be float32 [2048, 8], and scales must "
"be finite non-negative float16 [numel / 256]; IQ1_M carries that scale in the four "
Expand All @@ -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].");
}
Loading
Loading