combs_formats/
source.rs

1//! The format-agnostic [`ModelSource`] trait and tensor reader.
2
3use burn::tensor::{Device, Int, Tensor, TensorData, backend::Backend};
4
5use crate::metadata::ModelMetadata;
6use crate::tokenizer::TokenizerSpec;
7use crate::{FormatError, Result};
8
9/// The central adapter trait: a source of model weights + config, independent
10/// of the on-disk format (LiteRT-LM `ModelResources` equivalent).
11///
12/// Implementations must be cheap to query for metadata and names, and lazy /
13/// zero-copy (e.g. mmap-backed) when opening tensors.
14pub trait ModelSource: Send + Sync {
15    /// Architecture + hyperparameter metadata.
16    fn metadata(&self) -> &ModelMetadata;
17
18    /// Names of all tensors available in this source.
19    fn tensor_names(&self) -> Vec<String>;
20
21    /// Opens a tensor by name, returning a lazy reader over the raw bytes.
22    fn open_tensor(&self, name: &str) -> Result<TensorReader<'_>>;
23
24    /// Tokenizer specification (path to `tokenizer.json` + added tokens).
25    fn tokenizer(&self) -> Result<TokenizerSpec>;
26
27    /// Sampler defaults from `generation_config.json`, if present.
28    fn sampler_defaults(&self) -> Option<SamplerConfig>;
29
30    /// Raw *packed* quantized bytes for a tensor, when this source stores it
31    /// in a quant format that has a device kernel (GGUF Q4_0/Q4_K/Q6_K).
32    /// `None` means "no packed representation" — the caller falls back to
33    /// [`ModelSource::open_tensor`], which dequantizes to float. Sources
34    /// without packed formats keep this default.
35    fn open_tensor_quant(&self, _name: &str) -> Result<Option<QuantTensor<'_>>> {
36        Ok(None)
37    }
38}
39
40/// GGUF quant formats with a native device kernel (see `combs-models`).
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum QuantFormat {
43    /// 32-value blocks, f16 scale + 16 nibble bytes (18 B).
44    Q4_0,
45    /// 32-value blocks, f16 scale + u32 high bits + 16 nibble bytes (22 B).
46    Q5_0,
47    /// 32-value blocks, f16 scale + 32 i8 values (34 B).
48    Q8_0,
49    /// 256-value superblocks, 6-bit sub-scales + 4-bit quants (144 B).
50    Q4K,
51    /// 256-value superblocks, 6-bit sub-scales + 5-bit quants (176 B).
52    Q5K,
53    /// 256-value superblocks, i8 sub-scales + 6-bit quants (210 B).
54    Q6K,
55}
56
57/// A quantized tensor's packed bytes, exactly as stored in the file.
58pub struct QuantTensor<'a> {
59    /// Block format of `data`.
60    pub format: QuantFormat,
61    /// Logical shape, HF layout (`[out_features, in_features]` for weights).
62    pub shape: Vec<usize>,
63    /// The raw block stream — mmap-backed when served verbatim, owned when
64    /// the source had to reorder rows (GGUF RoPE de-permutation).
65    pub data: std::borrow::Cow<'a, [u8]>,
66}
67
68/// Element dtypes supported by the loaders.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum TensorDtype {
71    /// IEEE 64-bit float.
72    F64,
73    /// IEEE 32-bit float.
74    F32,
75    /// IEEE 16-bit half float.
76    F16,
77    /// bfloat16.
78    BF16,
79    /// Signed 64-bit integer.
80    I64,
81    /// Signed 32-bit integer.
82    I32,
83    /// Signed 16-bit integer.
84    I16,
85    /// Signed 8-bit integer.
86    I8,
87    /// Unsigned 64-bit integer.
88    U64,
89    /// Unsigned 32-bit integer.
90    U32,
91    /// Unsigned 16-bit integer.
92    U16,
93    /// Unsigned 8-bit integer (raw packed data, e.g. quantized weights).
94    U8,
95    /// Boolean (stored as one byte).
96    Bool,
97}
98
99impl TensorDtype {
100    /// Byte size of one element.
101    pub fn size(&self) -> usize {
102        match self {
103            TensorDtype::F64 | TensorDtype::I64 | TensorDtype::U64 => 8,
104            TensorDtype::F32 | TensorDtype::I32 | TensorDtype::U32 => 4,
105            TensorDtype::F16 | TensorDtype::BF16 | TensorDtype::I16 | TensorDtype::U16 => 2,
106            TensorDtype::I8 | TensorDtype::U8 | TensorDtype::Bool => 1,
107        }
108    }
109}
110
111impl std::fmt::Display for TensorDtype {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        match self {
114            TensorDtype::F64 => write!(f, "F64"),
115            TensorDtype::F32 => write!(f, "F32"),
116            TensorDtype::F16 => write!(f, "F16"),
117            TensorDtype::BF16 => write!(f, "BF16"),
118            TensorDtype::I64 => write!(f, "I64"),
119            TensorDtype::I32 => write!(f, "I32"),
120            TensorDtype::I16 => write!(f, "I16"),
121            TensorDtype::I8 => write!(f, "I8"),
122            TensorDtype::U64 => write!(f, "U64"),
123            TensorDtype::U32 => write!(f, "U32"),
124            TensorDtype::U16 => write!(f, "U16"),
125            TensorDtype::U8 => write!(f, "U8"),
126            TensorDtype::Bool => write!(f, "Bool"),
127        }
128    }
129}
130
131/// A lazy view over one tensor's raw bytes inside a [`ModelSource`].
132///
133/// The byte slice borrows from the source (e.g. an mmap region) — no copy is
134/// made until [`TensorReader::load_data`] is called. Format adapters that
135/// decode on open (e.g. GGUF quantization) use [`TensorReader::owned`].
136pub struct TensorReader<'a> {
137    name: String,
138    shape: Vec<usize>,
139    dtype: TensorDtype,
140    data: std::borrow::Cow<'a, [u8]>,
141}
142
143impl<'a> TensorReader<'a> {
144    /// Creates a reader from raw parts. `data` must be
145    /// `shape.iter().product::<usize>() * dtype.size()` little-endian bytes.
146    pub fn new(name: String, shape: Vec<usize>, dtype: TensorDtype, data: &'a [u8]) -> Self {
147        TensorReader {
148            name,
149            shape,
150            dtype,
151            data: std::borrow::Cow::Borrowed(data),
152        }
153    }
154
155    /// Creates a reader over owned (already-decoded) f32 bytes.
156    pub fn owned(name: String, shape: Vec<usize>, data: Vec<u8>) -> Self {
157        TensorReader {
158            name,
159            shape,
160            dtype: TensorDtype::F32,
161            data: std::borrow::Cow::Owned(data),
162        }
163    }
164
165    /// Creates a reader over owned bytes of an explicit dtype (used when a
166    /// passthrough tensor had to be row-reordered on load).
167    pub fn owned_with_dtype(
168        name: String,
169        shape: Vec<usize>,
170        dtype: TensorDtype,
171        data: Vec<u8>,
172    ) -> Self {
173        TensorReader {
174            name,
175            shape,
176            dtype,
177            data: std::borrow::Cow::Owned(data),
178        }
179    }
180
181    /// Element shape.
182    pub fn shape(&self) -> &[usize] {
183        &self.shape
184    }
185
186    /// On-disk dtype.
187    pub fn dtype(&self) -> TensorDtype {
188        self.dtype
189    }
190
191    /// Number of elements.
192    pub fn num_elements(&self) -> usize {
193        self.shape.iter().product()
194    }
195
196    /// Converts the raw bytes to f32 `TensorData`. Integer and boolean tensors
197    /// are cast to f32 so weight loaders never abort on buffer dtypes such as
198    /// I64 `position_ids`.
199    pub fn load_data(&self) -> Result<TensorData> {
200        let values: Vec<f32> = match self.dtype {
201            TensorDtype::F64 => self
202                .data
203                .chunks_exact(8)
204                .map(|c| f64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32)
205                .collect(),
206            TensorDtype::F32 => self
207                .data
208                .chunks_exact(4)
209                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
210                .collect(),
211            TensorDtype::F16 => self
212                .data
213                .chunks_exact(2)
214                .map(|c| half::f16::from_le_bytes([c[0], c[1]]).to_f32())
215                .collect(),
216            TensorDtype::BF16 => self
217                .data
218                .chunks_exact(2)
219                .map(|c| half::bf16::from_le_bytes([c[0], c[1]]).to_f32())
220                .collect(),
221            TensorDtype::I64 => self
222                .data
223                .chunks_exact(8)
224                .map(|c| i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32)
225                .collect(),
226            TensorDtype::I32 => self
227                .data
228                .chunks_exact(4)
229                .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f32)
230                .collect(),
231            TensorDtype::I16 => self
232                .data
233                .chunks_exact(2)
234                .map(|c| i16::from_le_bytes([c[0], c[1]]) as f32)
235                .collect(),
236            TensorDtype::I8 => self
237                .data
238                .iter()
239                .map(|&b| b as i8 as f32)
240                .collect(),
241            TensorDtype::U64 => self
242                .data
243                .chunks_exact(8)
244                .map(|c| u64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32)
245                .collect(),
246            TensorDtype::U32 => self
247                .data
248                .chunks_exact(4)
249                .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f32)
250                .collect(),
251            TensorDtype::U16 => self
252                .data
253                .chunks_exact(2)
254                .map(|c| u16::from_le_bytes([c[0], c[1]]) as f32)
255                .collect(),
256            TensorDtype::U8 => self.data.iter().map(|&b| b as f32).collect(),
257            TensorDtype::Bool => self.data.iter().map(|&b| (b != 0) as i32 as f32).collect(),
258        };
259        if values.len() != self.num_elements() {
260            return Err(FormatError::Safetensors(format!(
261                "tensor {}: expected {} elements, got {} (dtype {})",
262                self.name,
263                self.num_elements(),
264                values.len(),
265                self.dtype
266            )));
267        }
268        Ok(TensorData::new(values, self.shape.clone()))
269    }
270
271    /// Loads the tensor onto a backend device as an f32 tensor of rank `D`.
272    pub fn load_to_tensor<B: Backend, const D: usize>(
273        &self,
274        device: &Device<B>,
275    ) -> Result<Tensor<B, D>> {
276        let data = self.load_data()?;
277        if self.shape.len() != D {
278            return Err(FormatError::Safetensors(format!(
279                "tensor {}: expected rank {D}, got {}",
280                self.name,
281                self.shape.len()
282            )));
283        }
284        Ok(Tensor::from_data(data, device))
285    }
286
287    /// Loads raw unsigned-byte data onto a backend device as an i32 tensor
288    /// of rank `D` (values 0..=255). Only valid for `U8` tensors; used to
289    /// feed packed quantized weights to GPU-side dequantization.
290    pub fn load_int_tensor<B: Backend, const D: usize>(
291        &self,
292        device: &Device<B>,
293    ) -> Result<Tensor<B, D, Int>> {
294        if self.dtype != TensorDtype::U8 {
295            return Err(FormatError::Safetensors(format!(
296                "tensor {}: load_int_tensor requires U8, got {}",
297                self.name, self.dtype
298            )));
299        }
300        if self.shape.len() != D {
301            return Err(FormatError::Safetensors(format!(
302                "tensor {}: expected rank {D}, got {}",
303                self.name,
304                self.shape.len()
305            )));
306        }
307        let values: Vec<i32> = self.data.iter().map(|&b| b as i32).collect();
308        Ok(Tensor::from_data(
309            TensorData::new(values, self.shape.clone()),
310            device,
311        ))
312    }
313
314    /// The raw little-endian bytes, borrowing from the source (zero-copy).
315    pub fn raw_bytes(&self) -> &[u8] {
316        &self.data
317    }
318}
319
320/// Blanket forwarding so `Box<dyn ModelSource>` (returned by
321/// `open_model_source`) can be passed anywhere a `&dyn ModelSource` goes.
322impl<T: ModelSource + ?Sized> ModelSource for Box<T> {
323    fn metadata(&self) -> &crate::ModelMetadata {
324        (**self).metadata()
325    }
326    fn tensor_names(&self) -> Vec<String> {
327        (**self).tensor_names()
328    }
329    fn open_tensor(&self, name: &str) -> crate::Result<TensorReader<'_>> {
330        (**self).open_tensor(name)
331    }
332    fn tokenizer(&self) -> crate::Result<TokenizerSpec> {
333        (**self).tokenizer()
334    }
335    fn sampler_defaults(&self) -> Option<SamplerConfig> {
336        (**self).sampler_defaults()
337    }
338    // Every method must be forwarded, including defaulted ones: a missing
339    // forward silently pins callers of `Box<dyn ModelSource>` to the trait
340    // default (this bit `open_tensor_quant` — quantized GGUF weights fell
341    // back to dense for every CLI run while unit tests on the concrete
342    // type passed).
343    fn open_tensor_quant(&self, name: &str) -> Result<Option<QuantTensor<'_>>> {
344        (**self).open_tensor_quant(name)
345    }
346}
347
348/// Default sampler parameters, typically from `generation_config.json`.
349#[derive(Debug, Clone, Default)]
350pub struct SamplerConfig {    /// Sampling temperature (1.0 = neutral, 0.0 = greedy).
351    pub temperature: Option<f32>,
352    /// Top-p (nucleus) threshold.
353    pub top_p: Option<f32>,
354    /// Top-k cutoff.
355    pub top_k: Option<usize>,
356    /// Repetition penalty.
357    pub repetition_penalty: Option<f32>,
358    /// Suggested maximum new tokens.
359    pub max_new_tokens: Option<usize>,
360}
combs_formats/source.rs

combs_formats/
source.rs
