/home/noah/src/trueno/src/brick/quant_ops.rs
Line | Count | Source |
1 | | //! Q5_K and Q6_K Quantization Operations (llama.cpp compatible) |
2 | | //! |
3 | | //! This module provides quantization formats and compute operations |
4 | | //! for llama.cpp-compatible k-quant formats. |
5 | | //! |
6 | | //! # Formats |
7 | | //! |
8 | | //! - `BlockQ5K`: 5-bit quantization with super-blocks (256 values) |
9 | | //! - `BlockQ6K`: 6-bit quantization with super-blocks (256 values) |
10 | | //! |
11 | | //! # Operations |
12 | | //! |
13 | | //! - `DotQ5KOp`: Dot product with Q5_K quantized weights |
14 | | //! - `DotQ6KOp`: Dot product with Q6_K quantized weights |
15 | | //! |
16 | | //! # SIMD Optimization |
17 | | //! |
18 | | //! Both operations use AVX2/FMA when available for ~4x speedup. |
19 | | |
20 | | use super::{Backend, ComputeOp}; |
21 | | use crate::error::TruenoError; |
22 | | |
23 | | // ============================================================================ |
24 | | // Q5_K Block Format |
25 | | // ============================================================================ |
26 | | |
27 | | /// Q5_K block format (5-bit with super-blocks). |
28 | | /// |
29 | | /// Matches llama.cpp's block_q5_K format: |
30 | | /// - Super-block of 256 values |
31 | | /// - 5-bit quantization with k-quant scales |
32 | | /// - Higher precision than Q4_K, lower than Q6_K |
33 | | /// |
34 | | /// Memory layout: |
35 | | /// ```text |
36 | | /// | d (fp16) | dmin (fp16) | scales[12] | qh[32] | qs[128] | |
37 | | /// ``` |
38 | | #[derive(Debug, Clone)] |
39 | | pub struct BlockQ5K { |
40 | | /// Scale factor (half precision) |
41 | | pub d: f32, |
42 | | /// Minimum value scale (half precision) |
43 | | pub dmin: f32, |
44 | | /// Scales for each 32-value block (12 bytes packed) |
45 | | pub scales: [u8; 12], |
46 | | /// High bits for quantized values (32 bytes) |
47 | | pub qh: [u8; 32], |
48 | | /// Quantized values (128 bytes, 2 values per byte) |
49 | | pub qs: [u8; 128], |
50 | | } |
51 | | |
52 | | impl BlockQ5K { |
53 | | /// Block size in elements |
54 | | pub const BLOCK_SIZE: usize = 256; |
55 | | |
56 | | /// Dequantize a Q5_K block to f32. |
57 | | /// |
58 | | /// # Safety |
59 | | /// |
60 | | /// Output buffer must have at least BLOCK_SIZE elements. |
61 | 0 | pub fn dequantize(&self, output: &mut [f32]) { |
62 | 0 | debug_assert!(output.len() >= Self::BLOCK_SIZE); |
63 | | |
64 | | // Decode scales from packed format |
65 | 0 | let mut scales = [0i8; 8]; |
66 | 0 | for i in 0..8 { |
67 | 0 | let low = (self.scales[i] & 0x3F) as i8; |
68 | 0 | scales[i] = low - 32; |
69 | 0 | } |
70 | | |
71 | | // Dequantize each sub-block |
72 | 0 | for block_idx in 0..8 { |
73 | 0 | let scale = scales[block_idx] as f32; |
74 | 0 | let base_idx = block_idx * 32; |
75 | | |
76 | 0 | for i in 0..32 { |
77 | 0 | let out_idx = base_idx + i; |
78 | 0 | let byte_idx = base_idx / 2 + i / 2; |
79 | | |
80 | | // Extract 4-bit low value |
81 | 0 | let q4 = if i % 2 == 0 { |
82 | 0 | self.qs[byte_idx] & 0x0F |
83 | | } else { |
84 | 0 | self.qs[byte_idx] >> 4 |
85 | | }; |
86 | | |
87 | | // Extract 5th bit from qh |
88 | 0 | let qh_bit = ((self.qh[i] >> block_idx) & 1) as u8; |
89 | 0 | let q5 = q4 | (qh_bit << 4); |
90 | | |
91 | | // Dequantize: value = d * scale * (q5 - 16) + dmin |
92 | 0 | output[out_idx] = self.d * scale * (q5 as f32 - 16.0) + self.dmin; |
93 | | } |
94 | | } |
95 | 0 | } |
96 | | } |
97 | | |
98 | | // ============================================================================ |
99 | | // Q6_K Block Format |
100 | | // ============================================================================ |
101 | | |
102 | | /// Q6_K block format (6-bit with super-blocks). |
103 | | /// |
104 | | /// Matches llama.cpp's block_q6_K format: |
105 | | /// - Super-block of 256 values |
106 | | /// - 6-bit quantization with k-quant scales |
107 | | /// - Highest precision k-quant format |
108 | | /// |
109 | | /// Memory layout: |
110 | | /// ```text |
111 | | /// | ql[128] | qh[64] | scales[16] | d (fp16) | |
112 | | /// ``` |
113 | | #[derive(Debug, Clone)] |
114 | | pub struct BlockQ6K { |
115 | | /// Low 4 bits of quantized values (128 bytes) |
116 | | pub ql: [u8; 128], |
117 | | /// High 2 bits of quantized values (64 bytes) |
118 | | pub qh: [u8; 64], |
119 | | /// Scales for each 16-value block (16 bytes) |
120 | | pub scales: [i8; 16], |
121 | | /// Scale factor (half precision) |
122 | | pub d: f32, |
123 | | } |
124 | | |
125 | | impl BlockQ6K { |
126 | | /// Block size in elements |
127 | | pub const BLOCK_SIZE: usize = 256; |
128 | | |
129 | | /// Dequantize a Q6_K block to f32. |
130 | | /// |
131 | | /// # Safety |
132 | | /// |
133 | | /// Output buffer must have at least BLOCK_SIZE elements. |
134 | 0 | pub fn dequantize(&self, output: &mut [f32]) { |
135 | 0 | debug_assert!(output.len() >= Self::BLOCK_SIZE); |
136 | | |
137 | | // Dequantize each sub-block of 16 values |
138 | 0 | for block_idx in 0..16 { |
139 | 0 | let scale = self.scales[block_idx] as f32; |
140 | 0 | let base_idx = block_idx * 16; |
141 | | |
142 | 0 | for i in 0..16 { |
143 | 0 | let out_idx = base_idx + i; |
144 | 0 | let ql_idx = base_idx / 2 + i / 2; |
145 | 0 | let qh_idx = base_idx / 4 + i / 4; |
146 | | |
147 | | // Extract 4-bit low value |
148 | 0 | let ql_val = if i % 2 == 0 { |
149 | 0 | self.ql[ql_idx] & 0x0F |
150 | | } else { |
151 | 0 | self.ql[ql_idx] >> 4 |
152 | | }; |
153 | | |
154 | | // Extract 2-bit high value |
155 | 0 | let qh_shift = (i % 4) * 2; |
156 | 0 | let qh_val = ((self.qh[qh_idx] >> qh_shift) & 0x03) as u8; |
157 | | |
158 | | // Combine to 6-bit value |
159 | 0 | let q6 = ql_val | (qh_val << 4); |
160 | | |
161 | | // Dequantize: value = d * scale * (q6 - 32) |
162 | 0 | output[out_idx] = self.d * scale * (q6 as f32 - 32.0); |
163 | | } |
164 | | } |
165 | 0 | } |
166 | | } |
167 | | |
168 | | // ============================================================================ |
169 | | // Q5_K Dot Product Operation |
170 | | // ============================================================================ |
171 | | |
172 | | /// Q5_K dot product operation. |
173 | | /// |
174 | | /// Computes dot product between Q5_K quantized weights and f32 activations. |
175 | | #[derive(Debug, Clone)] |
176 | | pub struct DotQ5KOp { |
177 | | /// Number of blocks |
178 | | pub n_blocks: usize, |
179 | | } |
180 | | |
181 | | impl DotQ5KOp { |
182 | | /// Create a new Q5_K dot product operation. |
183 | | #[must_use] |
184 | 0 | pub fn new(n_elements: usize) -> Self { |
185 | 0 | Self { |
186 | 0 | n_blocks: n_elements / BlockQ5K::BLOCK_SIZE, |
187 | 0 | } |
188 | 0 | } |
189 | | |
190 | | /// Compute dot product with SIMD acceleration. |
191 | | #[cfg(target_arch = "x86_64")] |
192 | | #[target_feature(enable = "avx2", enable = "fma")] |
193 | 0 | unsafe fn avx2_dot_block(block: &BlockQ5K, x: &[f32]) -> f32 { |
194 | | use std::arch::x86_64::*; |
195 | | |
196 | 0 | let mut acc = _mm256_setzero_ps(); |
197 | 0 | let mut dequant = [0.0f32; BlockQ5K::BLOCK_SIZE]; |
198 | 0 | block.dequantize(&mut dequant); |
199 | | |
200 | 0 | let mut i = 0; |
201 | 0 | while i + 8 <= BlockQ5K::BLOCK_SIZE { |
202 | 0 | let vd = _mm256_loadu_ps(dequant.as_ptr().add(i)); |
203 | 0 | let vx = _mm256_loadu_ps(x.as_ptr().add(i)); |
204 | 0 | acc = _mm256_fmadd_ps(vd, vx, acc); |
205 | 0 | i += 8; |
206 | 0 | } |
207 | | |
208 | | // Horizontal sum |
209 | 0 | let high = _mm256_extractf128_ps(acc, 1); |
210 | 0 | let low = _mm256_castps256_ps128(acc); |
211 | 0 | let sum128 = _mm_add_ps(high, low); |
212 | 0 | let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); |
213 | 0 | let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1)); |
214 | 0 | _mm_cvtss_f32(sum32) |
215 | 0 | } |
216 | | } |
217 | | |
218 | | impl ComputeOp for DotQ5KOp { |
219 | | type Input = (Vec<BlockQ5K>, Vec<f32>); |
220 | | type Output = f32; |
221 | | |
222 | 0 | fn name(&self) -> &'static str { |
223 | 0 | "dot_q5k" |
224 | 0 | } |
225 | | |
226 | 0 | fn execute(&self, input: Self::Input, backend: Backend) -> Result<Self::Output, TruenoError> { |
227 | 0 | let (blocks, x) = input; |
228 | | |
229 | 0 | if blocks.is_empty() || x.is_empty() { |
230 | 0 | return Ok(0.0); |
231 | 0 | } |
232 | | |
233 | 0 | let mut sum = 0.0f32; |
234 | | |
235 | | #[cfg(target_arch = "x86_64")] |
236 | | { |
237 | 0 | if matches!(backend, Backend::Avx2 | Backend::Auto) && is_x86_feature_detected!("avx2") |
238 | | { |
239 | 0 | for (i, block) in blocks.iter().enumerate() { |
240 | 0 | let x_slice = &x[i * BlockQ5K::BLOCK_SIZE..]; |
241 | 0 | sum += unsafe { Self::avx2_dot_block(block, x_slice) }; |
242 | 0 | } |
243 | 0 | return Ok(sum); |
244 | 0 | } |
245 | | } |
246 | | |
247 | | // Scalar fallback |
248 | 0 | let mut dequant = [0.0f32; BlockQ5K::BLOCK_SIZE]; |
249 | 0 | for (i, block) in blocks.iter().enumerate() { |
250 | 0 | block.dequantize(&mut dequant); |
251 | 0 | let x_slice = &x[i * BlockQ5K::BLOCK_SIZE..]; |
252 | 0 | for j in 0..BlockQ5K::BLOCK_SIZE { |
253 | 0 | sum += dequant[j] * x_slice[j]; |
254 | 0 | } |
255 | | } |
256 | | |
257 | 0 | Ok(sum) |
258 | 0 | } |
259 | | |
260 | 0 | fn tokens(&self, _input: &Self::Input) -> usize { |
261 | 0 | self.n_blocks * BlockQ5K::BLOCK_SIZE |
262 | 0 | } |
263 | | } |
264 | | |
265 | | // ============================================================================ |
266 | | // Q6_K Dot Product Operation |
267 | | // ============================================================================ |
268 | | |
269 | | /// Q6_K dot product operation. |
270 | | /// |
271 | | /// Computes dot product between Q6_K quantized weights and f32 activations. |
272 | | #[derive(Debug, Clone)] |
273 | | pub struct DotQ6KOp { |
274 | | /// Number of blocks |
275 | | pub n_blocks: usize, |
276 | | } |
277 | | |
278 | | impl DotQ6KOp { |
279 | | /// Create a new Q6_K dot product operation. |
280 | | #[must_use] |
281 | 0 | pub fn new(n_elements: usize) -> Self { |
282 | 0 | Self { |
283 | 0 | n_blocks: n_elements / BlockQ6K::BLOCK_SIZE, |
284 | 0 | } |
285 | 0 | } |
286 | | |
287 | | /// Compute dot product with SIMD acceleration. |
288 | | #[cfg(target_arch = "x86_64")] |
289 | | #[target_feature(enable = "avx2", enable = "fma")] |
290 | 0 | unsafe fn avx2_dot_block(block: &BlockQ6K, x: &[f32]) -> f32 { |
291 | | use std::arch::x86_64::*; |
292 | | |
293 | 0 | let mut acc = _mm256_setzero_ps(); |
294 | 0 | let mut dequant = [0.0f32; BlockQ6K::BLOCK_SIZE]; |
295 | 0 | block.dequantize(&mut dequant); |
296 | | |
297 | 0 | let mut i = 0; |
298 | 0 | while i + 8 <= BlockQ6K::BLOCK_SIZE { |
299 | 0 | let vd = _mm256_loadu_ps(dequant.as_ptr().add(i)); |
300 | 0 | let vx = _mm256_loadu_ps(x.as_ptr().add(i)); |
301 | 0 | acc = _mm256_fmadd_ps(vd, vx, acc); |
302 | 0 | i += 8; |
303 | 0 | } |
304 | | |
305 | | // Horizontal sum |
306 | 0 | let high = _mm256_extractf128_ps(acc, 1); |
307 | 0 | let low = _mm256_castps256_ps128(acc); |
308 | 0 | let sum128 = _mm_add_ps(high, low); |
309 | 0 | let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); |
310 | 0 | let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1)); |
311 | 0 | _mm_cvtss_f32(sum32) |
312 | 0 | } |
313 | | } |
314 | | |
315 | | impl ComputeOp for DotQ6KOp { |
316 | | type Input = (Vec<BlockQ6K>, Vec<f32>); |
317 | | type Output = f32; |
318 | | |
319 | 0 | fn name(&self) -> &'static str { |
320 | 0 | "dot_q6k" |
321 | 0 | } |
322 | | |
323 | 0 | fn execute(&self, input: Self::Input, backend: Backend) -> Result<Self::Output, TruenoError> { |
324 | 0 | let (blocks, x) = input; |
325 | | |
326 | 0 | if blocks.is_empty() || x.is_empty() { |
327 | 0 | return Ok(0.0); |
328 | 0 | } |
329 | | |
330 | 0 | let mut sum = 0.0f32; |
331 | | |
332 | | #[cfg(target_arch = "x86_64")] |
333 | | { |
334 | 0 | if matches!(backend, Backend::Avx2 | Backend::Auto) && is_x86_feature_detected!("avx2") |
335 | | { |
336 | 0 | for (i, block) in blocks.iter().enumerate() { |
337 | 0 | let x_slice = &x[i * BlockQ6K::BLOCK_SIZE..]; |
338 | 0 | sum += unsafe { Self::avx2_dot_block(block, x_slice) }; |
339 | 0 | } |
340 | 0 | return Ok(sum); |
341 | 0 | } |
342 | | } |
343 | | |
344 | | // Scalar fallback |
345 | 0 | let mut dequant = [0.0f32; BlockQ6K::BLOCK_SIZE]; |
346 | 0 | for (i, block) in blocks.iter().enumerate() { |
347 | 0 | block.dequantize(&mut dequant); |
348 | 0 | let x_slice = &x[i * BlockQ6K::BLOCK_SIZE..]; |
349 | 0 | for j in 0..BlockQ6K::BLOCK_SIZE { |
350 | 0 | sum += dequant[j] * x_slice[j]; |
351 | 0 | } |
352 | | } |
353 | | |
354 | 0 | Ok(sum) |
355 | 0 | } |
356 | | |
357 | 0 | fn tokens(&self, _input: &Self::Input) -> usize { |
358 | 0 | self.n_blocks * BlockQ6K::BLOCK_SIZE |
359 | 0 | } |
360 | | } |
361 | | |
362 | | #[cfg(test)] |
363 | | mod tests { |
364 | | use super::*; |
365 | | |
366 | | // ===== BlockQ5K Tests ===== |
367 | | |
368 | | #[test] |
369 | | fn test_block_q5k_size() { |
370 | | assert_eq!(BlockQ5K::BLOCK_SIZE, 256); |
371 | | } |
372 | | |
373 | | #[test] |
374 | | fn test_block_q5k_dequantize_basic() { |
375 | | let block = BlockQ5K { |
376 | | d: 0.1, |
377 | | dmin: 0.0, |
378 | | scales: [32; 12], // Neutral scales (32 - 32 = 0) |
379 | | qh: [0; 32], // All high bits 0 |
380 | | qs: [0x88; 128], // 8,8 pattern (mid-range 4-bit) |
381 | | }; |
382 | | |
383 | | let mut output = [0.0f32; 256]; |
384 | | block.dequantize(&mut output); |
385 | | |
386 | | // With scale=0, all outputs should be dmin (0.0) |
387 | | for val in &output { |
388 | | assert!(val.abs() < 1.0, "Expected near-zero, got {}", val); |
389 | | } |
390 | | } |
391 | | |
392 | | #[test] |
393 | | fn test_block_q5k_dequantize_with_scale() { |
394 | | let block = BlockQ5K { |
395 | | d: 1.0, |
396 | | dmin: 0.5, |
397 | | scales: [33; 12], // Scale of 1 (33 - 32 = 1) |
398 | | qh: [0xFF; 32], // All high bits set |
399 | | qs: [0xFF; 128], // All low bits set (15,15) |
400 | | }; |
401 | | |
402 | | let mut output = [0.0f32; 256]; |
403 | | block.dequantize(&mut output); |
404 | | |
405 | | // Values should be non-zero with positive scale |
406 | | let non_zero_count = output.iter().filter(|&&v| v.abs() > 1e-6).count(); |
407 | | assert!(non_zero_count > 0, "Should have non-zero values"); |
408 | | } |
409 | | |
410 | | #[test] |
411 | | fn test_block_q5k_dequantize_alternating() { |
412 | | let block = BlockQ5K { |
413 | | d: 0.5, |
414 | | dmin: 0.1, |
415 | | scales: [34; 12], // Scale of 2 |
416 | | qh: [0xAA; 32], // Alternating bits |
417 | | qs: [0x55; 128], // Alternating nibbles (5,5) |
418 | | }; |
419 | | |
420 | | let mut output = [0.0f32; 256]; |
421 | | block.dequantize(&mut output); |
422 | | |
423 | | // All values should be finite |
424 | | for val in &output { |
425 | | assert!(val.is_finite(), "Value should be finite"); |
426 | | } |
427 | | } |
428 | | |
429 | | #[test] |
430 | | fn test_block_q5k_dequantize_odd_even_bytes() { |
431 | | // Test both even and odd index paths in dequantization |
432 | | let block = BlockQ5K { |
433 | | d: 1.0, |
434 | | dmin: 0.0, |
435 | | scales: [48; 12], // Scale of 16 (48 - 32 = 16) |
436 | | qh: [0; 32], |
437 | | qs: [0x12; 128], // Low nibble = 2, high nibble = 1 |
438 | | }; |
439 | | |
440 | | let mut output = [0.0f32; 256]; |
441 | | block.dequantize(&mut output); |
442 | | |
443 | | // Check that alternating values differ (even vs odd extraction) |
444 | | // Since qs[i] = 0x12, even indices extract 2, odd indices extract 1 |
445 | | // Note: the actual dequant formula is complex, but values should differ |
446 | | assert!(output[0] != output[1] || output[0].abs() < 1e-6); |
447 | | } |
448 | | |
449 | | // ===== BlockQ6K Tests ===== |
450 | | |
451 | | #[test] |
452 | | fn test_block_q6k_size() { |
453 | | assert_eq!(BlockQ6K::BLOCK_SIZE, 256); |
454 | | } |
455 | | |
456 | | #[test] |
457 | | fn test_block_q6k_dequantize_basic() { |
458 | | let block = BlockQ6K { |
459 | | ql: [0; 128], |
460 | | qh: [0; 64], |
461 | | scales: [0; 16], // Zero scales |
462 | | d: 0.1, |
463 | | }; |
464 | | |
465 | | let mut output = [0.0f32; 256]; |
466 | | block.dequantize(&mut output); |
467 | | |
468 | | // With scale=0, all outputs should be d * 0 * (q6 - 32) = 0 |
469 | | for val in &output { |
470 | | assert!(val.abs() < 1e-6, "Expected 0, got {}", val); |
471 | | } |
472 | | } |
473 | | |
474 | | #[test] |
475 | | fn test_block_q6k_dequantize_with_scale() { |
476 | | let block = BlockQ6K { |
477 | | ql: [0xFF; 128], // Max low bits |
478 | | qh: [0xFF; 64], // Max high bits |
479 | | scales: [1; 16], // Positive scale |
480 | | d: 0.5, |
481 | | }; |
482 | | |
483 | | let mut output = [0.0f32; 256]; |
484 | | block.dequantize(&mut output); |
485 | | |
486 | | // Values should be non-zero |
487 | | let non_zero = output.iter().any(|&v| v.abs() > 1e-6); |
488 | | assert!(non_zero, "Should have non-zero values"); |
489 | | } |
490 | | |
491 | | #[test] |
492 | | fn test_block_q6k_dequantize_negative_scale() { |
493 | | let block = BlockQ6K { |
494 | | ql: [0x88; 128], |
495 | | qh: [0x55; 64], |
496 | | scales: [-1; 16], // Negative scale |
497 | | d: 1.0, |
498 | | }; |
499 | | |
500 | | let mut output = [0.0f32; 256]; |
501 | | block.dequantize(&mut output); |
502 | | |
503 | | // All values should be finite |
504 | | for val in &output { |
505 | | assert!(val.is_finite()); |
506 | | } |
507 | | } |
508 | | |
509 | | #[test] |
510 | | fn test_block_q6k_dequantize_all_subblocks() { |
511 | | // Test that all 16 sub-blocks are processed |
512 | | let block = BlockQ6K { |
513 | | ql: [0x12; 128], |
514 | | qh: [0x03; 64], // Different pattern per position |
515 | | scales: [1, 2, 3, 4, 5, 6, 7, 8, -1, -2, -3, -4, -5, -6, -7, -8], |
516 | | d: 0.1, |
517 | | }; |
518 | | |
519 | | let mut output = [0.0f32; 256]; |
520 | | block.dequantize(&mut output); |
521 | | |
522 | | // Check values at different sub-block boundaries |
523 | | assert!(output[0].is_finite()); |
524 | | assert!(output[15].is_finite()); |
525 | | assert!(output[16].is_finite()); |
526 | | assert!(output[127].is_finite()); |
527 | | assert!(output[255].is_finite()); |
528 | | } |
529 | | |
530 | | #[test] |
531 | | fn test_block_q6k_qh_extraction() { |
532 | | // Test the 2-bit high value extraction logic |
533 | | // qh_shift cycles through 0, 2, 4, 6 for i % 4 = 0, 1, 2, 3 |
534 | | let block = BlockQ6K { |
535 | | ql: [0; 128], |
536 | | qh: [0b11_10_01_00; 64], // Pattern: 0,1,2,3 across 4 positions |
537 | | scales: [1; 16], |
538 | | d: 1.0, |
539 | | }; |
540 | | |
541 | | let mut output = [0.0f32; 256]; |
542 | | block.dequantize(&mut output); |
543 | | |
544 | | // Different qh values should produce different outputs |
545 | | // Position 0: qh_val = 0, Position 1: qh_val = 1, etc. |
546 | | // This tests the (i % 4) * 2 shift logic |
547 | | assert!(output[0].is_finite()); |
548 | | assert!(output[1].is_finite()); |
549 | | assert!(output[2].is_finite()); |
550 | | assert!(output[3].is_finite()); |
551 | | } |
552 | | |
553 | | // ===== DotQ5KOp Tests ===== |
554 | | |
555 | | #[test] |
556 | | fn test_dot_q5k_new() { |
557 | | let op = DotQ5KOp::new(512); |
558 | | assert_eq!(op.n_blocks, 2); |
559 | | } |
560 | | |
561 | | #[test] |
562 | | fn test_dot_q5k_name() { |
563 | | let op = DotQ5KOp::new(256); |
564 | | assert_eq!(op.name(), "dot_q5k"); |
565 | | } |
566 | | |
567 | | #[test] |
568 | | fn test_dot_q5k_empty() { |
569 | | let op = DotQ5KOp::new(256); |
570 | | let result = op.execute((vec![], vec![]), Backend::Scalar).unwrap(); |
571 | | assert!((result - 0.0).abs() < 1e-6); |
572 | | } |
573 | | |
574 | | #[test] |
575 | | fn test_dot_q5k_empty_activations() { |
576 | | let op = DotQ5KOp::new(256); |
577 | | let block = BlockQ5K { |
578 | | d: 1.0, |
579 | | dmin: 0.0, |
580 | | scales: [32; 12], |
581 | | qh: [0; 32], |
582 | | qs: [0; 128], |
583 | | }; |
584 | | let result = op.execute((vec![block], vec![]), Backend::Scalar).unwrap(); |
585 | | assert!((result - 0.0).abs() < 1e-6); |
586 | | } |
587 | | |
588 | | #[test] |
589 | | fn test_dot_q5k_tokens() { |
590 | | let op = DotQ5KOp::new(512); // 2 blocks |
591 | | let input = (vec![], vec![]); |
592 | | assert_eq!(op.tokens(&input), 512); |
593 | | } |
594 | | |
595 | | #[test] |
596 | | fn test_dot_q5k_scalar_execution() { |
597 | | let op = DotQ5KOp::new(256); |
598 | | let block = BlockQ5K { |
599 | | d: 1.0, |
600 | | dmin: 0.0, |
601 | | scales: [33; 12], // Scale = 1 |
602 | | qh: [0; 32], |
603 | | qs: [0x88; 128], // Mid-range values |
604 | | }; |
605 | | let x = vec![1.0f32; 256]; |
606 | | let result = op.execute((vec![block], x), Backend::Scalar).unwrap(); |
607 | | assert!(result.is_finite()); |
608 | | } |
609 | | |
610 | | #[test] |
611 | | fn test_dot_q5k_multiple_blocks() { |
612 | | let op = DotQ5KOp::new(512); |
613 | | let block = BlockQ5K { |
614 | | d: 0.5, |
615 | | dmin: 0.1, |
616 | | scales: [34; 12], |
617 | | qh: [0; 32], |
618 | | qs: [0x44; 128], |
619 | | }; |
620 | | let x = vec![0.5f32; 512]; |
621 | | let result = op.execute((vec![block.clone(), block], x), Backend::Scalar).unwrap(); |
622 | | assert!(result.is_finite()); |
623 | | } |
624 | | |
625 | | #[test] |
626 | | fn test_dot_q5k_auto_backend() { |
627 | | let op = DotQ5KOp::new(256); |
628 | | let block = BlockQ5K { |
629 | | d: 1.0, |
630 | | dmin: 0.0, |
631 | | scales: [32; 12], |
632 | | qh: [0; 32], |
633 | | qs: [0; 128], |
634 | | }; |
635 | | let x = vec![1.0f32; 256]; |
636 | | // Auto backend should work (may use AVX2 if available) |
637 | | let result = op.execute((vec![block], x), Backend::Auto).unwrap(); |
638 | | assert!(result.is_finite()); |
639 | | } |
640 | | |
641 | | #[test] |
642 | | fn test_dot_q5k_avx2_backend() { |
643 | | let op = DotQ5KOp::new(256); |
644 | | let block = BlockQ5K { |
645 | | d: 1.0, |
646 | | dmin: 0.0, |
647 | | scales: [33; 12], |
648 | | qh: [0; 32], |
649 | | qs: [0x11; 128], |
650 | | }; |
651 | | let x = vec![2.0f32; 256]; |
652 | | // Request AVX2, will fall back to scalar if not available |
653 | | let result = op.execute((vec![block], x), Backend::Avx2).unwrap(); |
654 | | assert!(result.is_finite()); |
655 | | } |
656 | | |
657 | | // ===== DotQ6KOp Tests ===== |
658 | | |
659 | | #[test] |
660 | | fn test_dot_q6k_new() { |
661 | | let op = DotQ6KOp::new(768); |
662 | | assert_eq!(op.n_blocks, 3); |
663 | | } |
664 | | |
665 | | #[test] |
666 | | fn test_dot_q6k_name() { |
667 | | let op = DotQ6KOp::new(256); |
668 | | assert_eq!(op.name(), "dot_q6k"); |
669 | | } |
670 | | |
671 | | #[test] |
672 | | fn test_dot_q6k_empty() { |
673 | | let op = DotQ6KOp::new(256); |
674 | | let result = op.execute((vec![], vec![]), Backend::Scalar).unwrap(); |
675 | | assert!((result - 0.0).abs() < 1e-6); |
676 | | } |
677 | | |
678 | | #[test] |
679 | | fn test_dot_q6k_empty_activations() { |
680 | | let op = DotQ6KOp::new(256); |
681 | | let block = BlockQ6K { |
682 | | ql: [0; 128], |
683 | | qh: [0; 64], |
684 | | scales: [0; 16], |
685 | | d: 1.0, |
686 | | }; |
687 | | let result = op.execute((vec![block], vec![]), Backend::Scalar).unwrap(); |
688 | | assert!((result - 0.0).abs() < 1e-6); |
689 | | } |
690 | | |
691 | | #[test] |
692 | | fn test_dot_q6k_tokens() { |
693 | | let op = DotQ6KOp::new(768); // 3 blocks |
694 | | let input = (vec![], vec![]); |
695 | | assert_eq!(op.tokens(&input), 768); |
696 | | } |
697 | | |
698 | | #[test] |
699 | | fn test_dot_q6k_scalar_execution() { |
700 | | let op = DotQ6KOp::new(256); |
701 | | let block = BlockQ6K { |
702 | | ql: [0x55; 128], |
703 | | qh: [0x55; 64], |
704 | | scales: [1; 16], |
705 | | d: 0.5, |
706 | | }; |
707 | | let x = vec![1.0f32; 256]; |
708 | | let result = op.execute((vec![block], x), Backend::Scalar).unwrap(); |
709 | | assert!(result.is_finite()); |
710 | | } |
711 | | |
712 | | #[test] |
713 | | fn test_dot_q6k_multiple_blocks() { |
714 | | let op = DotQ6KOp::new(512); |
715 | | let block = BlockQ6K { |
716 | | ql: [0x33; 128], |
717 | | qh: [0x33; 64], |
718 | | scales: [2; 16], |
719 | | d: 0.25, |
720 | | }; |
721 | | let x = vec![0.5f32; 512]; |
722 | | let result = op.execute((vec![block.clone(), block], x), Backend::Scalar).unwrap(); |
723 | | assert!(result.is_finite()); |
724 | | } |
725 | | |
726 | | #[test] |
727 | | fn test_dot_q6k_auto_backend() { |
728 | | let op = DotQ6KOp::new(256); |
729 | | let block = BlockQ6K { |
730 | | ql: [0; 128], |
731 | | qh: [0; 64], |
732 | | scales: [1; 16], |
733 | | d: 1.0, |
734 | | }; |
735 | | let x = vec![1.0f32; 256]; |
736 | | let result = op.execute((vec![block], x), Backend::Auto).unwrap(); |
737 | | assert!(result.is_finite()); |
738 | | } |
739 | | |
740 | | #[test] |
741 | | fn test_dot_q6k_avx2_backend() { |
742 | | let op = DotQ6KOp::new(256); |
743 | | let block = BlockQ6K { |
744 | | ql: [0xAA; 128], |
745 | | qh: [0xAA; 64], |
746 | | scales: [3; 16], |
747 | | d: 0.1, |
748 | | }; |
749 | | let x = vec![2.0f32; 256]; |
750 | | let result = op.execute((vec![block], x), Backend::Avx2).unwrap(); |
751 | | assert!(result.is_finite()); |
752 | | } |
753 | | |
754 | | // ===== Backend Equivalence Tests ===== |
755 | | |
756 | | #[test] |
757 | | fn test_q5k_backend_equivalence() { |
758 | | let op = DotQ5KOp::new(256); |
759 | | let block = BlockQ5K { |
760 | | d: 0.5, |
761 | | dmin: 0.1, |
762 | | scales: [35; 12], |
763 | | qh: [0x55; 32], |
764 | | qs: [0x77; 128], |
765 | | }; |
766 | | let x = vec![1.5f32; 256]; |
767 | | |
768 | | let scalar = op.execute((vec![block.clone()], x.clone()), Backend::Scalar).unwrap(); |
769 | | let auto = op.execute((vec![block], x), Backend::Auto).unwrap(); |
770 | | |
771 | | // Allow small FP differences due to SIMD operation ordering |
772 | | let rel_diff = (scalar - auto).abs() / scalar.abs().max(1e-6); |
773 | | assert!(rel_diff < 1e-4, "scalar={scalar}, auto={auto}, rel_diff={rel_diff}"); |
774 | | } |
775 | | |
776 | | #[test] |
777 | | fn test_q6k_backend_equivalence() { |
778 | | let op = DotQ6KOp::new(256); |
779 | | let block = BlockQ6K { |
780 | | ql: [0x66; 128], |
781 | | qh: [0x22; 64], |
782 | | scales: [4; 16], |
783 | | d: 0.2, |
784 | | }; |
785 | | let x = vec![1.5f32; 256]; |
786 | | |
787 | | let scalar = op.execute((vec![block.clone()], x.clone()), Backend::Scalar).unwrap(); |
788 | | let auto = op.execute((vec![block], x), Backend::Auto).unwrap(); |
789 | | |
790 | | // Allow small FP differences due to SIMD operation ordering |
791 | | let rel_diff = (scalar - auto).abs() / scalar.abs().max(1e-6); |
792 | | assert!(rel_diff < 1e-4, "scalar={scalar}, auto={auto}, rel_diff={rel_diff}"); |
793 | | } |
794 | | |
795 | | // ===== Clone/Debug Trait Tests ===== |
796 | | |
797 | | #[test] |
798 | | fn test_block_q5k_clone_debug() { |
799 | | let block = BlockQ5K { |
800 | | d: 1.0, |
801 | | dmin: 0.5, |
802 | | scales: [32; 12], |
803 | | qh: [0; 32], |
804 | | qs: [0; 128], |
805 | | }; |
806 | | let cloned = block.clone(); |
807 | | assert_eq!(format!("{:?}", block), format!("{:?}", cloned)); |
808 | | } |
809 | | |
810 | | #[test] |
811 | | fn test_block_q6k_clone_debug() { |
812 | | let block = BlockQ6K { |
813 | | ql: [0; 128], |
814 | | qh: [0; 64], |
815 | | scales: [0; 16], |
816 | | d: 1.0, |
817 | | }; |
818 | | let cloned = block.clone(); |
819 | | assert_eq!(format!("{:?}", block), format!("{:?}", cloned)); |
820 | | } |
821 | | |
822 | | #[test] |
823 | | fn test_dot_q5k_op_clone_debug() { |
824 | | let op = DotQ5KOp::new(256); |
825 | | let cloned = op.clone(); |
826 | | assert_eq!(format!("{:?}", op), format!("{:?}", cloned)); |
827 | | } |
828 | | |
829 | | #[test] |
830 | | fn test_dot_q6k_op_clone_debug() { |
831 | | let op = DotQ6KOp::new(256); |
832 | | let cloned = op.clone(); |
833 | | assert_eq!(format!("{:?}", op), format!("{:?}", cloned)); |
834 | | } |
835 | | } |