model_artifact/
gguf.rs

1use std::io::{Read, Seek, SeekFrom};
2use std::path::Path;
3
4mod kv_cache;
5
6pub use kv_cache::{GgufKvCacheQuant, GgufKvCacheType};
7
8const MAX_GGUF_STRING_BYTES: u64 = 1_000_000;
9const MAX_GGUF_ARRAY_ELEMENTS: u64 = 1_000_000;
10const MAX_GGUF_ARRAY_DEPTH: u32 = 64;
11const MAX_GGUF_TENSOR_DIMS: u32 = 8;
12const MAX_GGUF_HEADER_KV_COUNT: usize = 1_000_000;
13const MAX_GGUF_TENSOR_COUNT: usize = 1_000_000;
14
15/// GGUF value types (matching gguf.h enum).
16#[repr(u32)]
17#[derive(Debug, Clone, Copy, PartialEq)]
18enum GgufType {
19    Uint8 = 0,
20    Int8 = 1,
21    Uint16 = 2,
22    Int16 = 3,
23    Uint32 = 4,
24    Int32 = 5,
25    Float32 = 6,
26    Bool = 7,
27    String = 8,
28    Array = 9,
29    Uint64 = 10,
30    Int64 = 11,
31    Float64 = 12,
32}
33
34impl GgufType {
35    fn from_u32(v: u32) -> Option<Self> {
36        match v {
37            0 => Some(Self::Uint8),
38            1 => Some(Self::Int8),
39            2 => Some(Self::Uint16),
40            3 => Some(Self::Int16),
41            4 => Some(Self::Uint32),
42            5 => Some(Self::Int32),
43            6 => Some(Self::Float32),
44            7 => Some(Self::Bool),
45            8 => Some(Self::String),
46            9 => Some(Self::Array),
47            10 => Some(Self::Uint64),
48            11 => Some(Self::Int64),
49            12 => Some(Self::Float64),
50            _ => None,
51        }
52    }
53
54    fn fixed_size(self) -> Option<usize> {
55        match self {
56            Self::Uint8 | Self::Int8 | Self::Bool => Some(1),
57            Self::Uint16 | Self::Int16 => Some(2),
58            Self::Uint32 | Self::Int32 | Self::Float32 => Some(4),
59            Self::Uint64 | Self::Int64 | Self::Float64 => Some(8),
60            Self::String | Self::Array => None,
61        }
62    }
63}
64
65fn read_u32(f: &mut std::fs::File) -> std::io::Result<u32> {
66    let mut buf = [0u8; 4];
67    f.read_exact(&mut buf)?;
68    Ok(u32::from_le_bytes(buf))
69}
70
71fn read_u64(f: &mut std::fs::File) -> std::io::Result<u64> {
72    let mut buf = [0u8; 8];
73    f.read_exact(&mut buf)?;
74    Ok(u64::from_le_bytes(buf))
75}
76
77fn read_i32(f: &mut std::fs::File) -> std::io::Result<i32> {
78    let mut buf = [0u8; 4];
79    f.read_exact(&mut buf)?;
80    Ok(i32::from_le_bytes(buf))
81}
82
83fn read_i64(f: &mut std::fs::File) -> std::io::Result<i64> {
84    let mut buf = [0u8; 8];
85    f.read_exact(&mut buf)?;
86    Ok(i64::from_le_bytes(buf))
87}
88
89fn read_gguf_header_count(
90    f: &mut std::fs::File,
91    max: usize,
92    label: &str,
93) -> std::io::Result<usize> {
94    let value = read_i64(f)?;
95    let count = usize::try_from(value).map_err(|_| {
96        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("negative {label}"))
97    })?;
98    if count > max {
99        return Err(std::io::Error::new(
100            std::io::ErrorKind::InvalidData,
101            format!("{label} too large"),
102        ));
103    }
104    Ok(count)
105}
106
107fn read_bounded_len(f: &mut std::fs::File, max: u64, label: &str) -> std::io::Result<usize> {
108    let len = read_u64(f)?;
109    if len > max {
110        return Err(std::io::Error::new(
111            std::io::ErrorKind::InvalidData,
112            format!("{label} too long"),
113        ));
114    }
115    usize::try_from(len).map_err(|_| {
116        std::io::Error::new(
117            std::io::ErrorKind::InvalidData,
118            format!("{label} too large"),
119        )
120    })
121}
122
123fn read_gguf_string(f: &mut std::fs::File) -> std::io::Result<String> {
124    let len = read_bounded_len(f, MAX_GGUF_STRING_BYTES, "string")?;
125    let mut buf = vec![0u8; len];
126    f.read_exact(&mut buf)?;
127    String::from_utf8(buf).map_err(|_| {
128        std::io::Error::new(
129            std::io::ErrorKind::InvalidData,
130            "invalid UTF-8 in GGUF string",
131        )
132    })
133}
134
135fn skip_gguf_value(f: &mut std::fs::File, typ: GgufType) -> std::io::Result<()> {
136    skip_gguf_value_with_depth(f, typ, 0)
137}
138
139fn skip_gguf_value_with_depth(
140    f: &mut std::fs::File,
141    typ: GgufType,
142    depth: u32,
143) -> std::io::Result<()> {
144    match typ {
145        GgufType::String => {
146            let _ = read_gguf_string(f)?;
147        }
148        GgufType::Array => {
149            if depth >= MAX_GGUF_ARRAY_DEPTH {
150                return Err(std::io::Error::new(
151                    std::io::ErrorKind::InvalidData,
152                    "GGUF nesting too deep",
153                ));
154            }
155            let elem_type = GgufType::from_u32(read_u32(f)?).ok_or_else(|| {
156                std::io::Error::new(std::io::ErrorKind::InvalidData, "bad array type")
157            })?;
158            let count = read_bounded_len(f, MAX_GGUF_ARRAY_ELEMENTS, "array")?;
159            for _ in 0..count {
160                skip_gguf_value_with_depth(f, elem_type, depth + 1)?;
161            }
162        }
163        other => {
164            let size = other.fixed_size().unwrap_or(0);
165            f.seek(SeekFrom::Current(size as i64))?;
166        }
167    }
168    Ok(())
169}
170
171fn read_gguf_value_as_u32(f: &mut std::fs::File, typ: GgufType) -> std::io::Result<Option<u32>> {
172    match typ {
173        GgufType::Uint32 => Ok(Some(read_u32(f)?)),
174        GgufType::Int32 => {
175            let value = read_i32(f)?;
176            let value = u32::try_from(value).map_err(|_| {
177                std::io::Error::new(
178                    std::io::ErrorKind::InvalidData,
179                    "negative Int32 where unsigned GGUF value was expected",
180                )
181            })?;
182            Ok(Some(value))
183        }
184        GgufType::Uint16 => {
185            let mut buf = [0u8; 2];
186            f.read_exact(&mut buf)?;
187            Ok(Some(u16::from_le_bytes(buf) as u32))
188        }
189        GgufType::Uint8 => {
190            let mut buf = [0u8; 1];
191            f.read_exact(&mut buf)?;
192            Ok(Some(buf[0] as u32))
193        }
194        _ => {
195            skip_gguf_value(f, typ)?;
196            Ok(None)
197        }
198    }
199}
200
201fn read_gguf_value_as_f32(f: &mut std::fs::File, typ: GgufType) -> std::io::Result<Option<f32>> {
202    match typ {
203        GgufType::Float32 => {
204            let mut buf = [0u8; 4];
205            f.read_exact(&mut buf)?;
206            Ok(Some(f32::from_le_bytes(buf)))
207        }
208        _ => {
209            skip_gguf_value(f, typ)?;
210            Ok(None)
211        }
212    }
213}
214
215fn read_gguf_value_as_string_opt(
216    f: &mut std::fs::File,
217    typ: GgufType,
218) -> std::io::Result<Option<String>> {
219    match typ {
220        GgufType::String => Ok(Some(read_gguf_string(f)?)),
221        _ => {
222            skip_gguf_value(f, typ)?;
223            Ok(None)
224        }
225    }
226}
227
228#[derive(Clone, Debug, Default)]
229pub struct GgufCompactMeta {
230    pub architecture: String,
231    pub parameter_size: Option<String>,
232    pub context_length: u32,
233    pub vocab_size: u32,
234    pub embedding_size: u32,
235    pub head_count: u32,
236    pub kv_head_count: u32,
237    pub layer_count: u32,
238    pub feed_forward_length: u32,
239    pub key_length: u32,
240    pub value_length: u32,
241    pub kv_lora_rank: u32,
242    pub tokenizer_model_name: String,
243    pub rope_scale: f32,
244    pub rope_freq_base: f32,
245    pub expert_count: u32,
246    pub expert_used_count: u32,
247    pub nextn_predict_layers: u32,
248}
249
250impl GgufCompactMeta {
251    pub fn effective_kv_head_count(&self) -> Option<u32> {
252        if self.kv_head_count > 0 {
253            Some(self.kv_head_count)
254        } else if self.head_count > 0 {
255            Some(self.head_count)
256        } else {
257            None
258        }
259    }
260}
261
262#[derive(Clone, Debug, Default, Eq, PartialEq)]
263pub struct GgufTensorByteProfile {
264    pub expert_count: u32,
265    pub expert_used_count: u32,
266    pub full_model_bytes: u64,
267    pub base_resident_bytes: u64,
268    pub expert_tensor_bytes: u64,
269    pub file_overhead_bytes: u64,
270}
271
272#[derive(Clone, Debug)]
273struct GgufTensorInfo {
274    name: String,
275    offset: u64,
276}
277
278/// Scan a GGUF file header and return compact structural metadata.
279/// Reads only the KV section, never tensor data. Returns None on any parse failure.
280struct GgufHeader {
281    file: std::fs::File,
282    n_tensors: usize,
283    n_kv: usize,
284}
285
286/// Opens `path`, validates the GGUF magic/version, and reads the tensor/KV
287/// header counts shared by every GGUF scan entry point below.
288fn open_gguf_header(path: &Path) -> Option<GgufHeader> {
289    let mut f = std::fs::File::open(path).ok()?;
290
291    let mut magic = [0u8; 4];
292    f.read_exact(&mut magic).ok()?;
293    if &magic != b"GGUF" {
294        return None;
295    }
296    let version = read_u32(&mut f).ok()?;
297    if version < 2 {
298        return None;
299    }
300
301    let n_tensors = read_gguf_header_count(&mut f, MAX_GGUF_TENSOR_COUNT, "tensor count").ok()?;
302    let n_kv = read_gguf_header_count(&mut f, MAX_GGUF_HEADER_KV_COUNT, "KV count").ok()?;
303
304    Some(GgufHeader {
305        file: f,
306        n_tensors,
307        n_kv,
308    })
309}
310
311/// Skips every KV pair without inspecting keys/values, for scans that only
312/// need the tensor-info table.
313fn skip_all_kv_pairs(f: &mut std::fs::File, n_kv: usize) -> Option<()> {
314    for _ in 0..n_kv {
315        let _key = read_gguf_string(f).ok()?;
316        let vtype = GgufType::from_u32(read_u32(f).ok()?)?;
317        skip_gguf_value(f, vtype).ok()?;
318    }
319    Some(())
320}
321
322pub fn scan_gguf_compact_meta(path: &Path) -> Option<GgufCompactMeta> {
323    let GgufHeader {
324        file: mut f, n_kv, ..
325    } = open_gguf_header(path)?;
326
327    let mut meta = GgufCompactMeta::default();
328    for _ in 0..n_kv {
329        let key = read_gguf_string(&mut f).ok()?;
330        let vtype = GgufType::from_u32(read_u32(&mut f).ok()?)?;
331
332        if key == "general.architecture" {
333            meta.architecture = read_gguf_value_as_string_opt(&mut f, vtype).ok()??;
334        } else if key == "general.size_label" {
335            meta.parameter_size = read_gguf_value_as_string_opt(&mut f, vtype).ok()?;
336        } else if key == "tokenizer.ggml.model" {
337            meta.tokenizer_model_name = read_gguf_value_as_string_opt(&mut f, vtype).ok()??;
338        } else if key.ends_with(".context_length") {
339            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
340                meta.context_length = v;
341            }
342        } else if key.ends_with(".embedding_length") {
343            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
344                meta.embedding_size = v;
345            }
346        } else if key.ends_with(".head_count") && !key.ends_with("_kv") {
347            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
348                meta.head_count = v;
349            }
350        } else if key.ends_with(".attention.head_count_kv") {
351            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
352                meta.kv_head_count = v;
353            }
354        } else if key.ends_with(".block_count") {
355            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
356                meta.layer_count = v;
357            }
358        } else if key.ends_with(".feed_forward_length") {
359            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
360                meta.feed_forward_length = v;
361            }
362        } else if key.ends_with(".attention.key_length") {
363            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
364                meta.key_length = v;
365            }
366        } else if key.ends_with(".attention.value_length") {
367            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
368                meta.value_length = v;
369            }
370        } else if key.ends_with(".attention.kv_lora_rank") {
371            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
372                meta.kv_lora_rank = v;
373            }
374        } else if key.ends_with(".rope.scale") {
375            if let Ok(Some(v)) = read_gguf_value_as_f32(&mut f, vtype) {
376                meta.rope_scale = v;
377            }
378        } else if key.ends_with(".rope.freq_base") {
379            if let Ok(Some(v)) = read_gguf_value_as_f32(&mut f, vtype) {
380                meta.rope_freq_base = v;
381            }
382        } else if key.ends_with(".vocab_size") {
383            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
384                meta.vocab_size = v;
385            }
386        } else if key.ends_with(".expert_count") {
387            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
388                meta.expert_count = v;
389            }
390        } else if key.ends_with(".expert_used_count") {
391            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
392                meta.expert_used_count = v;
393            }
394        } else if key.ends_with(".nextn_predict_layers") {
395            if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) {
396                meta.nextn_predict_layers = v;
397            }
398        } else {
399            skip_gguf_value(&mut f, vtype).ok()?;
400        }
401    }
402
403    let derived_key_length = (meta.key_length == 0 && meta.head_count > 0)
404        .then(|| meta.embedding_size.checked_div(meta.head_count))
405        .flatten();
406    if let Some(key_length) = derived_key_length {
407        meta.key_length = key_length;
408    }
409    let derived_value_length = (meta.value_length == 0)
410        .then(|| {
411            meta.effective_kv_head_count()
412                .and_then(|effective_kv| meta.embedding_size.checked_div(effective_kv))
413        })
414        .flatten();
415    if let Some(value_length) = derived_value_length {
416        meta.value_length = value_length;
417    }
418
419    Some(meta)
420}
421
422fn align_offset(value: u64, alignment: u32) -> u64 {
423    let alignment = u64::from(alignment.max(1));
424    let remainder = value % alignment;
425    if remainder == 0 {
426        value
427    } else {
428        value + (alignment - remainder)
429    }
430}
431
432fn read_tensor_infos(
433    f: &mut std::fs::File,
434    n_tensors: usize,
435) -> std::io::Result<Vec<GgufTensorInfo>> {
436    let mut tensors = Vec::new();
437    tensors.try_reserve(n_tensors).map_err(|_| {
438        std::io::Error::new(
439            std::io::ErrorKind::InvalidData,
440            "GGUF tensor count requires too much memory",
441        )
442    })?;
443    for _ in 0..n_tensors {
444        let name = read_gguf_string(f)?;
445        let n_dims = read_u32(f)?;
446        if n_dims > MAX_GGUF_TENSOR_DIMS {
447            return Err(std::io::Error::new(
448                std::io::ErrorKind::InvalidData,
449                "too many GGUF tensor dimensions",
450            ));
451        }
452        for _ in 0..n_dims {
453            let _ = read_u64(f)?;
454        }
455        let _ = read_u32(f)?;
456        let offset = read_u64(f)?;
457        tensors.push(GgufTensorInfo { name, offset });
458    }
459    Ok(tensors)
460}
461
462/// Scan GGUF tensor names and return whether any tensor matches the predicate.
463/// Reads only the header and tensor-info table, never tensor data.
464pub fn scan_gguf_tensor_names_any(
465    path: &Path,
466    mut matches: impl FnMut(&str) -> bool,
467) -> Option<bool> {
468    let GgufHeader {
469        file: mut f,
470        n_tensors,
471        n_kv,
472    } = open_gguf_header(path)?;
473
474    skip_all_kv_pairs(&mut f, n_kv)?;
475
476    for _ in 0..n_tensors {
477        let name = read_gguf_string(&mut f).ok()?;
478        if matches(&name) {
479            return Some(true);
480        }
481        let n_dims = read_u32(&mut f).ok()?;
482        if n_dims > MAX_GGUF_TENSOR_DIMS {
483            return None;
484        }
485        for _ in 0..n_dims {
486            let _ = read_u64(&mut f).ok()?;
487        }
488        let _ggml_type = read_u32(&mut f).ok()?;
489        let _offset = read_u64(&mut f).ok()?;
490    }
491
492    Some(false)
493}
494
495fn is_expert_partitioned_tensor(name: &str) -> bool {
496    let lower = name.to_ascii_lowercase();
497    if lower.contains("shared_expert") || lower.contains("sharedexpert") || lower.contains("shexp")
498    {
499        return false;
500    }
501
502    lower.contains("ffn_gate_exps")
503        || lower.contains("ffn_up_exps")
504        || lower.contains("ffn_down_exps")
505        || lower.contains("exp_probs")
506        || lower.contains(".expert")
507        || lower.contains("_expert")
508}
509
510/// Scan GGUF tensor metadata and estimate which bytes are always resident versus
511/// expert-partitioned. Reads only the header and tensor-info tables.
512pub fn scan_gguf_tensor_byte_profile(path: &Path) -> Option<GgufTensorByteProfile> {
513    let GgufHeader {
514        file: mut f,
515        n_tensors,
516        n_kv,
517    } = open_gguf_header(path)?;
518    let file_len = f.metadata().ok()?.len();
519
520    let mut expert_count = 0u32;
521    let mut expert_used_count = 0u32;
522    let mut alignment = 32u32;
523
524    for _ in 0..n_kv {
525        let key = read_gguf_string(&mut f).ok()?;
526        let vtype = GgufType::from_u32(read_u32(&mut f).ok()?)?;
527
528        if key == "general.alignment" {
529            if let Ok(Some(value)) = read_gguf_value_as_u32(&mut f, vtype) {
530                alignment = value.max(1);
531            }
532        } else if key.ends_with(".expert_count") {
533            if let Ok(Some(value)) = read_gguf_value_as_u32(&mut f, vtype) {
534                expert_count = value;
535            }
536        } else if key.ends_with(".expert_used_count") {
537            if let Ok(Some(value)) = read_gguf_value_as_u32(&mut f, vtype) {
538                expert_used_count = value;
539            }
540        } else {
541            skip_gguf_value(&mut f, vtype).ok()?;
542        }
543    }
544
545    let mut tensors = read_tensor_infos(&mut f, n_tensors).ok()?;
546    if tensors.is_empty() {
547        return Some(GgufTensorByteProfile {
548            expert_count,
549            expert_used_count,
550            full_model_bytes: file_len,
551            base_resident_bytes: 0,
552            expert_tensor_bytes: 0,
553            file_overhead_bytes: file_len,
554        });
555    }
556
557    let tensor_info_end = f.stream_position().ok()?;
558    let data_start = align_offset(tensor_info_end, alignment);
559    if data_start > file_len {
560        return None;
561    }
562    let data_len = file_len - data_start;
563
564    tensors.sort_by_key(|tensor| tensor.offset);
565    if tensors.first()?.offset > data_len {
566        return None;
567    }
568
569    let mut base_resident_bytes = 0u64;
570    let mut expert_tensor_bytes = 0u64;
571    for (index, tensor) in tensors.iter().enumerate() {
572        let next_offset = tensors
573            .get(index + 1)
574            .map(|next| next.offset)
575            .unwrap_or(data_len);
576        if next_offset < tensor.offset || next_offset > data_len {
577            return None;
578        }
579        let tensor_bytes = next_offset - tensor.offset;
580        if is_expert_partitioned_tensor(&tensor.name) {
581            expert_tensor_bytes = expert_tensor_bytes.saturating_add(tensor_bytes);
582        } else {
583            base_resident_bytes = base_resident_bytes.saturating_add(tensor_bytes);
584        }
585    }
586
587    let file_overhead_bytes = file_len.saturating_sub(base_resident_bytes + expert_tensor_bytes);
588    Some(GgufTensorByteProfile {
589        expert_count,
590        expert_used_count,
591        full_model_bytes: file_len,
592        base_resident_bytes,
593        expert_tensor_bytes,
594        file_overhead_bytes,
595    })
596}
597
598#[cfg(test)]
599mod tests {
600    use super::*;
601    use std::io::Write;
602    use std::path::PathBuf;
603    use std::time::{SystemTime, UNIX_EPOCH};
604
605    fn temp_file_path(prefix: &str) -> PathBuf {
606        let unique = SystemTime::now()
607            .duration_since(UNIX_EPOCH)
608            .unwrap()
609            .as_nanos();
610        std::env::temp_dir().join(format!("{prefix}-{unique}.gguf"))
611    }
612
613    fn write_bytes(prefix: &str, bytes: &[u8]) -> PathBuf {
614        let path = temp_file_path(prefix);
615        let mut file = std::fs::File::create(&path).unwrap();
616        file.write_all(bytes).unwrap();
617        file.flush().unwrap();
618        path
619    }
620
621    fn push_array_header(bytes: &mut Vec<u8>, elem_type: GgufType, count: u64) {
622        bytes.extend_from_slice(&(elem_type as u32).to_le_bytes());
623        bytes.extend_from_slice(&count.to_le_bytes());
624    }
625
626    fn push_gguf_string(bytes: &mut Vec<u8>, value: &str) {
627        bytes.extend_from_slice(&(value.len() as u64).to_le_bytes());
628        bytes.extend_from_slice(value.as_bytes());
629    }
630
631    fn push_u32_kv(bytes: &mut Vec<u8>, key: &str, value: u32) {
632        push_gguf_string(bytes, key);
633        bytes.extend_from_slice(&(GgufType::Uint32 as u32).to_le_bytes());
634        bytes.extend_from_slice(&value.to_le_bytes());
635    }
636
637    fn push_tensor_info(bytes: &mut Vec<u8>, name: &str, offset: u64) {
638        push_gguf_string(bytes, name);
639        bytes.extend_from_slice(&1u32.to_le_bytes());
640        bytes.extend_from_slice(&16u64.to_le_bytes());
641        bytes.extend_from_slice(&(GgufType::Uint8 as u32).to_le_bytes());
642        bytes.extend_from_slice(&offset.to_le_bytes());
643    }
644
645    #[test]
646    fn skip_gguf_value_rejects_excessive_array_depth() {
647        let mut bytes = Vec::new();
648        for _ in 0..=MAX_GGUF_ARRAY_DEPTH {
649            push_array_header(&mut bytes, GgufType::Array, 1);
650        }
651        push_array_header(&mut bytes, GgufType::Uint8, 1);
652        bytes.push(0);
653
654        let path = write_bytes("model-artifact-gguf-depth", &bytes);
655        let mut file = std::fs::File::open(&path).unwrap();
656        let err = skip_gguf_value(&mut file, GgufType::Array).unwrap_err();
657        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
658        assert!(err.to_string().contains("nesting too deep"));
659        let _ = std::fs::remove_file(path);
660    }
661
662    #[test]
663    fn skip_gguf_value_rejects_excessive_array_count() {
664        let mut bytes = Vec::new();
665        push_array_header(&mut bytes, GgufType::Uint8, MAX_GGUF_ARRAY_ELEMENTS + 1);
666
667        let path = write_bytes("model-artifact-gguf-count", &bytes);
668        let mut file = std::fs::File::open(&path).unwrap();
669        let err = skip_gguf_value(&mut file, GgufType::Array).unwrap_err();
670        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
671        assert!(err.to_string().contains("array too long"));
672        let _ = std::fs::remove_file(path);
673    }
674
675    #[test]
676    fn scan_gguf_compact_meta_returns_none_on_malicious_nested_array() {
677        let mut bytes = Vec::new();
678        bytes.extend_from_slice(b"GGUF");
679        bytes.extend_from_slice(&2u32.to_le_bytes());
680        bytes.extend_from_slice(&0i64.to_le_bytes());
681        bytes.extend_from_slice(&1i64.to_le_bytes());
682        push_gguf_string(&mut bytes, "general.architecture");
683        bytes.extend_from_slice(&(GgufType::Array as u32).to_le_bytes());
684        for _ in 0..=MAX_GGUF_ARRAY_DEPTH {
685            push_array_header(&mut bytes, GgufType::Array, 1);
686        }
687        push_array_header(&mut bytes, GgufType::Uint8, 1);
688        bytes.push(0);
689
690        let path = write_bytes("model-artifact-gguf-malicious", &bytes);
691        assert!(scan_gguf_compact_meta(&path).is_none());
692        let _ = std::fs::remove_file(path);
693    }
694
695    #[test]
696    fn scan_gguf_compact_meta_derives_value_length_from_kv_heads_without_head_count() {
697        let mut bytes = Vec::new();
698        bytes.extend_from_slice(b"GGUF");
699        bytes.extend_from_slice(&2u32.to_le_bytes());
700        bytes.extend_from_slice(&0i64.to_le_bytes());
701        bytes.extend_from_slice(&2i64.to_le_bytes());
702        push_u32_kv(&mut bytes, "llama.embedding_length", 4096);
703        push_u32_kv(&mut bytes, "llama.attention.head_count_kv", 8);
704
705        let path = write_bytes("model-artifact-gguf-kv-heads", &bytes);
706        let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF");
707        assert_eq!(meta.head_count, 0);
708        assert_eq!(meta.kv_head_count, 8);
709        assert_eq!(meta.key_length, 0);
710        assert_eq!(meta.value_length, 512);
711        let _ = std::fs::remove_file(path);
712    }
713
714    #[test]
715    fn scan_gguf_compact_meta_preserves_kv_head_count() {
716        let mut bytes = Vec::new();
717        bytes.extend_from_slice(b"GGUF");
718        bytes.extend_from_slice(&2u32.to_le_bytes());
719        bytes.extend_from_slice(&0i64.to_le_bytes());
720        bytes.extend_from_slice(&6i64.to_le_bytes());
721        push_u32_kv(&mut bytes, "llama.embedding_length", 4096);
722        push_u32_kv(&mut bytes, "llama.attention.head_count", 32);
723        push_u32_kv(&mut bytes, "llama.attention.head_count_kv", 8);
724        push_u32_kv(&mut bytes, "llama.block_count", 24);
725        push_u32_kv(&mut bytes, "llama.attention.key_length", 128);
726        push_u32_kv(&mut bytes, "llama.attention.value_length", 128);
727
728        let path = write_bytes("model-artifact-gguf-kv-head-count", &bytes);
729        let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF");
730        assert_eq!(meta.head_count, 32);
731        assert_eq!(meta.kv_head_count, 8);
732        assert_eq!(meta.effective_kv_head_count(), Some(8));
733        assert_eq!(meta.k_cache_bytes_per_token_f16(), Some(49_152));
734        assert_eq!(meta.v_cache_bytes_per_token_f16(), Some(49_152));
735        let _ = std::fs::remove_file(path);
736    }
737
738    #[test]
739    fn scan_gguf_compact_meta_preserves_kv_lora_rank() {
740        let mut bytes = Vec::new();
741        bytes.extend_from_slice(b"GGUF");
742        bytes.extend_from_slice(&2u32.to_le_bytes());
743        bytes.extend_from_slice(&0i64.to_le_bytes());
744        bytes.extend_from_slice(&1i64.to_le_bytes());
745        push_u32_kv(&mut bytes, "glm-dsa.attention.kv_lora_rank", 512);
746
747        let path = write_bytes("model-artifact-gguf-kv-lora-rank", &bytes);
748        let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF");
749        assert_eq!(meta.kv_lora_rank, 512);
750        let _ = std::fs::remove_file(path);
751    }
752
753    #[test]
754    fn scan_gguf_compact_meta_preserves_nextn_predict_layers() {
755        let mut bytes = Vec::new();
756        bytes.extend_from_slice(b"GGUF");
757        bytes.extend_from_slice(&2u32.to_le_bytes());
758        bytes.extend_from_slice(&0i64.to_le_bytes());
759        bytes.extend_from_slice(&2i64.to_le_bytes());
760        push_gguf_string(&mut bytes, "general.architecture");
761        bytes.extend_from_slice(&(GgufType::String as u32).to_le_bytes());
762        push_gguf_string(&mut bytes, "deepseek2");
763        push_u32_kv(&mut bytes, "deepseek2.nextn_predict_layers", 1);
764
765        let path = write_bytes("model-artifact-gguf-nextn", &bytes);
766        let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF");
767        assert_eq!(meta.architecture, "deepseek2");
768        assert_eq!(meta.nextn_predict_layers, 1);
769        let _ = std::fs::remove_file(path);
770    }
771
772    #[test]
773    fn scan_gguf_tensor_names_any_detects_nextn_tensor_without_reading_data() {
774        let mut bytes = Vec::new();
775        bytes.extend_from_slice(b"GGUF");
776        bytes.extend_from_slice(&2u32.to_le_bytes());
777        bytes.extend_from_slice(&2i64.to_le_bytes());
778        bytes.extend_from_slice(&0i64.to_le_bytes());
779        push_tensor_info(&mut bytes, "blk.0.attn_q.weight", 0);
780        push_tensor_info(&mut bytes, "blk.23.nextn.eh_proj.weight", 64);
781
782        let path = write_bytes("model-artifact-gguf-nextn-tensor", &bytes);
783        let has_nextn = scan_gguf_tensor_names_any(&path, |name| name.contains(".nextn."))
784            .expect("tensor scan should parse");
785        assert!(has_nextn);
786        let _ = std::fs::remove_file(path);
787    }
788
789    #[test]
790    fn scan_gguf_compact_meta_rejects_negative_kv_count() {
791        let mut bytes = Vec::new();
792        bytes.extend_from_slice(b"GGUF");
793        bytes.extend_from_slice(&2u32.to_le_bytes());
794        bytes.extend_from_slice(&0i64.to_le_bytes());
795        bytes.extend_from_slice(&(-1i64).to_le_bytes());
796
797        let path = write_bytes("model-artifact-gguf-negative-kv", &bytes);
798        assert!(scan_gguf_compact_meta(&path).is_none());
799        let _ = std::fs::remove_file(path);
800    }
801
802    #[test]
803    fn scan_gguf_tensor_byte_profile_rejects_excessive_tensor_count() {
804        let mut bytes = Vec::new();
805        bytes.extend_from_slice(b"GGUF");
806        bytes.extend_from_slice(&2u32.to_le_bytes());
807        bytes.extend_from_slice(&((MAX_GGUF_TENSOR_COUNT as i64) + 1).to_le_bytes());
808        bytes.extend_from_slice(&0i64.to_le_bytes());
809
810        let path = write_bytes("model-artifact-gguf-too-many-tensors", &bytes);
811        assert!(scan_gguf_tensor_byte_profile(&path).is_none());
812        let _ = std::fs::remove_file(path);
813    }
814
815    #[test]
816    fn read_gguf_value_as_u32_rejects_negative_int32() {
817        let path = write_bytes("model-artifact-gguf-negative-int32", &(-1i32).to_le_bytes());
818        let mut file = std::fs::File::open(&path).unwrap();
819        let err = read_gguf_value_as_u32(&mut file, GgufType::Int32).unwrap_err();
820        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
821        assert!(
822            err.to_string()
823                .contains("negative Int32 where unsigned GGUF value was expected")
824        );
825        let _ = std::fs::remove_file(path);
826    }
827
828    #[test]
829    fn scan_gguf_tensor_byte_profile_splits_base_and_expert_bytes() {
830        let mut bytes = Vec::new();
831        bytes.extend_from_slice(b"GGUF");
832        bytes.extend_from_slice(&2u32.to_le_bytes());
833        bytes.extend_from_slice(&2i64.to_le_bytes());
834        bytes.extend_from_slice(&3i64.to_le_bytes());
835
836        push_u32_kv(&mut bytes, "general.alignment", 32);
837        push_u32_kv(&mut bytes, "llama.expert_count", 8);
838        push_u32_kv(&mut bytes, "llama.expert_used_count", 2);
839
840        push_tensor_info(&mut bytes, "blk.0.ffn_up_exps.weight", 0);
841        push_tensor_info(&mut bytes, "blk.0.attn_q.weight", 64);
842
843        let data_start = align_offset(bytes.len() as u64, 32) as usize;
844        bytes.resize(data_start, 0);
845        bytes.resize(data_start + 96, 0);
846
847        let path = write_bytes("model-artifact-gguf-tensors", &bytes);
848        let profile = scan_gguf_tensor_byte_profile(&path).unwrap();
849        assert_eq!(profile.expert_count, 8);
850        assert_eq!(profile.expert_used_count, 2);
851        assert_eq!(profile.expert_tensor_bytes, 64);
852        assert_eq!(profile.base_resident_bytes, 32);
853        assert_eq!(profile.full_model_bytes, bytes.len() as u64);
854        assert_eq!(
855            profile.full_model_bytes,
856            profile.base_resident_bytes + profile.expert_tensor_bytes + profile.file_overhead_bytes
857        );
858        let _ = std::fs::remove_file(path);
859    }
860}
model_artifact/gguf.rs

model_artifact/
gguf.rs
