/home/noah/src/realizar/src/gguf/inference/matmul.rs
Line | Count | Source |
1 | | //! Quantized matrix operations for OwnedQuantizedModel |
2 | | //! |
3 | | //! Contains embed, fused_matmul, qkv_matmul methods with real implementations |
4 | | //! for Q4_0, Q8_0, Q4_K, Q5_K, Q6_K quantization formats. |
5 | | |
6 | | use crate::error::{RealizarError, Result}; |
7 | | use crate::gguf::types::{ |
8 | | GGUF_TYPE_Q4_0, GGUF_TYPE_Q4_1, GGUF_TYPE_Q4_K, GGUF_TYPE_Q5_0, GGUF_TYPE_Q5_K, |
9 | | GGUF_TYPE_Q6_K, GGUF_TYPE_Q8_0, |
10 | | }; |
11 | | use crate::gguf::{ops, OwnedQKVWeights, OwnedQuantizedModel, OwnedQuantizedTensor}; |
12 | | |
13 | | impl OwnedQuantizedModel { |
14 | | /// Look up token embeddings (public for debugging PAR-001) |
15 | 718 | pub fn embed(&self, token_ids: &[u32]) -> Vec<f32> { |
16 | 718 | let hidden_dim = self.config.hidden_dim; |
17 | 718 | let mut embeddings = Vec::with_capacity(token_ids.len() * hidden_dim); |
18 | | |
19 | 1.48k | for &token_id764 in token_ids { |
20 | 764 | let start = (token_id as usize) * hidden_dim; |
21 | 764 | let end = start + hidden_dim; |
22 | 764 | if end <= self.token_embedding.len() { |
23 | 763 | embeddings.extend_from_slice(&self.token_embedding[start..end]); |
24 | 763 | } else { |
25 | 1 | embeddings.extend(std::iter::repeat_n(0.0, hidden_dim)); |
26 | 1 | } |
27 | | } |
28 | | |
29 | 718 | embeddings |
30 | 718 | } |
31 | | |
32 | | /// Look up single token embedding into pre-allocated buffer (IMP-131) |
33 | 2 | pub(crate) fn embed_into(&self, token_id: u32, output: &mut [f32]) { |
34 | 2 | let hidden_dim = self.config.hidden_dim; |
35 | 2 | let start = (token_id as usize) * hidden_dim; |
36 | 2 | let end = start + hidden_dim; |
37 | 2 | if end <= self.token_embedding.len() { |
38 | 1 | output[..hidden_dim].copy_from_slice(&self.token_embedding[start..end]); |
39 | 1 | } else { |
40 | 1 | output[..hidden_dim].iter_mut().for_each(|x| *x = 0.0); |
41 | | } |
42 | 2 | } |
43 | | |
44 | | /// Fused dequantize + matmul for quantized weights |
45 | | /// |
46 | | /// Supports Q4_0, Q8_0, Q4_1, Q5_0, Q4_K, Q5_K, Q6_K quantization formats. |
47 | | /// Uses SIMD-accelerated implementations for optimal performance. |
48 | 3.50k | pub(crate) fn fused_matmul( |
49 | 3.50k | &self, |
50 | 3.50k | input: &[f32], |
51 | 3.50k | weight: &OwnedQuantizedTensor, |
52 | 3.50k | ) -> Result<Vec<f32>> { |
53 | | use crate::quantize::{ |
54 | | dequantize_q4_1, dequantize_q5_0, fused_q4_0_q8_0_parallel_matvec, |
55 | | fused_q4k_parallel_matvec, fused_q5k_parallel_matvec, fused_q6k_parallel_matvec, |
56 | | fused_q8_0_q8_0_parallel_matvec, |
57 | | }; |
58 | | use trueno::{Matrix as TruenoMatrix, Vector as TruenoVector}; |
59 | | |
60 | 3.50k | let in_dim = weight.in_dim; |
61 | 3.50k | let out_dim = weight.out_dim; |
62 | 3.50k | let seq_len = input.len() / in_dim; |
63 | | |
64 | | // CUDA path when enabled |
65 | | #[cfg(feature = "cuda")] |
66 | | if let Some(ref executor_mutex) = self.cuda_executor { |
67 | | return self.fused_matmul_cuda(input, weight, executor_mutex); |
68 | | } |
69 | | |
70 | | // CPU path: For Q4_0, use fused Q8_0 integer SIMD matmul (llama.cpp parity) |
71 | 3.50k | if weight.qtype == GGUF_TYPE_Q4_0 { |
72 | 3 | if seq_len == 1 { |
73 | 2 | return fused_q4_0_q8_0_parallel_matvec(&weight.data, input, in_dim, out_dim); |
74 | 1 | } |
75 | 1 | let mut output = Vec::with_capacity(seq_len * out_dim); |
76 | 4 | for s in 0..seq_len1 { |
77 | 4 | let x = &input[s * in_dim..(s + 1) * in_dim]; |
78 | 4 | let row_output = fused_q4_0_q8_0_parallel_matvec(&weight.data, x, in_dim, out_dim)?0 ; |
79 | 4 | output.extend_from_slice(&row_output); |
80 | | } |
81 | 1 | return Ok(output); |
82 | 3.50k | } |
83 | | |
84 | | // CPU path: For Q8_0, use fused Q8_0 × Q8_0 integer SIMD matmul |
85 | 3.50k | if weight.qtype == GGUF_TYPE_Q8_0 { |
86 | 3 | if seq_len == 1 { |
87 | 2 | return fused_q8_0_q8_0_parallel_matvec(&weight.data, input, in_dim, out_dim); |
88 | 1 | } |
89 | 1 | let mut output = Vec::with_capacity(seq_len * out_dim); |
90 | 3 | for s in 0..seq_len1 { |
91 | 3 | let x = &input[s * in_dim..(s + 1) * in_dim]; |
92 | 3 | let row_output = fused_q8_0_q8_0_parallel_matvec(&weight.data, x, in_dim, out_dim)?0 ; |
93 | 3 | output.extend_from_slice(&row_output); |
94 | | } |
95 | 1 | return Ok(output); |
96 | 3.49k | } |
97 | | |
98 | | // CPU path: For Q4_1, use dequantize + SIMD matmul |
99 | 3.49k | if weight.qtype == GGUF_TYPE_Q4_1 { |
100 | 3 | let weights_f32 = dequantize_q4_1(&weight.data)?0 ; |
101 | | |
102 | 3 | let weight_matrix = match TruenoMatrix::from_vec(out_dim, in_dim, weights_f32) { |
103 | 3 | Ok(m) => m, |
104 | | Err(_) => { |
105 | 0 | return Err(RealizarError::InvalidShape { |
106 | 0 | reason: "Failed to create weight matrix for Q4_1".to_string(), |
107 | 0 | }); |
108 | | } |
109 | | }; |
110 | | |
111 | 3 | let mut output = Vec::with_capacity(seq_len * out_dim); |
112 | 4 | for s in 0..seq_len3 { |
113 | 4 | let x = &input[s * in_dim..(s + 1) * in_dim]; |
114 | 4 | let x_vec = TruenoVector::from_slice(x); |
115 | 4 | match weight_matrix.matvec(&x_vec) { |
116 | 4 | Ok(r) => output.extend_from_slice(r.as_slice()), |
117 | | Err(_) => { |
118 | 0 | return Err(RealizarError::InvalidShape { |
119 | 0 | reason: "SIMD matvec failed for Q4_1".to_string(), |
120 | 0 | }); |
121 | | } |
122 | | } |
123 | | } |
124 | 3 | return Ok(output); |
125 | 3.49k | } |
126 | | |
127 | | // CPU path: For Q5_0, use dequantize + SIMD matmul |
128 | 3.49k | if weight.qtype == GGUF_TYPE_Q5_0 { |
129 | 3 | let weights_f32 = dequantize_q5_0(&weight.data)?0 ; |
130 | | |
131 | 3 | let weight_matrix = match TruenoMatrix::from_vec(out_dim, in_dim, weights_f32) { |
132 | 3 | Ok(m) => m, |
133 | | Err(_) => { |
134 | 0 | return Err(RealizarError::InvalidShape { |
135 | 0 | reason: "Failed to create weight matrix for Q5_0".to_string(), |
136 | 0 | }); |
137 | | } |
138 | | }; |
139 | | |
140 | 3 | let mut output = Vec::with_capacity(seq_len * out_dim); |
141 | 4 | for s in 0..seq_len3 { |
142 | 4 | let x = &input[s * in_dim..(s + 1) * in_dim]; |
143 | 4 | let x_vec = TruenoVector::from_slice(x); |
144 | 4 | match weight_matrix.matvec(&x_vec) { |
145 | 4 | Ok(r) => output.extend_from_slice(r.as_slice()), |
146 | | Err(_) => { |
147 | 0 | return Err(RealizarError::InvalidShape { |
148 | 0 | reason: "SIMD matvec failed for Q5_0".to_string(), |
149 | 0 | }); |
150 | | } |
151 | | } |
152 | | } |
153 | 3 | return Ok(output); |
154 | 3.49k | } |
155 | | |
156 | | // CPU path: Process each position in sequence for Q4_K, Q5_K, Q6_K |
157 | 3.49k | if seq_len > 1 { |
158 | 18 | let mut output = Vec::with_capacity(seq_len * out_dim); |
159 | 70 | for s in 0..seq_len18 { |
160 | 70 | let x = &input[s * in_dim..(s + 1) * in_dim]; |
161 | 70 | let row_output = match weight.qtype { |
162 | | GGUF_TYPE_Q4_K => { |
163 | 66 | fused_q4k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0 |
164 | | } |
165 | | GGUF_TYPE_Q5_K => { |
166 | 2 | fused_q5k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0 |
167 | | } |
168 | | GGUF_TYPE_Q6_K => { |
169 | 2 | fused_q6k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0 |
170 | | } |
171 | | _ => { |
172 | 0 | return Err(RealizarError::UnsupportedOperation { |
173 | 0 | operation: "owned_fused_matmul".to_string(), |
174 | 0 | reason: format!( |
175 | 0 | "Fused matmul only supports Q4_0/Q4_1/Q5_0/Q8_0/Q4_K/Q5_K/Q6_K, got type {}", |
176 | 0 | weight.qtype |
177 | 0 | ), |
178 | 0 | }); |
179 | | } |
180 | | }; |
181 | 70 | output.extend_from_slice(&row_output); |
182 | | } |
183 | 18 | Ok(output) |
184 | | } else { |
185 | | // Single position - most common case in generation |
186 | 3.47k | match weight.qtype { |
187 | 3.46k | GGUF_TYPE_Q4_K => fused_q4k_parallel_matvec(&weight.data, input, in_dim, out_dim), |
188 | 2 | GGUF_TYPE_Q5_K => fused_q5k_parallel_matvec(&weight.data, input, in_dim, out_dim), |
189 | 2 | GGUF_TYPE_Q6_K => fused_q6k_parallel_matvec(&weight.data, input, in_dim, out_dim), |
190 | 1 | _ => Err(RealizarError::UnsupportedOperation { |
191 | 1 | operation: "owned_fused_matmul".to_string(), |
192 | 1 | reason: format!( |
193 | 1 | "Fused matmul only supports Q4_0/Q8_0/Q4_K/Q5_K/Q6_K, got type {}", |
194 | 1 | weight.qtype |
195 | 1 | ), |
196 | 1 | }), |
197 | | } |
198 | | } |
199 | 3.50k | } |
200 | | |
201 | | /// CUDA path for fused matmul |
202 | | #[cfg(feature = "cuda")] |
203 | | fn fused_matmul_cuda( |
204 | | &self, |
205 | | input: &[f32], |
206 | | weight: &OwnedQuantizedTensor, |
207 | | executor_mutex: &std::sync::Mutex<crate::cuda::CudaExecutor>, |
208 | | ) -> Result<Vec<f32>> { |
209 | | use tracing::info_span; |
210 | | |
211 | | let in_dim = weight.in_dim; |
212 | | let out_dim = weight.out_dim; |
213 | | let seq_len = input.len() / in_dim; |
214 | | let gemm_start = std::time::Instant::now(); |
215 | | let mut output = vec![0.0f32; seq_len * out_dim]; |
216 | | |
217 | | // Use native quantized GEMV kernels for single-token generation |
218 | | if seq_len == 1 { |
219 | | let cache_key = format!("{}_{:016x}", |
220 | | match weight.qtype { |
221 | | GGUF_TYPE_Q4_K => "q4k", |
222 | | GGUF_TYPE_Q5_K => "q5k", |
223 | | GGUF_TYPE_Q6_K => "q6k", |
224 | | _ => "unknown", |
225 | | }, |
226 | | weight.data.as_ptr() as usize |
227 | | ); |
228 | | |
229 | | if weight.qtype == GGUF_TYPE_Q4_K |
230 | | || weight.qtype == GGUF_TYPE_Q5_K |
231 | | || weight.qtype == GGUF_TYPE_Q6_K |
232 | | { |
233 | | let mut executor = |
234 | | executor_mutex |
235 | | .lock() |
236 | | .map_err(|e| RealizarError::UnsupportedOperation { |
237 | | operation: "cuda_lock".to_string(), |
238 | | reason: format!("Failed to acquire CUDA executor lock: {e}"), |
239 | | })?; |
240 | | |
241 | | executor |
242 | | .make_current() |
243 | | .map_err(|e| RealizarError::UnsupportedOperation { |
244 | | operation: "cuda_make_current".to_string(), |
245 | | reason: format!("Failed to set CUDA context current: {e}"), |
246 | | })?; |
247 | | |
248 | | if !executor.has_quantized_weights(&cache_key) { |
249 | | executor |
250 | | .load_quantized_weights(&cache_key, &weight.data) |
251 | | .map_err(|e| RealizarError::UnsupportedOperation { |
252 | | operation: "cuda_cache".to_string(), |
253 | | reason: format!("Failed to cache weights: {e}"), |
254 | | })?; |
255 | | } |
256 | | |
257 | | let result = match weight.qtype { |
258 | | GGUF_TYPE_Q4_K => executor.q4k_gemv_cached( |
259 | | &cache_key, |
260 | | input, |
261 | | &mut output, |
262 | | out_dim as u32, |
263 | | in_dim as u32, |
264 | | ), |
265 | | GGUF_TYPE_Q5_K => executor.q5k_gemv_cached( |
266 | | &cache_key, |
267 | | input, |
268 | | &mut output, |
269 | | out_dim as u32, |
270 | | in_dim as u32, |
271 | | ), |
272 | | GGUF_TYPE_Q6_K => executor.q6k_gemv_cached( |
273 | | &cache_key, |
274 | | input, |
275 | | &mut output, |
276 | | out_dim as u32, |
277 | | in_dim as u32, |
278 | | ), |
279 | | _ => unreachable!(), |
280 | | }; |
281 | | |
282 | | result.map_err(|e| RealizarError::UnsupportedOperation { |
283 | | operation: "cuda_gemv".to_string(), |
284 | | reason: format!("CUDA GEMV failed: {e}"), |
285 | | })?; |
286 | | |
287 | | let gemm_duration_us = gemm_start.elapsed().as_micros() as u64; |
288 | | let _span = info_span!( |
289 | | "gpu_kernel:gemv", |
290 | | gpu.backend = "cuda", |
291 | | gpu.dimensions.n = out_dim, |
292 | | gpu.dimensions.k = in_dim, |
293 | | duration_us = gemm_duration_us, |
294 | | ) |
295 | | .entered(); |
296 | | |
297 | | self.cuda_kernel_count |
298 | | .fetch_add(1, std::sync::atomic::Ordering::Relaxed); |
299 | | |
300 | | return Ok(output); |
301 | | } |
302 | | } |
303 | | |
304 | | // Fallback: Dequantize and use FP32 GEMM |
305 | | let dequant_weight = self.dequantize_weight_for_cuda(weight)?; |
306 | | |
307 | | { |
308 | | let mut executor = |
309 | | executor_mutex |
310 | | .lock() |
311 | | .map_err(|e| RealizarError::UnsupportedOperation { |
312 | | operation: "cuda_gemm_lock".to_string(), |
313 | | reason: format!("Failed to acquire CUDA executor lock: {e}"), |
314 | | })?; |
315 | | |
316 | | executor |
317 | | .make_current() |
318 | | .map_err(|e| RealizarError::UnsupportedOperation { |
319 | | operation: "cuda_make_current".to_string(), |
320 | | reason: format!("Failed to set CUDA context current: {e}"), |
321 | | })?; |
322 | | |
323 | | executor |
324 | | .gemm( |
325 | | input, |
326 | | &dequant_weight, |
327 | | &mut output, |
328 | | seq_len as u32, |
329 | | out_dim as u32, |
330 | | in_dim as u32, |
331 | | ) |
332 | | .map_err(|e| RealizarError::UnsupportedOperation { |
333 | | operation: "cuda_gemm".to_string(), |
334 | | reason: format!("CUDA GEMM failed: {e}"), |
335 | | })?; |
336 | | } |
337 | | |
338 | | self.cuda_kernel_count |
339 | | .fetch_add(1, std::sync::atomic::Ordering::Relaxed); |
340 | | |
341 | | Ok(output) |
342 | | } |
343 | | |
344 | | /// Fused matmul into pre-allocated output buffer |
345 | 0 | pub(crate) fn fused_matmul_into( |
346 | 0 | &self, |
347 | 0 | input: &[f32], |
348 | 0 | weight: &OwnedQuantizedTensor, |
349 | 0 | output: &mut [f32], |
350 | 0 | ) -> Result<()> { |
351 | | use crate::quantize::{ |
352 | | fused_q4_0_q8_0_parallel_matvec_into, fused_q4k_parallel_matvec_into, |
353 | | fused_q5k_parallel_matvec_into, fused_q6k_parallel_matvec_into, |
354 | | fused_q8_0_q8_0_parallel_matvec_into, |
355 | | }; |
356 | | |
357 | 0 | let in_dim = weight.in_dim; |
358 | 0 | let out_dim = weight.out_dim; |
359 | 0 | let seq_len = input.len() / in_dim; |
360 | | |
361 | | // Only support single-token case for now (most common in generation) |
362 | 0 | if seq_len != 1 { |
363 | 0 | let result = self.fused_matmul(input, weight)?; |
364 | 0 | output[..result.len()].copy_from_slice(&result); |
365 | 0 | return Ok(()); |
366 | 0 | } |
367 | | |
368 | 0 | debug_assert!( |
369 | 0 | output.len() >= out_dim, |
370 | 0 | "Output buffer too small: {} < {}", |
371 | 0 | output.len(), |
372 | | out_dim |
373 | | ); |
374 | | |
375 | 0 | match weight.qtype { |
376 | | GGUF_TYPE_Q4_0 => { |
377 | 0 | fused_q4_0_q8_0_parallel_matvec_into( |
378 | 0 | &weight.data, |
379 | 0 | input, |
380 | 0 | in_dim, |
381 | 0 | &mut output[..out_dim], |
382 | | ) |
383 | | } |
384 | 0 | GGUF_TYPE_Q8_0 => fused_q8_0_q8_0_parallel_matvec_into( |
385 | 0 | &weight.data, |
386 | 0 | input, |
387 | 0 | in_dim, |
388 | 0 | out_dim, |
389 | 0 | &mut output[..out_dim], |
390 | | ), |
391 | 0 | GGUF_TYPE_Q4_K => fused_q4k_parallel_matvec_into( |
392 | 0 | &weight.data, |
393 | 0 | input, |
394 | 0 | in_dim, |
395 | 0 | out_dim, |
396 | 0 | &mut output[..out_dim], |
397 | | ), |
398 | 0 | GGUF_TYPE_Q5_K => fused_q5k_parallel_matvec_into( |
399 | 0 | &weight.data, |
400 | 0 | input, |
401 | 0 | in_dim, |
402 | 0 | out_dim, |
403 | 0 | &mut output[..out_dim], |
404 | | ), |
405 | 0 | GGUF_TYPE_Q6_K => fused_q6k_parallel_matvec_into( |
406 | 0 | &weight.data, |
407 | 0 | input, |
408 | 0 | in_dim, |
409 | 0 | out_dim, |
410 | 0 | &mut output[..out_dim], |
411 | | ), |
412 | | _ => { |
413 | 0 | let result = self.fused_matmul(input, weight)?; |
414 | 0 | output[..result.len()].copy_from_slice(&result); |
415 | 0 | Ok(()) |
416 | | } |
417 | | } |
418 | 0 | } |
419 | | |
420 | | /// QKV projection matmul |
421 | 714 | pub fn qkv_matmul(&self, input: &[f32], qkv: &OwnedQKVWeights) -> Result<Vec<f32>> { |
422 | 714 | let hidden_dim = self.config.hidden_dim; |
423 | 714 | match qkv { |
424 | 714 | OwnedQKVWeights::Fused(ref weight) => self.fused_matmul(input, weight), |
425 | | OwnedQKVWeights::Separate { |
426 | 0 | ref q, |
427 | 0 | ref k, |
428 | 0 | ref v, |
429 | | } => { |
430 | 0 | let seq_len = input.len() / hidden_dim; |
431 | | |
432 | 0 | let q_out = self.fused_matmul(input, q)?; |
433 | 0 | let k_out = self.fused_matmul(input, k)?; |
434 | 0 | let v_out = self.fused_matmul(input, v)?; |
435 | | |
436 | | // Interleave Q, K, V for each position |
437 | 0 | let qkv_dim = q.out_dim + k.out_dim + v.out_dim; |
438 | 0 | let mut output = Vec::with_capacity(seq_len * qkv_dim); |
439 | 0 | for s in 0..seq_len { |
440 | 0 | output.extend_from_slice(&q_out[s * q.out_dim..(s + 1) * q.out_dim]); |
441 | 0 | output.extend_from_slice(&k_out[s * k.out_dim..(s + 1) * k.out_dim]); |
442 | 0 | output.extend_from_slice(&v_out[s * v.out_dim..(s + 1) * v.out_dim]); |
443 | 0 | } |
444 | 0 | Ok(output) |
445 | | } |
446 | | } |
447 | 714 | } |
448 | | |
449 | | /// QKV matmul into pre-allocated buffer |
450 | 0 | pub fn qkv_matmul_into( |
451 | 0 | &self, |
452 | 0 | input: &[f32], |
453 | 0 | qkv: &OwnedQKVWeights, |
454 | 0 | output: &mut [f32], |
455 | 0 | ) -> Result<()> { |
456 | 0 | match qkv { |
457 | 0 | OwnedQKVWeights::Fused(ref weight) => self.fused_matmul_into(input, weight, output), |
458 | | OwnedQKVWeights::Separate { |
459 | 0 | ref q, |
460 | 0 | ref k, |
461 | 0 | ref v, |
462 | | } => { |
463 | 0 | let q_dim = q.out_dim; |
464 | 0 | let k_dim = k.out_dim; |
465 | 0 | let v_dim = v.out_dim; |
466 | | |
467 | 0 | self.fused_matmul_into(input, q, &mut output[..q_dim])?; |
468 | 0 | self.fused_matmul_into(input, k, &mut output[q_dim..q_dim + k_dim])?; |
469 | 0 | self.fused_matmul_into( |
470 | 0 | input, |
471 | 0 | v, |
472 | 0 | &mut output[q_dim + k_dim..q_dim + k_dim + v_dim], |
473 | 0 | )?; |
474 | | |
475 | 0 | Ok(()) |
476 | | } |
477 | | } |
478 | 0 | } |
479 | | |
480 | | /// Layer normalization |
481 | 18 | pub fn layer_norm( |
482 | 18 | &self, |
483 | 18 | input: &[f32], |
484 | 18 | weight: &[f32], |
485 | 18 | bias: Option<&[f32]>, |
486 | 18 | eps: f32, |
487 | 18 | ) -> Vec<f32> { |
488 | 18 | ops::layer_norm(input, weight, bias, eps) |
489 | 18 | } |
490 | | |
491 | | /// Add bias to activations |
492 | 0 | pub fn add_bias(&self, input: &mut [f32], bias: &[f32]) { |
493 | 0 | for (x, b) in input.iter_mut().zip(bias.iter()) { |
494 | 0 | *x += b; |
495 | 0 | } |
496 | 0 | } |
497 | | |
498 | | /// GELU activation |
499 | 6 | pub fn gelu(&self, input: &mut [f32]) { |
500 | 2.30k | for x in input6 .iter_mut6 () { |
501 | 2.30k | *x = 0.5 * *x * (1.0 + (*x * 0.7978845608 * (1.0 + 0.044715 * *x * *x)).tanh()); |
502 | 2.30k | } |
503 | 6 | } |
504 | | |
505 | | /// Fused RMSNorm + matmul helper |
506 | 0 | fn fused_rmsnorm_matmul( |
507 | 0 | &self, |
508 | 0 | input: &[f32], |
509 | 0 | norm_weight: &[f32], |
510 | 0 | eps: f32, |
511 | 0 | weight: &OwnedQuantizedTensor, |
512 | 0 | ) -> Result<Vec<f32>> { |
513 | | use crate::quantize::fused_rmsnorm_q4_0_matmul; |
514 | | |
515 | | // Only use fused path for Q4_0 weights (most common) |
516 | 0 | if weight.qtype == GGUF_TYPE_Q4_0 && input.len() == weight.in_dim { |
517 | 0 | return fused_rmsnorm_q4_0_matmul( |
518 | 0 | input, |
519 | 0 | norm_weight, |
520 | 0 | eps, |
521 | 0 | &weight.data, |
522 | 0 | weight.in_dim, |
523 | 0 | weight.out_dim, |
524 | | ); |
525 | 0 | } |
526 | | |
527 | | // Fallback to separate RMSNorm + matmul for other types |
528 | 0 | let normed = ops::rms_norm(input, norm_weight, eps); |
529 | 0 | self.fused_matmul(&normed, weight) |
530 | 0 | } |
531 | | |
532 | | /// Fused RMSNorm + QKV matmul |
533 | 0 | pub fn fused_rmsnorm_qkv_matmul( |
534 | 0 | &self, |
535 | 0 | input: &[f32], |
536 | 0 | norm_weight: &[f32], |
537 | 0 | eps: f32, |
538 | 0 | qkv: &OwnedQKVWeights, |
539 | 0 | ) -> Result<Vec<f32>> { |
540 | 0 | match qkv { |
541 | 0 | OwnedQKVWeights::Fused(ref weight) => { |
542 | 0 | self.fused_rmsnorm_matmul(input, norm_weight, eps, weight) |
543 | | } |
544 | | OwnedQKVWeights::Separate { |
545 | 0 | ref q, |
546 | 0 | ref k, |
547 | 0 | ref v, |
548 | | } => { |
549 | | // For separate Q/K/V, normalize once and reuse |
550 | 0 | let normed = ops::rms_norm(input, norm_weight, eps); |
551 | | |
552 | 0 | let q_out = self.fused_matmul(&normed, q)?; |
553 | 0 | let k_out = self.fused_matmul(&normed, k)?; |
554 | 0 | let v_out = self.fused_matmul(&normed, v)?; |
555 | | |
556 | 0 | let qkv_dim = q.out_dim + k.out_dim + v.out_dim; |
557 | 0 | let mut output = Vec::with_capacity(qkv_dim); |
558 | 0 | output.extend_from_slice(&q_out); |
559 | 0 | output.extend_from_slice(&k_out); |
560 | 0 | output.extend_from_slice(&v_out); |
561 | 0 | Ok(output) |
562 | | } |
563 | | } |
564 | 0 | } |
565 | | |
566 | | /// Fused RMSNorm + LM head |
567 | 0 | pub fn fused_rmsnorm_lm_head(&self, input: &[f32]) -> Result<Vec<f32>> { |
568 | | use crate::quantize::fused_rmsnorm_q4_0_matmul; |
569 | | |
570 | | // Only use fused path for Q4_0 weights |
571 | 0 | if self.lm_head_weight.qtype == GGUF_TYPE_Q4_0 |
572 | 0 | && input.len() == self.lm_head_weight.in_dim |
573 | | { |
574 | 0 | return fused_rmsnorm_q4_0_matmul( |
575 | 0 | input, |
576 | 0 | &self.output_norm_weight, |
577 | 0 | self.config.eps, |
578 | 0 | &self.lm_head_weight.data, |
579 | 0 | self.lm_head_weight.in_dim, |
580 | 0 | self.lm_head_weight.out_dim, |
581 | | ); |
582 | 0 | } |
583 | | |
584 | | // Fallback to separate RMSNorm + matmul for other types |
585 | 0 | let normed = ops::rms_norm(input, &self.output_norm_weight, self.config.eps); |
586 | 0 | self.fused_matmul(&normed, &self.lm_head_weight) |
587 | 0 | } |
588 | | |
589 | | /// Fused RMSNorm + FFN up/gate projections for SwiGLU |
590 | 0 | pub fn fused_rmsnorm_ffn_up_gate( |
591 | 0 | &self, |
592 | 0 | input: &[f32], |
593 | 0 | norm_weight: &[f32], |
594 | 0 | eps: f32, |
595 | 0 | up_weight: &OwnedQuantizedTensor, |
596 | 0 | gate_weight: &OwnedQuantizedTensor, |
597 | 0 | ) -> Result<(Vec<f32>, Vec<f32>)> { |
598 | | use crate::quantize::fused_rmsnorm_ffn_up_gate; |
599 | | |
600 | | // Only use fused path for Q4_0 weights |
601 | 0 | if up_weight.qtype == GGUF_TYPE_Q4_0 |
602 | 0 | && gate_weight.qtype == GGUF_TYPE_Q4_0 |
603 | 0 | && input.len() == up_weight.in_dim |
604 | 0 | && up_weight.in_dim == gate_weight.in_dim |
605 | 0 | && up_weight.out_dim == gate_weight.out_dim |
606 | | { |
607 | 0 | return fused_rmsnorm_ffn_up_gate( |
608 | 0 | input, |
609 | 0 | norm_weight, |
610 | 0 | eps, |
611 | 0 | &up_weight.data, |
612 | 0 | &gate_weight.data, |
613 | 0 | up_weight.in_dim, |
614 | 0 | up_weight.out_dim, |
615 | | ); |
616 | 0 | } |
617 | | |
618 | | // Fallback to separate RMSNorm + matmuls for other types |
619 | 0 | let normed = ops::rms_norm(input, norm_weight, eps); |
620 | 0 | let up_out = self.fused_matmul(&normed, up_weight)?; |
621 | 0 | let gate_out = self.fused_matmul(&normed, gate_weight)?; |
622 | 0 | Ok((up_out, gate_out)) |
623 | 0 | } |
624 | | |
625 | | /// Q8K QKV matmul into buffer |
626 | | /// |
627 | | /// Uses pre-quantized Q8K activations for faster matmul with Q4K weights. |
628 | | #[allow(unused_variables)] |
629 | 0 | pub fn qkv_matmul_q8k_into( |
630 | 0 | &self, |
631 | 0 | input: &[f32], |
632 | 0 | qkv: &OwnedQKVWeights, |
633 | 0 | output: &mut [f32], |
634 | 0 | scales: &[f32], |
635 | 0 | quants: &[i8], |
636 | 0 | ) -> Result<()> { |
637 | | // Fall back to regular qkv_matmul_into for now |
638 | | // TODO: Implement Q8K-accelerated path using pre-quantized activations |
639 | 0 | self.qkv_matmul_into(input, qkv, output) |
640 | 0 | } |
641 | | |
642 | | /// Helper to dequantize weights for CUDA GEMM |
643 | | #[cfg(feature = "cuda")] |
644 | | fn dequantize_weight_for_cuda(&self, weight: &OwnedQuantizedTensor) -> Result<Vec<f32>> { |
645 | | use crate::quantize::{ |
646 | | dequantize_q4_0, dequantize_q4_1, dequantize_q4_k, dequantize_q5_0, dequantize_q5_k, |
647 | | dequantize_q6_k, dequantize_q8_0, |
648 | | }; |
649 | | |
650 | | match weight.qtype { |
651 | | GGUF_TYPE_Q4_0 => dequantize_q4_0(&weight.data), |
652 | | GGUF_TYPE_Q4_1 => dequantize_q4_1(&weight.data), |
653 | | GGUF_TYPE_Q5_0 => dequantize_q5_0(&weight.data), |
654 | | GGUF_TYPE_Q8_0 => dequantize_q8_0(&weight.data), |
655 | | GGUF_TYPE_Q4_K => dequantize_q4_k(&weight.data), |
656 | | GGUF_TYPE_Q5_K => dequantize_q5_k(&weight.data), |
657 | | GGUF_TYPE_Q6_K => dequantize_q6_k(&weight.data), |
658 | | _ => Err(RealizarError::UnsupportedOperation { |
659 | | operation: "dequantize_weight_for_cuda".to_string(), |
660 | | reason: format!("Unsupported quantization type: {}", weight.qtype), |
661 | | }), |
662 | | } |
663 | | } |
664 | | } |