teeny_quant/quant/
pack4.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! INT4 nibble packing.
18//!
19//! safetensors 0.7 has no native `I4`/`U4` dtype, so packed INT4 tensors are stored as plain
20//! `U8`. **Packing layout** (this crate's own convention -- not bit-for-bit compatible with any
21//! particular GPTQ/AWQ int32-packing scheme, which vary across implementations and versions):
22//! two consecutive elements in row-major order share one byte, the first in the low nibble and
23//! the second in the high nibble, each a 4-bit two's-complement value in `-8..=7`. A tensor with
24//! an odd element count gets one trailing byte whose high nibble is unused padding (`0`). The
25//! packed tensor's logical shape and element count are recorded separately in the
26//! `quantization_config` metadata (see [`crate::format`]) since the packed `U8` tensor's own
27//! shape (`[ceil(n / 2)]`) doesn't reflect it.
28
29/// Packs 4-bit two's-complement values (each expected to be in `-8..=7`; out-of-range bits above
30/// the low nibble are silently dropped) two-per-byte, low nibble first.
31pub fn pack_i4(values: &[i8]) -> Vec<u8> {
32    let mut out = Vec::with_capacity(values.len().div_ceil(2));
33    for pair in values.chunks(2) {
34        let lo = (pair[0] as u8) & 0x0F;
35        let hi = pair.get(1).map(|&v| (v as u8) & 0x0F).unwrap_or(0);
36        out.push(lo | (hi << 4));
37    }
38    out
39}
40
41/// Inverse of [`pack_i4`]: unpacks `n` sign-extended 4-bit values from `bytes`.
42pub fn unpack_i4(bytes: &[u8], n: usize) -> Vec<i8> {
43    fn sign_extend(nibble: u8) -> i8 {
44        if nibble >= 8 {
45            nibble as i8 - 16
46        } else {
47            nibble as i8
48        }
49    }
50
51    let mut out = Vec::with_capacity(n);
52    for &byte in bytes {
53        if out.len() >= n {
54            break;
55        }
56        out.push(sign_extend(byte & 0x0F));
57        if out.len() < n {
58            out.push(sign_extend((byte >> 4) & 0x0F));
59        }
60    }
61    out
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn round_trips_full_int4_range() {
70        let values: Vec<i8> = (-8..=7).collect();
71        let packed = pack_i4(&values);
72        assert_eq!(packed.len(), values.len().div_ceil(2));
73        let unpacked = unpack_i4(&packed, values.len());
74        assert_eq!(values, unpacked);
75    }
76
77    #[test]
78    fn odd_length_pads_trailing_nibble() {
79        let values = vec![-8i8, 7, -1];
80        let packed = pack_i4(&values);
81        assert_eq!(packed.len(), 2);
82        let unpacked = unpack_i4(&packed, values.len());
83        assert_eq!(values, unpacked);
84    }
85
86    #[test]
87    fn empty_round_trips() {
88        let values: Vec<i8> = vec![];
89        let packed = pack_i4(&values);
90        assert!(packed.is_empty());
91        assert!(unpack_i4(&packed, 0).is_empty());
92    }
93}
teeny_quant/quant/pack4.rs

teeny_quant/quant/
pack4.rs
