zsfm_gguf/
reader.rs

1//! Minimal GGUF v2/v3 reader for the tensor types this crate's writer emits
2//! (F32, F16, BF16, Q8_0).
3//!
4//! candle's `gguf_file` reader covers many more quantization formats but does
5//! not know BF16 (ggml dtype 30), so BF16 files written here would otherwise be
6//! unreadable by our own tooling. Callers should prefer candle's reader and
7//! fall back to this one.
8
9use std::io::{Read, Seek, SeekFrom};
10
11use anyhow::{bail, Context, Result};
12use byteorder::{LittleEndian, ReadBytesExt};
13
14use super::types::{GGMLType, GGUFMetaValue};
15
16const GGUF_MAGIC: &[u8; 4] = b"GGUF";
17const DEFAULT_ALIGNMENT: u64 = 32;
18
19/// Capacity every model crate should wrap its `BufReader<File>` with before parsing a GGUF
20/// file. Header/metadata parsing is hundreds to thousands of few-byte `read_u32`/`read_u64`
21/// calls (one per tensor dim, name length, metadata value, …); the default 8 KiB `BufReader`
22/// capacity is fine for small files but forces a syscall every few tensors on models with
23/// hundreds of them. 1 MiB comfortably covers any GGUF header in one underlying read while
24/// staying small relative to the model file itself.
25pub const READ_BUF_CAPACITY: usize = 1 << 20;
26
27pub struct GGUFTensorInfo {
28    pub name: String,
29    /// Dims exactly as stored in the file (GGUF order — reversed row-major).
30    pub shape: Vec<u64>,
31    pub dtype: GGMLType,
32    /// Byte offset relative to the start of the data section.
33    pub offset: u64,
34}
35
36impl GGUFTensorInfo {
37    pub fn n_elems(&self) -> u64 {
38        self.shape.iter().product()
39    }
40
41    pub fn n_bytes(&self) -> u64 {
42        match self.dtype {
43            GGMLType::F32 => self.n_elems() * 4,
44            GGMLType::F16 | GGMLType::BF16 => self.n_elems() * 2,
45            GGMLType::Q8_0 => self.n_elems() / 32 * 34,
46        }
47    }
48}
49
50pub struct GGUFFile {
51    pub metadata: Vec<(String, GGUFMetaValue)>,
52    pub tensors: Vec<GGUFTensorInfo>,
53    /// Absolute file offset of the tensor data section.
54    pub data_start: u64,
55}
56
57impl GGUFFile {
58    pub fn read(reader: &mut (impl Read + Seek)) -> Result<Self> {
59        let mut magic = [0u8; 4];
60        reader.read_exact(&mut magic).context("read GGUF magic")?;
61        if &magic != GGUF_MAGIC {
62            bail!("not a GGUF file (bad magic)");
63        }
64        let version = reader.read_u32::<LittleEndian>()?;
65        if !(2..=3).contains(&version) {
66            bail!("unsupported GGUF version {version}");
67        }
68        let tensor_count = reader.read_u64::<LittleEndian>()?;
69        let kv_count = reader.read_u64::<LittleEndian>()?;
70
71        let mut metadata = Vec::with_capacity(kv_count as usize);
72        for _ in 0..kv_count {
73            let key = read_string(reader)?;
74            let vtype = reader.read_u32::<LittleEndian>()?;
75            let value = read_value(reader, vtype)?;
76            if let Some(v) = value {
77                metadata.push((key, v));
78            }
79        }
80
81        let alignment = metadata
82            .iter()
83            .find(|(k, _)| k == "general.alignment")
84            .and_then(|(_, v)| match v {
85                GGUFMetaValue::Uint32(a) => Some(*a as u64),
86                GGUFMetaValue::Uint64(a) => Some(*a),
87                _ => None,
88            })
89            .unwrap_or(DEFAULT_ALIGNMENT);
90
91        let mut tensors = Vec::with_capacity(tensor_count as usize);
92        for _ in 0..tensor_count {
93            let name = read_string(reader)?;
94            let n_dims = reader.read_u32::<LittleEndian>()?;
95            let mut shape = Vec::with_capacity(n_dims as usize);
96            for _ in 0..n_dims {
97                shape.push(reader.read_u64::<LittleEndian>()?);
98            }
99            let dtype_raw = reader.read_u32::<LittleEndian>()?;
100            let dtype = match dtype_raw {
101                0 => GGMLType::F32,
102                1 => GGMLType::F16,
103                8 => GGMLType::Q8_0,
104                30 => GGMLType::BF16,
105                other => bail!("tensor {name}: ggml dtype {other} not supported by this reader"),
106            };
107            let offset = reader.read_u64::<LittleEndian>()?;
108            tensors.push(GGUFTensorInfo { name, shape, dtype, offset });
109        }
110
111        let pos = reader.stream_position()?;
112        let data_start = pos.div_ceil(alignment) * alignment;
113
114        Ok(GGUFFile { metadata, tensors, data_start })
115    }
116
117    /// Raw stored bytes of one tensor.
118    pub fn tensor_bytes(
119        &self,
120        reader: &mut (impl Read + Seek),
121        info: &GGUFTensorInfo,
122    ) -> Result<Vec<u8>> {
123        reader.seek(SeekFrom::Start(self.data_start + info.offset))?;
124        let mut buf = vec![0u8; info.n_bytes() as usize];
125        reader
126            .read_exact(&mut buf)
127            .with_context(|| format!("read tensor {} data", info.name))?;
128        Ok(buf)
129    }
130
131    /// Tensor decoded to f32 (dequantizing Q8_0).
132    pub fn tensor_f32(
133        &self,
134        reader: &mut (impl Read + Seek),
135        info: &GGUFTensorInfo,
136    ) -> Result<Vec<f32>> {
137        let bytes = self.tensor_bytes(reader, info)?;
138        match info.dtype {
139            GGMLType::F32 => Ok(bytes
140                .chunks_exact(4)
141                .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
142                .collect()),
143            GGMLType::F16 => Ok(bytes
144                .chunks_exact(2)
145                .map(|c| f16_bits_to_f32(u16::from_le_bytes([c[0], c[1]])))
146                .collect()),
147            GGMLType::BF16 => Ok(bytes
148                .chunks_exact(2)
149                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
150                .collect()),
151            GGMLType::Q8_0 => {
152                let mut out = Vec::with_capacity(info.n_elems() as usize);
153                for block in bytes.chunks_exact(34) {
154                    let d = f16_bits_to_f32(u16::from_le_bytes([block[0], block[1]]));
155                    for &q in &block[2..34] {
156                        out.push(d * (q as i8) as f32);
157                    }
158                }
159                Ok(out)
160            }
161        }
162    }
163}
164
165fn read_string(reader: &mut impl Read) -> Result<String> {
166    let len = reader.read_u64::<LittleEndian>()? as usize;
167    let mut buf = vec![0u8; len];
168    reader.read_exact(&mut buf)?;
169    String::from_utf8(buf).context("GGUF string is not UTF-8")
170}
171
172/// Parse one metadata value. Returns `None` for value types we don't model
173/// (after consuming the correct number of bytes so the stream stays aligned).
174fn read_value(reader: &mut impl Read, vtype: u32) -> Result<Option<GGUFMetaValue>> {
175    Ok(match vtype {
176        0 => Some(GGUFMetaValue::Uint8(reader.read_u8()?)),
177        1 => Some(GGUFMetaValue::Int8(reader.read_i8()?)),
178        2 => Some(GGUFMetaValue::Uint16(reader.read_u16::<LittleEndian>()?)),
179        3 => Some(GGUFMetaValue::Int16(reader.read_i16::<LittleEndian>()?)),
180        4 => Some(GGUFMetaValue::Uint32(reader.read_u32::<LittleEndian>()?)),
181        5 => Some(GGUFMetaValue::Int32(reader.read_i32::<LittleEndian>()?)),
182        6 => Some(GGUFMetaValue::Float32(reader.read_f32::<LittleEndian>()?)),
183        7 => Some(GGUFMetaValue::Bool(reader.read_u8()? != 0)),
184        8 => Some(GGUFMetaValue::String(read_string(reader)?)),
185        9 => {
186            let elem_type = reader.read_u32::<LittleEndian>()?;
187            let count = reader.read_u64::<LittleEndian>()?;
188            match elem_type {
189                4 => {
190                    let mut v = Vec::with_capacity(count as usize);
191                    for _ in 0..count {
192                        v.push(reader.read_u32::<LittleEndian>()?);
193                    }
194                    Some(GGUFMetaValue::ArrayUint32(v))
195                }
196                6 => {
197                    let mut v = Vec::with_capacity(count as usize);
198                    for _ in 0..count {
199                        v.push(reader.read_f32::<LittleEndian>()?);
200                    }
201                    Some(GGUFMetaValue::ArrayFloat32(v))
202                }
203                8 => {
204                    let mut v = Vec::with_capacity(count as usize);
205                    for _ in 0..count {
206                        v.push(read_string(reader)?);
207                    }
208                    Some(GGUFMetaValue::ArrayString(v))
209                }
210                other => {
211                    for _ in 0..count {
212                        read_value(reader, other)?;
213                    }
214                    None
215                }
216            }
217        }
218        10 => Some(GGUFMetaValue::Uint64(reader.read_u64::<LittleEndian>()?)),
219        11 => Some(GGUFMetaValue::Int64(reader.read_i64::<LittleEndian>()?)),
220        12 => Some(GGUFMetaValue::Float64(reader.read_f64::<LittleEndian>()?)),
221        other => bail!("unsupported GGUF metadata value type {other}"),
222    })
223}
224
225fn f16_bits_to_f32(bits: u16) -> f32 {
226    let sign = ((bits & 0x8000) as u32) << 16;
227    let exp = ((bits >> 10) & 0x1F) as i32;
228    let mantissa = (bits & 0x03FF) as u32;
229    let f32_bits = if exp == 0 {
230        if mantissa == 0 {
231            sign
232        } else {
233            let mut m = mantissa;
234            let mut e = 0i32;
235            while m & 0x0400 == 0 {
236                m <<= 1;
237                e += 1;
238            }
239            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
240        }
241    } else if exp == 31 {
242        sign | 0x7F80_0000 | (mantissa << 13)
243    } else {
244        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
245    };
246    f32::from_bits(f32_bits)
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use crate::writer::GGUFWriter;
253    use std::io::Cursor;
254
255    #[test]
256    fn writer_reader_roundtrip_all_dtypes() {
257        let mut w = GGUFWriter::new();
258        w.add_metadata("general.architecture", GGUFMetaValue::String("test".into()));
259        w.add_metadata("test.count", GGUFMetaValue::Uint32(7));
260        w.add_metadata(
261            "test.floats",
262            GGUFMetaValue::ArrayFloat32(vec![0.5, 1.5]),
263        );
264
265        let f32_vals: Vec<f32> = (0..32).map(|i| i as f32 / 4.0).collect();
266        let f32_bytes: Vec<u8> = f32_vals.iter().flat_map(|v| v.to_le_bytes()).collect();
267        w.add_tensor("t.f32", vec![32], GGMLType::F32, f32_bytes);
268
269        // 1.25 is exactly representable in bf16.
270        let bf16_bytes: Vec<u8> = (0..16)
271            .flat_map(|_| ((1.25f32.to_bits() >> 16) as u16).to_le_bytes())
272            .collect();
273        w.add_tensor("t.bf16", vec![4, 4], GGMLType::BF16, bf16_bytes);
274
275        let mut buf = Cursor::new(Vec::new());
276        w.write_to(&mut buf).unwrap();
277
278        buf.set_position(0);
279        let f = GGUFFile::read(&mut buf).unwrap();
280        assert_eq!(f.metadata.len(), 3);
281        assert_eq!(f.tensors.len(), 2);
282
283        let tb = f.tensors.iter().find(|t| t.name == "t.bf16").unwrap();
284        assert_eq!(tb.dtype, GGMLType::BF16);
285        assert_eq!(tb.shape, vec![4, 4]);
286        let vals = f.tensor_f32(&mut buf, tb).unwrap();
287        assert!(vals.iter().all(|&v| v == 1.25));
288
289        let tf = f.tensors.iter().find(|t| t.name == "t.f32").unwrap();
290        let vals = f.tensor_f32(&mut buf, tf).unwrap();
291        assert_eq!(vals, f32_vals);
292    }
293
294    #[test]
295    fn q8_0_dequant_roundtrip() {
296        // d = 1.0, quants = -127..-96 → values exactly d*q.
297        let mut block = vec![0u8; 34];
298        block[0..2].copy_from_slice(&0x3C00u16.to_le_bytes()); // f16 1.0
299        for i in 0..32 {
300            block[2 + i] = (-(127 - i as i8)) as u8;
301        }
302        let mut w = GGUFWriter::new();
303        w.add_tensor("q", vec![32], GGMLType::Q8_0, block);
304        let mut buf = Cursor::new(Vec::new());
305        w.write_to(&mut buf).unwrap();
306        buf.set_position(0);
307        let f = GGUFFile::read(&mut buf).unwrap();
308        let vals = f.tensor_f32(&mut buf, &f.tensors[0]).unwrap();
309        assert_eq!(vals[0], -127.0);
310        assert_eq!(vals[31], -96.0);
311    }
312}
zsfm_gguf/reader.rs

zsfm_gguf/
reader.rs
