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

zsfm_gguf/
reader.rs
