/home/noah/src/realizar/src/gguf/types.rs
Line | Count | Source |
1 | | //! GGUF type definitions and constants |
2 | | //! |
3 | | //! Core types used throughout the GGUF parser and model loading. |
4 | | //! |
5 | | //! This module defines the foundational types for GGUF file parsing: |
6 | | //! - Magic numbers and version constants |
7 | | //! - Quantization type constants (Q4_0, Q4_K, Q6_K, etc.) |
8 | | //! - Buffer types with small buffer optimization |
9 | | //! - Core structs: GGUFValue, GGUFHeader, TensorInfo, GGUFModel |
10 | | |
11 | | use std::collections::HashMap; |
12 | | |
13 | | // ============================================================================ |
14 | | // GGUF Magic and Version Constants |
15 | | // ============================================================================ |
16 | | |
17 | | /// GGUF magic number: "GGUF" in little-endian |
18 | | pub const GGUF_MAGIC: u32 = 0x4655_4747; |
19 | | |
20 | | /// Supported GGUF versions |
21 | | pub const GGUF_VERSION_V3: u32 = 3; |
22 | | |
23 | | // ============================================================================ |
24 | | // Quantization Type Constants |
25 | | // ============================================================================ |
26 | | |
27 | | /// GGUF quantization type: F32 (unquantized float32) |
28 | | pub const GGUF_TYPE_F32: u32 = 0; |
29 | | |
30 | | /// GGUF quantization type: F16 (half precision float16) |
31 | | pub const GGUF_TYPE_F16: u32 = 1; |
32 | | |
33 | | /// GGUF quantization type: `Q4_0` (4-bit quantization, block size 32) |
34 | | pub const GGUF_TYPE_Q4_0: u32 = 2; |
35 | | |
36 | | /// GGUF quantization type: `Q4_1` (4-bit quantization with min, block size 32) |
37 | | pub const GGUF_TYPE_Q4_1: u32 = 3; |
38 | | |
39 | | /// GGUF quantization type: `Q5_0` (5-bit quantization, block size 32) |
40 | | pub const GGUF_TYPE_Q5_0: u32 = 6; |
41 | | |
42 | | /// GGUF quantization type: `Q5_1` (5-bit quantization with min, block size 32) |
43 | | pub const GGUF_TYPE_Q5_1: u32 = 7; |
44 | | |
45 | | /// GGUF quantization type: `Q8_0` (8-bit quantization, block size 32) |
46 | | pub const GGUF_TYPE_Q8_0: u32 = 8; |
47 | | |
48 | | /// GGUF quantization type: `Q2_K` (2-bit K-quantization, super-block size 256) |
49 | | pub const GGUF_TYPE_Q2_K: u32 = 10; |
50 | | |
51 | | /// GGUF quantization type: `Q3_K` (3-bit K-quantization, super-block size 256) |
52 | | pub const GGUF_TYPE_Q3_K: u32 = 11; |
53 | | |
54 | | /// GGUF quantization type: `Q4_K` (4-bit K-quantization, super-block size 256) |
55 | | pub const GGUF_TYPE_Q4_K: u32 = 12; |
56 | | |
57 | | /// GGUF quantization type: `Q5_K` (5-bit K-quantization, super-block size 256) |
58 | | pub const GGUF_TYPE_Q5_K: u32 = 13; |
59 | | |
60 | | /// GGUF quantization type: `Q6_K` (6-bit K-quantization, super-block size 256) |
61 | | pub const GGUF_TYPE_Q6_K: u32 = 14; |
62 | | |
63 | | // ============================================================================ |
64 | | // IMP-117: Small Buffer Optimization Constants (per spec Section 4.1-4.2) |
65 | | // ============================================================================ |
66 | | |
67 | | /// Small buffer inline capacity for token IDs (IMP-117) |
68 | | /// Most prompts are < 32 tokens, avoiding heap allocation |
69 | | pub const TOKEN_BUFFER_INLINE_CAP: usize = 32; |
70 | | |
71 | | /// Small buffer inline capacity for attention scores (IMP-117) |
72 | | /// Stack-allocated for short sequences (per-head, small context) |
73 | | pub const ATTENTION_BUFFER_INLINE_CAP: usize = 64; |
74 | | |
75 | | /// Small buffer inline capacity for hidden states (IMP-117) |
76 | | /// Inline storage for small models (hidden_dim <= 128) |
77 | | pub const HIDDEN_BUFFER_INLINE_CAP: usize = 128; |
78 | | |
79 | | /// Buffer watermark: Low mark for inline/stack allocation |
80 | | pub const BUFFER_LW_SIZE: usize = 1024; |
81 | | |
82 | | /// Buffer watermark: High mark for pooled allocations |
83 | | pub const BUFFER_HW_SIZE: usize = 8 * 1024; |
84 | | |
85 | | /// Buffer watermark: Maximum before chunking |
86 | | pub const BUFFER_MAX_SIZE: usize = 32 * 1024; |
87 | | |
88 | | // ============================================================================ |
89 | | // Buffer Type Aliases with Small Buffer Optimization |
90 | | // ============================================================================ |
91 | | |
92 | | /// Token buffer with inline storage (IMP-117) |
93 | | /// Uses SmallVec for stack allocation when size <= TOKEN_BUFFER_INLINE_CAP |
94 | | pub type TokenBuffer = smallvec::SmallVec<[u32; TOKEN_BUFFER_INLINE_CAP]>; |
95 | | |
96 | | /// Attention score buffer with inline storage (IMP-117) |
97 | | /// Uses SmallVec for stack allocation when size <= ATTENTION_BUFFER_INLINE_CAP |
98 | | pub type AttentionBuffer = smallvec::SmallVec<[f32; ATTENTION_BUFFER_INLINE_CAP]>; |
99 | | |
100 | | /// Hidden state buffer with inline storage (IMP-117) |
101 | | /// Uses SmallVec for stack allocation when size <= HIDDEN_BUFFER_INLINE_CAP |
102 | | pub type HiddenBuffer = smallvec::SmallVec<[f32; HIDDEN_BUFFER_INLINE_CAP]>; |
103 | | |
104 | | // ============================================================================ |
105 | | // Core GGUF Types |
106 | | // ============================================================================ |
107 | | |
108 | | /// GGUF alignment requirement (32 bytes) |
109 | | pub const GGUF_ALIGNMENT: usize = 32; |
110 | | |
111 | | /// GGUF metadata value types |
112 | | #[derive(Debug, Clone, PartialEq)] |
113 | | pub enum GGUFValue { |
114 | | /// Unsigned 8-bit integer |
115 | | UInt8(u8), |
116 | | /// Signed 8-bit integer |
117 | | Int8(i8), |
118 | | /// Unsigned 16-bit integer |
119 | | UInt16(u16), |
120 | | /// Signed 16-bit integer |
121 | | Int16(i16), |
122 | | /// Unsigned 32-bit integer |
123 | | UInt32(u32), |
124 | | /// Signed 32-bit integer |
125 | | Int32(i32), |
126 | | /// 32-bit floating point |
127 | | Float32(f32), |
128 | | /// Boolean |
129 | | Bool(bool), |
130 | | /// UTF-8 string |
131 | | String(String), |
132 | | /// Array of values |
133 | | Array(Vec<GGUFValue>), |
134 | | /// Unsigned 64-bit integer |
135 | | UInt64(u64), |
136 | | /// Signed 64-bit integer |
137 | | Int64(i64), |
138 | | /// 64-bit floating point |
139 | | Float64(f64), |
140 | | } |
141 | | |
142 | | /// GGUF file header |
143 | | #[derive(Debug, Clone, PartialEq)] |
144 | | pub struct GGUFHeader { |
145 | | /// Magic number (must be `GGUF_MAGIC`) |
146 | | pub magic: u32, |
147 | | /// Format version |
148 | | pub version: u32, |
149 | | /// Number of tensors in the file |
150 | | pub tensor_count: u64, |
151 | | /// Number of metadata key-value pairs |
152 | | pub metadata_count: u64, |
153 | | } |
154 | | |
155 | | /// Tensor information |
156 | | #[derive(Debug, Clone, PartialEq)] |
157 | | pub struct TensorInfo { |
158 | | /// Tensor name |
159 | | pub name: String, |
160 | | /// Number of dimensions |
161 | | pub n_dims: u32, |
162 | | /// Dimensions (shape) |
163 | | pub dims: Vec<u64>, |
164 | | /// Quantization type |
165 | | pub qtype: u32, |
166 | | /// Offset in the file where tensor data starts |
167 | | pub offset: u64, |
168 | | } |
169 | | |
170 | | /// GGUF model container |
171 | | #[derive(Debug, Clone)] |
172 | | pub struct GGUFModel { |
173 | | /// File header |
174 | | pub header: GGUFHeader, |
175 | | /// Metadata key-value pairs |
176 | | pub metadata: HashMap<String, GGUFValue>, |
177 | | /// Tensor information |
178 | | pub tensors: Vec<TensorInfo>, |
179 | | /// Offset where tensor data starts (after header/metadata/tensor_info + alignment) |
180 | | pub tensor_data_start: usize, |
181 | | } |
182 | | |
183 | | #[cfg(test)] |
184 | | mod tests { |
185 | | use super::*; |
186 | | |
187 | | #[test] |
188 | 1 | fn test_magic_constant() { |
189 | 1 | assert_eq!(GGUF_MAGIC, 0x4655_4747); |
190 | 1 | } |
191 | | |
192 | | #[test] |
193 | 1 | fn test_quantization_constants() { |
194 | 1 | assert_eq!(GGUF_TYPE_F32, 0); |
195 | 1 | assert_eq!(GGUF_TYPE_F16, 1); |
196 | 1 | assert_eq!(GGUF_TYPE_Q4_0, 2); |
197 | 1 | assert_eq!(GGUF_TYPE_Q8_0, 8); |
198 | 1 | assert_eq!(GGUF_TYPE_Q4_K, 12); |
199 | 1 | assert_eq!(GGUF_TYPE_Q6_K, 14); |
200 | 1 | } |
201 | | |
202 | | #[test] |
203 | 1 | fn test_buffer_constants() { |
204 | 1 | assert_eq!(TOKEN_BUFFER_INLINE_CAP, 32); |
205 | 1 | assert_eq!(ATTENTION_BUFFER_INLINE_CAP, 64); |
206 | 1 | assert_eq!(HIDDEN_BUFFER_INLINE_CAP, 128); |
207 | 1 | } |
208 | | |
209 | | #[test] |
210 | 1 | fn test_buffer_watermarks() { |
211 | 1 | assert_eq!(BUFFER_LW_SIZE, 1024); |
212 | 1 | assert_eq!(BUFFER_HW_SIZE, 8 * 1024); |
213 | 1 | assert_eq!(BUFFER_MAX_SIZE, 32 * 1024); |
214 | 1 | } |
215 | | |
216 | | #[test] |
217 | 1 | fn test_version_constant() { |
218 | 1 | assert_eq!(GGUF_VERSION_V3, 3); |
219 | 1 | } |
220 | | |
221 | | #[test] |
222 | 1 | fn test_gguf_value_variants() { |
223 | 1 | let uint8 = GGUFValue::UInt8(255); |
224 | 1 | let string = GGUFValue::String("test".to_string()); |
225 | 1 | let array = GGUFValue::Array(vec![GGUFValue::UInt32(1), GGUFValue::UInt32(2)]); |
226 | | |
227 | 1 | assert_eq!(uint8, GGUFValue::UInt8(255)); |
228 | 1 | assert_eq!(string, GGUFValue::String("test".to_string())); |
229 | 1 | assert!(matches!0 (array, GGUFValue::Array(_))); |
230 | 1 | } |
231 | | |
232 | | #[test] |
233 | 1 | fn test_gguf_header() { |
234 | 1 | let header = GGUFHeader { |
235 | 1 | magic: GGUF_MAGIC, |
236 | 1 | version: GGUF_VERSION_V3, |
237 | 1 | tensor_count: 100, |
238 | 1 | metadata_count: 50, |
239 | 1 | }; |
240 | | |
241 | 1 | assert_eq!(header.magic, 0x4655_4747); |
242 | 1 | assert_eq!(header.version, 3); |
243 | 1 | assert_eq!(header.tensor_count, 100); |
244 | 1 | assert_eq!(header.metadata_count, 50); |
245 | 1 | } |
246 | | |
247 | | #[test] |
248 | 1 | fn test_tensor_info() { |
249 | 1 | let info = TensorInfo { |
250 | 1 | name: "model.layers.0.attn.wq".to_string(), |
251 | 1 | n_dims: 2, |
252 | 1 | dims: vec![4096, 4096], |
253 | 1 | qtype: GGUF_TYPE_Q4_K, |
254 | 1 | offset: 1024, |
255 | 1 | }; |
256 | | |
257 | 1 | assert_eq!(info.name, "model.layers.0.attn.wq"); |
258 | 1 | assert_eq!(info.n_dims, 2); |
259 | 1 | assert_eq!(info.dims, vec![4096, 4096]); |
260 | 1 | assert_eq!(info.qtype, GGUF_TYPE_Q4_K); |
261 | 1 | assert_eq!(info.offset, 1024); |
262 | 1 | } |
263 | | |
264 | | #[test] |
265 | 1 | fn test_gguf_model() { |
266 | 1 | let model = GGUFModel { |
267 | 1 | header: GGUFHeader { |
268 | 1 | magic: GGUF_MAGIC, |
269 | 1 | version: GGUF_VERSION_V3, |
270 | 1 | tensor_count: 1, |
271 | 1 | metadata_count: 0, |
272 | 1 | }, |
273 | 1 | metadata: HashMap::new(), |
274 | 1 | tensors: vec![], |
275 | 1 | tensor_data_start: 128, |
276 | 1 | }; |
277 | | |
278 | 1 | assert_eq!(model.header.magic, GGUF_MAGIC); |
279 | 1 | assert!(model.tensors.is_empty()); |
280 | 1 | assert_eq!(model.tensor_data_start, 128); |
281 | 1 | } |
282 | | |
283 | | #[test] |
284 | 1 | fn test_alignment_constant() { |
285 | 1 | assert_eq!(GGUF_ALIGNMENT, 32); |
286 | | // Verify it's a power of 2 |
287 | 1 | assert_eq!(GGUF_ALIGNMENT & (GGUF_ALIGNMENT - 1), 0); |
288 | 1 | } |
289 | | } |