Skip to main content

ferrum_quantization/gguf/
source.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use ferrum_interfaces::vnext::{
5    ElementType, RetainedHostMemoryRegion, StableHostMemory, VNextError, WeightComponentPayload,
6    WeightComponentSource, WeightComponentSpec, WeightEncoding,
7};
8use ferrum_types::{FerrumError, Result};
9
10use super::{GgmlDType, GgufFile};
11use crate::safetensors_archive::transcode_dense_bytes;
12
13/// Schema-addressed, mmap-backed GGUF source for vNext static weights.
14/// Fixed-block payloads borrow the immutable file mapping without
15/// dequantization or repacking. Dense floating-point payloads are borrowed
16/// when their type matches the schema and materialized on a cold-path source
17/// request when the typed execution plan requires another floating-point type.
18pub struct GgufWeightComponentSource {
19    file: Arc<GgufFile>,
20    source_file: String,
21}
22
23impl GgufWeightComponentSource {
24    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
25        let path = path.as_ref();
26        let source_file = path
27            .file_name()
28            .and_then(|name| name.to_str())
29            .filter(|name| is_portable_source_file(name))
30            .ok_or_else(|| {
31                FerrumError::model(format!(
32                    "GGUF path must end in one portable UTF-8 file name: {}",
33                    path.display()
34                ))
35            })?
36            .to_owned();
37        let file = GgufFile::open(path).map_err(|error| {
38            FerrumError::model(format!(
39                "open vNext GGUF source {}: {error}",
40                path.display()
41            ))
42        })?;
43        Ok(Self {
44            file: Arc::new(file),
45            source_file,
46        })
47    }
48
49    pub fn file(&self) -> &GgufFile {
50        &self.file
51    }
52
53    pub fn source_file(&self) -> &str {
54        &self.source_file
55    }
56}
57
58impl WeightComponentSource for GgufWeightComponentSource {
59    fn component<'source>(
60        &'source self,
61        component: &WeightComponentSpec,
62    ) -> std::result::Result<WeightComponentPayload<'source>, VNextError> {
63        let [external_name] = component.external_names.as_slice() else {
64            return Err(invalid_component(
65                component,
66                "GGUF components must bind exactly one physical tensor; combine tensors in PhysicalWeightLayout instead",
67            ));
68        };
69        let info = self.file.tensor_info(external_name).ok_or_else(|| {
70            invalid_component(
71                component,
72                format!("GGUF tensor {external_name:?} is absent"),
73            )
74        })?;
75        let bytes = self.file.tensor_byte_slice(external_name).ok_or_else(|| {
76            invalid_component(
77                component,
78                format!("GGUF tensor {external_name:?} has an invalid block or byte range"),
79            )
80        })?;
81
82        let (element_type, payload_bytes) = match &component.encoding {
83            WeightEncoding::Dense { element_type } => {
84                let actual = dense_element_type(info.ggml_dtype).ok_or_else(|| {
85                    invalid_component(
86                        component,
87                        format!(
88                            "GGUF dtype {:?} is quantized but the schema declares dense bytes",
89                            info.ggml_dtype
90                        ),
91                    )
92                })?;
93                let dimensions = info
94                    .shape
95                    .dims()
96                    .iter()
97                    .map(|dimension| *dimension as u64)
98                    .collect::<Vec<_>>();
99                if dimensions != component.dimensions {
100                    return Err(invalid_component(
101                        component,
102                        format!(
103                            "GGUF dense tensor dimensions differ: source_dtype={:?} dimensions={dimensions:?}",
104                            info.ggml_dtype,
105                        ),
106                    ));
107                }
108                let materialized =
109                    transcode_dense_bytes(bytes, actual, *element_type, external_name, None)?;
110                (*element_type, materialized)
111            }
112            WeightEncoding::BlockQuantized(spec) => {
113                spec.validate()?;
114                let actual_format =
115                    block_quantization_format(info.ggml_dtype).ok_or_else(|| {
116                        invalid_component(
117                            component,
118                            format!(
119                            "GGUF dtype {:?} is not a supported fixed-block quantization format",
120                            info.ggml_dtype
121                        ),
122                        )
123                    })?;
124                let logical_elements = u64::try_from(info.shape.elem_count()).map_err(|_| {
125                    invalid_component(component, "GGUF tensor element count exceeds u64")
126                })?;
127                let block_width = u64::from(spec.logical_values_per_block);
128                let mut physical_dimensions = info
129                    .shape
130                    .dims()
131                    .iter()
132                    .map(|dimension| *dimension as u64)
133                    .collect::<Vec<_>>();
134                let innermost = physical_dimensions.last_mut().ok_or_else(|| {
135                    invalid_component(
136                        component,
137                        "GGUF quantized tensor must have at least one axis",
138                    )
139                })?;
140                if !innermost.is_multiple_of(block_width) {
141                    return Err(invalid_component(
142                        component,
143                        format!(
144                            "GGUF innermost dimension {innermost} is not divisible by block width {block_width}"
145                        ),
146                    ));
147                }
148                *innermost /= block_width;
149                if actual_format != spec.format_id.as_str()
150                    || info.ggml_dtype.block_size() != spec.logical_values_per_block as usize
151                    || info.ggml_dtype.type_size() != spec.bytes_per_block as usize
152                    || !logical_elements.is_multiple_of(block_width)
153                    || physical_dimensions != component.dimensions
154                {
155                    return Err(invalid_component(
156                        component,
157                        format!(
158                            "GGUF block ABI differs: dtype={:?} format={actual_format} values_per_block={} bytes_per_block={} logical_elements={logical_elements} physical_dimensions={physical_dimensions:?}",
159                            info.ggml_dtype,
160                            info.ggml_dtype.block_size(),
161                            info.ggml_dtype.type_size(),
162                        ),
163                    ));
164                }
165                (ElementType::U8, bytes.into())
166            }
167            WeightEncoding::DenseAffine { .. } => {
168                return Err(invalid_component(
169                    component,
170                    "GGUF source values are already transformed and cannot apply a dense affine source transform",
171                ));
172            }
173            WeightEncoding::Quantized(_) => {
174                return Err(invalid_component(
175                    component,
176                    "GGUF fixed-block bytes cannot satisfy a separate-component quantization encoding",
177                ));
178            }
179        };
180
181        let retained_host_memory =
182            if payload_bytes.as_ptr() == bytes.as_ptr() && payload_bytes.len() == bytes.len() {
183                let (offset_bytes, length_bytes) =
184                    self.file.tensor_byte_range(external_name).ok_or_else(|| {
185                        invalid_component(component, "GGUF tensor byte range is invalid")
186                    })?;
187                Some(RetainedHostMemoryRegion::new(
188                    Arc::clone(&self.file),
189                    offset_bytes,
190                    length_bytes,
191                )?)
192            } else {
193                None
194            };
195        let payload = WeightComponentPayload::new(
196            component,
197            external_name.clone(),
198            self.source_file.clone(),
199            component.dimensions.clone(),
200            element_type,
201            payload_bytes,
202        )?;
203        match retained_host_memory {
204            Some(retained) => payload.with_retained_host_memory(retained),
205            None => Ok(payload),
206        }
207    }
208}
209
210// SAFETY: GgufFile owns an immutable Mmap whose address and length do not
211// change during its lifetime.
212unsafe impl StableHostMemory for GgufFile {
213    fn stable_bytes(&self) -> &[u8] {
214        self.mmap_bytes()
215    }
216}
217
218fn dense_element_type(dtype: GgmlDType) -> Option<ElementType> {
219    match dtype {
220        GgmlDType::F16 => Some(ElementType::F16),
221        GgmlDType::BF16 => Some(ElementType::Bf16),
222        GgmlDType::F32 => Some(ElementType::F32),
223        _ => None,
224    }
225}
226
227pub fn block_quantization_format(dtype: GgmlDType) -> Option<&'static str> {
228    match dtype {
229        GgmlDType::Q4_0 => Some("quantization.gguf.q4-0"),
230        GgmlDType::Q4_1 => Some("quantization.gguf.q4-1"),
231        GgmlDType::Q5_0 => Some("quantization.gguf.q5-0"),
232        GgmlDType::Q5_1 => Some("quantization.gguf.q5-1"),
233        GgmlDType::Q8_0 => Some("quantization.gguf.q8-0"),
234        GgmlDType::Q8_1 => Some("quantization.gguf.q8-1"),
235        GgmlDType::Q2K => Some("quantization.gguf.q2-k"),
236        GgmlDType::Q3K => Some("quantization.gguf.q3-k"),
237        GgmlDType::Q4K => Some("quantization.gguf.q4-k"),
238        GgmlDType::Q5K => Some("quantization.gguf.q5-k"),
239        GgmlDType::Q6K => Some("quantization.gguf.q6-k"),
240        GgmlDType::Q8K => Some("quantization.gguf.q8-k"),
241        GgmlDType::F16 | GgmlDType::BF16 | GgmlDType::F32 => None,
242    }
243}
244
245fn is_portable_source_file(name: &str) -> bool {
246    !name.is_empty()
247        && !name.contains(['/', '\\'])
248        && !matches!(name, "." | "..")
249        && !name.bytes().any(|byte| byte.is_ascii_control())
250}
251
252fn invalid_component(component: &WeightComponentSpec, reason: impl Into<String>) -> VNextError {
253    VNextError::InvalidExecutionPlan {
254        reason: format!(
255            "GGUF component `{}` does not match its typed source: {}",
256            component.id,
257            reason.into()
258        ),
259    }
260}