Pack each IQ weight once and decode the IQ formats on CUDA

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

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<uint32_t>(bytes[0]) | (static_cast<uint32_t>(bytes[1]) << 8);
+}
+
+__device__ __forceinline__ uint32_t load_u32(const uint8_t *bytes) {
+  return load_u16(bytes) | (load_u16(bytes + 2) << 16);
+}
+
+__device__ __forceinline__ float half_bits_to_float(uint32_t bits) {
+  return __half2float(__ushort_as_half(static_cast<unsigned short>(bits)));
+}
+
+// IQ2_XS and IQ2_XXS store seven sign bits; the eighth makes the count of negatives even.
+__device__ __forceinline__ uint32_t with_parity_bit(uint32_t sign_index) {
+  return sign_index | ((__popc(sign_index) & 1u) << 7);
+}
+
+// x * scale, with coordinate j negated where bit j of sign_mask is set.
+__device__ __forceinline__ void signed_scaled(const float *q, uint32_t sign_mask, float scale,
+                                              float (&values)[kVectorSize]) {
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j)
+    values[j] = __fmul_rn((sign_mask >> j) & 1 ? -q[j] : q[j], scale);
+}
+
+// (q + delta) * scale, the IQ1 formats' shifted ternary grid.
+__device__ __forceinline__ void shifted_scaled(const float *q, float delta, float scale,
+                                               float (&values)[kVectorSize]) {
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j)
+    values[j] = __fmul_rn(__fadd_rn(q[j], delta), scale);
+}
+
+template <typename Format, typename out_t>
+__global__ void decode_vectors(const uint8_t *packed, int64_t num_vectors, const float *grid,
+                               out_t *output) {
+  const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= num_vectors)
+    return;
+  float values[kVectorSize];
+  Format::decode(packed + (index / kVectorsPerBlock) * Format::kPayloadBytes,
+                 static_cast<int>(index % kVectorsPerBlock), grid, values);
+  out_t *out = output + index * kVectorSize;
+#pragma unroll
+  for (int j = 0; j < kVectorSize; ++j)
+    out[j] = static_cast<out_t>(values[j]);
+}
+
+// Decodes uint8 [blocks, kPayloadBytes] into [blocks, 256] of dtype on the payload's device.
+template <typename Format>
+at::Tensor decode_blocks(const char *format, const at::Tensor &packed, const at::Tensor &grid,
+                         at::ScalarType dtype) {
+  constexpr int kPayloadBytes = Format::kPayloadBytes;
+  constexpr int kEntries = Format::kEntries;
+  TORCH_CHECK(packed.is_cuda() && grid.is_cuda(), format, " decoding requires CUDA tensors");
+  TORCH_CHECK(packed.get_device() == grid.get_device(), "payload and grid must share a device");
+  TORCH_CHECK(packed.scalar_type() == at::kByte && packed.dim() == 2 &&
+                  packed.size(1) == kPayloadBytes,
+              format, " payloads must be uint8 [blocks, ", kPayloadBytes, "]");
+  TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.dim() == 2 && grid.size(0) == kEntries &&
+                  grid.size(1) == kVectorSize,
+              format, " grid must be float32 [", kEntries, ", ", kVectorSize, "]");
+  TORCH_CHECK(at::isFloatingType(dtype), format, " decodes to a floating-point dtype");
+  c10::cuda::CUDAGuard guard(packed.device());
+  const auto payload = packed.contiguous();
+  const auto table = grid.contiguous();
+  auto output = at::empty({packed.size(0), kBlockSize}, packed.options().dtype(dtype));
+  const int64_t num_vectors = packed.size(0) * kVectorsPerBlock;
+  if (num_vectors == 0)
+    return output;
+  const auto launch_blocks = static_cast<unsigned>((num_vectors + kThreads - 1) / kThreads);
+  const auto stream = c10::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, dtype, "ggml_decode", [&] {
+        decode_vectors<Format, scalar_t><<<launch_blocks, kThreads, 0, stream>>>(
+            payload.data_ptr<uint8_t>(), num_vectors, table.data_ptr<float>(),
+            output.data_ptr<scalar_t>());
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+      });
+  return output;
+}
+
 #endif // __CUDACC__
 
 } // namespace modelopt::ggml

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].");
 }

modelopt/torch/kernels/quantization/ggml/iq1_m.cu

@@ -78,6 +78,26 @@ struct Format {
       payload[kScaleWordOffset + 2 * tid + 1] = static_cast<uint8_t>(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<float>(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<Format>(input, grid, scales);
 }
+
+at::Tensor iq1_m_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) {
+  return decode_blocks<Format>("IQ1_M", packed, grid, dtype);
+}

modelopt/torch/kernels/quantization/ggml/iq1_s.cu

@@ -54,6 +54,19 @@ struct Format {
       payload[kMetadataOffset + 2 * tid + 1] = static_cast<uint8_t>(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<float>(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<Format>(values, grid, scales);
 }
+
+at::Tensor iq1_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) {
+  return decode_blocks<Format>("IQ1_S", packed, grid, dtype);
+}

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<float>(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<Format>(input, grid, scales);
 }
+
+at::Tensor iq2_s_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) {
+  return decode_blocks<Format>("IQ2_S", packed, grid, dtype);
+}

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<float>(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<Format>(input, grid, scales);
 }
+
+at::Tensor iq2_xs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) {
+  return decode_blocks<Format>("IQ2_XS", packed, grid, dtype);
+}

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<uint8_t>(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<float>(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<Format>(input, grid, scales);
 }
+
+at::Tensor iq2_xxs_unpack_cuda(at::Tensor packed, at::Tensor grid, at::ScalarType dtype) {
+  return decode_blocks<Format>("IQ2_XXS", packed, grid, dtype);
+}

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 ``<format>_unpack(payload, grid, dtype)``.
     """
     if not hasattr(get_cuda_ext_ggml, "extension") or (
         raise_if_failed and get_cuda_ext_ggml.extension is None