/home/noah/src/realizar/src/gguf/quantized.rs
Line | Count | Source |
1 | | //! Quantized tensor types for GGUF models |
2 | | //! |
3 | | //! This module contains the fundamental quantized tensor types that form |
4 | | //! the backbone of efficient LLM inference: |
5 | | //! |
6 | | //! - `QuantizedTensorRef`: Reference to quantized data in memory-mapped file |
7 | | //! - `OwnedQuantizedTensor`: Owned copy of quantized data |
8 | | //! - `QKVWeights`: Fused or separate QKV projection weights (borrowed) |
9 | | //! - `OwnedQKVWeights`: Fused or separate QKV projection weights (owned) |
10 | | //! |
11 | | //! Per Wulf & McKee (1995) "Hitting the Memory Wall", memory bandwidth is the |
12 | | //! bottleneck for LLM inference. These types support 8x bandwidth reduction |
13 | | //! via Q4_K quantization. |
14 | | |
15 | | // ============================================================================ |
16 | | // QuantizedTensorRef - Reference to quantized data in mmap |
17 | | // ============================================================================ |
18 | | |
19 | | /// Reference to quantized tensor data in memory-mapped file |
20 | | /// |
21 | | /// Per Wulf & McKee (1995) "Hitting the Memory Wall", memory bandwidth is the |
22 | | /// bottleneck for LLM inference. By keeping weights in quantized form and |
23 | | /// dequantizing inline during computation, we achieve 8x memory bandwidth |
24 | | /// reduction for Q4_K format. |
25 | | #[derive(Debug, Clone)] |
26 | | pub struct QuantizedTensorRef { |
27 | | /// Byte offset in file where tensor data starts |
28 | | pub offset: usize, |
29 | | /// Size in bytes of the quantized data |
30 | | pub byte_size: usize, |
31 | | /// Number of elements after dequantization |
32 | | pub num_elements: usize, |
33 | | /// Quantization type (GGUF_TYPE_Q4_K, GGUF_TYPE_Q6_K, etc.) |
34 | | pub qtype: u32, |
35 | | } |
36 | | |
37 | | // ============================================================================ |
38 | | // QKVWeights - Borrowed QKV weight storage |
39 | | // ============================================================================ |
40 | | |
41 | | /// QKV weight storage - supports both fused (phi-2) and separate (llama) formats |
42 | | /// |
43 | | /// Five Whys Root Cause Fix: TinyLlama and other LLaMA-style models use separate |
44 | | /// Q, K, V tensors while phi-2 style models use fused QKV. This enum supports both. |
45 | | #[derive(Clone)] |
46 | | pub enum QKVWeights { |
47 | | /// Fused QKV tensor (phi-2 style): single [hidden_dim, 3*hidden_dim] tensor |
48 | | Fused(QuantizedTensorRef), |
49 | | /// Separate Q, K, V tensors (llama style): three separate tensors |
50 | | Separate { |
51 | | /// Query projection [hidden_dim, hidden_dim] |
52 | | q: QuantizedTensorRef, |
53 | | /// Key projection [hidden_dim, kv_dim] (may differ for GQA) |
54 | | k: QuantizedTensorRef, |
55 | | /// Value projection [hidden_dim, kv_dim] |
56 | | v: QuantizedTensorRef, |
57 | | }, |
58 | | } |
59 | | |
60 | | impl QKVWeights { |
61 | | /// Calculate the output dimension per position (q_dim + k_dim + v_dim) |
62 | 2 | pub fn out_dim(&self, hidden_dim: usize) -> usize { |
63 | 2 | match self { |
64 | 1 | Self::Fused(ref weight) => weight.num_elements / hidden_dim, |
65 | | Self::Separate { |
66 | 1 | ref q, |
67 | 1 | ref k, |
68 | 1 | ref v, |
69 | | } => { |
70 | 1 | let q_dim = q.num_elements / hidden_dim; |
71 | 1 | let k_dim = k.num_elements / hidden_dim; |
72 | 1 | let v_dim = v.num_elements / hidden_dim; |
73 | 1 | q_dim + k_dim + v_dim |
74 | | }, |
75 | | } |
76 | 2 | } |
77 | | |
78 | | /// Get the Q dimension (query projection output dimension) |
79 | 2 | pub fn q_dim(&self, hidden_dim: usize) -> usize { |
80 | 2 | match self { |
81 | 1 | Self::Fused(ref weight) => weight.num_elements / hidden_dim / 3, |
82 | 1 | Self::Separate { ref q, .. } => q.num_elements / hidden_dim, |
83 | | } |
84 | 2 | } |
85 | | } |
86 | | |
87 | | // ============================================================================ |
88 | | // OwnedQuantizedTensor - Owned copy of quantized data |
89 | | // ============================================================================ |
90 | | |
91 | | /// Owned quantized tensor - copies data to avoid lifetime issues |
92 | | /// |
93 | | /// IMP-100: This allows storing quantized models in AppState with 'static lifetime |
94 | | #[derive(Debug, Clone)] |
95 | | pub struct OwnedQuantizedTensor { |
96 | | /// Raw quantized data (owned copy) |
97 | | pub data: Vec<u8>, |
98 | | /// Input dimension |
99 | | pub in_dim: usize, |
100 | | /// Output dimension |
101 | | pub out_dim: usize, |
102 | | /// Quantization type |
103 | | pub qtype: u32, |
104 | | } |
105 | | |
106 | | impl OwnedQuantizedTensor { |
107 | | /// Create owned tensor from a tensor reference and data slice with explicit dimensions |
108 | | #[must_use] |
109 | 3 | pub fn from_ref_with_dims( |
110 | 3 | tensor_ref: &QuantizedTensorRef, |
111 | 3 | data: &[u8], |
112 | 3 | in_dim: usize, |
113 | 3 | out_dim: usize, |
114 | 3 | ) -> Self { |
115 | 3 | let start = tensor_ref.offset; |
116 | 3 | let end = start + tensor_ref.byte_size; |
117 | 3 | let tensor_data = if end <= data.len() { |
118 | 2 | data[start..end].to_vec() |
119 | | } else { |
120 | 1 | Vec::new() |
121 | | }; |
122 | | |
123 | 3 | Self { |
124 | 3 | data: tensor_data, |
125 | 3 | in_dim, |
126 | 3 | out_dim, |
127 | 3 | qtype: tensor_ref.qtype, |
128 | 3 | } |
129 | 3 | } |
130 | | } |
131 | | |
132 | | // ============================================================================ |
133 | | // OwnedQKVWeights - Owned QKV weight storage |
134 | | // ============================================================================ |
135 | | |
136 | | /// Owned QKV weight storage - supports both fused (phi-2) and separate (llama) formats |
137 | | #[derive(Debug, Clone)] |
138 | | pub enum OwnedQKVWeights { |
139 | | /// Fused QKV tensor (phi-2 style) |
140 | | Fused(OwnedQuantizedTensor), |
141 | | /// Separate Q, K, V tensors (llama style) |
142 | | Separate { |
143 | | /// Query projection weights |
144 | | q: OwnedQuantizedTensor, |
145 | | /// Key projection weights |
146 | | k: OwnedQuantizedTensor, |
147 | | /// Value projection weights |
148 | | v: OwnedQuantizedTensor, |
149 | | }, |
150 | | } |
151 | | |
152 | | impl OwnedQKVWeights { |
153 | | /// Create from borrowed QKVWeights |
154 | | #[must_use] |
155 | 1 | pub fn from_borrowed(qkv: &QKVWeights, data: &[u8], hidden_dim: usize) -> Self { |
156 | 1 | match qkv { |
157 | 1 | QKVWeights::Fused(ref tensor) => { |
158 | 1 | let qkv_dim = 3 * hidden_dim; |
159 | 1 | OwnedQKVWeights::Fused(OwnedQuantizedTensor::from_ref_with_dims( |
160 | 1 | tensor, data, hidden_dim, qkv_dim, |
161 | 1 | )) |
162 | | }, |
163 | | QKVWeights::Separate { |
164 | 0 | ref q, |
165 | 0 | ref k, |
166 | 0 | ref v, |
167 | | } => { |
168 | 0 | let q_dim = q.num_elements / hidden_dim; |
169 | 0 | let k_dim = k.num_elements / hidden_dim; |
170 | 0 | let v_dim = v.num_elements / hidden_dim; |
171 | 0 | OwnedQKVWeights::Separate { |
172 | 0 | q: OwnedQuantizedTensor::from_ref_with_dims(q, data, hidden_dim, q_dim), |
173 | 0 | k: OwnedQuantizedTensor::from_ref_with_dims(k, data, hidden_dim, k_dim), |
174 | 0 | v: OwnedQuantizedTensor::from_ref_with_dims(v, data, hidden_dim, v_dim), |
175 | 0 | } |
176 | | }, |
177 | | } |
178 | 1 | } |
179 | | |
180 | | /// Get the output dimension (total Q+K+V dim) |
181 | | #[must_use] |
182 | 21 | pub fn out_dim(&self) -> usize { |
183 | 21 | match self { |
184 | 21 | OwnedQKVWeights::Fused(t) => t.out_dim, |
185 | 0 | OwnedQKVWeights::Separate { q, k, v } => q.out_dim + k.out_dim + v.out_dim, |
186 | | } |
187 | 21 | } |
188 | | |
189 | | /// Get the Q dimension (query projection output dimension) |
190 | | #[must_use] |
191 | 21 | pub fn q_dim(&self) -> usize { |
192 | 21 | match self { |
193 | 21 | OwnedQKVWeights::Fused(t) => t.out_dim / 3, |
194 | 0 | OwnedQKVWeights::Separate { q, .. } => q.out_dim, |
195 | | } |
196 | 21 | } |
197 | | } |
198 | | |
199 | | // ============================================================================ |
200 | | // OwnedQuantizedLayer - Owned transformer layer weights |
201 | | // ============================================================================ |
202 | | |
203 | | /// Owned quantized transformer layer - copies all weight data |
204 | | /// |
205 | | /// IMP-100: Allows storing in Arc without lifetime parameters |
206 | | #[derive(Debug, Clone)] |
207 | | pub struct OwnedQuantizedLayer { |
208 | | /// Attention norm weight (f32, small) |
209 | | pub attn_norm_weight: Vec<f32>, |
210 | | /// Attention norm bias (optional) |
211 | | pub attn_norm_bias: Option<Vec<f32>>, |
212 | | /// QKV projection weights (owned quantized data) - supports fused or separate |
213 | | pub qkv_weight: OwnedQKVWeights, |
214 | | /// QKV bias (optional, f32) |
215 | | pub qkv_bias: Option<Vec<f32>>, |
216 | | /// Attention output projection weights |
217 | | pub attn_output_weight: OwnedQuantizedTensor, |
218 | | /// Attention output bias (optional) |
219 | | pub attn_output_bias: Option<Vec<f32>>, |
220 | | /// FFN up projection weights |
221 | | pub ffn_up_weight: OwnedQuantizedTensor, |
222 | | /// FFN up bias (optional) |
223 | | pub ffn_up_bias: Option<Vec<f32>>, |
224 | | /// FFN down projection weights |
225 | | pub ffn_down_weight: OwnedQuantizedTensor, |
226 | | /// FFN down bias (optional) |
227 | | pub ffn_down_bias: Option<Vec<f32>>, |
228 | | /// FFN gate projection weights (SwiGLU models like LLaMA) |
229 | | pub ffn_gate_weight: Option<OwnedQuantizedTensor>, |
230 | | /// FFN gate bias (optional) |
231 | | pub ffn_gate_bias: Option<Vec<f32>>, |
232 | | /// FFN norm weight (pre-FFN layer norm, LLaMA-style) |
233 | | pub ffn_norm_weight: Option<Vec<f32>>, |
234 | | /// FFN norm bias (optional) |
235 | | pub ffn_norm_bias: Option<Vec<f32>>, |
236 | | } |
237 | | |
238 | | impl OwnedQuantizedLayer { |
239 | | /// Convert from borrowed layer with data reference and model config |
240 | | #[must_use] |
241 | 0 | pub fn from_borrowed( |
242 | 0 | layer: &crate::gguf::QuantizedGGUFTransformerLayer, |
243 | 0 | data: &[u8], |
244 | 0 | config: &crate::gguf::GGUFConfig, |
245 | 0 | ) -> Self { |
246 | 0 | let hidden_dim = config.hidden_dim; |
247 | 0 | let intermediate_dim = config.intermediate_dim; |
248 | | |
249 | | Self { |
250 | 0 | attn_norm_weight: layer.attn_norm_weight.clone(), |
251 | 0 | attn_norm_bias: layer.attn_norm_bias.clone(), |
252 | 0 | qkv_weight: OwnedQKVWeights::from_borrowed(&layer.qkv_weight, data, hidden_dim), |
253 | 0 | qkv_bias: layer.qkv_bias.clone(), |
254 | 0 | attn_output_weight: OwnedQuantizedTensor::from_ref_with_dims( |
255 | 0 | &layer.attn_output_weight, |
256 | 0 | data, |
257 | 0 | hidden_dim, |
258 | 0 | hidden_dim, |
259 | | ), |
260 | 0 | attn_output_bias: layer.attn_output_bias.clone(), |
261 | 0 | ffn_up_weight: OwnedQuantizedTensor::from_ref_with_dims( |
262 | 0 | &layer.ffn_up_weight, |
263 | 0 | data, |
264 | 0 | hidden_dim, |
265 | 0 | intermediate_dim, |
266 | | ), |
267 | 0 | ffn_up_bias: layer.ffn_up_bias.clone(), |
268 | 0 | ffn_down_weight: OwnedQuantizedTensor::from_ref_with_dims( |
269 | 0 | &layer.ffn_down_weight, |
270 | 0 | data, |
271 | 0 | intermediate_dim, |
272 | 0 | hidden_dim, |
273 | | ), |
274 | 0 | ffn_down_bias: layer.ffn_down_bias.clone(), |
275 | 0 | ffn_gate_weight: layer.ffn_gate_weight.as_ref().map(|gate_ref| { |
276 | 0 | OwnedQuantizedTensor::from_ref_with_dims( |
277 | 0 | gate_ref, |
278 | 0 | data, |
279 | 0 | hidden_dim, |
280 | 0 | intermediate_dim, |
281 | | ) |
282 | 0 | }), |
283 | 0 | ffn_gate_bias: layer.ffn_gate_bias.clone(), |
284 | 0 | ffn_norm_weight: layer.ffn_norm_weight.clone(), |
285 | 0 | ffn_norm_bias: layer.ffn_norm_bias.clone(), |
286 | | } |
287 | 0 | } |
288 | | } |
289 | | |
290 | | #[cfg(test)] |
291 | | mod tests { |
292 | | use super::*; |
293 | | use crate::gguf::types::GGUF_TYPE_Q4_K; |
294 | | |
295 | | #[test] |
296 | 1 | fn test_quantized_tensor_ref() { |
297 | 1 | let tensor = QuantizedTensorRef { |
298 | 1 | offset: 1024, |
299 | 1 | byte_size: 4096, |
300 | 1 | num_elements: 8192, |
301 | 1 | qtype: GGUF_TYPE_Q4_K, |
302 | 1 | }; |
303 | | |
304 | 1 | assert_eq!(tensor.offset, 1024); |
305 | 1 | assert_eq!(tensor.byte_size, 4096); |
306 | 1 | assert_eq!(tensor.num_elements, 8192); |
307 | 1 | assert_eq!(tensor.qtype, GGUF_TYPE_Q4_K); |
308 | 1 | } |
309 | | |
310 | | #[test] |
311 | 1 | fn test_qkv_weights_fused() { |
312 | 1 | let tensor = QuantizedTensorRef { |
313 | 1 | offset: 0, |
314 | 1 | byte_size: 1024, |
315 | 1 | num_elements: 4096 * 3, // 3 * hidden_dim |
316 | 1 | qtype: GGUF_TYPE_Q4_K, |
317 | 1 | }; |
318 | 1 | let qkv = QKVWeights::Fused(tensor); |
319 | | |
320 | 1 | assert_eq!(qkv.out_dim(4096), 3); // 12288 / 4096 = 3 |
321 | 1 | assert_eq!(qkv.q_dim(4096), 1); // 3 / 3 = 1 |
322 | 1 | } |
323 | | |
324 | | #[test] |
325 | 1 | fn test_qkv_weights_separate() { |
326 | 1 | let q = QuantizedTensorRef { |
327 | 1 | offset: 0, |
328 | 1 | byte_size: 1024, |
329 | 1 | num_elements: 4096 * 4096, // hidden_dim * hidden_dim |
330 | 1 | qtype: GGUF_TYPE_Q4_K, |
331 | 1 | }; |
332 | 1 | let k = QuantizedTensorRef { |
333 | 1 | offset: 1024, |
334 | 1 | byte_size: 256, |
335 | 1 | num_elements: 4096 * 512, // hidden_dim * kv_dim |
336 | 1 | qtype: GGUF_TYPE_Q4_K, |
337 | 1 | }; |
338 | 1 | let v = QuantizedTensorRef { |
339 | 1 | offset: 1280, |
340 | 1 | byte_size: 256, |
341 | 1 | num_elements: 4096 * 512, |
342 | 1 | qtype: GGUF_TYPE_Q4_K, |
343 | 1 | }; |
344 | | |
345 | 1 | let qkv = QKVWeights::Separate { q, k, v }; |
346 | | |
347 | 1 | assert_eq!(qkv.out_dim(4096), 4096 + 512 + 512); |
348 | 1 | assert_eq!(qkv.q_dim(4096), 4096); |
349 | 1 | } |
350 | | |
351 | | #[test] |
352 | 1 | fn test_owned_quantized_tensor() { |
353 | 1 | let tensor_ref = QuantizedTensorRef { |
354 | 1 | offset: 0, |
355 | 1 | byte_size: 8, |
356 | 1 | num_elements: 16, |
357 | 1 | qtype: GGUF_TYPE_Q4_K, |
358 | 1 | }; |
359 | 1 | let data = vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10]; |
360 | | |
361 | 1 | let owned = OwnedQuantizedTensor::from_ref_with_dims(&tensor_ref, &data, 4, 4); |
362 | | |
363 | 1 | assert_eq!(owned.data, &[1, 2, 3, 4, 5, 6, 7, 8]); |
364 | 1 | assert_eq!(owned.in_dim, 4); |
365 | 1 | assert_eq!(owned.out_dim, 4); |
366 | 1 | assert_eq!(owned.qtype, GGUF_TYPE_Q4_K); |
367 | 1 | } |
368 | | |
369 | | #[test] |
370 | 1 | fn test_owned_qkv_weights() { |
371 | 1 | let tensor = QuantizedTensorRef { |
372 | 1 | offset: 0, |
373 | 1 | byte_size: 12, |
374 | 1 | num_elements: 12, // 4 * 3 |
375 | 1 | qtype: GGUF_TYPE_Q4_K, |
376 | 1 | }; |
377 | 1 | let qkv_borrowed = QKVWeights::Fused(tensor); |
378 | 1 | let data = vec![0u8; 20]; |
379 | | |
380 | 1 | let owned = OwnedQKVWeights::from_borrowed(&qkv_borrowed, &data, 4); |
381 | | |
382 | 1 | assert_eq!(owned.out_dim(), 12); // 3 * 4 |
383 | 1 | assert_eq!(owned.q_dim(), 4); // 12 / 3 |
384 | 1 | } |
385 | | |
386 | | #[test] |
387 | 1 | fn test_owned_quantized_tensor_bounds() { |
388 | 1 | let tensor_ref = QuantizedTensorRef { |
389 | 1 | offset: 100, |
390 | 1 | byte_size: 50, |
391 | 1 | num_elements: 100, |
392 | 1 | qtype: GGUF_TYPE_Q4_K, |
393 | 1 | }; |
394 | | // Data too small - offset 100, needs 50 bytes |
395 | 1 | let data = vec![0u8; 50]; |
396 | | |
397 | 1 | let owned = OwnedQuantizedTensor::from_ref_with_dims(&tensor_ref, &data, 10, 10); |
398 | | |
399 | | // Should return empty data when out of bounds |
400 | 1 | assert!(owned.data.is_empty()); |
401 | 1 | } |
402 | | } |