Skip to main content

ferrum_quantization/gguf/
native.rs

1//! vNext physical GGUF source. The validated native descriptor table retains
2//! file type IDs, including IQ types absent from Candle's QTensor enum. Model
3//! semantics and tokenizer metadata remain in their independently bound sources.
4
5use 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
27/// Describes physical bytes only. An accepted encoding does not imply that an
28/// operation provider exists for any particular device or model tensor role.
29pub 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        // SAFETY: The caller keeps the model artifact immutable while it is
67        // loaded. This handle exposes only a read-only file mapping.
68        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}