Skip to main content

cherenkov_model_data/
gguf.rs

1//! Header-only GGUF inspection. Tensor codecs are identified, never treated as
2//! interchangeable with MLX affine Q4. Execution/conversion is not implemented.
3
4use super::*;
5const BYTES_PER_MIB: usize = 1024 * 1024;
6use anyhow::{Context, Result, ensure};
7use serde_json::{Value, json};
8
9pub(crate) fn inspect(map: &[u8]) -> Result<Inventory> {
10    let mut reader = Header { bytes: map, pos: 0 };
11
12    ensure!(reader.take(4)? == b"GGUF", "expected a GGUF file");
13
14    let version = reader.u32()?;
15
16    ensure!(
17        matches!(version, 2 | 3),
18        "unsupported GGUF version {version}"
19    );
20
21    let tensor_count = reader.count()?;
22    let metadata_count = reader.count()?;
23    let mut metadata = serde_json::Map::new();
24    let mut alignment = 32_u64;
25
26    for _ in 0..metadata_count {
27        let key = reader.string()?.to_owned();
28        let kind = reader.u32()?;
29        let value = reader.value(kind, 0)?;
30
31        if key == "general.alignment" {
32            ensure!(kind == 4, "GGUF alignment must be uint32");
33
34            alignment = value.as_u64().context("invalid GGUF alignment")?;
35        }
36
37        ensure!(
38            !metadata.contains_key(&key),
39            "duplicate GGUF metadata {key}"
40        );
41        metadata.insert(key, json!({"type": kind, "value": value}));
42    }
43
44    let mut tensors = Vec::new();
45
46    for _ in 0..tensor_count {
47        let name = reader.string()?.to_owned();
48        let rank = reader.u32()?;
49
50        ensure!((1..=4).contains(&rank), "invalid GGUF tensor rank");
51
52        let mut shape = (0..rank)
53            .map(|_| reader.u64())
54            .collect::<Result<Vec<_>>>()?;
55
56        shape.reverse(); // GGUF lists the fastest-varying dimension first.
57
58        let encoding = codec(reader.u32()?);
59        let offset = reader.u64()?;
60
61        tensors.push((name, shape, encoding, offset));
62    }
63
64    ensure!(
65        alignment > 0 && alignment.is_power_of_two(),
66        "invalid GGUF alignment"
67    );
68
69    let data_start = (reader.pos as u64)
70        .checked_add(alignment - 1)
71        .context("GGUF alignment overflow")?
72        / alignment
73        * alignment;
74
75    ensure!(data_start <= map.len() as u64, "truncated GGUF padding");
76    tensors.sort_by_key(|t| t.3);
77
78    let mut described = Vec::new();
79
80    for (i, (name, shape, encoding, relative)) in tensors.iter().enumerate() {
81        let offset = data_start
82            .checked_add(*relative)
83            .context("GGUF tensor offset overflow")?;
84        let end = match tensors.get(i + 1) {
85            Some(next) => data_start
86                .checked_add(next.3)
87                .context("GGUF tensor offset overflow")?,
88            None => map.len() as u64,
89        };
90
91        ensure!(
92            relative.is_multiple_of(alignment) && offset <= end && end <= map.len() as u64,
93            "GGUF tensor range out of bounds"
94        );
95
96        let length = encoded_len(*encoding, shape)?.unwrap_or(end - offset);
97
98        ensure!(length <= end - offset, "truncated GGUF tensor {name}");
99        described.push(Tensor::new(
100            name.clone(),
101            TensorRole::Opaque,
102            Some(shape.clone()),
103            TensorEncoding::Ggml {
104                encoding: *encoding,
105                data: DataSpan {
106                    object: ObjectId(0),
107                    offset,
108                    length,
109                },
110            },
111        ));
112    }
113
114    Ok(Inventory {
115        format: ContainerFormat::Gguf { version },
116        metadata: Value::Object(metadata),
117        tensors: described,
118    })
119}
120
121fn encoded_len(encoding: GgmlEncoding, shape: &[u64]) -> Result<Option<u64>> {
122    let (block, bytes) = match encoding {
123        GgmlEncoding::F32 => (1, 4),
124        GgmlEncoding::F16 | GgmlEncoding::Bf16 => (1, 2),
125        GgmlEncoding::Q4_0 => (32, 18),
126        GgmlEncoding::Q4_1 => (32, 20),
127        GgmlEncoding::Q8_0 => (32, 34),
128        GgmlEncoding::Q4K => (256, 144),
129        GgmlEncoding::Q5K => (256, 176),
130        GgmlEncoding::Q6K => (256, 210),
131        GgmlEncoding::Opaque { .. } => return Ok(None),
132    };
133    let width = *shape.last().context("GGUF tensor has no dimensions")?;
134
135    ensure!(
136        width > 0 && width.is_multiple_of(block),
137        "GGUF row does not fit its quantization block"
138    );
139
140    let elements = shape.iter().try_fold(1_u64, |size, &dim| {
141        ensure!(dim > 0, "empty GGUF tensor dimension");
142
143        size.checked_mul(dim).context("GGUF tensor size overflow")
144    })?;
145
146    Ok(Some(
147        (elements / block)
148            .checked_mul(bytes)
149            .context("GGUF tensor size overflow")?,
150    ))
151}
152
153fn codec(code: u32) -> GgmlEncoding {
154    match code {
155        0 => GgmlEncoding::F32,
156        1 => GgmlEncoding::F16,
157        2 => GgmlEncoding::Q4_0,
158        3 => GgmlEncoding::Q4_1,
159        8 => GgmlEncoding::Q8_0,
160        12 => GgmlEncoding::Q4K,
161        13 => GgmlEncoding::Q5K,
162        14 => GgmlEncoding::Q6K,
163        30 => GgmlEncoding::Bf16,
164        code => GgmlEncoding::Opaque { code },
165    }
166}
167
168struct Header<'a> {
169    bytes: &'a [u8],
170    pos: usize,
171}
172
173impl<'a> Header<'a> {
174    fn take(&mut self, count: usize) -> Result<&'a [u8]> {
175        let end = self
176            .pos
177            .checked_add(count)
178            .context("GGUF header overflow")?;
179
180        ensure!(end <= 64 * BYTES_PER_MIB, "GGUF header exceeds 64 MiB");
181
182        let bytes = self
183            .bytes
184            .get(self.pos..end)
185            .context("truncated GGUF header")?;
186        self.pos = end;
187
188        Ok(bytes)
189    }
190
191    fn u32(&mut self) -> Result<u32> {
192        Ok(u32::from_le_bytes(self.take(4)?.try_into()?))
193    }
194
195    fn u64(&mut self) -> Result<u64> {
196        Ok(u64::from_le_bytes(self.take(8)?.try_into()?))
197    }
198
199    fn count(&mut self) -> Result<u64> {
200        let count = self.u64()?;
201
202        ensure!(
203            count <= 1_000_000,
204            "GGUF entry count exceeds inspection limit"
205        );
206
207        Ok(count)
208    }
209
210    fn string(&mut self) -> Result<&'a str> {
211        let len = usize::try_from(self.u64()?).context("GGUF string size overflow")?;
212
213        std::str::from_utf8(self.take(len)?).context("invalid UTF-8 in GGUF header")
214    }
215
216    fn value(&mut self, kind: u32, depth: usize) -> Result<Value> {
217        ensure!(depth <= 8, "GGUF metadata nesting limit exceeded");
218
219        Ok(match kind {
220            0 => json!(self.take(1)?[0]),
221            1 => json!(self.take(1)?[0] as i8),
222            2 => json!(u16::from_le_bytes(self.take(2)?.try_into()?)),
223            3 => json!(i16::from_le_bytes(self.take(2)?.try_into()?)),
224            4 => json!(self.u32()?),
225            5 => json!(self.u32()? as i32),
226            6 => {
227                let bits = self.u32()?;
228
229                float_value(f64::from(f32::from_bits(bits)), u64::from(bits))
230            }
231            7 => {
232                let value = self.take(1)?[0];
233
234                ensure!(value <= 1, "invalid GGUF boolean");
235
236                json!(value == 1)
237            }
238            8 => json!(self.string()?),
239            9 => {
240                let element = self.u32()?;
241                let count = self.count()?;
242                let values = (0..count)
243                    .map(|_| self.value(element, depth + 1))
244                    .collect::<Result<Vec<_>>>()?;
245
246                json!({"element_type": element, "values": values})
247            }
248            10 => json!(self.u64()?),
249            11 => json!(self.u64()? as i64),
250            12 => {
251                let bits = self.u64()?;
252
253                float_value(f64::from_bits(bits), bits)
254            }
255            _ => anyhow::bail!("unknown GGUF metadata type {kind}"),
256        })
257    }
258}
259
260fn float_value(value: f64, bits: u64) -> Value {
261    serde_json::Number::from_f64(value)
262        .map(Value::Number)
263        .unwrap_or_else(|| json!({"nonfinite_bits": bits}))
264}