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