From 49f266328870415c5362b7b9f4b9c4b8180a8c3a Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 30 Sep 2026 20:41:48 +0000 Subject: [PATCH 1/3] Export IQ weights from the payload fake quant already packed hf_ptq ran the IQ search twice for every weight: once on the first forward after quantization, where fake quant packs the weight and caches the payload, and again in export, which called the format's encoder on the same unchanged weight. Instrumenting the TinyLlama example showed 154 encodes from fake quant and 154 more from export for 154 weights, in all five formats. On a 27B model at IQ1_M's 318 M elem/s the second pass is about 85 seconds of repeated work. IQFormat.pack now returns the quantizer's cached payload when it was packed from the weight being exported, and runs the search only otherwise. Fake quant is handed the weight reshaped into 256-value blocks, so the cache is keyed on that view rather than on the weight export holds. pack therefore also accepts another contiguous view of the same storage, version and length -- the same values in the same order, hence the same GGML blocks -- and reshapes the payload to the weight's layout. The version counter still makes a weight edited after its forward pack afresh. Reusing the payload also means the checkpoint holds exactly the bytes the evaluated model decoded. The HF exporter calls it; the Megatron exporter still packs itself, since it can remap or slice a weight before packing it. The cache key is now a NamedTuple so the view comparison can name its fields. The reuse test uses a 512-wide weight so the blocked view really differs from the weight, as it does on real models; with an exact-key match only, it fails for every format. The IQ payload export test now covers all five formats rather than two. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/unified_export_hf.py | 5 +- modelopt/torch/quantization/ggml/common.py | 99 ++++++++++++++++--- tests/unit/torch/export/test_export_weight.py | 60 ++++++++++- 3 files changed, 144 insertions(+), 20 deletions(-) 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/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/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__() From c898758f0d3f8490c61ca1956bf712e97dec8b00 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 30 Sep 2026 20:41:59 +0000 Subject: [PATCH 2/3] Decode the IQ formats on CUDA Every format had a CUDA encoder but decoded with PyTorch ops, which was fine while decoding looked like a one-off. Fake quant decodes every weight on every forward, though, so once packing was fast and cached, decoding became the cost: in the TinyLlama hf_ptq example it was 32 to 46 seconds of a 55 to 70 second run, 3 ms per weight per forward, for the 100 preview tokens. On a 27B model it would be about 10 seconds per forward. Each format now has a CUDA unpacker beside its packer. One thread decodes one 8-value vector: it reads that vector's index, sign or delta bits and local scale from the block, and writes the eight values in the requested dtype. common.cuh holds the launcher, the bit readers and the two value forms, sign-flipped for IQ2 and delta-shifted for IQ1. Every float operation is explicitly rounded and follows the PyTorch decoder's order, so the compiler cannot fuse a multiply into an add, and the output is bit-identical: across random payloads, including block scales that decode to inf or NaN, and real encodings, in float32, bfloat16, float16 and float64. The CUDA path also reproduces llama.cpp's own values on the conformance blocks. dequantize_ uses it for CUDA payloads and keeps the PyTorch path otherwise. Decoding a 5632x2048 weight drops from 3.3 to 5.2 ms to 0.06 to 0.14 ms. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../kernels/quantization/ggml/common.cuh | 86 +++++++++++++++++++ .../torch/kernels/quantization/ggml/ggml.cpp | 27 +++++- .../torch/kernels/quantization/ggml/iq1_m.cu | 25 ++++++ .../torch/kernels/quantization/ggml/iq1_s.cu | 18 ++++ .../torch/kernels/quantization/ggml/iq2_s.cu | 19 ++++ .../torch/kernels/quantization/ggml/iq2_xs.cu | 18 ++++ .../kernels/quantization/ggml/iq2_xxs.cu | 20 +++++ modelopt/torch/quantization/extensions.py | 2 +- modelopt/torch/quantization/ggml/iq1_m.py | 5 ++ modelopt/torch/quantization/ggml/iq1_s.py | 5 ++ modelopt/torch/quantization/ggml/iq2_s.py | 5 ++ modelopt/torch/quantization/ggml/iq2_xs.py | 5 ++ modelopt/torch/quantization/ggml/iq2_xxs.py | 5 ++ .../quantization/test_iq_formats_cuda.py | 37 ++++++++ 14 files changed, 275 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 5eccec6cd03..bdb7b1adbde 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -229,6 +229,92 @@ __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 supplies a Decoder +// with decode(block, vector, grid, values) and binds decode_blocks for its layout. +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]; + Decoder::decode(packed + (index / kVectorsPerBlock) * 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) { + 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..ba51e6e947f 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -17,7 +17,7 @@ // 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. +// translation unit and exposes a packer and an unpacker. #include "common.cuh" @@ -27,6 +27,11 @@ 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 { @@ -155,4 +160,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_m.cu b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu index b5eadc85677..138b9ec708e 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq1_m.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq1_m.cu @@ -178,6 +178,27 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g } } +// 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. +struct Decoder { + __device__ static void decode(const uint8_t *block, int vector, const float *grid, + float (&values)[kVectorSize]) { + uint32_t words[kScaleWords]; +#pragma unroll + 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) ? -kDelta : kDelta, scale, values); + } +}; + } // namespace at::Tensor iq1_m_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { @@ -203,3 +224,7 @@ at::Tensor iq1_m_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) }); 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..8ce40cbd3e7 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu @@ -161,6 +161,20 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g } } +// 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). +struct Decoder { + __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) ? -kDelta : kDelta, scale, values); + } +}; + } // namespace at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) { @@ -185,3 +199,7 @@ at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) { }); return output; } + +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_s.cu b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu index 36f9d546e78..3013001fdb0 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_s.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_s.cu @@ -183,6 +183,21 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } +// 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. +struct Decoder { + __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); + } +}; + } // namespace at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { @@ -208,3 +223,7 @@ at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) }); 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..c2c8683098a 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu @@ -197,6 +197,20 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } +// 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. +struct Decoder { + __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); + } +}; + } // namespace at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { @@ -222,3 +236,7 @@ at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales }); 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..687d61aac6b 100644 --- a/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu @@ -213,6 +213,22 @@ __global__ void encode(const scalar_t *input, int64_t num_blocks, const float *g } } +// 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. +struct Decoder { + __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); + } +}; + } // namespace at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { @@ -238,3 +254,7 @@ at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scale }); 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/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) From 7a32325f94407ec54d15ba696f010dab261e2c01 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 1 Oct 2026 05:53:57 +0000 Subject: [PATCH 3/3] Share one CUDA encoder per IQ family The five IQ encoders were five copies of the same search. IQ2_XS and IQ2_XXS shared 139 of their roughly 150 lines of encoder and launcher, and IQ2_S 113 of them; IQ1_S and IQ1_M had the same structure with a different choice space. Every scaled packer also validated its scales twice, in the pybind wrapper and again in the CUDA entry point. iq2_family.cuh now holds the one IQ2 encoder: the grid in shared memory, the 16 local scales scored per group, then the best entry per vector under the chosen scale. A format supplies its group shape, whether it stores seven sign bits and recovers the eighth from parity, and a store() that writes the chosen entries, sign masks and local scales into its layout. iq1_family.cuh holds the one IQ1 encoder over the shared ternary grid. Each group picks one of kChoices options; with kSharedShift the option also fixes the +/-1/8 delta (IQ1_S: shift * 8 + local), otherwise each vector picks its own (IQ1_M). IQ1_S's scale kernel now writes FP16 scales, so both IQ1 formats take the same input. Each format file is now one Format struct -- layout constants, store() and the existing decode() -- plus its two entry points. Validation lives once in common.cuh, as check_pack_inputs and check_scaled_pack_inputs, and ggml.cpp binds the CUDA entry points directly, as it already did the unpackers. Kernel sources drop from 1,747 to 1,445 lines. Nothing changes in the output. Every encoder and decoder hashes the same as before on float32, bfloat16, float16 and float64 inputs with zero, tiny, oversized and non-finite blocks, and encode speed is unchanged within 1% on an idle RTX PRO 6000. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../kernels/quantization/ggml/common.cuh | 40 ++- .../torch/kernels/quantization/ggml/ggml.cpp | 89 +----- .../kernels/quantization/ggml/iq1_family.cuh | 191 +++++++++++++ .../torch/kernels/quantization/ggml/iq1_m.cu | 228 ++++------------ .../torch/kernels/quantization/ggml/iq1_s.cu | 220 +++++---------- .../kernels/quantization/ggml/iq2_family.cuh | 256 ++++++++++++++++++ .../torch/kernels/quantization/ggml/iq2_s.cu | 215 +++------------ .../torch/kernels/quantization/ggml/iq2_xs.cu | 224 ++------------- .../kernels/quantization/ggml/iq2_xxs.cu | 237 ++-------------- 9 files changed, 699 insertions(+), 1001 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/ggml/iq1_family.cuh create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_family.cuh diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index bdb7b1adbde..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 @@ -231,8 +249,8 @@ __device__ __forceinline__ bool store_block_scale(uint8_t *payload, uint16_t d_b // 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 supplies a Decoder -// with decode(block, vector, grid, values) and binds decode_blocks for its layout. +// 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) { @@ -268,15 +286,15 @@ __device__ __forceinline__ void shifted_scaled(const float *q, float delta, floa values[j] = __fmul_rn(__fadd_rn(q[j], delta), scale); } -template +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]; - Decoder::decode(packed + (index / kVectorsPerBlock) * kPayloadBytes, - static_cast(index % kVectorsPerBlock), grid, values); + 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) @@ -284,9 +302,11 @@ __global__ void decode_vectors(const uint8_t *packed, int64_t num_vectors, const } // Decodes uint8 [blocks, kPayloadBytes] into [blocks, 256] of dtype on the payload's device. -template +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 && @@ -307,7 +327,7 @@ at::Tensor decode_blocks(const char *format, const at::Tensor &packed, const at: const auto stream = c10::cuda::getCurrentCUDAStream(); AT_DISPATCH_FLOATING_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, dtype, "ggml_decode", [&] { - decode_vectors<<>>( + decode_vectors<<>>( payload.data_ptr(), num_vectors, table.data_ptr(), output.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index ba51e6e947f..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 packer and an unpacker. +// 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" @@ -35,80 +36,6 @@ at::Tensor iq2_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType 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); @@ -118,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 " @@ -131,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 " @@ -139,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 " @@ -147,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 " 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 138b9ec708e..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,164 +24,64 @@ using namespace modelopt::ggml; // The IQ1_M packed payload layout and format constants below follow the GGML // definition at: // https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h -// IQ1_M shares the IQ1_S ternary grid. At 64 KiB it is past the 48 KiB static shared-memory -// limit, so both kernels read it from global memory and rely on the cache. -constexpr int kEntries = kIq1sEntries; -constexpr int kGroups = 16; // one 3-bit local scale per 16 values -constexpr int kVectorsPerGroup = 2; -constexpr int kVectors = kGroups * kVectorsPerGroup; -constexpr int kLocalScales = 8; -constexpr int kSubBlocks = 8; -constexpr int kScaleWords = 4; -constexpr int kLowOffset = 0; // no leading block scale in this format -constexpr int kHighOffset = kLowOffset + kVectors; // 16 bytes, two vectors per byte -constexpr int kScaleWordOffset = kHighOffset + 2 * kSubBlocks; -constexpr int kPayloadBytes = kScaleWordOffset + 2 * kScaleWords; -constexpr float kDelta = 0.125f; -constexpr int kEntryBits = 11; // 2048 entries; the shift bit sits above them in the sort key - -static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); -static_assert(kPayloadBytes == 56, "IQ1_M blocks are 56 bytes"); - -template -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - - const int tid = threadIdx.x; - const int64_t block = blockIdx.x; - if (block >= num_blocks) - return; - - const scalar_t *source = input + block * kBlockSize; - uint8_t *payload = output + block * kPayloadBytes; - const __half d_half = scales[block]; - const uint16_t d_bits = __half_as_ushort(d_half); - const float d = __half2float(d_half); - // IQ1_M has no leading scale field, so the shared helper does not apply: a zero block is all - // zero bytes, which decodes to a zero scale and therefore zero values. - if ((d_bits & 0x7FFF) == 0) { - if (tid < kPayloadBytes) - payload[tid] = 0; - return; +struct Format { + static constexpr int kEntries = kIq1sEntries; // IQ1_M shares the IQ1_S ternary grid + static constexpr int kGroups = 16; // one 3-bit local scale per 16 values + static constexpr int kVectorsPerGroup = 2; + static constexpr int kChoices = 8; // the local scale; each vector picks its own shift + static constexpr bool kSharedShift = false; + static constexpr int kSubBlocks = 8; + static constexpr int kScaleWords = 4; + static constexpr int kLowOffset = 0; // no leading block scale + static constexpr int kHighOffset = kLowOffset + kVectorsPerBlock; // two vectors per byte + static constexpr int kScaleWordOffset = kHighOffset + 2 * kSubBlocks; + static constexpr int kPayloadBytes = kScaleWordOffset + 2 * kScaleWords; + + // IQ1_M has no leading scale field: a zero block is all zero bytes, which decodes to a zero + // scale and therefore zero values. + __device__ static bool begin(uint8_t *payload, uint16_t d_bits) { + if ((d_bits & 0x7FFF) != 0) + return true; + if (threadIdx.x < kPayloadBytes) + payload[threadIdx.x] = 0; + return false; } - if (tid < kHighOffset + 2 * kSubBlocks) - payload[tid] = 0; // index-high and shift nibbles are OR-ed into below - __syncthreads(); - -#pragma unroll 1 - for (int group = 0; group < kGroups; ++group) { - if (tid < kLocalScales) - group_error[tid] = 0.0f; - __syncthreads(); - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm, xsum; - load_vector(source + group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize, x, - xnorm, xsum); - float local_best[kLocalScales]; -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) - local_best[local] = FLT_MAX; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - float dot, qnorm, qsum; - grid_terms(x, grid + entry * kVectorSize, dot, qnorm, qsum); - // The shift is free per vector here, unlike IQ1_S where it is shared, so take the - // better of the two before the local scale is chosen. -#pragma unroll - for (int shift = 0; shift < 2; ++shift) { - const float delta = shift ? -kDelta : kDelta; -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1); - local_best[local] = fminf(local_best[local], - shifted_error(xnorm, xsum, dot, qnorm, qsum, scale, delta)); - } - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; + // A low index byte per vector; a nibble per vector, low nibble first, holding three index-high + // bits and the shift; then four scale words, each with two sub-blocks' 3-bit local scales in + // bits 0..11 and one nibble of the FP16 block scale in bits 12..15. + __device__ static void store(uint8_t *payload, const uint16_t *picks, const uint8_t *choices, + uint16_t d_bits) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) + payload[kLowOffset + tid] = static_cast(picks[tid]); + if (tid < kVectorsPerBlock / 2) { + uint32_t byte = 0; #pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } + for (int half = 0; half < 2; ++half) { + const uint32_t pick = picks[2 * tid + half]; + byte |= (((pick >> 8) & 0x7u) | (((pick >> kIq1EntryBits) & 0x1u) << 3)) << (4 * half); } - locals[group] = static_cast(selected_local); + payload[kHighOffset + tid] = static_cast(byte); } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1); - -#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); + if (tid < kScaleWords) { + uint32_t word = 0; #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; - } + 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); } - 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; -#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); - } -} - -// 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. -struct Decoder { + // 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]; @@ -195,36 +95,20 @@ struct Decoder { 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) ? -kDelta : kDelta, scale, values); + 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(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_m_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ1_M", input, grid, Format::kEntries, scales); + return iq1_encode_blocks(input, grid, scales); } at::Tensor iq1_m_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { - return decode_blocks("IQ1_M", packed, grid, 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 8ce40cbd3e7..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,146 +24,39 @@ 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; - -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"); - -template -__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { - const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (block >= num_blocks) - return; - - float amax = 0.0f; - const scalar_t *values = input + block * kBlockSize; -#pragma unroll 1 - for (int i = 0; i < kBlockSize; ++i) - amax = fmaxf(amax, fabsf(load_float(values + i))); - // 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); +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 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); - } + 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); } - - 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(); } -} -// 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). -struct Decoder { + // 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)); @@ -171,35 +64,54 @@ struct Decoder { 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) ? -kDelta : kDelta, scale, values); + shifted_scaled(grid + entry * kVectorSize, (qh & 0x8000) ? -kIq1Delta : kIq1Delta, scale, + values); } }; +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, __half *scales) { + const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) + return; + + float amax = 0.0f; + const scalar_t *values = input + block * kBlockSize; +#pragma unroll 1 + for (int i = 0; i < kBlockSize; ++i) + amax = fmaxf(amax, fabsf(load_float(values + i))); + 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); + 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 3013001fdb0..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,168 +24,40 @@ 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; -#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]; + if (tid < kSubBlocks) { + uint32_t high = 0; #pragma unroll - for (int local = 0; local < kLocalScales; ++local) - local_best[local] = FLT_MAX; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float dot = magnitude_dot(x, shared_grid + entry * kVectorSize); -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; -#pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } - } - locals[group] = static_cast(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - } - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float error = - clamped_quant_error(xnorm, magnitude_dot(x, shared_grid + entry * kVectorSize), - grid_norm[entry], selected_scale); - const unsigned long long candidate = error_key(error, entry); - key = candidate < key ? candidate : key; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - int sign_mask = 0; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) - sign_mask |= static_cast(x[j] < 0.0f) << j; - const int slot = group * kVectorsPerGroup + vector; // 0..31, GGML's qs index - payload[kLowOffset + slot] = static_cast(entry & 0xFF); - payload[kSignOffset + slot] = static_cast(sign_mask); - // Two high index bits per vector, packed four to a byte in vector order. - payload[kHighOffset + slot / 4] |= - static_cast(((entry >> 8) & 0x3) << (2 * (slot % 4))); - } - __syncthreads(); + for (int k = 0; k < 4; ++k) + high |= ((entries[4 * tid + k] >> 8) & 0x3u) << (2 * k); + payload[kHighOffset + tid] = static_cast(high); } + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } - if (tid < kGroups / 2) - payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); -} - -// 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. -struct Decoder { + // 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] | @@ -198,32 +70,15 @@ struct Decoder { } }; +static_assert(Format::kPayloadBytes == 82, "IQ2_S blocks are 82 bytes"); + } // namespace at::Tensor iq2_s_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(), - "inputs must be contiguous"); - check_pack_inputs("IQ2_S", input, grid, kEntries); - const int64_t num_blocks = input.numel() / kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_s_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_S", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } at::Tensor iq2_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { - return decode_blocks("IQ2_S", packed, grid, 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 c2c8683098a..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,182 +24,31 @@ using namespace modelopt::ggml; // The IQ2_XS packed payload layout and format constants below follow the GGML // definition at: // https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h -constexpr int kEntries = kIq2xsEntries; -constexpr int kGroups = 16; -constexpr int kVectorsPerGroup = 2; -constexpr int kLocalScales = 16; -constexpr int kCodeOffset = kScaleBytes; -constexpr int kCodeBytes = 2 * (kBlockSize / kVectorSize); -constexpr int kLocalScaleOffset = kCodeOffset + kCodeBytes; -constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; -constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. - -static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); -static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two"); - -// Dot product of |x| against one codebook vector, under the format's even-parity sign rule. The -// grid must hold non-negative magnitudes: the signs live in the packed 7-bit field, and the -// eighth sign is recovered from the parity of the other seven during decoding. For odd parity, -// flip the coordinate with the smallest |x| * q penalty. -__device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) { - float dot = 0.0f; - float weakest = FLT_MAX; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - const float term = fabsf(x[j]) * q[j]; - dot += term; - weakest = fminf(weakest, term); - } - return odd_parity ? dot - 2.0f * weakest : dot; -} - -template -__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, - const __half *scales, uint8_t *output) { - __shared__ float shared_grid[kEntries * kVectorSize]; - __shared__ float grid_norm[kEntries]; - __shared__ float warp_best[kWarps * kLocalScales]; - __shared__ float group_error[kLocalScales]; - __shared__ unsigned long long warp_keys[kWarps]; - __shared__ int selected_local; - __shared__ uint8_t locals[kGroups]; - - const int tid = threadIdx.x; - const int64_t block = blockIdx.x; - if (block >= num_blocks) - return; - - for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x) - shared_grid[i] = grid[i]; - __syncthreads(); - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - float norm = 0.0f; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - const float q = shared_grid[entry * kVectorSize + j]; - norm = fmaf(q, q, norm); - } - grid_norm[entry] = norm; - } - __syncthreads(); - - const scalar_t *source = input + block * kBlockSize; - uint8_t *payload = output + block * kPayloadBytes; - const __half d_half = scales[block]; - const uint16_t d_bits = __half_as_ushort(d_half); - const float d = __half2float(d_half); - if (!store_block_scale(payload, d_bits)) - return; - -#pragma unroll 1 - for (int group = 0; group < kGroups; ++group) { - if (tid < kLocalScales) - group_error[tid] = 0.0f; - __syncthreads(); - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - negative_count += x[j] < 0.0f; - } - const bool odd_parity = (negative_count & 1) != 0; - float local_best[kLocalScales]; -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) - local_best[local] = FLT_MAX; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float *q = shared_grid + entry * kVectorSize; - const float dot = even_parity_dot(x, q, odd_parity); -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; -#pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } - } - locals[group] = static_cast(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - negative_count += x[j] < 0.0f; - } - const bool odd_parity = (negative_count & 1) != 0; - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float error = clamped_quant_error( - xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), - grid_norm[entry], selected_scale); - const unsigned long long candidate = error_key(error, entry); - key = candidate < key ? candidate : key; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - const float *q = shared_grid + entry * kVectorSize; - int flip_index = 0; - float weakest = fabsf(x[0]) * q[0]; -#pragma unroll - for (int j = 1; j < kVectorSize; ++j) { - const float term = fabsf(x[j]) * q[j]; - if (term < weakest) { - weakest = term; - flip_index = j; - } - } - int sign_mask = 0; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - bool is_negative = x[j] < 0.0f; - if (odd_parity && j == flip_index) - is_negative = !is_negative; - sign_mask |= static_cast(is_negative) << j; - } - const uint16_t code = static_cast(entry | ((sign_mask & 0x7f) << 9)); - const int code_offset = kCodeOffset + 2 * (group * kVectorsPerGroup + vector); - payload[code_offset] = static_cast(code); - payload[code_offset + 1] = static_cast(code >> 8); - } +struct Format { + static constexpr int kEntries = kIq2xsEntries; + static constexpr int kGroups = 16; // one 4-bit local scale per 16 values + static constexpr int kVectorsPerGroup = 2; + static constexpr bool kParitySigns = true; + static constexpr int kCodeOffset = kScaleBytes; + static constexpr int kLocalScaleOffset = kCodeOffset + 2 * kVectorsPerBlock; + static constexpr int kPayloadBytes = kLocalScaleOffset + kGroups / 2; + + // One uint16 code per vector, a 9-bit entry under seven sign bits, then the local scales as + // nibbles, two per byte. + __device__ static void store(uint8_t *payload, const uint16_t *entries, const uint8_t *signs, + const uint8_t *locals) { + const int tid = threadIdx.x; + if (tid < kVectorsPerBlock) { + const uint32_t code = entries[tid] | ((signs[tid] & 0x7Fu) << 9); + payload[kCodeOffset + 2 * tid] = static_cast(code); + payload[kCodeOffset + 2 * tid + 1] = static_cast(code >> 8); } + if (tid < kGroups / 2) + payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); } - if (tid < kGroups / 2) - payload[kLocalScaleOffset + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); -} - -// 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. -struct Decoder { + // 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); @@ -211,32 +60,15 @@ struct Decoder { } }; +static_assert(Format::kPayloadBytes == 74, "IQ2_XS blocks are 74 bytes"); + } // namespace at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(), - "inputs must be contiguous"); - check_pack_inputs("IQ2_XS", input, grid, kEntries); - const int64_t num_blocks = input.numel() / kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_XS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } at::Tensor iq2_xs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { - return decode_blocks("IQ2_XS", packed, grid, 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 687d61aac6b..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,198 +24,36 @@ 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]; +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) - local_best[local] = FLT_MAX; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float *q = shared_grid + entry * kVectorSize; - const float dot = even_parity_dot(x, q, odd_parity); -#pragma unroll - for (int local = 0; local < kLocalScales; ++local) { - const float scale = d * (2 * local + 1) * kLocalScaleStep; - local_best[local] = - fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); - } - } - block_min_accumulate(local_best, warp_best, group_error); - } - - if (tid == 0) { - selected_local = 0; - float best = group_error[0]; -#pragma unroll - for (int local = 1; local < kLocalScales; ++local) { - if (group_error[local] < best) { - best = group_error[local]; - selected_local = local; - } - } - locals[group] = static_cast(selected_local); - } - __syncthreads(); - const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; - -#pragma unroll - for (int vector = 0; vector < kVectorsPerGroup; ++vector) { - float x[kVectorSize]; - float xnorm = 0.0f; - int negative_count = 0; - const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - x[j] = load_float(source + offset + j); - xnorm = fmaf(x[j], x[j], xnorm); - negative_count += x[j] < 0.0f; - } - const bool odd_parity = (negative_count & 1) != 0; - unsigned long long key = ~0ULL; - for (int entry = tid; entry < kEntries; entry += blockDim.x) { - const float error = clamped_quant_error( - xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), - grid_norm[entry], selected_scale); - const unsigned long long candidate = error_key(error, entry); - key = candidate < key ? candidate : key; - } - key = block_min_key(key, warp_keys); - if (tid == 0) { - const int entry = static_cast(key & (kEntries - 1)); - const float *q = shared_grid + entry * kVectorSize; - int flip_index = 0; - float weakest = fabsf(x[0]) * q[0]; -#pragma unroll - for (int j = 1; j < kVectorSize; ++j) { - const float term = fabsf(x[j]) * q[j]; - if (term < weakest) { - weakest = term; - flip_index = j; - } - } - int sign_mask = 0; -#pragma unroll - for (int j = 0; j < kVectorSize; ++j) { - bool is_negative = x[j] < 0.0f; - if (odd_parity && j == flip_index) - is_negative = !is_negative; - sign_mask |= static_cast(is_negative) << j; - } - const int slot = group * kVectorsPerGroup + vector; - entry_bytes[slot] = static_cast(entry); - sign_bits[slot] = static_cast(sign_mask & 0x7f); - } - __syncthreads(); + for (int j = 0; j < kVectorsPerGroup; ++j) { + record[j] = static_cast(entries[group * kVectorsPerGroup + j]); + aux |= (signs[group * kVectorsPerGroup + j] & 0x7Fu) << (7 * j); } - } - - // 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. -struct Decoder { + // 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); @@ -229,32 +67,15 @@ struct Decoder { } }; +static_assert(Format::kPayloadBytes == 66, "IQ2_XXS blocks are 66 bytes"); + } // namespace at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { - TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(), - "inputs must be contiguous"); - check_pack_inputs("IQ2_XXS", input, grid, kEntries); - const int64_t num_blocks = input.numel() / kBlockSize; - TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && - scales.numel() == num_blocks, - "scales must be float16 [numel / 256]"); - TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); - c10::cuda::CUDAGuard guard(input.device()); - auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); - const auto stream = c10::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xxs_pack", [&] { - encode<<(num_blocks), kThreads, 0, stream>>>( - input.data_ptr(), num_blocks, grid.data_ptr(), - reinterpret_cast(scales.data_ptr()), - output.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); - return output; + check_scaled_pack_inputs("IQ2_XXS", input, grid, Format::kEntries, scales); + return iq2_encode_blocks(input, grid, scales); } at::Tensor iq2_xxs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) { - return decode_blocks("IQ2_XXS", packed, grid, dtype); + return decode_blocks("IQ2_XXS", packed, grid, dtype); }