ferrum_quantization/gguf/
native.rs1use std::collections::BTreeMap;
6use std::fs::File;
7use std::io::Cursor;
8use std::path::Path;
9
10use candle_core::{Error, Result};
11use ferrum_interfaces::vnext::{BlockQuantizationSpec, ElementType, WeightEncoding};
12use memmap2::Mmap;
13
14use super::inventory::{block_abi, GgufInventory};
15use super::GgufHadamard;
16
17#[derive(Debug, Clone)]
18pub struct NativeGgufTensor {
19 pub ggml_type: u32,
20 pub dimensions: Vec<u64>,
21 pub encoding: WeightEncoding,
22 pub elements: u64,
23 offset: usize,
24 bytes: usize,
25}
26
27pub fn gguf_weight_encoding(ggml_type: u32) -> Result<WeightEncoding> {
30 let abi = block_abi(ggml_type)?;
31 if let Some(format) = abi.format {
32 let spec = BlockQuantizationSpec {
33 format_id: format.to_owned().try_into().map_err(Error::wrap)?,
34 logical_values_per_block: u32::try_from(abi.values).map_err(Error::wrap)?,
35 bytes_per_block: u32::try_from(abi.bytes).map_err(Error::wrap)?,
36 };
37 spec.validate().map_err(Error::wrap)?;
38 Ok(WeightEncoding::BlockQuantized(spec))
39 } else {
40 let element_type = match ggml_type {
41 0 => ElementType::F32,
42 1 => ElementType::F16,
43 30 => ElementType::Bf16,
44 _ => {
45 return Err(Error::Msg(format!(
46 "GGUF type {ggml_type} has no declared dense encoding"
47 )))
48 }
49 };
50 Ok(WeightEncoding::Dense { element_type })
51 }
52}
53
54#[derive(Debug)]
55pub struct NativeGgufFile {
56 mmap: Mmap,
57 architecture: String,
58 quantization_version: Option<u64>,
59 hadamard: Option<GgufHadamard>,
60 tensors: BTreeMap<String, NativeGgufTensor>,
61}
62
63impl NativeGgufFile {
64 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
65 let file = File::open(path.as_ref()).map_err(Error::wrap)?;
66 let mmap = unsafe { Mmap::map(&file) }.map_err(Error::wrap)?;
69 let inventory = GgufInventory::read(&mut Cursor::new(&mmap[..]), mmap.len() as u64)?;
70 if inventory
71 .split
72 .as_ref()
73 .is_some_and(|split| split.count != 1)
74 {
75 return Err(Error::Msg("a GGUF shard cannot be loaded as a complete weight artifact; all declared shards are required".into()));
76 }
77 let tensors = inventory
78 .tensors
79 .into_iter()
80 .map(|tensor| {
81 let offset = usize::try_from(tensor.absolute_offset).map_err(Error::wrap)?;
82 let bytes = usize::try_from(tensor.bytes).map_err(Error::wrap)?;
83 let encoding = gguf_weight_encoding(tensor.ggml_type)?;
84 Ok((
85 tensor.name,
86 NativeGgufTensor {
87 ggml_type: tensor.ggml_type,
88 dimensions: tensor.dimensions,
89 encoding,
90 elements: tensor.elements,
91 offset,
92 bytes,
93 },
94 ))
95 })
96 .collect::<Result<BTreeMap<_, _>>>()?;
97 Ok(Self {
98 mmap,
99 architecture: inventory.architecture,
100 quantization_version: inventory.quantization_version,
101 hadamard: inventory.hadamard,
102 tensors,
103 })
104 }
105
106 pub fn architecture(&self) -> Result<&str> {
107 Ok(&self.architecture)
108 }
109 pub fn quantization_version(&self) -> Option<u64> {
110 self.quantization_version
111 }
112 pub fn hadamard(&self) -> Option<&GgufHadamard> {
113 self.hadamard.as_ref()
114 }
115 pub fn tensor_count(&self) -> usize {
116 self.tensors.len()
117 }
118 pub fn tensor_names(&self) -> impl Iterator<Item = &str> {
119 self.tensors.keys().map(String::as_str)
120 }
121 pub fn tensor_info(&self, name: &str) -> Option<&NativeGgufTensor> {
122 self.tensors.get(name)
123 }
124 pub fn has_tensor(&self, name: &str) -> bool {
125 self.tensors.contains_key(name)
126 }
127 pub fn mmap_bytes(&self) -> &[u8] {
128 &self.mmap
129 }
130
131 pub fn tensor_byte_slice(&self, name: &str) -> Option<&[u8]> {
132 let (offset, bytes) = self.tensor_byte_range(name)?;
133 self.mmap.get(offset..offset.checked_add(bytes)?)
134 }
135
136 pub fn tensor_byte_range(&self, name: &str) -> Option<(usize, usize)> {
137 let info = self.tensor_info(name)?;
138 Some((info.offset, info.bytes))
139 }
140}