gguf_rs_lib/tensor/
quantization.rs

1//! Quantization format structures and utilities
2
3use crate::format::types::GGUFTensorType as TensorType;
4#[cfg(feature = "std")]
5use serde::{Deserialize, Serialize};
6#[cfg(feature = "std")]
7use std::cmp;
8
9#[cfg(all(not(feature = "std"), feature = "alloc"))]
10extern crate alloc;
11#[cfg(all(not(feature = "std"), feature = "alloc"))]
12use alloc::{vec, vec::Vec};
13#[cfg(feature = "std")]
14use std::vec::Vec;
15
16// Import core modules for no_std compatibility
17#[cfg(not(feature = "std"))]
18use core::{cmp, fmt};
19
20/// Exact-size opaque views of GGML quantization blocks.
21///
22/// The encoded byte layouts are stable, but many formats use packed bitfields
23/// whose semantic field representation depends on GGML implementation details.
24/// Keeping those blocks opaque prevents this crate from promising incorrect
25/// field layouts while still making their canonical sizes available.
26pub mod blocks {
27    macro_rules! opaque_block {
28        ($(#[$meta:meta])* $name:ident, $size:expr) => {
29            $(#[$meta])*
30            #[repr(transparent)]
31            #[allow(non_camel_case_types)]
32            #[derive(Debug, Clone, Copy, PartialEq, Eq)]
33            pub struct $name {
34                /// The canonical encoded GGML block bytes.
35                pub bytes: [u8; $size],
36            }
37
38            impl $name {
39                /// Encoded size of one block in bytes.
40                pub const SIZE: usize = $size;
41
42                /// Construct a block from its exact encoded representation.
43                pub const fn from_bytes(bytes: [u8; $size]) -> Self {
44                    Self { bytes }
45                }
46
47                /// Return the exact encoded representation.
48                pub const fn into_bytes(self) -> [u8; $size] {
49                    self.bytes
50                }
51            }
52        };
53    }
54
55    opaque_block!(/// Q4_0 block: 32 weights in 18 bytes.
56        Q4_0Block, 18);
57    opaque_block!(/// Q4_1 block: 32 weights in 20 bytes.
58        Q4_1Block, 20);
59    opaque_block!(/// Q5_0 block: 32 weights in 22 bytes.
60        Q5_0Block, 22);
61    opaque_block!(/// Q5_1 block: 32 weights in 24 bytes.
62        Q5_1Block, 24);
63    opaque_block!(/// Q8_0 block: 32 weights in 34 bytes.
64        Q8_0Block, 34);
65    opaque_block!(/// Q8_1 block: 32 weights in 36 bytes.
66        Q8_1Block, 36);
67    opaque_block!(/// Q2_K super-block: 256 weights in 84 bytes.
68        Q2_KBlock, 84);
69    opaque_block!(/// Q3_K super-block: 256 weights in 110 bytes.
70        Q3_KBlock, 110);
71    opaque_block!(/// Q4_K super-block: 256 weights in 144 bytes.
72        Q4_KBlock, 144);
73    opaque_block!(/// Q5_K super-block: 256 weights in 176 bytes.
74        Q5_KBlock, 176);
75    opaque_block!(/// Q6_K super-block: 256 weights in 210 bytes.
76        Q6_KBlock, 210);
77    opaque_block!(/// Q8_K super-block: 256 weights in 292 bytes.
78        Q8_KBlock, 292);
79    opaque_block!(/// IQ2_XXS super-block: 256 weights in 66 bytes.
80        IQ2_XXSBlock, 66);
81    opaque_block!(/// IQ2_XS super-block: 256 weights in 74 bytes.
82        IQ2_XSBlock, 74);
83    opaque_block!(/// IQ3_XXS super-block: 256 weights in 98 bytes.
84        IQ3_XXSBlock, 98);
85    opaque_block!(/// IQ1_S super-block: 256 weights in 50 bytes.
86        IQ1_SBlock, 50);
87    opaque_block!(/// IQ4_NL block: 32 weights in 18 bytes.
88        IQ4_NLBlock, 18);
89    opaque_block!(/// IQ3_S super-block: 256 weights in 110 bytes.
90        IQ3_SBlock, 110);
91    opaque_block!(/// IQ2_S super-block: 256 weights in 82 bytes.
92        IQ2_SBlock, 82);
93    opaque_block!(/// IQ4_XS super-block: 256 weights in 136 bytes.
94        IQ4_XSBlock, 136);
95    opaque_block!(/// IQ1_M super-block: 256 weights in 56 bytes.
96        IQ1_MBlock, 56);
97    opaque_block!(/// TQ1_0 block: 256 weights in 54 bytes.
98        TQ1_0Block, 54);
99    opaque_block!(/// TQ2_0 block: 256 weights in 66 bytes.
100        TQ2_0Block, 66);
101    opaque_block!(/// MXFP4 block: 32 weights in 17 bytes.
102        MXFP4Block, 17);
103    opaque_block!(/// NVFP4 block: 64 weights in 36 bytes.
104        NVFP4Block, 36);
105    opaque_block!(/// Q1_0 block: 128 weights in 18 bytes.
106        Q1_0Block, 18);
107    opaque_block!(/// Q2_0 block: 64 weights in 18 bytes.
108        Q2_0Block, 18);
109}
110
111/// Quantization parameters for different tensor types
112#[derive(Debug, Clone, PartialEq)]
113#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
114pub struct QuantizationParams {
115    /// The tensor type this applies to
116    pub tensor_type: TensorType,
117    /// Block size (number of elements per quantization block)
118    pub block_size: usize,
119    /// Physical storage bits per weight, including block metadata overhead
120    pub bits_per_weight: f32,
121    /// Whether this format supports scales
122    pub has_scales: bool,
123    /// Whether this format supports minimum values
124    pub has_min: bool,
125    /// Whether this format has high-bit extensions
126    pub has_high_bits: bool,
127    /// Size of each block in bytes
128    pub block_size_bytes: usize,
129}
130
131impl QuantizationParams {
132    /// Get quantization parameters for a tensor type
133    pub fn for_type(tensor_type: TensorType) -> Self {
134        let block_size = tensor_type.block_size();
135        let block_size_bytes = tensor_type.block_size_bytes().unwrap_or(0);
136        let bits_per_weight = tensor_type.storage_bits_per_weight().unwrap_or(0.0);
137        let has_scales = tensor_type.is_quantized();
138        let has_min = matches!(
139            tensor_type,
140            TensorType::Q4_1
141                | TensorType::Q5_1
142                | TensorType::Q2_K
143                | TensorType::Q4_K
144                | TensorType::Q5_K
145        );
146        let has_high_bits = matches!(
147            tensor_type,
148            TensorType::Q5_0
149                | TensorType::Q5_1
150                | TensorType::Q3_K
151                | TensorType::Q5_K
152                | TensorType::Q6_K
153        );
154
155        Self {
156            tensor_type,
157            block_size,
158            bits_per_weight,
159            has_scales,
160            has_min,
161            has_high_bits,
162            block_size_bytes,
163        }
164    }
165
166    /// Whether GGUF readers and writers support this tensor type.
167    pub fn is_supported(&self) -> bool {
168        self.block_size > 0 && self.block_size_bytes > 0
169    }
170
171    /// Calculate the storage size for a given number of elements.
172    ///
173    /// Returns `None` for unsupported tensor types and on arithmetic overflow.
174    pub fn calculate_storage_size(&self, element_count: u64) -> Option<u64> {
175        self.tensor_type.calculate_size(element_count)
176    }
177
178    /// Calculate the number of blocks needed for a given element count
179    pub fn calculate_num_blocks(&self, element_count: u64) -> u64 {
180        if self.block_size == 0 {
181            return 0;
182        }
183        if self.block_size <= 1 {
184            return element_count;
185        }
186
187        element_count.div_ceil(self.block_size as u64)
188    }
189}
190
191/// Utilities for working with quantization
192pub struct QuantizationUtils;
193
194impl QuantizationUtils {
195    /// Get all supported quantization types
196    #[cfg(any(feature = "std", feature = "alloc"))]
197    pub fn all_quantized_types() -> Vec<TensorType> {
198        vec![
199            TensorType::Q4_0,
200            TensorType::Q4_1,
201            TensorType::Q5_0,
202            TensorType::Q5_1,
203            TensorType::Q8_0,
204            TensorType::Q8_1,
205            TensorType::Q2_K,
206            TensorType::Q3_K,
207            TensorType::Q4_K,
208            TensorType::Q5_K,
209            TensorType::Q6_K,
210            TensorType::Q8_K,
211            TensorType::IQ2_XXS,
212            TensorType::IQ2_XS,
213            TensorType::IQ3_XXS,
214            TensorType::IQ1_S,
215            TensorType::IQ4_NL,
216            TensorType::IQ3_S,
217            TensorType::IQ2_S,
218            TensorType::IQ4_XS,
219            TensorType::IQ1_M,
220            TensorType::TQ1_0,
221            TensorType::TQ2_0,
222            TensorType::MXFP4,
223            TensorType::NVFP4,
224            TensorType::Q1_0,
225            TensorType::Q2_0,
226        ]
227    }
228
229    /// Compare two quantization formats
230    pub fn compare_formats(type_a: TensorType, type_b: TensorType) -> cmp::Ordering {
231        let params_a = QuantizationParams::for_type(type_a);
232        let params_b = QuantizationParams::for_type(type_b);
233
234        // Compare by physical storage rate. This does not attempt to rank
235        // perceptual quality across different quantization families.
236        params_a
237            .bits_per_weight
238            .partial_cmp(&params_b.bits_per_weight)
239            .unwrap_or(cmp::Ordering::Equal)
240    }
241
242    /// Get the most similar quantization to a target bit rate
243    pub fn find_closest_quantization(target_bits: f32) -> TensorType {
244        let all_types = Self::all_quantized_types();
245        let mut best_type = TensorType::Q4_0;
246        let mut best_diff = f32::INFINITY;
247
248        for tensor_type in all_types {
249            let params = QuantizationParams::for_type(tensor_type);
250            let diff = (params.bits_per_weight - target_bits).abs();
251
252            if diff < best_diff {
253                best_diff = diff;
254                best_type = tensor_type;
255            }
256        }
257
258        best_type
259    }
260
261    /// Get the storage family for a supported tensor type.
262    ///
263    /// Removed GGML types and the SDK-only `IQ4_UNI` variant return `None`.
264    pub fn get_quantization_family(tensor_type: TensorType) -> Option<&'static str> {
265        tensor_type.block_size_bytes()?;
266        if matches!(tensor_type, TensorType::TQ1_0 | TensorType::TQ2_0) {
267            Some("ternary")
268        } else if matches!(tensor_type, TensorType::MXFP4 | TensorType::NVFP4) {
269            Some("microscaling")
270        } else if matches!(tensor_type, TensorType::Q1_0 | TensorType::Q2_0) {
271            Some("block")
272        } else if tensor_type.is_k_quant() {
273            Some("k-quant")
274        } else if tensor_type.is_iq_quant() {
275            Some("i-quant")
276        } else if tensor_type.is_quantized() {
277            Some("legacy")
278        } else {
279            Some("unquantized")
280        }
281    }
282}
283
284#[cfg(feature = "std")]
285impl std::fmt::Display for QuantizationParams {
286    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
287        write!(
288            f,
289            "{} (block_size: {}, {:.3} storage bits/weight, {} bytes/block)",
290            self.tensor_type.name(),
291            self.block_size,
292            self.bits_per_weight,
293            self.block_size_bytes
294        )
295    }
296}
297
298#[cfg(not(feature = "std"))]
299impl fmt::Display for QuantizationParams {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        write!(
302            f,
303            "QuantizationParams {{ type: {:?}, block_size: {}, block_size_bytes: {}, storage_bits_per_weight: {} }}",
304            self.tensor_type,
305            self.block_size,
306            self.block_size_bytes,
307            self.bits_per_weight
308        )
309    }
310}
311
312#[cfg(all(test, feature = "std"))]
313mod tests {
314    use super::*;
315    use core::mem;
316
317    const QUANT_GEOMETRY: &[(TensorType, usize, usize)] = &[
318        (TensorType::Q4_0, 32, 18),
319        (TensorType::Q4_1, 32, 20),
320        (TensorType::Q5_0, 32, 22),
321        (TensorType::Q5_1, 32, 24),
322        (TensorType::Q8_0, 32, 34),
323        (TensorType::Q8_1, 32, 36),
324        (TensorType::Q2_K, 256, 84),
325        (TensorType::Q3_K, 256, 110),
326        (TensorType::Q4_K, 256, 144),
327        (TensorType::Q5_K, 256, 176),
328        (TensorType::Q6_K, 256, 210),
329        (TensorType::Q8_K, 256, 292),
330        (TensorType::IQ2_XXS, 256, 66),
331        (TensorType::IQ2_XS, 256, 74),
332        (TensorType::IQ3_XXS, 256, 98),
333        (TensorType::IQ1_S, 256, 50),
334        (TensorType::IQ4_NL, 32, 18),
335        (TensorType::IQ3_S, 256, 110),
336        (TensorType::IQ2_S, 256, 82),
337        (TensorType::IQ4_XS, 256, 136),
338        (TensorType::IQ1_M, 256, 56),
339        (TensorType::TQ1_0, 256, 54),
340        (TensorType::TQ2_0, 256, 66),
341        (TensorType::MXFP4, 32, 17),
342        (TensorType::NVFP4, 64, 36),
343        (TensorType::Q1_0, 128, 18),
344        (TensorType::Q2_0, 64, 18),
345    ];
346
347    #[test]
348    fn test_quantization_params_q4_0() {
349        let params = QuantizationParams::for_type(TensorType::Q4_0);
350        assert_eq!(params.tensor_type, TensorType::Q4_0);
351        assert_eq!(params.block_size, 32);
352        assert_eq!(params.bits_per_weight, 4.5);
353        assert!(params.has_scales);
354        assert!(!params.has_min);
355        assert!(!params.has_high_bits);
356        assert_eq!(params.block_size_bytes, 18); // 2 bytes scale + 16 bytes data
357    }
358
359    #[test]
360    fn test_quantization_params_q5_1() {
361        let params = QuantizationParams::for_type(TensorType::Q5_1);
362        assert_eq!(params.block_size, 32);
363        assert_eq!(params.bits_per_weight, 6.0);
364        assert!(params.has_scales);
365        assert!(params.has_min);
366        assert!(params.has_high_bits);
367    }
368
369    #[test]
370    fn test_quantization_params_k_quant() {
371        let params = QuantizationParams::for_type(TensorType::Q4_K);
372        assert_eq!(params.block_size, 256); // K-quants use larger blocks
373        assert_eq!(params.bits_per_weight, 4.5);
374    }
375
376    #[test]
377    fn test_k_quant_minimum_flags() {
378        for tensor_type in [TensorType::Q2_K, TensorType::Q4_K, TensorType::Q5_K] {
379            assert!(QuantizationParams::for_type(tensor_type).has_min);
380        }
381        for tensor_type in [TensorType::Q3_K, TensorType::Q6_K, TensorType::Q8_K] {
382            assert!(!QuantizationParams::for_type(tensor_type).has_min);
383        }
384    }
385
386    #[test]
387    fn test_storage_size_calculation() {
388        let params = QuantizationParams::for_type(TensorType::Q4_0);
389
390        // One block worth of elements
391        let size_32 = params.calculate_storage_size(32);
392        assert_eq!(size_32, Some(params.block_size_bytes as u64));
393
394        // Two blocks worth
395        let size_64 = params.calculate_storage_size(64);
396        assert_eq!(size_64, Some(2 * params.block_size_bytes as u64));
397
398        // Partial block (should round up)
399        let size_33 = params.calculate_storage_size(33);
400        assert_eq!(size_33, Some(2 * params.block_size_bytes as u64));
401    }
402
403    #[test]
404    fn test_quantization_utils() {
405        let all_quantized = QuantizationUtils::all_quantized_types();
406        let expected: Vec<_> =
407            QUANT_GEOMETRY.iter().map(|(tensor_type, _, _)| *tensor_type).collect();
408        assert_eq!(all_quantized, expected);
409        assert!(!all_quantized.contains(&TensorType::IQ4_UNI));
410        assert!(!all_quantized.contains(&TensorType::Q4_2));
411    }
412
413    #[test]
414    fn test_format_comparison() {
415        use cmp::Ordering;
416
417        let cmp = QuantizationUtils::compare_formats(TensorType::Q8_0, TensorType::Q4_0);
418        assert_eq!(cmp, Ordering::Greater); // Q8_0 has more bits, so it's "greater"
419
420        let cmp = QuantizationUtils::compare_formats(TensorType::Q4_0, TensorType::Q4_K);
421        // Both have 4 bits, so should be equal
422        assert_eq!(cmp, Ordering::Equal);
423    }
424
425    #[test]
426    fn test_closest_quantization() {
427        let closest_to_4_5 = QuantizationUtils::find_closest_quantization(4.5);
428        let params = QuantizationParams::for_type(closest_to_4_5);
429        assert_eq!(params.bits_per_weight, 4.5);
430
431        let closest_to_6 = QuantizationUtils::find_closest_quantization(6.0);
432        let params = QuantizationParams::for_type(closest_to_6);
433        assert!((params.bits_per_weight - 6.0).abs() <= 1.0); // Should be close
434    }
435
436    #[test]
437    fn test_quantization_families() {
438        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::Q4_0), Some("legacy"));
439        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::Q4_K), Some("k-quant"));
440        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::IQ2_XS), Some("i-quant"));
441        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::TQ1_0), Some("ternary"));
442        assert_eq!(
443            QuantizationUtils::get_quantization_family(TensorType::MXFP4),
444            Some("microscaling")
445        );
446        assert_eq!(
447            QuantizationUtils::get_quantization_family(TensorType::NVFP4),
448            Some("microscaling")
449        );
450        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::Q1_0), Some("block"));
451        assert_eq!(
452            QuantizationUtils::get_quantization_family(TensorType::F32),
453            Some("unquantized")
454        );
455        assert_eq!(QuantizationUtils::get_quantization_family(TensorType::IQ4_UNI), None);
456    }
457
458    #[test]
459    fn test_block_struct_sizes() {
460        use blocks::*;
461
462        macro_rules! assert_block_size {
463            ($block:ty, $size:expr) => {{
464                assert_eq!(mem::size_of::<$block>(), $size);
465                assert_eq!(<$block>::SIZE, $size);
466            }};
467        }
468
469        assert_block_size!(Q4_0Block, 18);
470        assert_block_size!(Q4_1Block, 20);
471        assert_block_size!(Q5_0Block, 22);
472        assert_block_size!(Q5_1Block, 24);
473        assert_block_size!(Q8_0Block, 34);
474        assert_block_size!(Q8_1Block, 36);
475        assert_block_size!(Q2_KBlock, 84);
476        assert_block_size!(Q3_KBlock, 110);
477        assert_block_size!(Q4_KBlock, 144);
478        assert_block_size!(Q5_KBlock, 176);
479        assert_block_size!(Q6_KBlock, 210);
480        assert_block_size!(Q8_KBlock, 292);
481        assert_block_size!(IQ2_XXSBlock, 66);
482        assert_block_size!(IQ2_XSBlock, 74);
483        assert_block_size!(IQ3_XXSBlock, 98);
484        assert_block_size!(IQ1_SBlock, 50);
485        assert_block_size!(IQ4_NLBlock, 18);
486        assert_block_size!(IQ3_SBlock, 110);
487        assert_block_size!(IQ2_SBlock, 82);
488        assert_block_size!(IQ4_XSBlock, 136);
489        assert_block_size!(IQ1_MBlock, 56);
490        assert_block_size!(TQ1_0Block, 54);
491        assert_block_size!(TQ2_0Block, 66);
492        assert_block_size!(MXFP4Block, 17);
493        assert_block_size!(NVFP4Block, 36);
494        assert_block_size!(Q1_0Block, 18);
495        assert_block_size!(Q2_0Block, 18);
496    }
497
498    #[test]
499    fn test_all_supported_quantization_geometry_is_consistent() {
500        for &(tensor_type, block_size, block_bytes) in QUANT_GEOMETRY {
501            let params = QuantizationParams::for_type(tensor_type);
502            assert!(params.is_supported(), "{}", tensor_type.name());
503            assert!(tensor_type.is_quantized(), "{}", tensor_type.name());
504            assert_eq!(tensor_type.element_size(), None, "{}", tensor_type.name());
505            assert_eq!(params.block_size, block_size, "{}", tensor_type.name());
506            assert_eq!(params.block_size_bytes, block_bytes, "{}", tensor_type.name());
507            assert_eq!(tensor_type.calculate_size(block_size as u64), Some(block_bytes as u64));
508            let expected_bpw = block_bytes as f32 * 8.0 / block_size as f32;
509            assert!((params.bits_per_weight - expected_bpw).abs() < f32::EPSILON);
510        }
511
512        for unsupported in [TensorType::Q4_2, TensorType::Q4_3, TensorType::IQ4_UNI] {
513            let params = QuantizationParams::for_type(unsupported);
514            assert!(!params.is_supported());
515            assert_eq!(params.block_size, 0);
516            assert_eq!(params.block_size_bytes, 0);
517            assert_eq!(params.bits_per_weight, 0.0);
518        }
519    }
520
521    #[test]
522    fn test_params_display() {
523        let params = QuantizationParams::for_type(TensorType::Q4_K);
524        let display = format!("{}", params);
525        assert!(display.contains("Q4_K"));
526        assert!(display.contains("256"));
527        assert!(display.contains("4.500"));
528    }
529}
gguf_rs_lib/tensor/quantization.rs

gguf_rs_lib/tensor/
quantization.rs
