1pub mod encode;
16pub use encode::q4_k::{encode_block_q4_k, encode_row_q4_k};
17pub use encode::q5_k::{encode_block_q5_k, encode_row_q5_k};
18pub use encode::q6_k::{encode_block_q6_k, encode_row_q6_k, probe_q6_k_group};
19pub use encode::{encode_block_q8_0, encode_row_q8_0};
20
21pub mod iq_tables;
22#[cfg(test)]
26mod iq_tier_goldens;
27pub mod repack;
28
29pub use repack::{
30 gemm_q4_0x4_group, gemm_q4_0x4_group_x4, gemm_q4_0x4_group_x4_on, gemm_q4_kx8_group,
31 gemm_q4_kx8_group_x4, gemm_q4_kx8_group_x4_on, gemm_q5_kx8_group, gemm_q5_kx8_group_x4,
32 gemm_q5_kx8_group_x4_on, gemm_q6_kx8_group, gemm_q6_kx8_group_x4, gemm_q6_kx8_group_x4_on,
33 gemm_q8_0x4_group, gemm_q8_0x4_group_x4, gemm_q8_0x4_group_x4_on, gemv_q4_0x4_group,
34 gemv_q4_kx8_group, gemv_q4_kx8_q8_k, gemv_q5_kx8_group, gemv_q5_kx8_q8_k, gemv_q6_kx8_group,
35 gemv_q6_kx8_q8_k, gemv_q8_0x4_group, gemv_q8_0x4_q8_0, interleaved_gemm_is_accelerated,
36 make_block_q4_0x4, make_block_q4_kx8, make_block_q5_kx8, make_block_q6_kx8, make_block_q8_0x4,
37 pack_q4_0_matrix_x4, pack_q4_k_matrix_x8, pack_q5_k_matrix_x8, pack_q6_k_matrix_x8,
38 pack_q8_0_matrix_x4, preferred_interleave, prepare_q8_acts_x4, prepare_q8_k_acts_x4,
39 q4_0x4_gemm_uses_acts_x4, q4_0x4_interleave, q4_kx8_gemm_uses_acts_x4, q4_kx8_interleave,
40 q5_kx8_gemm_uses_acts_x4, q5_kx8_interleave, q6_kx8_gemm_uses_acts_x4, q6_kx8_interleave,
41 q8_0x4_gemm_uses_acts_x4, q8_0x4_interleave, AccelX4, Q8ActsX4, Q8KActsX4, Q4_0X4_BLOCK_BYTES,
42 Q4_0X4_GEMM_NC, Q4_0X4_INTERLEAVE, Q4_0X4_NROWS, Q4_KX8_BLOCK_BYTES, Q4_KX8_GEMM_NC,
43 Q4_KX8_NROWS, Q5_KX8_BLOCK_BYTES, Q5_KX8_GEMM_NC, Q5_KX8_NROWS, Q6_KX8_BLOCK_BYTES,
44 Q6_KX8_GEMM_NC, Q6_KX8_NROWS, Q8K_ACTS_X4_NC, Q8_0X4_BLOCK_BYTES, Q8_0X4_GEMM_NC,
45 Q8_0X4_INTERLEAVE, Q8_0X4_NROWS,
46};
47
48use half::f16;
49
50pub const Q8_0_BLOCK_BYTES: usize = 34;
52pub const Q8_0_BLOCK_ELEMS: usize = 32;
53
54pub const Q4_0_BLOCK_BYTES: usize = 18;
56pub const Q4_0_BLOCK_ELEMS: usize = 32;
57
58pub const Q4_1_BLOCK_BYTES: usize = 20;
63pub const Q4_1_BLOCK_ELEMS: usize = 32;
64
65pub const Q5_0_BLOCK_BYTES: usize = 22;
69pub const Q5_0_BLOCK_ELEMS: usize = 32;
70
71pub const Q5_1_BLOCK_BYTES: usize = 24;
75pub const Q5_1_BLOCK_ELEMS: usize = 32;
76
77pub const Q8_1_BLOCK_BYTES: usize = 36;
83pub const Q8_1_BLOCK_ELEMS: usize = 32;
84
85pub const TURBO4_KV_GROUP: usize = 32;
87pub const TURBO4_KV_BLOCK_BYTES: usize = 18;
88
89pub const FP8_KV_GROUP: usize = 32;
92pub const FP8_KV_BLOCK_BYTES: usize = 34;
93
94pub fn pack_turbo4_kv_blocks(x: &[f32]) -> Vec<u8> {
96 assert_eq!(x.len() % TURBO4_KV_GROUP, 0);
97 let n_blocks = x.len() / TURBO4_KV_GROUP;
98 let mut out = vec![0u8; n_blocks * TURBO4_KV_BLOCK_BYTES];
99 for b in 0..n_blocks {
100 let chunk = &x[b * TURBO4_KV_GROUP..(b + 1) * TURBO4_KV_GROUP];
101 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
102 let scale = if amax > 0.0 { amax / 7.0 } else { 0.0 };
103 let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
104 let bits = f16::from_f32(scale).to_le_bytes();
105 let dst = &mut out[b * TURBO4_KV_BLOCK_BYTES..(b + 1) * TURBO4_KV_BLOCK_BYTES];
106 dst[0] = bits[0];
107 dst[1] = bits[1];
108 for i in 0..16 {
109 let q0 = (chunk[i * 2] * inv).round().clamp(-8.0, 7.0) as i8;
110 let q1 = (chunk[i * 2 + 1] * inv).round().clamp(-8.0, 7.0) as i8;
111 dst[2 + i] = ((q0 as u8) & 0x0f) | (((q1 as u8) & 0x0f) << 4);
112 }
113 }
114 out
115}
116
117pub fn unpack_turbo4_kv_blocks(bytes: &[u8]) -> Result<Vec<f32>, QuantError> {
119 if !bytes.len().is_multiple_of(TURBO4_KV_BLOCK_BYTES) {
120 return Err(QuantError::Misaligned(bytes.len(), TURBO4_KV_BLOCK_BYTES));
121 }
122 let n_blocks = bytes.len() / TURBO4_KV_BLOCK_BYTES;
123 let mut out = Vec::with_capacity(n_blocks * TURBO4_KV_GROUP);
124 for b in 0..n_blocks {
125 let block = &bytes[b * TURBO4_KV_BLOCK_BYTES..(b + 1) * TURBO4_KV_BLOCK_BYTES];
126 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
127 for i in 0..16 {
128 let byte = block[2 + i];
129 let q0 = ((byte & 0x0f) as i8) << 4 >> 4;
130 let q1 = ((byte >> 4) as i8) << 4 >> 4;
131 out.push(q0 as f32 * scale);
132 out.push(q1 as f32 * scale);
133 }
134 }
135 Ok(out)
136}
137
138pub fn pack_fp8_kv_blocks(x: &[f32]) -> Vec<u8> {
140 quantize_q8_0(x)
142}
143
144pub fn unpack_fp8_kv_blocks(bytes: &[u8]) -> Result<Vec<f32>, QuantError> {
146 dequant_q8_0(bytes)
147}
148
149pub const Q4_K_BLOCK_BYTES: usize = 144;
154pub const Q4_K_BLOCK_ELEMS: usize = 256;
155const Q4_K_SCALE_BYTES: usize = 12;
156
157pub const Q5_K_BLOCK_BYTES: usize = 176;
163pub const Q5_K_BLOCK_ELEMS: usize = 256;
164
165pub const Q6_K_BLOCK_BYTES: usize = 210;
170pub const Q6_K_BLOCK_ELEMS: usize = 256;
171
172pub const Q2_K_BLOCK_BYTES: usize = 84;
182pub const Q2_K_BLOCK_ELEMS: usize = 256;
183const Q2_K_SCALE_BYTES: usize = 16;
184
185pub const Q3_K_BLOCK_BYTES: usize = 110;
194pub const Q3_K_BLOCK_ELEMS: usize = 256;
195const Q3_K_SCALE_BYTES: usize = 12;
196
197#[derive(Debug, thiserror::Error)]
198pub enum QuantError {
199 #[error("buffer length {0} is not a multiple of the block size {1}")]
200 Misaligned(usize, usize),
201 #[error("MXFP4 packed buffer is {0} bytes but scales buffer implies {1} bytes ({1} = scales.len() * MXFP4_GROUP_SIZE / 2)")]
202 Mxfp4RowMismatch(usize, usize),
203}
204
205pub fn dequant_bf16(src: &[u8]) -> Result<Vec<f32>, QuantError> {
214 if !src.len().is_multiple_of(2) {
215 return Err(QuantError::Misaligned(src.len(), 2));
216 }
217 Ok(src
218 .as_chunks::<2>()
219 .0
220 .iter()
221 .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
222 .collect())
223}
224
225pub fn dequant_f16(src: &[u8]) -> Result<Vec<f32>, QuantError> {
232 if !src.len().is_multiple_of(2) {
233 return Err(QuantError::Misaligned(src.len(), 2));
234 }
235 Ok(src
236 .as_chunks::<2>()
237 .0
238 .iter()
239 .map(|c| f16::from_le_bytes([c[0], c[1]]).to_f32())
240 .collect())
241}
242
243pub fn dequant_q8_0(src: &[u8]) -> Result<Vec<f32>, QuantError> {
245 if !src.len().is_multiple_of(Q8_0_BLOCK_BYTES) {
246 return Err(QuantError::Misaligned(src.len(), Q8_0_BLOCK_BYTES));
247 }
248 let n_blocks = src.len() / Q8_0_BLOCK_BYTES;
249 let mut out = Vec::with_capacity(n_blocks * Q8_0_BLOCK_ELEMS);
250 for b in 0..n_blocks {
251 let block = &src[b * Q8_0_BLOCK_BYTES..(b + 1) * Q8_0_BLOCK_BYTES];
252 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
253 for i in 0..Q8_0_BLOCK_ELEMS {
254 let q = block[2 + i] as i8;
255 out.push(q as f32 * scale);
256 }
257 }
258 Ok(out)
259}
260
261pub fn dequant_q4_0(src: &[u8]) -> Result<Vec<f32>, QuantError> {
265 if !src.len().is_multiple_of(Q4_0_BLOCK_BYTES) {
266 return Err(QuantError::Misaligned(src.len(), Q4_0_BLOCK_BYTES));
267 }
268 let n_blocks = src.len() / Q4_0_BLOCK_BYTES;
269 let mut out = vec![0f32; n_blocks * Q4_0_BLOCK_ELEMS];
270 for b in 0..n_blocks {
271 let block = &src[b * Q4_0_BLOCK_BYTES..(b + 1) * Q4_0_BLOCK_BYTES];
272 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
273 let nibbles = &block[2..18];
274 let base = b * Q4_0_BLOCK_ELEMS;
275 for i in 0..16 {
276 let byte = nibbles[i];
277 let lo = (byte & 0x0F) as i32 - 8;
278 let hi = ((byte >> 4) & 0x0F) as i32 - 8;
279 out[base + i] = lo as f32 * scale;
280 out[base + i + 16] = hi as f32 * scale;
281 }
282 }
283 Ok(out)
284}
285
286fn q4_k_scale_min(j: usize, scales: &[u8; Q4_K_SCALE_BYTES]) -> (u8, u8) {
295 if j < 4 {
296 (scales[j] & 63, scales[j + 4] & 63)
297 } else {
298 (
299 (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4),
300 (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4),
301 )
302 }
303}
304
305pub fn dequant_q4_k(src: &[u8]) -> Result<Vec<f32>, QuantError> {
308 if !src.len().is_multiple_of(Q4_K_BLOCK_BYTES) {
309 return Err(QuantError::Misaligned(src.len(), Q4_K_BLOCK_BYTES));
310 }
311 let n_blocks = src.len() / Q4_K_BLOCK_BYTES;
312 let mut out = Vec::with_capacity(n_blocks * Q4_K_BLOCK_ELEMS);
313 for block in src.as_chunks::<Q4_K_BLOCK_BYTES>().0 {
314 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
315 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
316 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
317 let qs = &block[16..144];
318
319 let mut is = 0usize;
320 let mut q_off = 0usize;
321 for _ in 0..4 {
322 let (sc1, m1) = q4_k_scale_min(is, &scales);
323 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
324 let (d1, min1) = (d * sc1 as f32, dmin * m1 as f32);
325 let (d2, min2) = (d * sc2 as f32, dmin * m2 as f32);
326 for l in 0..32 {
327 out.push(d1 * (qs[q_off + l] & 0x0F) as f32 - min1);
328 }
329 for l in 0..32 {
330 out.push(d2 * (qs[q_off + l] >> 4) as f32 - min2);
331 }
332 q_off += 32;
333 is += 2;
334 }
335 }
336 Ok(out)
337}
338
339pub fn dot_q4_k_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
344 #[cfg(target_arch = "x86_64")]
345 {
346 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
347 return unsafe { simd_x86::dot_q4_k_f32_avx2(row_bytes, x) };
348 }
349 }
350 #[cfg(target_arch = "aarch64")]
351 {
352 if std::arch::is_aarch64_feature_detected!("neon") {
353 return unsafe { simd_aarch64::dot_q4_k_f32_neon(row_bytes, x) };
354 }
355 }
356 dot_q4_k_f32_scalar(row_bytes, x)
357}
358
359pub fn dot_q4_k_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
360 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
361 let mut acc = 0f32;
362 let mut base = 0usize;
363 for block in row_bytes.as_chunks::<Q4_K_BLOCK_BYTES>().0 {
364 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
365 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
366 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
367 let qs = &block[16..144];
368
369 let mut is = 0usize;
370 let mut q_off = 0usize;
371 for _ in 0..4 {
372 let (sc1, m1) = q4_k_scale_min(is, &scales);
373 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
374 let (d1, min1) = (d * sc1 as f32, dmin * m1 as f32);
375 let (d2, min2) = (d * sc2 as f32, dmin * m2 as f32);
376 for l in 0..32 {
377 acc += (d1 * (qs[q_off + l] & 0x0F) as f32 - min1) * x[base + l];
378 }
379 for l in 0..32 {
380 acc += (d2 * (qs[q_off + l] >> 4) as f32 - min2) * x[base + 32 + l];
381 }
382 q_off += 32;
383 base += 64;
384 is += 2;
385 }
386 }
387 acc
388}
389
390pub fn dequant_q5_k(src: &[u8]) -> Result<Vec<f32>, QuantError> {
398 if !src.len().is_multiple_of(Q5_K_BLOCK_BYTES) {
399 return Err(QuantError::Misaligned(src.len(), Q5_K_BLOCK_BYTES));
400 }
401 let n_blocks = src.len() / Q5_K_BLOCK_BYTES;
402 let mut out = Vec::with_capacity(n_blocks * Q5_K_BLOCK_ELEMS);
403 for block in src.as_chunks::<Q5_K_BLOCK_BYTES>().0 {
404 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
405 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
406 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
407 let qh = &block[16..48];
408 let qs = &block[48..176];
409
410 let mut is = 0usize;
411 let (mut u1, mut u2) = (1u8, 2u8);
412 for oi in 0..4 {
413 let (sc1, m1) = q4_k_scale_min(is, &scales);
414 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
415 let (d1, min1) = (d * sc1 as f32, dmin * m1 as f32);
416 let (d2, min2) = (d * sc2 as f32, dmin * m2 as f32);
417 let ql = &qs[oi * 32..oi * 32 + 32];
418 for l in 0..32 {
419 let hi = if qh[l] & u1 != 0 { 16 } else { 0 };
420 out.push(d1 * ((ql[l] & 0x0F) + hi) as f32 - min1);
421 }
422 for l in 0..32 {
423 let hi = if qh[l] & u2 != 0 { 16 } else { 0 };
424 out.push(d2 * ((ql[l] >> 4) + hi) as f32 - min2);
425 }
426 is += 2;
427 u1 <<= 2;
428 u2 <<= 2;
429 }
430 }
431 Ok(out)
432}
433
434pub fn dot_q5_k_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
439 #[cfg(target_arch = "x86_64")]
440 {
441 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
442 return unsafe { simd_x86::dot_q5_k_f32_avx2(row_bytes, x) };
443 }
444 }
445 #[cfg(target_arch = "aarch64")]
446 {
447 if std::arch::is_aarch64_feature_detected!("neon") {
448 return unsafe { simd_aarch64::dot_q5_k_f32_neon(row_bytes, x) };
449 }
450 }
451 dot_q5_k_f32_scalar(row_bytes, x)
452}
453
454pub fn dot_q5_k_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
455 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
456 let mut acc = 0f32;
457 let mut base = 0usize;
458 for block in row_bytes.as_chunks::<Q5_K_BLOCK_BYTES>().0 {
459 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
460 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
461 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
462 let qh = &block[16..48];
463 let qs = &block[48..176];
464
465 let mut is = 0usize;
466 let (mut u1, mut u2) = (1u8, 2u8);
467 for oi in 0..4 {
468 let (sc1, m1) = q4_k_scale_min(is, &scales);
469 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
470 let (d1, min1) = (d * sc1 as f32, dmin * m1 as f32);
471 let (d2, min2) = (d * sc2 as f32, dmin * m2 as f32);
472 let ql = &qs[oi * 32..oi * 32 + 32];
473 for l in 0..32 {
474 let hi = if qh[l] & u1 != 0 { 16 } else { 0 };
475 acc += (d1 * ((ql[l] & 0x0F) + hi) as f32 - min1) * x[base + l];
476 }
477 for l in 0..32 {
478 let hi = if qh[l] & u2 != 0 { 16 } else { 0 };
479 acc += (d2 * ((ql[l] >> 4) + hi) as f32 - min2) * x[base + 32 + l];
480 }
481 base += 64;
482 is += 2;
483 u1 <<= 2;
484 u2 <<= 2;
485 }
486 }
487 acc
488}
489
490pub fn dequant_q6_k(src: &[u8]) -> Result<Vec<f32>, QuantError> {
493 if !src.len().is_multiple_of(Q6_K_BLOCK_BYTES) {
494 return Err(QuantError::Misaligned(src.len(), Q6_K_BLOCK_BYTES));
495 }
496 let n_blocks = src.len() / Q6_K_BLOCK_BYTES;
497 let mut out = vec![0f32; n_blocks * Q6_K_BLOCK_ELEMS];
498 for (b, block) in src.as_chunks::<Q6_K_BLOCK_BYTES>().0.iter().enumerate() {
499 let ql_full = &block[0..128];
500 let qh_full = &block[128..192];
501 let sc_full = &block[192..208];
502 let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
503 let out_base = b * Q6_K_BLOCK_ELEMS;
504
505 for half in 0..2 {
506 let ql = &ql_full[half * 64..half * 64 + 64];
507 let qh = &qh_full[half * 32..half * 32 + 32];
508 let sc = &sc_full[half * 8..half * 8 + 8];
509 let y = &mut out[out_base + half * 128..out_base + half * 128 + 128];
510
511 for l in 0..32 {
512 let is = l / 16;
513 let q1 = ((ql[l] & 0x0F) | ((qh[l] & 3) << 4)) as i8 - 32;
514 let q2 = ((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) as i8 - 32;
515 let q3 = ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) as i8 - 32;
516 let q4 = ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) as i8 - 32;
517 y[l] = d * (sc[is] as i8 as f32) * (q1 as f32);
518 y[l + 32] = d * (sc[is + 2] as i8 as f32) * (q2 as f32);
519 y[l + 64] = d * (sc[is + 4] as i8 as f32) * (q3 as f32);
520 y[l + 96] = d * (sc[is + 6] as i8 as f32) * (q4 as f32);
521 }
522 }
523 }
524 Ok(out)
525}
526
527pub fn dot_q6_k_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
532 #[cfg(target_arch = "x86_64")]
533 {
534 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
535 return unsafe { simd_x86::dot_q6_k_f32_avx2(row_bytes, x) };
536 }
537 }
538 #[cfg(target_arch = "aarch64")]
539 {
540 if std::arch::is_aarch64_feature_detected!("neon") {
541 return unsafe { simd_aarch64::dot_q6_k_f32_neon(row_bytes, x) };
542 }
543 }
544 dot_q6_k_f32_scalar(row_bytes, x)
545}
546
547pub fn dot_q6_k_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
548 debug_assert_eq!(row_bytes.len() % Q6_K_BLOCK_BYTES, 0);
549 let mut acc = 0f32;
550 let mut x_base = 0usize;
551 for block in row_bytes.as_chunks::<Q6_K_BLOCK_BYTES>().0 {
552 let ql_full = &block[0..128];
553 let qh_full = &block[128..192];
554 let sc_full = &block[192..208];
555 let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
556
557 for half in 0..2 {
558 let ql = &ql_full[half * 64..half * 64 + 64];
559 let qh = &qh_full[half * 32..half * 32 + 32];
560 let sc = &sc_full[half * 8..half * 8 + 8];
561 let xh = &x[x_base..x_base + 128];
562
563 for l in 0..32 {
564 let is = l / 16;
565 let q1 = ((ql[l] & 0x0F) | ((qh[l] & 3) << 4)) as i8 - 32;
566 let q2 = ((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) as i8 - 32;
567 let q3 = ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) as i8 - 32;
568 let q4 = ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) as i8 - 32;
569 acc += d * (sc[is] as i8 as f32) * (q1 as f32) * xh[l];
570 acc += d * (sc[is + 2] as i8 as f32) * (q2 as f32) * xh[l + 32];
571 acc += d * (sc[is + 4] as i8 as f32) * (q3 as f32) * xh[l + 64];
572 acc += d * (sc[is + 6] as i8 as f32) * (q4 as f32) * xh[l + 96];
573 }
574 x_base += 128;
575 }
576 }
577 acc
578}
579
580pub fn quantize_q8_0(src: &[f32]) -> Vec<u8> {
598 let mut out = Vec::with_capacity(src.len().div_ceil(Q8_0_BLOCK_ELEMS) * Q8_0_BLOCK_BYTES);
599 for chunk in src.chunks(Q8_0_BLOCK_ELEMS) {
600 let mut block = [0f32; Q8_0_BLOCK_ELEMS];
601 block[..chunk.len()].copy_from_slice(chunk);
602 encode::encode_block_q8_0(&block, &mut out);
603 }
604 out
605}
606
607pub fn dot_q8_0_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
621 #[cfg(target_arch = "x86_64")]
622 {
623 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
624 return unsafe { simd_x86::dot_q8_0_f32_avx2(row_bytes, x) };
625 }
626 }
627 #[cfg(target_arch = "aarch64")]
628 {
629 if std::arch::is_aarch64_feature_detected!("neon") {
630 return unsafe { simd_aarch64::dot_q8_0_f32_neon(row_bytes, x) };
631 }
632 }
633 dot_q8_0_f32_scalar(row_bytes, x)
634}
635
636pub fn dot_q8_0_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
637 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
638 debug_assert_eq!(
639 row_bytes.len() / Q8_0_BLOCK_BYTES * Q8_0_BLOCK_ELEMS,
640 x.len()
641 );
642 let mut acc = 0f32;
643 for (b, block) in row_bytes
644 .as_chunks::<Q8_0_BLOCK_BYTES>()
645 .0
646 .iter()
647 .enumerate()
648 {
649 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
650 let base = b * Q8_0_BLOCK_ELEMS;
651 let mut block_acc = 0f32;
652 for i in 0..Q8_0_BLOCK_ELEMS {
653 let q = block[2 + i] as i8;
654 block_acc += (q as f32) * x[base + i];
655 }
656 acc += block_acc * scale;
657 }
658 acc
659}
660
661#[derive(Clone, Debug)]
670pub struct Q8Activations {
671 pub q: Vec<i8>,
673 pub d: Vec<f32>,
675}
676
677impl Q8Activations {
678 pub fn n_blocks(&self) -> usize {
679 self.d.len()
680 }
681}
682
683#[derive(Clone, Debug)]
686pub struct Q8KActivations {
687 pub q: Vec<i8>,
688 pub d: Vec<f32>,
689 pub bsums: Vec<i16>,
691}
692
693impl Q8KActivations {
694 pub fn n_blocks(&self) -> usize {
695 self.d.len()
696 }
697}
698
699pub fn quantize_activations_q8_k(x: &[f32]) -> Q8KActivations {
703 debug_assert_eq!(x.len() % Q4_K_BLOCK_ELEMS, 0);
704 let n_blocks = x.len() / Q4_K_BLOCK_ELEMS;
705 let mut q = vec![0i8; n_blocks * Q4_K_BLOCK_ELEMS];
706 let mut d = vec![0f32; n_blocks];
707 let mut bsums = vec![0i16; n_blocks * 16];
708 let quant_one =
709 |(q_slot, d_slot, bsum_slot, chunk): (&mut [i8], &mut f32, &mut [i16], &[f32])| {
710 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
711 let scale = amax / 127.0;
712 let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
713 *d_slot = scale;
714 for (i, &v) in chunk.iter().enumerate() {
715 let qi = (v * inv).round();
716 q_slot[i] = qi.clamp(-127.0, 127.0) as i8;
717 }
718 for (slot, group) in bsum_slot.iter_mut().zip(q_slot.as_chunks::<16>().0) {
719 *slot = group.iter().map(|&q| q as i32).sum::<i32>() as i16;
720 }
721 };
722 for (b, chunk) in x.as_chunks::<Q4_K_BLOCK_ELEMS>().0.iter().enumerate() {
728 quant_one((
729 &mut q[b * Q4_K_BLOCK_ELEMS..(b + 1) * Q4_K_BLOCK_ELEMS],
730 &mut d[b],
731 &mut bsums[b * 16..(b + 1) * 16],
732 chunk,
733 ));
734 }
735 Q8KActivations { q, d, bsums }
736}
737
738pub fn quantize_activations_q8(x: &[f32]) -> Q8Activations {
742 debug_assert_eq!(x.len() % Q8_0_BLOCK_ELEMS, 0);
743 let n_blocks = x.len() / Q8_0_BLOCK_ELEMS;
744 let mut q = vec![0i8; n_blocks * Q8_0_BLOCK_ELEMS];
745 let mut d = vec![0f32; n_blocks];
746 let quant_one = |(q_slot, d_slot, chunk): (&mut [i8], &mut f32, &[f32])| {
747 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
748 let scale = amax / 127.0;
749 let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
750 *d_slot = scale;
751 for (i, &v) in chunk.iter().enumerate() {
752 let qi = (v * inv).round();
754 q_slot[i] = qi.clamp(-127.0, 127.0) as i8;
755 }
756 };
757 for (b, chunk) in x.as_chunks::<Q8_0_BLOCK_ELEMS>().0.iter().enumerate() {
761 quant_one((
762 &mut q[b * Q8_0_BLOCK_ELEMS..(b + 1) * Q8_0_BLOCK_ELEMS],
763 &mut d[b],
764 chunk,
765 ));
766 }
767 Q8Activations { q, d }
768}
769
770pub fn dot_q8_0_q8(row_bytes: &[u8], act: &Q8Activations) -> f32 {
775 #[cfg(target_arch = "x86_64")]
776 {
777 if is_x86_feature_detected!("avx2") {
778 return unsafe { simd_x86::dot_q8_0_q8_avx2(row_bytes, act) };
779 }
780 }
781 #[cfg(target_arch = "aarch64")]
782 {
783 if std::arch::is_aarch64_feature_detected!("dotprod") {
784 return unsafe { simd_aarch64::dot_q8_0_q8_neon_sdot(row_bytes, act) };
785 }
786 if std::arch::is_aarch64_feature_detected!("neon") {
787 return unsafe { simd_aarch64::dot_q8_0_q8_neon(row_bytes, act) };
788 }
789 }
790 dot_q8_0_q8_scalar(row_bytes, act)
791}
792
793pub fn dot_q8_0_q8_scalar(row_bytes: &[u8], act: &Q8Activations) -> f32 {
794 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
795 let n_blocks = row_bytes.len() / Q8_0_BLOCK_BYTES;
796 debug_assert_eq!(n_blocks, act.n_blocks());
797 let mut acc = 0f32;
798 for (b, block) in row_bytes
799 .as_chunks::<Q8_0_BLOCK_BYTES>()
800 .0
801 .iter()
802 .enumerate()
803 {
804 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
805 let base = b * Q8_0_BLOCK_ELEMS;
806 let mut isum = 0i32;
807 for i in 0..Q8_0_BLOCK_ELEMS {
808 let qw = block[2 + i] as i8 as i32;
809 let qa = act.q[base + i] as i32;
810 isum += qw * qa;
811 }
812 acc += dw * act.d[b] * isum as f32;
813 }
814 acc
815}
816
817pub fn dot_q4_0_q8(row_bytes: &[u8], act: &Q8Activations) -> f32 {
821 #[cfg(target_arch = "x86_64")]
822 {
823 if is_x86_feature_detected!("avx2") {
824 return unsafe { simd_x86::dot_q4_0_q8_avx2(row_bytes, act) };
825 }
826 }
827 #[cfg(target_arch = "aarch64")]
828 {
829 if std::arch::is_aarch64_feature_detected!("dotprod") {
830 return unsafe { simd_aarch64::dot_q4_0_q8_neon_sdot(row_bytes, act) };
831 }
832 if std::arch::is_aarch64_feature_detected!("neon") {
833 return unsafe { simd_aarch64::dot_q4_0_q8_neon(row_bytes, act) };
834 }
835 }
836 dot_q4_0_q8_scalar(row_bytes, act)
837}
838
839pub fn dot_q4_0_q8_2row(row0: &[u8], row1: &[u8], act: &Q8Activations) -> (f32, f32) {
842 #[cfg(target_arch = "aarch64")]
843 {
844 if std::arch::is_aarch64_feature_detected!("dotprod")
845 && row0.len() == row1.len()
846 && row0.len().is_multiple_of(Q4_0_BLOCK_BYTES)
847 {
848 return unsafe { simd_aarch64::dot_q4_0_q8_neon_sdot_2row(row0, row1, act) };
849 }
850 }
851 (dot_q4_0_q8(row0, act), dot_q4_0_q8(row1, act))
852}
853
854pub fn dot_q4_0_q8_scalar(row_bytes: &[u8], act: &Q8Activations) -> f32 {
855 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
856 let n_blocks = row_bytes.len() / Q4_0_BLOCK_BYTES;
857 debug_assert_eq!(n_blocks, act.n_blocks());
858 let mut acc = 0f32;
859 for (b, block) in row_bytes
860 .as_chunks::<Q4_0_BLOCK_BYTES>()
861 .0
862 .iter()
863 .enumerate()
864 {
865 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
866 let base = b * Q4_0_BLOCK_ELEMS;
867 let mut isum = 0i32;
868 for i in 0..16 {
869 let qs = block[2 + i];
870 let q0 = (qs & 0x0F) as i32 - 8;
871 let q1 = (qs >> 4) as i32 - 8;
872 isum += q0 * act.q[base + i] as i32;
873 isum += q1 * act.q[base + 16 + i] as i32;
874 }
875 acc += dw * act.d[b] * isum as f32;
876 }
877 acc
878}
879
880pub fn dot_q4_k_q8(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
883 #[cfg(target_arch = "x86_64")]
884 {
885 if is_x86_feature_detected!("avx2") {
886 return unsafe { simd_x86::dot_q4_k_q8_avx2(row_bytes, act) };
887 }
888 }
889 #[cfg(target_arch = "aarch64")]
890 {
891 if std::arch::is_aarch64_feature_detected!("i8mm") {
892 return unsafe { simd_aarch64::dot_q4_k_q8_neon_i8mm(row_bytes, act) };
893 }
894 if std::arch::is_aarch64_feature_detected!("dotprod") {
895 return unsafe { simd_aarch64::dot_q4_k_q8_neon_sdot(row_bytes, act) };
896 }
897 if std::arch::is_aarch64_feature_detected!("neon") {
898 return unsafe { simd_aarch64::dot_q4_k_q8_neon(row_bytes, act) };
899 }
900 }
901 dot_q4_k_q8_scalar(row_bytes, act)
902}
903
904pub fn dot_q4_k_q8_scalar(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
905 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
906 let n_blocks = row_bytes.len() / Q4_K_BLOCK_BYTES;
907 debug_assert_eq!(n_blocks, act.n_blocks());
908 let mut acc = 0f32;
909 for (b, block) in row_bytes
910 .as_chunks::<Q4_K_BLOCK_BYTES>()
911 .0
912 .iter()
913 .enumerate()
914 {
915 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
916 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
917 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
918 let qs = &block[16..144];
919 let da = act.d[b];
920 let q8 = &act.q[b * Q4_K_BLOCK_ELEMS..(b + 1) * Q4_K_BLOCK_ELEMS];
921 let bsums = &act.bsums[b * 16..(b + 1) * 16];
922
923 let mut sum_min = 0i32;
924 for i in 0..8 {
925 let (_, m) = q4_k_scale_min(i, &scales);
926 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
927 }
928 acc -= dmin * da * sum_min as f32;
929
930 let mut q_off = 0usize;
931 let mut base = 0usize;
932 let mut is = 0usize;
933 for _ in 0..4 {
934 let (sc1, _) = q4_k_scale_min(is, &scales);
935 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
936 let mut isum1 = 0i32;
937 let mut isum2 = 0i32;
938 for l in 0..32 {
939 isum1 += (qs[q_off + l] & 0x0F) as i32 * q8[base + l] as i32;
940 }
941 for l in 0..32 {
942 isum2 += (qs[q_off + l] >> 4) as i32 * q8[base + 32 + l] as i32;
943 }
944 acc += d * da * (sc1 as f32 * isum1 as f32 + sc2 as f32 * isum2 as f32);
945 q_off += 32;
946 base += 64;
947 is += 2;
948 }
949 }
950 acc
951}
952
953pub fn dot_q5_k_q8(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
956 #[cfg(target_arch = "aarch64")]
957 {
958 if std::arch::is_aarch64_feature_detected!("dotprod") {
959 return unsafe { simd_aarch64::dot_q5_k_q8_neon_sdot(row_bytes, act) };
960 }
961 if std::arch::is_aarch64_feature_detected!("neon") {
962 return unsafe { simd_aarch64::dot_q5_k_q8_neon(row_bytes, act) };
963 }
964 }
965 dot_q5_k_q8_scalar(row_bytes, act)
966}
967
968pub fn dot_q5_k_q8_scalar(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
969 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
970 let n_blocks = row_bytes.len() / Q5_K_BLOCK_BYTES;
971 debug_assert_eq!(n_blocks, act.n_blocks());
972 let mut acc = 0f32;
973 for (b, block) in row_bytes
974 .as_chunks::<Q5_K_BLOCK_BYTES>()
975 .0
976 .iter()
977 .enumerate()
978 {
979 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
980 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
981 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
982 let qh = &block[16..48];
983 let qs = &block[48..176];
984 let da = act.d[b];
985 let q8 = &act.q[b * Q5_K_BLOCK_ELEMS..(b + 1) * Q5_K_BLOCK_ELEMS];
986 let bsums = &act.bsums[b * 16..(b + 1) * 16];
987
988 let mut sum_min = 0i32;
989 for i in 0..8 {
990 let (_, m) = q4_k_scale_min(i, &scales);
991 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
992 }
993 acc -= dmin * da * sum_min as f32;
994
995 let mut q_off = 0usize;
996 let mut base = 0usize;
997 let mut is = 0usize;
998 let (mut u1, mut u2) = (1u8, 2u8);
999 for _ in 0..4 {
1000 let (sc1, _) = q4_k_scale_min(is, &scales);
1001 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
1002 let mut isum1 = 0i32;
1003 let mut isum2 = 0i32;
1004 for l in 0..32 {
1005 let hi = if qh[l] & u1 != 0 { 16 } else { 0 };
1006 isum1 += ((qs[q_off + l] & 0x0F) + hi) as i32 * q8[base + l] as i32;
1007 }
1008 for l in 0..32 {
1009 let hi = if qh[l] & u2 != 0 { 16 } else { 0 };
1010 isum2 += ((qs[q_off + l] >> 4) + hi) as i32 * q8[base + 32 + l] as i32;
1011 }
1012 acc += d * da * (sc1 as f32 * isum1 as f32 + sc2 as f32 * isum2 as f32);
1013 q_off += 32;
1014 base += 64;
1015 is += 2;
1016 u1 <<= 2;
1017 u2 <<= 2;
1018 }
1019 }
1020 acc
1021}
1022
1023pub const Q5_K_GEMM_NC: usize = 4;
1027pub const Q6_K_GEMM_NC: usize = 4;
1028
1029pub fn gemm_q5_k_q8_row(row_bytes: &[u8], acts: &[Q8KActivations], out: &mut [f32]) {
1035 assert_eq!(out.len(), acts.len());
1036 if acts.is_empty() {
1037 return;
1038 }
1039 #[cfg(target_arch = "aarch64")]
1040 {
1041 if acts.len() <= Q5_K_GEMM_NC && std::arch::is_aarch64_feature_detected!("dotprod") {
1042 unsafe {
1043 simd_aarch64::gemm_q5_k_q8_neon_sdot(row_bytes, acts, out);
1044 }
1045 return;
1046 }
1047 }
1048 gemm_q5_k_q8_row_scalar(row_bytes, acts, out);
1049}
1050
1051pub fn gemm_q5_k_q8_row_scalar(row_bytes: &[u8], acts: &[Q8KActivations], out: &mut [f32]) {
1052 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
1053 out.fill(0.0);
1054 let n_blocks = row_bytes.len() / Q5_K_BLOCK_BYTES;
1055 for act in acts {
1056 debug_assert_eq!(n_blocks, act.n_blocks());
1057 }
1058 for (b, block) in row_bytes
1059 .as_chunks::<Q5_K_BLOCK_BYTES>()
1060 .0
1061 .iter()
1062 .enumerate()
1063 {
1064 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1065 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
1066 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
1067 let qh = &block[16..48];
1068 let qs = &block[48..176];
1069 let mut mins = [0u8; 8];
1070 let mut sc_only = [0u8; 8];
1071 for i in 0..8 {
1072 let (s, m) = q4_k_scale_min(i, &scales);
1073 sc_only[i] = s;
1074 mins[i] = m;
1075 }
1076 for (j, act) in acts.iter().enumerate() {
1077 let da = act.d[b];
1078 let q8 = &act.q[b * Q5_K_BLOCK_ELEMS..(b + 1) * Q5_K_BLOCK_ELEMS];
1079 let bsums = &act.bsums[b * 16..(b + 1) * 16];
1080 let mut sum_min = 0i32;
1081 for i in 0..8 {
1082 sum_min += mins[i] as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
1083 }
1084 out[j] -= dmin * da * sum_min as f32;
1085
1086 let mut q_off = 0usize;
1087 let mut base = 0usize;
1088 let mut is = 0usize;
1089 let (mut u1, mut u2) = (1u8, 2u8);
1090 for _ in 0..4 {
1091 let sc1 = sc_only[is];
1092 let sc2 = sc_only[is + 1];
1093 let mut isum1 = 0i32;
1094 let mut isum2 = 0i32;
1095 for l in 0..32 {
1096 let hi = if qh[l] & u1 != 0 { 16 } else { 0 };
1097 isum1 += ((qs[q_off + l] & 0x0F) + hi) as i32 * q8[base + l] as i32;
1098 }
1099 for l in 0..32 {
1100 let hi = if qh[l] & u2 != 0 { 16 } else { 0 };
1101 isum2 += ((qs[q_off + l] >> 4) + hi) as i32 * q8[base + 32 + l] as i32;
1102 }
1103 out[j] += d * da * (sc1 as f32 * isum1 as f32 + sc2 as f32 * isum2 as f32);
1104 q_off += 32;
1105 base += 64;
1106 is += 2;
1107 u1 <<= 2;
1108 u2 <<= 2;
1109 }
1110 }
1111 }
1112}
1113
1114pub fn gemm_q6_k_q8_row(row_bytes: &[u8], acts: &[Q8KActivations], out: &mut [f32]) {
1116 assert_eq!(out.len(), acts.len());
1117 if acts.is_empty() {
1118 return;
1119 }
1120 #[cfg(target_arch = "aarch64")]
1121 {
1122 if acts.len() <= Q6_K_GEMM_NC && std::arch::is_aarch64_feature_detected!("dotprod") {
1123 unsafe {
1124 simd_aarch64::gemm_q6_k_q8_neon_sdot(row_bytes, acts, out);
1125 }
1126 return;
1127 }
1128 }
1129 gemm_q6_k_q8_row_scalar(row_bytes, acts, out);
1130}
1131
1132pub fn gemm_q6_k_q8_row_scalar(row_bytes: &[u8], acts: &[Q8KActivations], out: &mut [f32]) {
1133 out.fill(0.0);
1134 for (j, act) in acts.iter().enumerate() {
1135 out[j] = dot_q6_k_q8_scalar(row_bytes, act);
1136 }
1137}
1138
1139pub fn dot_q6_k_q8(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
1142 #[cfg(target_arch = "aarch64")]
1143 {
1144 if std::arch::is_aarch64_feature_detected!("dotprod") {
1145 return unsafe { simd_aarch64::dot_q6_k_q8_neon_sdot(row_bytes, act) };
1146 }
1147 }
1148 dot_q6_k_q8_scalar(row_bytes, act)
1149}
1150
1151pub fn dot_q6_k_q8_scalar(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
1152 debug_assert_eq!(row_bytes.len() % Q6_K_BLOCK_BYTES, 0);
1153 let n_blocks = row_bytes.len() / Q6_K_BLOCK_BYTES;
1154 debug_assert_eq!(n_blocks, act.n_blocks());
1155 debug_assert_eq!(Q6_K_BLOCK_ELEMS, Q4_K_BLOCK_ELEMS);
1157 let mut acc = 0f32;
1158 for (b, block) in row_bytes
1159 .as_chunks::<Q6_K_BLOCK_BYTES>()
1160 .0
1161 .iter()
1162 .enumerate()
1163 {
1164 let ql_full = &block[0..128];
1165 let qh_full = &block[128..192];
1166 let sc_full = &block[192..208];
1167 let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
1168 let da = act.d[b];
1169 let q8 = &act.q[b * Q6_K_BLOCK_ELEMS..(b + 1) * Q6_K_BLOCK_ELEMS];
1170 let mut isum = 0i32;
1171
1172 for half in 0..2 {
1173 let ql = &ql_full[half * 64..half * 64 + 64];
1174 let qh = &qh_full[half * 32..half * 32 + 32];
1175 let sc = &sc_full[half * 8..half * 8 + 8];
1176 let q8h = &q8[half * 128..half * 128 + 128];
1177 for l in 0..32 {
1178 let is = l / 16;
1179 let q1 = ((ql[l] & 0x0F) | ((qh[l] & 3) << 4)) as i8 as i32 - 32;
1180 let q2 = ((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) as i8 as i32 - 32;
1181 let q3 = ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) as i8 as i32 - 32;
1182 let q4 = ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) as i8 as i32 - 32;
1183 isum += (sc[is] as i8 as i32) * q1 * (q8h[l] as i32);
1184 isum += (sc[is + 2] as i8 as i32) * q2 * (q8h[l + 32] as i32);
1185 isum += (sc[is + 4] as i8 as i32) * q3 * (q8h[l + 64] as i32);
1186 isum += (sc[is + 6] as i8 as i32) * q4 * (q8h[l + 96] as i32);
1187 }
1188 }
1189 acc += d * da * isum as f32;
1190 }
1191 acc
1192}
1193
1194#[cfg(target_arch = "x86_64")]
1195mod simd_x86 {
1196 use super::{
1197 e8m0_scale, q3_k_unpack_scales, q4_k_scale_min, q5_fifth_bits, Q8Activations,
1198 Q8KActivations, IQ4_NL_BLOCK_BYTES, IQ4_NL_BLOCK_ELEMS, IQ4_XS_BLOCK_BYTES, KVALUES_IQ4NL,
1199 MXFP4_GROUP_SIZE, Q2_K_BLOCK_BYTES, Q2_K_SCALE_BYTES, Q3_K_BLOCK_BYTES, Q3_K_SCALE_BYTES,
1200 Q4_0_BLOCK_BYTES, Q4_0_BLOCK_ELEMS, Q4_1_BLOCK_BYTES, Q4_1_BLOCK_ELEMS, Q4_K_BLOCK_BYTES,
1201 Q4_K_BLOCK_ELEMS, Q4_K_SCALE_BYTES, Q5_0_BLOCK_BYTES, Q5_0_BLOCK_ELEMS, Q5_1_BLOCK_BYTES,
1202 Q5_1_BLOCK_ELEMS, Q5_K_BLOCK_BYTES, Q6_K_BLOCK_BYTES, Q6_K_BLOCK_ELEMS, Q8_0_BLOCK_BYTES,
1203 Q8_0_BLOCK_ELEMS, Q8_1_BLOCK_BYTES, Q8_1_BLOCK_ELEMS,
1204 };
1205 use half::f16;
1206 use std::arch::x86_64::*;
1207
1208 #[target_feature(enable = "avx2,fma")]
1218 pub unsafe fn dot_q8_0_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1219 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
1220 debug_assert_eq!(
1221 row_bytes.len() / Q8_0_BLOCK_BYTES * Q8_0_BLOCK_ELEMS,
1222 x.len()
1223 );
1224 let mut acc = 0f32;
1225 for (b, block) in row_bytes
1226 .as_chunks::<Q8_0_BLOCK_BYTES>()
1227 .0
1228 .iter()
1229 .enumerate()
1230 {
1231 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
1232 let base = b * Q8_0_BLOCK_ELEMS;
1233 let qs = &block[2..34];
1234
1235 let mut block_acc = _mm256_setzero_ps();
1236 for g in 0..4 {
1237 let raw8 = _mm_loadl_epi64(qs.as_ptr().add(g * 8) as *const __m128i);
1238 let i32x8 = _mm256_cvtepi8_epi32(raw8);
1239 let f32x8 = _mm256_cvtepi32_ps(i32x8);
1240 let xv = _mm256_loadu_ps(x.as_ptr().add(base + g * 8));
1241 block_acc = _mm256_fmadd_ps(f32x8, xv, block_acc);
1242 }
1243 acc += hsum256_ps(block_acc) * scale;
1244 }
1245 acc
1246 }
1247
1248 #[target_feature(enable = "avx2")]
1254 pub unsafe fn dot_q8_0_q8_avx2(row_bytes: &[u8], act: &Q8Activations) -> f32 {
1255 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
1256 debug_assert_eq!(row_bytes.len() / Q8_0_BLOCK_BYTES, act.n_blocks());
1257 let mut acc = 0f32;
1258 for (b, block) in row_bytes
1259 .as_chunks::<Q8_0_BLOCK_BYTES>()
1260 .0
1261 .iter()
1262 .enumerate()
1263 {
1264 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
1265 let base = b * Q8_0_BLOCK_ELEMS;
1266 let w = _mm256_loadu_si256(block.as_ptr().add(2) as *const __m256i);
1267 let a = _mm256_loadu_si256(act.q.as_ptr().add(base) as *const __m256i);
1268 let w_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(w));
1269 let w_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(w, 1));
1270 let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(a));
1271 let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(a, 1));
1272 let prod =
1273 _mm256_add_epi32(_mm256_madd_epi16(w_lo, a_lo), _mm256_madd_epi16(w_hi, a_hi));
1274 let hi128 = _mm256_extracti128_si256(prod, 1);
1276 let lo128 = _mm256_castsi256_si128(prod);
1277 let mut sum128 = _mm_add_epi32(lo128, hi128);
1278 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b01_00_11_10));
1279 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b00_00_00_01));
1280 let isum = _mm_cvtsi128_si32(sum128);
1281 acc += dw * act.d[b] * isum as f32;
1282 }
1283 acc
1284 }
1285
1286 #[target_feature(enable = "avx2")]
1290 pub unsafe fn dot_q4_0_q8_avx2(row_bytes: &[u8], act: &Q8Activations) -> f32 {
1291 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
1292 debug_assert_eq!(row_bytes.len() / Q4_0_BLOCK_BYTES, act.n_blocks());
1293 let low_mask = _mm_set1_epi8(0x0F);
1294 let bias = _mm_set1_epi8(8);
1295 let mut acc = 0f32;
1296 for (b, block) in row_bytes
1297 .as_chunks::<Q4_0_BLOCK_BYTES>()
1298 .0
1299 .iter()
1300 .enumerate()
1301 {
1302 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
1303 let base = b * Q4_0_BLOCK_ELEMS;
1304 let qs = _mm_loadu_si128(block.as_ptr().add(2) as *const __m128i);
1305 let lo = _mm_sub_epi8(_mm_and_si128(qs, low_mask), bias);
1306 let hi = _mm_sub_epi8(_mm_and_si128(_mm_srli_epi16(qs, 4), low_mask), bias);
1307 let w = _mm256_set_m128i(hi, lo);
1309 let a = _mm256_loadu_si256(act.q.as_ptr().add(base) as *const __m256i);
1310 let w_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(w));
1311 let w_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(w, 1));
1312 let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(a));
1313 let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(a, 1));
1314 let prod =
1315 _mm256_add_epi32(_mm256_madd_epi16(w_lo, a_lo), _mm256_madd_epi16(w_hi, a_hi));
1316 let hi128 = _mm256_extracti128_si256(prod, 1);
1317 let lo128 = _mm256_castsi256_si128(prod);
1318 let mut sum128 = _mm_add_epi32(lo128, hi128);
1319 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b01_00_11_10));
1320 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b00_00_00_01));
1321 let isum = _mm_cvtsi128_si32(sum128);
1322 acc += dw * act.d[b] * isum as f32;
1323 }
1324 acc
1325 }
1326
1327 #[target_feature(enable = "avx2")]
1329 pub unsafe fn dot_q4_k_q8_avx2(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
1330 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
1331 debug_assert_eq!(row_bytes.len() / Q4_K_BLOCK_BYTES, act.n_blocks());
1332 let low_mask = _mm256_set1_epi8(0x0F_u8 as i8);
1333 let mut acc = 0f32;
1334 for (b, block) in row_bytes
1335 .as_chunks::<Q4_K_BLOCK_BYTES>()
1336 .0
1337 .iter()
1338 .enumerate()
1339 {
1340 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1341 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
1342 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
1343 let qs = &block[16..144];
1344 let da = act.d[b];
1345 let q8 = act.q.as_ptr().add(b * Q4_K_BLOCK_ELEMS);
1346 let bsums = &act.bsums[b * 16..(b + 1) * 16];
1347
1348 let mut sum_min = 0i32;
1349 for i in 0..8 {
1350 let (_, m) = q4_k_scale_min(i, &scales);
1351 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
1352 }
1353 acc -= dmin * da * sum_min as f32;
1354
1355 let mut q_off = 0usize;
1356 let mut base = 0usize;
1357 let mut is = 0usize;
1358 for _ in 0..4 {
1359 let (sc1, _) = q4_k_scale_min(is, &scales);
1360 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
1361 let packed = _mm256_loadu_si256(qs.as_ptr().add(q_off) as *const __m256i);
1362 let lo = _mm256_and_si256(packed, low_mask);
1363 let hi = _mm256_and_si256(_mm256_srli_epi16(packed, 4), low_mask);
1364 let a0 = _mm256_loadu_si256(q8.add(base) as *const __m256i);
1365 let a1 = _mm256_loadu_si256(q8.add(base + 32) as *const __m256i);
1366 let isum1 = madd_i8_avx2(lo, a0);
1367 let isum2 = madd_i8_avx2(hi, a1);
1368 acc += d * da * (sc1 as f32 * isum1 as f32 + sc2 as f32 * isum2 as f32);
1369 q_off += 32;
1370 base += 64;
1371 is += 2;
1372 }
1373 }
1374 acc
1375 }
1376
1377 #[target_feature(enable = "avx2")]
1378 unsafe fn madd_i8_avx2(w: __m256i, a: __m256i) -> i32 {
1379 let w_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(w));
1380 let w_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(w, 1));
1381 let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(a));
1382 let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(a, 1));
1383 let prod = _mm256_add_epi32(_mm256_madd_epi16(w_lo, a_lo), _mm256_madd_epi16(w_hi, a_hi));
1384 let hi128 = _mm256_extracti128_si256(prod, 1);
1385 let lo128 = _mm256_castsi256_si128(prod);
1386 let mut sum128 = _mm_add_epi32(lo128, hi128);
1387 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b01_00_11_10));
1388 sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b00_00_00_01));
1389 _mm_cvtsi128_si32(sum128)
1390 }
1391
1392 #[target_feature(enable = "avx2,fma")]
1401 pub unsafe fn dot_q4_0_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1402 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
1403 let bias = _mm_set1_epi8(8);
1404 let low_mask = _mm_set1_epi8(0x0F);
1405
1406 let mut acc = 0f32;
1407 for (b, block) in row_bytes
1408 .as_chunks::<Q4_0_BLOCK_BYTES>()
1409 .0
1410 .iter()
1411 .enumerate()
1412 {
1413 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
1414 let base = b * Q4_0_BLOCK_ELEMS;
1415 let nibbles = _mm_loadu_si128(block.as_ptr().add(2) as *const __m128i);
1416
1417 let lo_nibbles = _mm_sub_epi8(_mm_and_si128(nibbles, low_mask), bias);
1418 let hi_nibbles =
1419 _mm_sub_epi8(_mm_and_si128(_mm_srli_epi16(nibbles, 4), low_mask), bias);
1420
1421 let mut block_acc = _mm256_setzero_ps();
1422 for (group_idx, half) in [
1424 (0usize, lo_nibbles),
1425 (1usize, _mm_srli_si128(lo_nibbles, 8)),
1426 (2usize, hi_nibbles),
1427 (3usize, _mm_srli_si128(hi_nibbles, 8)),
1428 ] {
1429 let i32x8 = _mm256_cvtepi8_epi32(half);
1430 let f32x8 = _mm256_cvtepi32_ps(i32x8);
1431 let elem_base = base + group_idx * 8;
1432 let xv = _mm256_loadu_ps(x.as_ptr().add(elem_base));
1433 block_acc = _mm256_fmadd_ps(f32x8, xv, block_acc);
1434 }
1435 acc += hsum256_ps(block_acc) * scale;
1436 }
1437 acc
1438 }
1439
1440 #[inline]
1441 #[target_feature(enable = "avx2")]
1442 unsafe fn hsum256_ps(v: __m256) -> f32 {
1443 let hi = _mm256_extractf128_ps(v, 1);
1444 let lo = _mm256_castps256_ps128(v);
1445 let sum128 = _mm_add_ps(hi, lo);
1446 let shuf = _mm_movehdup_ps(sum128);
1447 let sums = _mm_add_ps(sum128, shuf);
1448 let shuf2 = _mm_movehl_ps(shuf, sums);
1449 let sums2 = _mm_add_ss(sums, shuf2);
1450 _mm_cvtss_f32(sums2)
1451 }
1452
1453 #[inline]
1467 #[target_feature(enable = "avx2,fma")]
1468 unsafe fn fma_affine8(
1469 part: __m128i,
1470 d: f32,
1471 min: f32,
1472 x: &[f32],
1473 x_base: usize,
1474 acc: __m256,
1475 ) -> __m256 {
1476 let i32x8 = _mm256_cvtepu8_epi32(part);
1477 let f32x8 = _mm256_cvtepi32_ps(i32x8);
1478 let weight = _mm256_fmsub_ps(f32x8, _mm256_set1_ps(d), _mm256_set1_ps(min));
1479 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base));
1480 _mm256_fmadd_ps(weight, xv, acc)
1481 }
1482
1483 #[target_feature(enable = "avx2,fma")]
1494 pub unsafe fn dot_q4_k_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1495 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
1496 let low_mask = _mm_set1_epi8(0x0F);
1497 let mut acc = 0f32;
1498 let mut x_base = 0usize;
1499 for block in row_bytes.as_chunks::<Q4_K_BLOCK_BYTES>().0 {
1500 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1501 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
1502 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
1503 let qs = &block[16..144];
1504
1505 let mut is = 0usize;
1506 let mut q_off = 0usize;
1507 for _ in 0..4 {
1508 let (sc1, m1) = q4_k_scale_min(is, &scales);
1509 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
1510 let d1 = d * sc1 as f32;
1511 let min1 = dmin * m1 as f32;
1512 let d2 = d * sc2 as f32;
1513 let min2 = dmin * m2 as f32;
1514
1515 let mut lo_acc = _mm256_setzero_ps();
1516 let mut hi_acc = _mm256_setzero_ps();
1517 for g in 0..2 {
1518 let raw16 = _mm_loadu_si128(qs.as_ptr().add(q_off + g * 16) as *const __m128i);
1519 let lo_nib = _mm_and_si128(raw16, low_mask);
1520 let hi_nib = _mm_and_si128(_mm_srli_epi16(raw16, 4), low_mask);
1521
1522 for (part_idx, part) in
1523 [lo_nib, _mm_srli_si128(lo_nib, 8)].into_iter().enumerate()
1524 {
1525 lo_acc =
1526 fma_affine8(part, d1, min1, x, x_base + g * 16 + part_idx * 8, lo_acc);
1527 }
1528 for (part_idx, part) in
1529 [hi_nib, _mm_srli_si128(hi_nib, 8)].into_iter().enumerate()
1530 {
1531 hi_acc = fma_affine8(
1532 part,
1533 d2,
1534 min2,
1535 x,
1536 x_base + 32 + g * 16 + part_idx * 8,
1537 hi_acc,
1538 );
1539 }
1540 }
1541 acc += hsum256_ps(lo_acc) + hsum256_ps(hi_acc);
1542 q_off += 32;
1543 x_base += 64;
1544 is += 2;
1545 }
1546 }
1547 acc
1548 }
1549
1550 #[target_feature(enable = "avx2,fma")]
1563 pub unsafe fn dot_q5_k_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1564 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
1565 let low_mask = _mm_set1_epi8(0x0F);
1566 let zero = _mm_setzero_si128();
1567 let sixteen = _mm_set1_epi8(16);
1568 let mut acc = 0f32;
1569 let mut x_base = 0usize;
1570 for block in row_bytes.as_chunks::<Q5_K_BLOCK_BYTES>().0 {
1571 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1572 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
1573 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
1574 let qh = &block[16..48];
1575 let qs = &block[48..176];
1576
1577 let mut is = 0usize;
1578 let (mut u1, mut u2) = (1u8, 2u8);
1579 for _oi in 0..4 {
1580 let (sc1, m1) = q4_k_scale_min(is, &scales);
1581 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
1582 let d1 = d * sc1 as f32;
1583 let min1 = dmin * m1 as f32;
1584 let d2 = d * sc2 as f32;
1585 let min2 = dmin * m2 as f32;
1586 let ql = &qs[is / 2 * 32..is / 2 * 32 + 32];
1587 let u1_vec = _mm_set1_epi8(u1 as i8);
1588 let u2_vec = _mm_set1_epi8(u2 as i8);
1589
1590 let mut lo_acc = _mm256_setzero_ps();
1591 let mut hi_acc = _mm256_setzero_ps();
1592 for g in 0..2 {
1593 let raw16 = _mm_loadu_si128(ql.as_ptr().add(g * 16) as *const __m128i);
1594 let qh16 = _mm_loadu_si128(qh.as_ptr().add(g * 16) as *const __m128i);
1595
1596 let lo_nib = _mm_and_si128(raw16, low_mask);
1597 let hi_nib = _mm_and_si128(_mm_srli_epi16(raw16, 4), low_mask);
1598
1599 let is_zero1 = _mm_cmpeq_epi8(_mm_and_si128(qh16, u1_vec), zero);
1600 let hi_bit1 = _mm_andnot_si128(is_zero1, sixteen);
1601 let is_zero2 = _mm_cmpeq_epi8(_mm_and_si128(qh16, u2_vec), zero);
1602 let hi_bit2 = _mm_andnot_si128(is_zero2, sixteen);
1603
1604 let lo_full = _mm_or_si128(lo_nib, hi_bit1);
1605 let hi_full = _mm_or_si128(hi_nib, hi_bit2);
1606
1607 for (part_idx, part) in [lo_full, _mm_srli_si128(lo_full, 8)]
1608 .into_iter()
1609 .enumerate()
1610 {
1611 lo_acc =
1612 fma_affine8(part, d1, min1, x, x_base + g * 16 + part_idx * 8, lo_acc);
1613 }
1614 for (part_idx, part) in [hi_full, _mm_srli_si128(hi_full, 8)]
1615 .into_iter()
1616 .enumerate()
1617 {
1618 hi_acc = fma_affine8(
1619 part,
1620 d2,
1621 min2,
1622 x,
1623 x_base + 32 + g * 16 + part_idx * 8,
1624 hi_acc,
1625 );
1626 }
1627 }
1628 acc += hsum256_ps(lo_acc) + hsum256_ps(hi_acc);
1629 x_base += 64;
1630 is += 2;
1631 u1 <<= 2;
1632 u2 <<= 2;
1633 }
1634 }
1635 acc
1636 }
1637
1638 #[inline]
1659 #[target_feature(enable = "avx2,fma")]
1660 #[allow(clippy::too_many_arguments)]
1661 unsafe fn q6_k_group_avx2<const QH_SHIFT: i32, const HI_NIBBLE: bool>(
1662 ql: &[u8],
1663 ql_off: usize,
1664 qh: &[u8],
1665 sc: &[u8],
1666 sc_base: usize,
1667 d: f32,
1668 x: &[f32],
1669 x_base: usize,
1670 out_off: usize,
1671 low_mask: __m128i,
1672 two_bit_mask: __m128i,
1673 bias: __m256,
1674 ) -> f32 {
1675 let mut acc = 0f32;
1676 for sub in 0..2usize {
1677 let byte_off = sub * 16;
1678 let ql_raw = _mm_loadu_si128(ql.as_ptr().add(ql_off + byte_off) as *const __m128i);
1679 let qh_raw = _mm_loadu_si128(qh.as_ptr().add(byte_off) as *const __m128i);
1680
1681 let nib = if HI_NIBBLE {
1682 _mm_and_si128(_mm_srli_epi16(ql_raw, 4), low_mask)
1683 } else {
1684 _mm_and_si128(ql_raw, low_mask)
1685 };
1686 let qh_field = _mm_and_si128(_mm_srli_epi16(qh_raw, QH_SHIFT), two_bit_mask);
1687 let raw6 = _mm_or_si128(nib, _mm_slli_epi16(qh_field, 4));
1688
1689 let scale = d * (sc[sc_base + sub] as i8) as f32;
1690 let elem_base = x_base + out_off + sub * 16;
1691 for (part_idx, part) in [raw6, _mm_srli_si128(raw6, 8)].into_iter().enumerate() {
1692 let i32x8 = _mm256_cvtepu8_epi32(part);
1693 let f32x8 = _mm256_sub_ps(_mm256_cvtepi32_ps(i32x8), bias);
1694 let xv = _mm256_loadu_ps(x.as_ptr().add(elem_base + part_idx * 8));
1695 let weighted = _mm256_mul_ps(f32x8, _mm256_set1_ps(scale));
1696 acc += hsum256_ps(_mm256_mul_ps(weighted, xv));
1697 }
1698 }
1699 acc
1700 }
1701
1702 #[target_feature(enable = "avx2,fma")]
1708 pub unsafe fn dot_q6_k_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1709 debug_assert_eq!(row_bytes.len() % Q6_K_BLOCK_BYTES, 0);
1710 debug_assert_eq!(
1711 row_bytes.len() / Q6_K_BLOCK_BYTES * Q6_K_BLOCK_ELEMS,
1712 x.len()
1713 );
1714 let low_mask = _mm_set1_epi8(0x0F);
1715 let two_bit_mask = _mm_set1_epi8(0x03);
1716 let bias = _mm256_set1_ps(32.0);
1717
1718 let mut acc = 0f32;
1719 let mut x_base = 0usize;
1720 for block in row_bytes.as_chunks::<Q6_K_BLOCK_BYTES>().0 {
1721 let ql_full = &block[0..128];
1722 let qh_full = &block[128..192];
1723 let sc_full = &block[192..208];
1724 let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
1725
1726 for half in 0..2 {
1727 let ql = &ql_full[half * 64..half * 64 + 64];
1728 let qh = &qh_full[half * 32..half * 32 + 32];
1729 let sc = &sc_full[half * 8..half * 8 + 8];
1730 let half_base = x_base + half * 128;
1731
1732 acc += q6_k_group_avx2::<0, false>(
1733 ql,
1734 0,
1735 qh,
1736 sc,
1737 0,
1738 d,
1739 x,
1740 half_base,
1741 0,
1742 low_mask,
1743 two_bit_mask,
1744 bias,
1745 );
1746 acc += q6_k_group_avx2::<2, false>(
1747 ql,
1748 32,
1749 qh,
1750 sc,
1751 2,
1752 d,
1753 x,
1754 half_base,
1755 32,
1756 low_mask,
1757 two_bit_mask,
1758 bias,
1759 );
1760 acc += q6_k_group_avx2::<4, true>(
1761 ql,
1762 0,
1763 qh,
1764 sc,
1765 4,
1766 d,
1767 x,
1768 half_base,
1769 64,
1770 low_mask,
1771 two_bit_mask,
1772 bias,
1773 );
1774 acc += q6_k_group_avx2::<6, true>(
1775 ql,
1776 32,
1777 qh,
1778 sc,
1779 6,
1780 d,
1781 x,
1782 half_base,
1783 96,
1784 low_mask,
1785 two_bit_mask,
1786 bias,
1787 );
1788 }
1789 x_base += Q6_K_BLOCK_ELEMS;
1790 }
1791 acc
1792 }
1793
1794 #[inline]
1803 #[target_feature(enable = "avx2,fma")]
1804 unsafe fn mxfp4_nibbles_to_f32x8(nib: __m128i) -> __m256 {
1805 let sign_bit = _mm_and_si128(nib, _mm_set1_epi8(0x8));
1806 let e = _mm_and_si128(_mm_srli_epi16(nib, 1), _mm_set1_epi8(0x3));
1807 let m = _mm_and_si128(nib, _mm_set1_epi8(0x1));
1808
1809 let pow2_table = _mm_setr_epi8(1, 1, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
1810 let bias_table = _mm_setr_epi8(0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
1811 let pow2_u8 = _mm_shuffle_epi8(pow2_table, e);
1812 let bias_u8 = _mm_shuffle_epi8(bias_table, e);
1813
1814 let pow2_f = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(pow2_u8));
1815 let bias_f = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(bias_u8));
1816 let m_f = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(m));
1817 let sign_f = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(sign_bit));
1818
1819 let magnitude = _mm256_mul_ps(pow2_f, _mm256_fmadd_ps(m_f, _mm256_set1_ps(0.5), bias_f));
1821 let sign_mul = _mm256_fnmadd_ps(sign_f, _mm256_set1_ps(0.25), _mm256_set1_ps(1.0));
1822 _mm256_mul_ps(magnitude, sign_mul)
1823 }
1824
1825 pub unsafe fn dot_mxfp4_row_f32_avx2(packed: &[u8], scales: &[u8], x: &[f32]) -> f32 {
1835 debug_assert_eq!(packed.len(), scales.len() * (MXFP4_GROUP_SIZE / 2));
1836 let low_mask = _mm_set1_epi8(0x0F);
1837 let mut acc = 0f32;
1838 let mut x_base = 0usize;
1839 for (g, &e_byte) in scales.iter().enumerate() {
1840 let d = e8m0_scale(e_byte);
1841 let group = &packed[g * 16..(g + 1) * 16];
1842 let bytes = _mm_loadu_si128(group.as_ptr() as *const __m128i);
1843 let lo_nib = _mm_and_si128(bytes, low_mask);
1844 let hi_nib = _mm_and_si128(_mm_srli_epi16(bytes, 4), low_mask);
1845
1846 let mut block_acc = _mm256_setzero_ps();
1847 for (half_idx, nib) in [
1848 (0usize, lo_nib),
1849 (1usize, _mm_srli_si128(lo_nib, 8)),
1850 (2usize, hi_nib),
1851 (3usize, _mm_srli_si128(hi_nib, 8)),
1852 ] {
1853 let vals = mxfp4_nibbles_to_f32x8(nib);
1854 let elem_base = x_base + half_idx * 8;
1855 let xv = _mm256_loadu_ps(x.as_ptr().add(elem_base));
1856 block_acc = _mm256_fmadd_ps(vals, xv, block_acc);
1857 }
1858 acc += hsum256_ps(block_acc) * d;
1859 x_base += MXFP4_GROUP_SIZE;
1860 }
1861 acc
1862 }
1863
1864 #[target_feature(enable = "avx2,fma")]
1870 pub unsafe fn dot_q8_1_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1871 debug_assert_eq!(row_bytes.len() % Q8_1_BLOCK_BYTES, 0);
1872 let mut acc = 0f32;
1873 for (b, block) in row_bytes
1874 .as_chunks::<Q8_1_BLOCK_BYTES>()
1875 .0
1876 .iter()
1877 .enumerate()
1878 {
1879 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1880 let base = b * Q8_1_BLOCK_ELEMS;
1881 let qs = &block[4..36];
1882
1883 let mut block_acc = _mm256_setzero_ps();
1884 for g in 0..4 {
1885 let raw8 = _mm_loadl_epi64(qs.as_ptr().add(g * 8) as *const __m128i);
1886 let i32x8 = _mm256_cvtepi8_epi32(raw8);
1887 let f32x8 = _mm256_cvtepi32_ps(i32x8);
1888 let xv = _mm256_loadu_ps(x.as_ptr().add(base + g * 8));
1889 block_acc = _mm256_fmadd_ps(f32x8, xv, block_acc);
1890 }
1891 acc += hsum256_ps(block_acc) * d;
1892 }
1893 acc
1894 }
1895
1896 #[target_feature(enable = "avx2,fma")]
1902 pub unsafe fn dot_q4_1_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1903 debug_assert_eq!(row_bytes.len() % Q4_1_BLOCK_BYTES, 0);
1904 let low_mask = _mm_set1_epi8(0x0F);
1905 let mut acc = 0f32;
1906 for (b, block) in row_bytes
1907 .as_chunks::<Q4_1_BLOCK_BYTES>()
1908 .0
1909 .iter()
1910 .enumerate()
1911 {
1912 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1913 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
1914 let base = b * Q4_1_BLOCK_ELEMS;
1915 let nibbles = _mm_loadu_si128(block.as_ptr().add(4) as *const __m128i);
1916
1917 let lo_nibbles = _mm_and_si128(nibbles, low_mask);
1918 let hi_nibbles = _mm_and_si128(_mm_srli_epi16(nibbles, 4), low_mask);
1919
1920 let mut lo_acc = _mm256_setzero_ps();
1921 let mut hi_acc = _mm256_setzero_ps();
1922 for (part_idx, part) in [lo_nibbles, _mm_srli_si128(lo_nibbles, 8)]
1923 .into_iter()
1924 .enumerate()
1925 {
1926 lo_acc = fma_affine8(part, d, -m, x, base + part_idx * 8, lo_acc);
1927 }
1928 for (part_idx, part) in [hi_nibbles, _mm_srli_si128(hi_nibbles, 8)]
1929 .into_iter()
1930 .enumerate()
1931 {
1932 hi_acc = fma_affine8(part, d, -m, x, base + 16 + part_idx * 8, hi_acc);
1933 }
1934 acc += hsum256_ps(lo_acc) + hsum256_ps(hi_acc);
1935 }
1936 acc
1937 }
1938
1939 #[target_feature(enable = "avx2,fma")]
1951 pub unsafe fn dot_q5_0_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1952 debug_assert_eq!(row_bytes.len() % Q5_0_BLOCK_BYTES, 0);
1953 let mut acc = 0f32;
1954 for (b, block) in row_bytes
1955 .as_chunks::<Q5_0_BLOCK_BYTES>()
1956 .0
1957 .iter()
1958 .enumerate()
1959 {
1960 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
1961 let qh = u32::from_le_bytes(block[2..6].try_into().unwrap());
1962 let qs = &block[6..22];
1963 let base = b * Q5_0_BLOCK_ELEMS;
1964
1965 let mut vals = [0i8; 32];
1966 for j in 0..16 {
1967 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
1968 vals[j] = (((qs[j] & 0x0F) | xh_0) as i32 - 16) as i8;
1969 vals[j + 16] = (((qs[j] >> 4) | xh_1) as i32 - 16) as i8;
1970 }
1971
1972 let mut block_acc = _mm256_setzero_ps();
1973 for g in 0..4 {
1974 let raw8 = _mm_loadl_epi64(vals.as_ptr().add(g * 8) as *const __m128i);
1975 let i32x8 = _mm256_cvtepi8_epi32(raw8);
1976 let f32x8 = _mm256_cvtepi32_ps(i32x8);
1977 let xv = _mm256_loadu_ps(x.as_ptr().add(base + g * 8));
1978 block_acc = _mm256_fmadd_ps(f32x8, xv, block_acc);
1979 }
1980 acc += hsum256_ps(block_acc) * d;
1981 }
1982 acc
1983 }
1984
1985 #[target_feature(enable = "avx2,fma")]
1991 pub unsafe fn dot_q5_1_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
1992 debug_assert_eq!(row_bytes.len() % Q5_1_BLOCK_BYTES, 0);
1993 let mut acc = 0f32;
1994 for (b, block) in row_bytes
1995 .as_chunks::<Q5_1_BLOCK_BYTES>()
1996 .0
1997 .iter()
1998 .enumerate()
1999 {
2000 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2001 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
2002 let qh = u32::from_le_bytes(block[4..8].try_into().unwrap());
2003 let qs = &block[8..24];
2004 let base = b * Q5_1_BLOCK_ELEMS;
2005
2006 let mut vals = [0u8; 32];
2007 for j in 0..16 {
2008 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
2009 vals[j] = (qs[j] & 0x0F) | xh_0;
2010 vals[j + 16] = (qs[j] >> 4) | xh_1;
2011 }
2012
2013 let mut block_acc = _mm256_setzero_ps();
2014 for g in 0..4 {
2015 let raw8 = _mm_loadl_epi64(vals.as_ptr().add(g * 8) as *const __m128i);
2016 let i32x8 = _mm256_cvtepu8_epi32(raw8);
2017 let f32x8 = _mm256_cvtepi32_ps(i32x8);
2018 let weight = _mm256_fmadd_ps(f32x8, _mm256_set1_ps(d), _mm256_set1_ps(m));
2019 let xv = _mm256_loadu_ps(x.as_ptr().add(base + g * 8));
2020 block_acc = _mm256_fmadd_ps(weight, xv, block_acc);
2021 }
2022 acc += hsum256_ps(block_acc);
2023 }
2024 acc
2025 }
2026
2027 #[target_feature(enable = "avx2,fma")]
2043 pub unsafe fn dot_q2_k_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2044 debug_assert_eq!(row_bytes.len() % Q2_K_BLOCK_BYTES, 0);
2045 let two_bit_mask = _mm_set1_epi8(3);
2046 let mut acc = 0f32;
2047 let mut x_base = 0usize;
2048
2049 macro_rules! q2_k_sub_block {
2050 ($shift:literal, $q:expr, $scales:expr, $is:expr, $d:expr, $dmin:expr, $x:expr, $x_base:expr, $acc:expr) => {{
2051 let sc1 = $scales[$is];
2052 $is += 1;
2053 let dl1 = $d * (sc1 & 0x0F) as f32;
2054 let ml1 = $dmin * (sc1 >> 4) as f32;
2055 let sc2 = $scales[$is];
2056 $is += 1;
2057 let dl2 = $d * (sc2 & 0x0F) as f32;
2058 let ml2 = $dmin * (sc2 >> 4) as f32;
2059
2060 let lo16 = _mm_loadu_si128($q.as_ptr() as *const __m128i);
2061 let hi16 = _mm_loadu_si128($q.as_ptr().add(16) as *const __m128i);
2062 let lo2 = _mm_and_si128(_mm_srli_epi16(lo16, $shift), two_bit_mask);
2063 let hi2 = _mm_and_si128(_mm_srli_epi16(hi16, $shift), two_bit_mask);
2064
2065 let mut lo_acc = _mm256_setzero_ps();
2066 let mut hi_acc = _mm256_setzero_ps();
2067 for (part_idx, part) in [lo2, _mm_srli_si128(lo2, 8)].into_iter().enumerate() {
2068 lo_acc = fma_affine8(part, dl1, ml1, $x, $x_base + part_idx * 8, lo_acc);
2069 }
2070 for (part_idx, part) in [hi2, _mm_srli_si128(hi2, 8)].into_iter().enumerate() {
2071 hi_acc = fma_affine8(part, dl2, ml2, $x, $x_base + 16 + part_idx * 8, hi_acc);
2072 }
2073 $acc += hsum256_ps(lo_acc) + hsum256_ps(hi_acc);
2074 $x_base += 32;
2075 }};
2076 }
2077
2078 for block in row_bytes.as_chunks::<Q2_K_BLOCK_BYTES>().0 {
2079 let scales: &[u8; Q2_K_SCALE_BYTES] = block[0..16].try_into().unwrap();
2080 let qs = &block[16..80];
2081 let d = f16::from_le_bytes([block[80], block[81]]).to_f32();
2082 let dmin = f16::from_le_bytes([block[82], block[83]]).to_f32();
2083
2084 let mut is = 0usize;
2085 for n in 0..2 {
2086 let q = &qs[n * 32..n * 32 + 32];
2087 q2_k_sub_block!(0, q, scales, is, d, dmin, x, x_base, acc);
2088 q2_k_sub_block!(2, q, scales, is, d, dmin, x, x_base, acc);
2089 q2_k_sub_block!(4, q, scales, is, d, dmin, x, x_base, acc);
2090 q2_k_sub_block!(6, q, scales, is, d, dmin, x, x_base, acc);
2091 }
2092 }
2093 acc
2094 }
2095
2096 #[target_feature(enable = "avx2,fma")]
2111 pub unsafe fn dot_q3_k_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2112 debug_assert_eq!(row_bytes.len() % Q3_K_BLOCK_BYTES, 0);
2113 let two_bit_mask = _mm_set1_epi8(3);
2114 let zero = _mm_setzero_si128();
2115 let four = _mm_set1_epi8(4);
2116 let mut acc = 0f32;
2117 let mut x_base = 0usize;
2118
2119 macro_rules! q3_k_sub_block {
2120 ($shift:literal, $q:expr, $hmask:expr, $m_vec:expr, $dl1:expr, $dl2:expr, $x:expr, $x_base:expr, $acc:expr) => {{
2121 let lo16 = _mm_loadu_si128($q.as_ptr() as *const __m128i);
2122 let hi16 = _mm_loadu_si128($q.as_ptr().add(16) as *const __m128i);
2123 let lo2 = _mm_and_si128(_mm_srli_epi16(lo16, $shift), two_bit_mask);
2124 let hi2 = _mm_and_si128(_mm_srli_epi16(hi16, $shift), two_bit_mask);
2125
2126 let hmask_lo = _mm_loadu_si128($hmask.as_ptr() as *const __m128i);
2127 let hmask_hi = _mm_loadu_si128($hmask.as_ptr().add(16) as *const __m128i);
2128 let bit_clear_lo = _mm_cmpeq_epi8(_mm_and_si128(hmask_lo, $m_vec), zero);
2132 let bit_clear_hi = _mm_cmpeq_epi8(_mm_and_si128(hmask_hi, $m_vec), zero);
2133 let bias_lo = _mm_and_si128(bit_clear_lo, four);
2134 let bias_hi = _mm_and_si128(bit_clear_hi, four);
2135 let raw_lo = _mm_sub_epi8(lo2, bias_lo);
2136 let raw_hi = _mm_sub_epi8(hi2, bias_hi);
2137
2138 let mut lo_acc = _mm256_setzero_ps();
2139 let mut hi_acc = _mm256_setzero_ps();
2140 for (part_idx, part) in [raw_lo, _mm_srli_si128(raw_lo, 8)].into_iter().enumerate()
2141 {
2142 let i32x8 = _mm256_cvtepi8_epi32(part);
2143 let f32x8 = _mm256_cvtepi32_ps(i32x8);
2144 let xv = _mm256_loadu_ps($x.as_ptr().add($x_base + part_idx * 8));
2145 lo_acc = _mm256_fmadd_ps(f32x8, xv, lo_acc);
2146 }
2147 for (part_idx, part) in [raw_hi, _mm_srli_si128(raw_hi, 8)].into_iter().enumerate()
2148 {
2149 let i32x8 = _mm256_cvtepi8_epi32(part);
2150 let f32x8 = _mm256_cvtepi32_ps(i32x8);
2151 let xv = _mm256_loadu_ps($x.as_ptr().add($x_base + 16 + part_idx * 8));
2152 hi_acc = _mm256_fmadd_ps(f32x8, xv, hi_acc);
2153 }
2154 $acc += hsum256_ps(lo_acc) * $dl1 + hsum256_ps(hi_acc) * $dl2;
2155 $x_base += 32;
2156 }};
2157 }
2158
2159 for block in row_bytes.as_chunks::<Q3_K_BLOCK_BYTES>().0 {
2160 let hmask = &block[0..32];
2161 let qs = &block[32..96];
2162 let scales_raw: &[u8; Q3_K_SCALE_BYTES] = block[96..108].try_into().unwrap();
2163 let d_all = f16::from_le_bytes([block[108], block[109]]).to_f32();
2164 let scales = q3_k_unpack_scales(scales_raw);
2165
2166 let mut is = 0usize;
2167 let mut m = 1u8;
2168 for n in 0..2 {
2169 let q = &qs[n * 32..n * 32 + 32];
2170 for shift in [0u32, 2, 4, 6] {
2171 let dl1 = d_all * (scales[is] as f32 - 32.0);
2172 let dl2 = d_all * (scales[is + 1] as f32 - 32.0);
2173 is += 2;
2174 let m_vec = _mm_set1_epi8(m as i8);
2175 match shift {
2176 0 => q3_k_sub_block!(0, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
2177 2 => q3_k_sub_block!(2, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
2178 4 => q3_k_sub_block!(4, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
2179 6 => q3_k_sub_block!(6, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
2180 _ => unreachable!(),
2181 }
2182 m <<= 1;
2183 }
2184 }
2185 }
2186 acc
2187 }
2188
2189 #[target_feature(enable = "avx2,fma")]
2197 pub unsafe fn dot_iq4_nl_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2198 debug_assert_eq!(row_bytes.len() % IQ4_NL_BLOCK_BYTES, 0);
2199 let low_mask = _mm_set1_epi8(0x0F);
2200 let codebook = _mm_loadu_si128(KVALUES_IQ4NL.as_ptr() as *const __m128i);
2201 let mut acc = 0f32;
2202 let mut x_base = 0usize;
2203 for block in row_bytes.as_chunks::<IQ4_NL_BLOCK_BYTES>().0 {
2204 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2205 let qs = &block[2..18];
2206 let bytes = _mm_loadu_si128(qs.as_ptr() as *const __m128i);
2207 let lo_idx = _mm_and_si128(bytes, low_mask);
2208 let hi_idx = _mm_and_si128(_mm_srli_epi16(bytes, 4), low_mask);
2209 let lo_vals = _mm_shuffle_epi8(codebook, lo_idx);
2210 let hi_vals = _mm_shuffle_epi8(codebook, hi_idx);
2211
2212 let mut block_acc = _mm256_setzero_ps();
2213 for (half_idx, vals) in [
2214 (0usize, lo_vals),
2215 (1usize, _mm_srli_si128(lo_vals, 8)),
2216 (2usize, hi_vals),
2217 (3usize, _mm_srli_si128(hi_vals, 8)),
2218 ] {
2219 let i32x8 = _mm256_cvtepi8_epi32(vals);
2220 let f32x8 = _mm256_cvtepi32_ps(i32x8);
2221 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base + half_idx * 8));
2222 block_acc = _mm256_fmadd_ps(f32x8, xv, block_acc);
2223 }
2224 acc += hsum256_ps(block_acc) * d;
2225 x_base += IQ4_NL_BLOCK_ELEMS;
2226 }
2227 acc
2228 }
2229
2230 #[target_feature(enable = "avx2,fma")]
2236 pub unsafe fn dot_iq4_xs_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2237 debug_assert_eq!(row_bytes.len() % IQ4_XS_BLOCK_BYTES, 0);
2238 let low_mask = _mm_set1_epi8(0x0F);
2239 let codebook = _mm_loadu_si128(KVALUES_IQ4NL.as_ptr() as *const __m128i);
2240 let mut acc = 0f32;
2241 let mut x_base = 0usize;
2242 for block in row_bytes.as_chunks::<IQ4_XS_BLOCK_BYTES>().0 {
2243 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2244 let scales_h = u16::from_le_bytes([block[2], block[3]]);
2245 let scales_l = &block[4..8];
2246 let qs = &block[8..136];
2247
2248 for ib in 0..8 {
2249 let ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0xf)
2250 | (((scales_h >> (2 * ib)) & 3) as u8) << 4;
2251 let dl = d * (ls as f32 - 32.0);
2252 let sub = &qs[ib * 16..ib * 16 + 16];
2253 let bytes = _mm_loadu_si128(sub.as_ptr() as *const __m128i);
2254 let lo_idx = _mm_and_si128(bytes, low_mask);
2255 let hi_idx = _mm_and_si128(_mm_srli_epi16(bytes, 4), low_mask);
2256 let lo_vals = _mm_shuffle_epi8(codebook, lo_idx);
2257 let hi_vals = _mm_shuffle_epi8(codebook, hi_idx);
2258
2259 let mut sub_acc = _mm256_setzero_ps();
2260 for (half_idx, vals) in [
2261 (0usize, lo_vals),
2262 (1usize, _mm_srli_si128(lo_vals, 8)),
2263 (2usize, hi_vals),
2264 (3usize, _mm_srli_si128(hi_vals, 8)),
2265 ] {
2266 let i32x8 = _mm256_cvtepi8_epi32(vals);
2267 let f32x8 = _mm256_cvtepi32_ps(i32x8);
2268 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base + half_idx * 8));
2269 sub_acc = _mm256_fmadd_ps(f32x8, xv, sub_acc);
2270 }
2271 acc += hsum256_ps(sub_acc) * dl;
2272 x_base += 32;
2273 }
2274 }
2275 acc
2276 }
2277
2278 #[inline]
2286 #[target_feature(enable = "avx2", enable = "fma")]
2287 unsafe fn iq_grid_row_signed_f32(row_le: u64, signs: u8) -> __m256 {
2288 let mags = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_set_epi64x(0, row_le as i64)));
2289 let bit_mask = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
2290 let bits = _mm256_and_si256(_mm256_set1_epi32(signs as i32), bit_mask);
2291 let neg = _mm256_cmpeq_epi32(bits, bit_mask);
2292 let sign_bit = _mm256_and_si256(neg, _mm256_set1_epi32(0x8000_0000_u32 as i32));
2293 _mm256_xor_ps(mags, _mm256_castsi256_ps(sign_bit))
2294 }
2295
2296 #[target_feature(enable = "avx2", enable = "fma")]
2303 pub unsafe fn dot_iq1_s_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2304 debug_assert_eq!(row_bytes.len() % crate::IQ1_S_BLOCK_BYTES, 0);
2305 let mut acc = _mm256_setzero_ps();
2306 let mut x_base = 0usize;
2307 for block in row_bytes.as_chunks::<{ crate::IQ1_S_BLOCK_BYTES }>().0 {
2308 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2309 let qs = &block[2..34];
2310 let qh = &block[34..50];
2311 for ib in 0..8 {
2312 let h = u16::from_le_bytes([qh[2 * ib], qh[2 * ib + 1]]);
2313 let dl = d * (2.0 * ((h >> 12) & 7) as f32 + 1.0);
2314 let delta = if h & 0x8000 != 0 {
2315 -crate::IQ1S_DELTA
2316 } else {
2317 crate::IQ1S_DELTA
2318 };
2319 let dl_v = _mm256_set1_ps(dl);
2320 let delta_v = _mm256_set1_ps(delta);
2321 for l in 0..4 {
2322 let idx = qs[4 * ib + l] as usize | ((((h >> (3 * l)) & 7) as usize) << 8);
2323 let row = crate::iq_tables::IQ1S_GRID[idx];
2324 let g = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_set_epi64x(0, row as i64)));
2325 let vals = _mm256_mul_ps(dl_v, _mm256_add_ps(g, delta_v));
2326 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base));
2327 acc = _mm256_fmadd_ps(vals, xv, acc);
2328 x_base += 8;
2329 }
2330 }
2331 }
2332 hsum256_ps(acc)
2333 }
2334
2335 #[target_feature(enable = "avx2", enable = "fma")]
2340 pub unsafe fn dot_iq2_xxs_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2341 debug_assert_eq!(row_bytes.len() % crate::IQ2_XXS_BLOCK_BYTES, 0);
2342 let mut acc = _mm256_setzero_ps();
2343 let mut x_base = 0usize;
2344 for block in row_bytes.as_chunks::<{ crate::IQ2_XXS_BLOCK_BYTES }>().0 {
2345 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2346 for ib32 in 0..8 {
2347 let g0 = u16::from_le_bytes([block[2 + 8 * ib32], block[3 + 8 * ib32]]);
2348 let g1 = u16::from_le_bytes([block[4 + 8 * ib32], block[5 + 8 * ib32]]);
2349 let g2 = u16::from_le_bytes([block[6 + 8 * ib32], block[7 + 8 * ib32]]);
2350 let g3 = u16::from_le_bytes([block[8 + 8 * ib32], block[9 + 8 * ib32]]);
2351 let aux32_1 = g2 as u32 | ((g3 as u32) << 16);
2352 let db = _mm256_set1_ps(d * (0.5 + (aux32_1 >> 28) as f32) * 0.25);
2353 let aux8 = [
2354 (g0 & 0xFF) as usize,
2355 (g0 >> 8) as usize,
2356 (g1 & 0xFF) as usize,
2357 (g1 >> 8) as usize,
2358 ];
2359 for (l, &code) in aux8.iter().enumerate() {
2360 let signs =
2361 crate::iq_tables::KSIGNS_IQ2XS[((aux32_1 >> (7 * l)) & 127) as usize];
2362 let vals = iq_grid_row_signed_f32(crate::iq_tables::IQ2XXS_GRID[code], signs);
2363 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base));
2364 acc = _mm256_fmadd_ps(_mm256_mul_ps(db, vals), xv, acc);
2365 x_base += 8;
2366 }
2367 }
2368 }
2369 hsum256_ps(acc)
2370 }
2371
2372 #[target_feature(enable = "avx2", enable = "fma")]
2377 pub unsafe fn dot_iq3_xxs_f32_avx2(row_bytes: &[u8], x: &[f32]) -> f32 {
2378 debug_assert_eq!(row_bytes.len() % crate::IQ3_XXS_BLOCK_BYTES, 0);
2379 let mut acc = _mm256_setzero_ps();
2380 let mut x_base = 0usize;
2381 for block in row_bytes.as_chunks::<{ crate::IQ3_XXS_BLOCK_BYTES }>().0 {
2382 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2383 let qs = &block[2..66];
2384 let sas = &block[66..98];
2385 for ib32 in 0..8 {
2386 let aux32 = u32::from_le_bytes([
2387 sas[4 * ib32],
2388 sas[4 * ib32 + 1],
2389 sas[4 * ib32 + 2],
2390 sas[4 * ib32 + 3],
2391 ]);
2392 let db = _mm256_set1_ps(d * (0.5 + (aux32 >> 28) as f32) * 0.5);
2393 for l in 0..4 {
2394 let signs = crate::iq_tables::KSIGNS_IQ2XS[((aux32 >> (7 * l)) & 127) as usize];
2395 let r1 = crate::iq_tables::IQ3XXS_GRID[qs[8 * ib32 + 2 * l] as usize];
2396 let r2 = crate::iq_tables::IQ3XXS_GRID[qs[8 * ib32 + 2 * l + 1] as usize];
2397 let row = (r1 as u64) | ((r2 as u64) << 32);
2398 let vals = iq_grid_row_signed_f32(row, signs);
2399 let xv = _mm256_loadu_ps(x.as_ptr().add(x_base));
2400 acc = _mm256_fmadd_ps(_mm256_mul_ps(db, vals), xv, acc);
2401 x_base += 8;
2402 }
2403 }
2404 }
2405 hsum256_ps(acc)
2406 }
2407}
2408
2409#[cfg(target_arch = "aarch64")]
2422mod simd_aarch64 {
2423 use super::{
2424 e8m0_scale, q3_k_unpack_scales, q4_k_scale_min, q5_fifth_bits, Q8Activations,
2425 Q8KActivations, IQ4_NL_BLOCK_BYTES, IQ4_NL_BLOCK_ELEMS, IQ4_XS_BLOCK_BYTES, KVALUES_IQ4NL,
2426 MXFP4_GROUP_SIZE, Q2_K_BLOCK_BYTES, Q2_K_SCALE_BYTES, Q3_K_BLOCK_BYTES, Q3_K_SCALE_BYTES,
2427 Q4_0_BLOCK_BYTES, Q4_0_BLOCK_ELEMS, Q4_1_BLOCK_BYTES, Q4_1_BLOCK_ELEMS, Q4_K_BLOCK_BYTES,
2428 Q4_K_BLOCK_ELEMS, Q4_K_SCALE_BYTES, Q5_0_BLOCK_BYTES, Q5_0_BLOCK_ELEMS, Q5_1_BLOCK_BYTES,
2429 Q5_1_BLOCK_ELEMS, Q5_K_BLOCK_BYTES, Q5_K_BLOCK_ELEMS, Q6_K_BLOCK_BYTES, Q6_K_BLOCK_ELEMS,
2430 Q8_0_BLOCK_BYTES, Q8_0_BLOCK_ELEMS, Q8_1_BLOCK_BYTES, Q8_1_BLOCK_ELEMS,
2431 };
2432 use half::f16;
2433 use std::arch::aarch64::*;
2434
2435 #[target_feature(enable = "neon")]
2448 pub unsafe fn dot_q8_0_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
2449 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
2450 debug_assert_eq!(
2451 row_bytes.len() / Q8_0_BLOCK_BYTES * Q8_0_BLOCK_ELEMS,
2452 x.len()
2453 );
2454 let mut acc = 0f32;
2455 for (b, block) in row_bytes
2456 .as_chunks::<Q8_0_BLOCK_BYTES>()
2457 .0
2458 .iter()
2459 .enumerate()
2460 {
2461 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
2462 let base = b * Q8_0_BLOCK_ELEMS;
2463 let qs = &block[2..34];
2464
2465 let mut block_acc = vdupq_n_f32(0.0);
2466 for g in 0..2 {
2467 let raw16 = vld1q_s8(qs.as_ptr().add(g * 16) as *const i8);
2468 let lo16 = vmovl_s8(vget_low_s8(raw16));
2469 let hi16 = vmovl_s8(vget_high_s8(raw16));
2470 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
2471 let lo32 = vmovl_s16(vget_low_s16(half16));
2472 let hi32 = vmovl_s16(vget_high_s16(half16));
2473 let f_lo = vcvtq_f32_s32(lo32);
2474 let f_hi = vcvtq_f32_s32(hi32);
2475 let elem_base = base + g * 16 + half_idx * 8;
2476 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
2477 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
2478 block_acc = vfmaq_f32(block_acc, f_lo, x_lo);
2479 block_acc = vfmaq_f32(block_acc, f_hi, x_hi);
2480 }
2481 }
2482 acc += vaddvq_f32(block_acc) * scale;
2483 }
2484 acc
2485 }
2486
2487 #[target_feature(enable = "neon")]
2490 pub unsafe fn dot_q8_0_q8_neon(row_bytes: &[u8], act: &Q8Activations) -> f32 {
2491 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
2492 debug_assert_eq!(row_bytes.len() / Q8_0_BLOCK_BYTES, act.n_blocks());
2493 let mut acc = 0f32;
2494 for (b, block) in row_bytes
2495 .as_chunks::<Q8_0_BLOCK_BYTES>()
2496 .0
2497 .iter()
2498 .enumerate()
2499 {
2500 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
2501 let base = b * Q8_0_BLOCK_ELEMS;
2502 let mut isum = vdupq_n_s32(0);
2503 for g in 0..2 {
2504 let w = vld1q_s8(block.as_ptr().add(2 + g * 16) as *const i8);
2505 let a = vld1q_s8(act.q.as_ptr().add(base + g * 16));
2506 let prod_lo = vmull_s8(vget_low_s8(w), vget_low_s8(a));
2507 let prod_hi = vmull_s8(vget_high_s8(w), vget_high_s8(a));
2508 isum = vpadalq_s16(isum, prod_lo);
2509 isum = vpadalq_s16(isum, prod_hi);
2510 }
2511 acc += dw * act.d[b] * vaddvq_s32(isum) as f32;
2512 }
2513 acc
2514 }
2515
2516 #[target_feature(enable = "neon,dotprod")]
2518 unsafe fn neon_sdot(mut acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
2519 std::arch::asm!(
2520 "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
2521 acc = inout(vreg) acc,
2522 a = in(vreg) a,
2523 b = in(vreg) b,
2524 options(pure, nomem, nostack),
2525 );
2526 acc
2527 }
2528
2529 #[target_feature(enable = "neon,dotprod")]
2532 pub unsafe fn dot_q8_0_q8_neon_sdot(row_bytes: &[u8], act: &Q8Activations) -> f32 {
2533 debug_assert_eq!(row_bytes.len() % Q8_0_BLOCK_BYTES, 0);
2534 debug_assert_eq!(row_bytes.len() / Q8_0_BLOCK_BYTES, act.n_blocks());
2535 let nb = row_bytes.len() / Q8_0_BLOCK_BYTES;
2536 let mut sumv0 = vdupq_n_f32(0.0);
2537 let mut sumv1 = vdupq_n_f32(0.0);
2538 let mut b = 0usize;
2539 while b + 1 < nb {
2540 let block0 = row_bytes.as_ptr().add(b * Q8_0_BLOCK_BYTES);
2541 let block1 = row_bytes.as_ptr().add((b + 1) * Q8_0_BLOCK_BYTES);
2542 let dw0 = f16::from_le_bytes([*block0, *block0.add(1)]).to_f32();
2543 let dw1 = f16::from_le_bytes([*block1, *block1.add(1)]).to_f32();
2544 let base0 = b * Q8_0_BLOCK_ELEMS;
2545 let base1 = (b + 1) * Q8_0_BLOCK_ELEMS;
2546 let mut isum0 = vdupq_n_s32(0);
2547 let mut isum1 = vdupq_n_s32(0);
2548 for g in 0..2 {
2549 let w0 = vld1q_s8(block0.add(2 + g * 16) as *const i8);
2550 let w1 = vld1q_s8(block1.add(2 + g * 16) as *const i8);
2551 let a0 = vld1q_s8(act.q.as_ptr().add(base0 + g * 16));
2552 let a1 = vld1q_s8(act.q.as_ptr().add(base1 + g * 16));
2553 isum0 = neon_sdot(isum0, w0, a0);
2554 isum1 = neon_sdot(isum1, w1, a1);
2555 }
2556 sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(isum0), dw0 * act.d[b]);
2557 sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(isum1), dw1 * act.d[b + 1]);
2558 b += 2;
2559 }
2560 let mut acc = vaddvq_f32(sumv0) + vaddvq_f32(sumv1);
2561 if b < nb {
2562 let block = row_bytes.as_ptr().add(b * Q8_0_BLOCK_BYTES);
2563 let dw = f16::from_le_bytes([*block, *block.add(1)]).to_f32();
2564 let base = b * Q8_0_BLOCK_ELEMS;
2565 let mut isum = vdupq_n_s32(0);
2566 for g in 0..2 {
2567 let w = vld1q_s8(block.add(2 + g * 16) as *const i8);
2568 let a = vld1q_s8(act.q.as_ptr().add(base + g * 16));
2569 isum = neon_sdot(isum, w, a);
2570 }
2571 acc += dw * act.d[b] * vaddvq_s32(isum) as f32;
2572 }
2573 acc
2574 }
2575
2576 #[target_feature(enable = "neon")]
2580 pub unsafe fn dot_q4_0_q8_neon(row_bytes: &[u8], act: &Q8Activations) -> f32 {
2581 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
2582 debug_assert_eq!(row_bytes.len() / Q4_0_BLOCK_BYTES, act.n_blocks());
2583 let bias = vdupq_n_s8(8);
2584 let low_mask = vdupq_n_u8(0x0F);
2585 let mut acc = 0f32;
2586 for (b, block) in row_bytes
2587 .as_chunks::<Q4_0_BLOCK_BYTES>()
2588 .0
2589 .iter()
2590 .enumerate()
2591 {
2592 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
2593 let base = b * Q4_0_BLOCK_ELEMS;
2594 let nibbles = vld1q_u8(block.as_ptr().add(2));
2595 let lo = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nibbles, low_mask)), bias);
2596 let hi = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nibbles, 4)), bias);
2597 let mut isum = vdupq_n_s32(0);
2598 let a0 = vld1q_s8(act.q.as_ptr().add(base));
2600 let a1 = vld1q_s8(act.q.as_ptr().add(base + 16));
2601 let p0_lo = vmull_s8(vget_low_s8(lo), vget_low_s8(a0));
2602 let p0_hi = vmull_s8(vget_high_s8(lo), vget_high_s8(a0));
2603 let p1_lo = vmull_s8(vget_low_s8(hi), vget_low_s8(a1));
2604 let p1_hi = vmull_s8(vget_high_s8(hi), vget_high_s8(a1));
2605 isum = vpadalq_s16(isum, p0_lo);
2606 isum = vpadalq_s16(isum, p0_hi);
2607 isum = vpadalq_s16(isum, p1_lo);
2608 isum = vpadalq_s16(isum, p1_hi);
2609 acc += dw * act.d[b] * vaddvq_s32(isum) as f32;
2610 }
2611 acc
2612 }
2613
2614 #[target_feature(enable = "neon,dotprod")]
2616 pub unsafe fn dot_q4_0_q8_neon_sdot_2row(
2617 row0: &[u8],
2618 row1: &[u8],
2619 act: &Q8Activations,
2620 ) -> (f32, f32) {
2621 debug_assert_eq!(row0.len(), row1.len());
2622 debug_assert_eq!(row0.len() % Q4_0_BLOCK_BYTES, 0);
2623 let bias = vdupq_n_s8(8);
2624 let low_mask = vdupq_n_u8(0x0F);
2625 let nb = row0.len() / Q4_0_BLOCK_BYTES;
2626 let mut sum0 = vdupq_n_f32(0.0);
2627 let mut sum1 = vdupq_n_f32(0.0);
2628 for b in 0..nb {
2629 let p0 = row0.as_ptr().add(b * Q4_0_BLOCK_BYTES);
2630 let p1 = row1.as_ptr().add(b * Q4_0_BLOCK_BYTES);
2631 let dw0 = f16::from_le_bytes([*p0, *p0.add(1)]).to_f32();
2632 let dw1 = f16::from_le_bytes([*p1, *p1.add(1)]).to_f32();
2633 let base = b * Q4_0_BLOCK_ELEMS;
2634 let a_lo = vld1q_s8(act.q.as_ptr().add(base));
2635 let a_hi = vld1q_s8(act.q.as_ptr().add(base + 16));
2636 let nib0 = vld1q_u8(p0.add(2));
2637 let nib1 = vld1q_u8(p1.add(2));
2638 let lo0 = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nib0, low_mask)), bias);
2639 let hi0 = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nib0, 4)), bias);
2640 let lo1 = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nib1, low_mask)), bias);
2641 let hi1 = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nib1, 4)), bias);
2642 let mut is0 = neon_sdot(vdupq_n_s32(0), lo0, a_lo);
2643 is0 = neon_sdot(is0, hi0, a_hi);
2644 let mut is1 = neon_sdot(vdupq_n_s32(0), lo1, a_lo);
2645 is1 = neon_sdot(is1, hi1, a_hi);
2646 let scale = act.d[b];
2647 sum0 = vmlaq_n_f32(sum0, vcvtq_f32_s32(is0), dw0 * scale);
2648 sum1 = vmlaq_n_f32(sum1, vcvtq_f32_s32(is1), dw1 * scale);
2649 }
2650 (vaddvq_f32(sum0), vaddvq_f32(sum1))
2651 }
2652
2653 #[target_feature(enable = "neon,dotprod")]
2655 pub unsafe fn dot_q4_0_q8_neon_sdot(row_bytes: &[u8], act: &Q8Activations) -> f32 {
2656 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
2657 debug_assert_eq!(row_bytes.len() / Q4_0_BLOCK_BYTES, act.n_blocks());
2658 let bias = vdupq_n_s8(8);
2659 let low_mask = vdupq_n_u8(0x0F);
2660 let nb = row_bytes.len() / Q4_0_BLOCK_BYTES;
2661 let mut sumv0 = vdupq_n_f32(0.0);
2662 let mut sumv1 = vdupq_n_f32(0.0);
2663 let mut b = 0usize;
2664 while b + 1 < nb {
2665 let block0 = row_bytes.as_ptr().add(b * Q4_0_BLOCK_BYTES);
2666 let block1 = row_bytes.as_ptr().add((b + 1) * Q4_0_BLOCK_BYTES);
2667 let dw0 = f16::from_le_bytes([*block0, *block0.add(1)]).to_f32();
2668 let dw1 = f16::from_le_bytes([*block1, *block1.add(1)]).to_f32();
2669 let base0 = b * Q4_0_BLOCK_ELEMS;
2670 let base1 = (b + 1) * Q4_0_BLOCK_ELEMS;
2671 let nib0 = vld1q_u8(block0.add(2));
2672 let nib1 = vld1q_u8(block1.add(2));
2673 let lo0 = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nib0, low_mask)), bias);
2674 let hi0 = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nib0, 4)), bias);
2675 let lo1 = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nib1, low_mask)), bias);
2676 let hi1 = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nib1, 4)), bias);
2677 let mut isum0 = neon_sdot(vdupq_n_s32(0), lo0, vld1q_s8(act.q.as_ptr().add(base0)));
2678 isum0 = neon_sdot(isum0, hi0, vld1q_s8(act.q.as_ptr().add(base0 + 16)));
2679 let mut isum1 = neon_sdot(vdupq_n_s32(0), lo1, vld1q_s8(act.q.as_ptr().add(base1)));
2680 isum1 = neon_sdot(isum1, hi1, vld1q_s8(act.q.as_ptr().add(base1 + 16)));
2681 sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(isum0), dw0 * act.d[b]);
2682 sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(isum1), dw1 * act.d[b + 1]);
2683 b += 2;
2684 }
2685 let mut acc = vaddvq_f32(sumv0) + vaddvq_f32(sumv1);
2686 if b < nb {
2687 let block = &row_bytes[b * Q4_0_BLOCK_BYTES..(b + 1) * Q4_0_BLOCK_BYTES];
2688 let dw = f16::from_le_bytes([block[0], block[1]]).to_f32();
2689 let base = b * Q4_0_BLOCK_ELEMS;
2690 let nibbles = vld1q_u8(block.as_ptr().add(2));
2691 let lo = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(nibbles, low_mask)), bias);
2692 let hi = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(nibbles, 4)), bias);
2693 let mut isum = neon_sdot(vdupq_n_s32(0), lo, vld1q_s8(act.q.as_ptr().add(base)));
2694 isum = neon_sdot(isum, hi, vld1q_s8(act.q.as_ptr().add(base + 16)));
2695 acc += dw * act.d[b] * vaddvq_s32(isum) as f32;
2696 }
2697 acc
2698 }
2699
2700 #[target_feature(enable = "neon")]
2701 unsafe fn neon_i8_dot_widen(mut isum: int32x4_t, w: int8x16_t, a: int8x16_t) -> int32x4_t {
2702 let prod_lo = vmull_s8(vget_low_s8(w), vget_low_s8(a));
2703 let prod_hi = vmull_s8(vget_high_s8(w), vget_high_s8(a));
2704 isum = vpadalq_s16(isum, prod_lo);
2705 vpadalq_s16(isum, prod_hi)
2706 }
2707
2708 #[target_feature(enable = "neon")]
2710 pub unsafe fn dot_q4_k_q8_neon(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
2711 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
2712 debug_assert_eq!(row_bytes.len() / Q4_K_BLOCK_BYTES, act.n_blocks());
2713 let low_mask = vdupq_n_u8(0x0F);
2714 let mut acc = 0f32;
2715 for (b, block) in row_bytes
2716 .as_chunks::<Q4_K_BLOCK_BYTES>()
2717 .0
2718 .iter()
2719 .enumerate()
2720 {
2721 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2722 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
2723 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
2724 let qs = &block[16..144];
2725 let da = act.d[b];
2726 let q8 = act.q.as_ptr().add(b * Q4_K_BLOCK_ELEMS);
2727 let bsums = &act.bsums[b * 16..(b + 1) * 16];
2728
2729 let mut sum_min = 0i32;
2730 for i in 0..8 {
2731 let (_, m) = q4_k_scale_min(i, &scales);
2732 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
2733 }
2734 acc -= dmin * da * sum_min as f32;
2735
2736 let mut q_off = 0usize;
2737 let mut base = 0usize;
2738 let mut is = 0usize;
2739 for _ in 0..4 {
2740 let (sc1, _) = q4_k_scale_min(is, &scales);
2741 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
2742 let mut isum1 = vdupq_n_s32(0);
2743 let mut isum2 = vdupq_n_s32(0);
2744 for g in 0..2 {
2745 let packed = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
2746 let lo = vreinterpretq_s8_u8(vandq_u8(packed, low_mask));
2747 let hi = vreinterpretq_s8_u8(vshrq_n_u8(packed, 4));
2748 let a0 = vld1q_s8(q8.add(base + g * 16));
2749 let a1 = vld1q_s8(q8.add(base + 32 + g * 16));
2750 isum1 = neon_i8_dot_widen(isum1, lo, a0);
2751 isum2 = neon_i8_dot_widen(isum2, hi, a1);
2752 }
2753 acc += d
2754 * da
2755 * (sc1 as f32 * vaddvq_s32(isum1) as f32
2756 + sc2 as f32 * vaddvq_s32(isum2) as f32);
2757 q_off += 32;
2758 base += 64;
2759 is += 2;
2760 }
2761 }
2762 acc
2763 }
2764
2765 #[target_feature(enable = "neon,i8mm")]
2771 pub unsafe fn dot_q4_k_q8_neon_i8mm(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
2772 dot_q4_k_q8_neon_sdot(row_bytes, act)
2773 }
2774
2775 #[target_feature(enable = "neon,dotprod")]
2777 pub unsafe fn dot_q4_k_q8_neon_sdot(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
2778 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
2779 debug_assert_eq!(row_bytes.len() / Q4_K_BLOCK_BYTES, act.n_blocks());
2780 let low_mask = vdupq_n_u8(0x0F);
2781 let mut acc = 0f32;
2782 for (b, block) in row_bytes
2783 .as_chunks::<Q4_K_BLOCK_BYTES>()
2784 .0
2785 .iter()
2786 .enumerate()
2787 {
2788 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2789 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
2790 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
2791 let qs = &block[16..144];
2792 let da = act.d[b];
2793 let q8 = act.q.as_ptr().add(b * Q4_K_BLOCK_ELEMS);
2794 let bsums = &act.bsums[b * 16..(b + 1) * 16];
2795
2796 let mut sum_min = 0i32;
2797 for i in 0..8 {
2798 let (_, m) = q4_k_scale_min(i, &scales);
2799 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
2800 }
2801 acc -= dmin * da * sum_min as f32;
2802
2803 let mut q_off = 0usize;
2804 let mut base = 0usize;
2805 let mut is = 0usize;
2806 for _ in 0..4 {
2807 let (sc1, _) = q4_k_scale_min(is, &scales);
2808 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
2809 let mut isum1 = vdupq_n_s32(0);
2810 let mut isum2 = vdupq_n_s32(0);
2811 for g in 0..2 {
2812 let packed = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
2813 let lo = vreinterpretq_s8_u8(vandq_u8(packed, low_mask));
2814 let hi = vreinterpretq_s8_u8(vshrq_n_u8(packed, 4));
2815 let a0 = vld1q_s8(q8.add(base + g * 16));
2816 let a1 = vld1q_s8(q8.add(base + 32 + g * 16));
2817 isum1 = neon_sdot(isum1, lo, a0);
2818 isum2 = neon_sdot(isum2, hi, a1);
2819 }
2820 acc += d
2821 * da
2822 * (sc1 as f32 * vaddvq_s32(isum1) as f32
2823 + sc2 as f32 * vaddvq_s32(isum2) as f32);
2824 q_off += 32;
2825 base += 64;
2826 is += 2;
2827 }
2828 }
2829 acc
2830 }
2831
2832 #[target_feature(enable = "neon")]
2834 pub unsafe fn dot_q5_k_q8_neon(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
2835 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
2836 debug_assert_eq!(row_bytes.len() / Q5_K_BLOCK_BYTES, act.n_blocks());
2837 let low_mask = vdupq_n_u8(0x0F);
2838 let sixteen = vdupq_n_u8(16);
2839 let mut acc = 0f32;
2840 for (b, block) in row_bytes
2841 .as_chunks::<Q5_K_BLOCK_BYTES>()
2842 .0
2843 .iter()
2844 .enumerate()
2845 {
2846 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2847 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
2848 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
2849 let qh = block.as_ptr().add(16);
2850 let qs = &block[48..176];
2851 let da = act.d[b];
2852 let q8 = act.q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
2853 let bsums = &act.bsums[b * 16..(b + 1) * 16];
2854
2855 let mut sum_min = 0i32;
2856 for i in 0..8 {
2857 let (_, m) = q4_k_scale_min(i, &scales);
2858 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
2859 }
2860 acc -= dmin * da * sum_min as f32;
2861
2862 let mut q_off = 0usize;
2863 let mut base = 0usize;
2864 let mut is = 0usize;
2865 let (mut u1, mut u2) = (1u8, 2u8);
2866 for _ in 0..4 {
2867 let (sc1, _) = q4_k_scale_min(is, &scales);
2868 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
2869 let mut isum1 = vdupq_n_s32(0);
2870 let mut isum2 = vdupq_n_s32(0);
2871 let u1_vec = vdupq_n_u8(u1);
2872 let u2_vec = vdupq_n_u8(u2);
2873 for g in 0..2 {
2874 let packed = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
2875 let qh16 = vld1q_u8(qh.add(g * 16));
2876 let lo_nib = vandq_u8(packed, low_mask);
2877 let hi_nib = vshrq_n_u8(packed, 4);
2878 let hi_bit1 = vandq_u8(vtstq_u8(qh16, u1_vec), sixteen);
2879 let hi_bit2 = vandq_u8(vtstq_u8(qh16, u2_vec), sixteen);
2880 let lo = vreinterpretq_s8_u8(vorrq_u8(lo_nib, hi_bit1));
2881 let hi = vreinterpretq_s8_u8(vorrq_u8(hi_nib, hi_bit2));
2882 let a0 = vld1q_s8(q8.add(base + g * 16));
2883 let a1 = vld1q_s8(q8.add(base + 32 + g * 16));
2884 isum1 = neon_i8_dot_widen(isum1, lo, a0);
2885 isum2 = neon_i8_dot_widen(isum2, hi, a1);
2886 }
2887 acc += d
2888 * da
2889 * (sc1 as f32 * vaddvq_s32(isum1) as f32
2890 + sc2 as f32 * vaddvq_s32(isum2) as f32);
2891 q_off += 32;
2892 base += 64;
2893 is += 2;
2894 u1 <<= 2;
2895 u2 <<= 2;
2896 }
2897 }
2898 acc
2899 }
2900
2901 #[target_feature(enable = "neon,dotprod")]
2903 pub unsafe fn dot_q5_k_q8_neon_sdot(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
2904 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
2905 debug_assert_eq!(row_bytes.len() / Q5_K_BLOCK_BYTES, act.n_blocks());
2906 let low_mask = vdupq_n_u8(0x0F);
2907 let sixteen = vdupq_n_u8(16);
2908 let mut acc = 0f32;
2909 for (b, block) in row_bytes
2910 .as_chunks::<Q5_K_BLOCK_BYTES>()
2911 .0
2912 .iter()
2913 .enumerate()
2914 {
2915 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2916 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
2917 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
2918 let qh = block.as_ptr().add(16);
2919 let qs = &block[48..176];
2920 let da = act.d[b];
2921 let q8 = act.q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
2922 let bsums = &act.bsums[b * 16..(b + 1) * 16];
2923
2924 let mut sum_min = 0i32;
2925 for i in 0..8 {
2926 let (_, m) = q4_k_scale_min(i, &scales);
2927 sum_min += m as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
2928 }
2929 acc -= dmin * da * sum_min as f32;
2930
2931 let mut q_off = 0usize;
2932 let mut base = 0usize;
2933 let mut is = 0usize;
2934 let (mut u1, mut u2) = (1u8, 2u8);
2935 for _ in 0..4 {
2936 let (sc1, _) = q4_k_scale_min(is, &scales);
2937 let (sc2, _) = q4_k_scale_min(is + 1, &scales);
2938 let mut isum1 = vdupq_n_s32(0);
2939 let mut isum2 = vdupq_n_s32(0);
2940 let u1_vec = vdupq_n_u8(u1);
2941 let u2_vec = vdupq_n_u8(u2);
2942 for g in 0..2 {
2943 let packed = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
2944 let qh16 = vld1q_u8(qh.add(g * 16));
2945 let lo_nib = vandq_u8(packed, low_mask);
2946 let hi_nib = vshrq_n_u8(packed, 4);
2947 let hi_bit1 = vandq_u8(vtstq_u8(qh16, u1_vec), sixteen);
2948 let hi_bit2 = vandq_u8(vtstq_u8(qh16, u2_vec), sixteen);
2949 let lo = vreinterpretq_s8_u8(vorrq_u8(lo_nib, hi_bit1));
2950 let hi = vreinterpretq_s8_u8(vorrq_u8(hi_nib, hi_bit2));
2951 let a0 = vld1q_s8(q8.add(base + g * 16));
2952 let a1 = vld1q_s8(q8.add(base + 32 + g * 16));
2953 isum1 = neon_sdot(isum1, lo, a0);
2954 isum2 = neon_sdot(isum2, hi, a1);
2955 }
2956 acc += d
2957 * da
2958 * (sc1 as f32 * vaddvq_s32(isum1) as f32
2959 + sc2 as f32 * vaddvq_s32(isum2) as f32);
2960 q_off += 32;
2961 base += 64;
2962 is += 2;
2963 u1 <<= 2;
2964 u2 <<= 2;
2965 }
2966 }
2967 acc
2968 }
2969
2970 #[target_feature(enable = "neon,dotprod")]
2972 pub unsafe fn gemm_q5_k_q8_neon_sdot(
2973 row_bytes: &[u8],
2974 acts: &[Q8KActivations],
2975 out: &mut [f32],
2976 ) {
2977 debug_assert_eq!(out.len(), acts.len());
2978 debug_assert!(acts.len() <= super::Q5_K_GEMM_NC);
2979 out.fill(0.0);
2980 if acts.is_empty() {
2981 return;
2982 }
2983 let low_mask = vdupq_n_u8(0x0F);
2984 let sixteen = vdupq_n_u8(16);
2985 let n = acts.len();
2986 for (b, block) in row_bytes
2987 .as_chunks::<Q5_K_BLOCK_BYTES>()
2988 .0
2989 .iter()
2990 .enumerate()
2991 {
2992 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
2993 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
2994 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
2995 let qh = block.as_ptr().add(16);
2996 let qs = &block[48..176];
2997 let mut mins = [0u8; 8];
2998 let mut sc_only = [0u8; 8];
2999 for i in 0..8 {
3000 let (s, m) = q4_k_scale_min(i, &scales);
3001 sc_only[i] = s;
3002 mins[i] = m;
3003 }
3004 for j in 0..n {
3005 let act = &acts[j];
3006 let da = act.d[b];
3007 let bsums = &act.bsums[b * 16..(b + 1) * 16];
3008 let mut sum_min = 0i32;
3009 for i in 0..8 {
3010 sum_min += mins[i] as i32 * (bsums[2 * i] as i32 + bsums[2 * i + 1] as i32);
3011 }
3012 out[j] -= dmin * da * sum_min as f32;
3013 }
3014 let mut q_off = 0usize;
3015 let mut base = 0usize;
3016 let mut is = 0usize;
3017 let (mut u1, mut u2) = (1u8, 2u8);
3018 for _ in 0..4 {
3019 let sc1 = sc_only[is];
3020 let sc2 = sc_only[is + 1];
3021 let u1_vec = vdupq_n_u8(u1);
3022 let u2_vec = vdupq_n_u8(u2);
3023 let mut lo_cols = [vreinterpretq_s8_u8(vdupq_n_u8(0)); 2];
3025 let mut hi_cols = [vreinterpretq_s8_u8(vdupq_n_u8(0)); 2];
3026 for g in 0..2 {
3027 let packed = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
3028 let qh16 = vld1q_u8(qh.add(g * 16));
3029 let lo_nib = vandq_u8(packed, low_mask);
3030 let hi_nib = vshrq_n_u8(packed, 4);
3031 let hi_bit1 = vandq_u8(vtstq_u8(qh16, u1_vec), sixteen);
3032 let hi_bit2 = vandq_u8(vtstq_u8(qh16, u2_vec), sixteen);
3033 lo_cols[g] = vreinterpretq_s8_u8(vorrq_u8(lo_nib, hi_bit1));
3034 hi_cols[g] = vreinterpretq_s8_u8(vorrq_u8(hi_nib, hi_bit2));
3035 }
3036 for j in 0..n {
3037 let q8 = acts[j].q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
3038 let da = acts[j].d[b];
3039 let mut isum1 = vdupq_n_s32(0);
3040 let mut isum2 = vdupq_n_s32(0);
3041 for g in 0..2 {
3042 let a0 = vld1q_s8(q8.add(base + g * 16));
3043 let a1 = vld1q_s8(q8.add(base + 32 + g * 16));
3044 isum1 = neon_sdot(isum1, lo_cols[g], a0);
3045 isum2 = neon_sdot(isum2, hi_cols[g], a1);
3046 }
3047 out[j] += d
3048 * da
3049 * (sc1 as f32 * vaddvq_s32(isum1) as f32
3050 + sc2 as f32 * vaddvq_s32(isum2) as f32);
3051 }
3052 q_off += 32;
3053 base += 64;
3054 is += 2;
3055 u1 <<= 2;
3056 u2 <<= 2;
3057 }
3058 }
3059 }
3060
3061 #[target_feature(enable = "neon,dotprod")]
3064 pub unsafe fn gemm_q6_k_q8_neon_sdot(
3065 row_bytes: &[u8],
3066 acts: &[Q8KActivations],
3067 out: &mut [f32],
3068 ) {
3069 debug_assert_eq!(out.len(), acts.len());
3070 debug_assert!(acts.len() <= super::Q6_K_GEMM_NC);
3071 out.fill(0.0);
3072 let n = acts.len();
3073 if n == 0 {
3074 return;
3075 }
3076 let m4b = vdupq_n_u8(0x0F);
3077 let mone = vdupq_n_u8(3);
3078 for (b, block) in row_bytes
3079 .as_chunks::<Q6_K_BLOCK_BYTES>()
3080 .0
3081 .iter()
3082 .enumerate()
3083 {
3084 let d_all = f16::from_le_bytes([block[208], block[209]]).to_f32();
3085 let ql = block.as_ptr();
3086 let qh = block.as_ptr().add(128);
3087 let scale = block.as_ptr().add(192) as *const i8;
3088 let scales = vld1q_s8(scale);
3089 let q6scales0 = vmovl_s8(vget_low_s8(scales));
3090 let q6scales1 = vmovl_s8(vget_high_s8(scales));
3091
3092 let mut isum_mins = [0i32; 4];
3093 let mut isums = [0i32; 4];
3094 for j in 0..n {
3095 let bsums = acts[j].bsums.as_ptr().add(b * 16);
3096 let q8sums0 = vld1q_s16(bsums);
3097 let q8sums1 = vld1q_s16(bsums.add(8));
3098 let prod = vaddq_s32(
3099 vaddq_s32(
3100 vmull_s16(vget_low_s16(q8sums0), vget_low_s16(q6scales0)),
3101 vmull_s16(vget_high_s16(q8sums0), vget_high_s16(q6scales0)),
3102 ),
3103 vaddq_s32(
3104 vmull_s16(vget_low_s16(q8sums1), vget_low_s16(q6scales1)),
3105 vmull_s16(vget_high_s16(q8sums1), vget_high_s16(q6scales1)),
3106 ),
3107 );
3108 isum_mins[j] = vaddvq_s32(prod);
3109 }
3110
3111 for half in 0..2usize {
3112 let q6 = ql.add(half * 64);
3113 let qhp = qh.add(half * 32);
3114 let sc = scale.add(half * 8);
3115 let act_off = half * 128;
3116
3117 let qh0 = vld1q_u8(qhp);
3118 let qh1 = vld1q_u8(qhp.add(16));
3119 let q6_0 = vld1q_u8(q6);
3120 let q6_1 = vld1q_u8(q6.add(16));
3121 let q6_2 = vld1q_u8(q6.add(32));
3122 let q6_3 = vld1q_u8(q6.add(48));
3123
3124 let h0 = vshlq_n_u8(vandq_u8(mone, qh0), 4);
3125 let h1 = vshlq_n_u8(vandq_u8(mone, qh1), 4);
3126 let mut shifted = vshrq_n_u8(qh0, 2);
3127 let h2 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3128 shifted = vshrq_n_u8(qh1, 2);
3129 let h3 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3130 let wb0 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_0, m4b), h0));
3131 let wb1 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_1, m4b), h1));
3132 let wb2 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_2, m4b), h2));
3133 let wb3 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_3, m4b), h3));
3134 let sc0 = *sc.add(0) as i32;
3135 let sc1 = *sc.add(1) as i32;
3136 let sc2 = *sc.add(2) as i32;
3137 let sc3 = *sc.add(3) as i32;
3138 let z = vdupq_n_s32(0);
3139 for j in 0..n {
3140 let q8p = acts[j].q.as_ptr().add(b * Q6_K_BLOCK_ELEMS + act_off);
3141 isums[j] += vaddvq_s32(neon_sdot(z, wb0, vld1q_s8(q8p))) * sc0
3142 + vaddvq_s32(neon_sdot(z, wb1, vld1q_s8(q8p.add(16)))) * sc1
3143 + vaddvq_s32(neon_sdot(z, wb2, vld1q_s8(q8p.add(32)))) * sc2
3144 + vaddvq_s32(neon_sdot(z, wb3, vld1q_s8(q8p.add(48)))) * sc3;
3145 }
3146
3147 shifted = vshrq_n_u8(qh0, 4);
3148 let h0 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3149 shifted = vshrq_n_u8(qh1, 4);
3150 let h1 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3151 shifted = vshrq_n_u8(qh0, 6);
3152 let h2 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3153 shifted = vshrq_n_u8(qh1, 6);
3154 let h3 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3155 let wb0 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_0, 4), h0));
3156 let wb1 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_1, 4), h1));
3157 let wb2 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_2, 4), h2));
3158 let wb3 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_3, 4), h3));
3159 let sc0 = *sc.add(4) as i32;
3160 let sc1 = *sc.add(5) as i32;
3161 let sc2 = *sc.add(6) as i32;
3162 let sc3 = *sc.add(7) as i32;
3163 for j in 0..n {
3164 let q8p = acts[j].q.as_ptr().add(b * Q6_K_BLOCK_ELEMS + act_off + 64);
3165 isums[j] += vaddvq_s32(neon_sdot(z, wb0, vld1q_s8(q8p))) * sc0
3166 + vaddvq_s32(neon_sdot(z, wb1, vld1q_s8(q8p.add(16)))) * sc1
3167 + vaddvq_s32(neon_sdot(z, wb2, vld1q_s8(q8p.add(32)))) * sc2
3168 + vaddvq_s32(neon_sdot(z, wb3, vld1q_s8(q8p.add(48)))) * sc3;
3169 }
3170 }
3171 for j in 0..n {
3172 out[j] += d_all * acts[j].d[b] * (isums[j] - 32 * isum_mins[j]) as f32;
3173 }
3174 }
3175 }
3176
3177 #[target_feature(enable = "neon,dotprod")]
3181 pub unsafe fn dot_q6_k_q8_neon_sdot(row_bytes: &[u8], act: &Q8KActivations) -> f32 {
3182 debug_assert_eq!(row_bytes.len() % Q6_K_BLOCK_BYTES, 0);
3183 debug_assert_eq!(row_bytes.len() / Q6_K_BLOCK_BYTES, act.n_blocks());
3184 let m4b = vdupq_n_u8(0x0F);
3185 let mone = vdupq_n_u8(3);
3186 let mut acc = 0f32;
3187 for (b, block) in row_bytes
3188 .as_chunks::<Q6_K_BLOCK_BYTES>()
3189 .0
3190 .iter()
3191 .enumerate()
3192 {
3193 let d_all = f16::from_le_bytes([block[208], block[209]]).to_f32();
3194 let da = act.d[b];
3195 let ql = block.as_ptr();
3196 let qh = block.as_ptr().add(128);
3197 let scale = block.as_ptr().add(192) as *const i8;
3198 let q8 = act.q.as_ptr().add(b * Q6_K_BLOCK_ELEMS);
3199 let bsums = act.bsums.as_ptr().add(b * 16);
3200
3201 let scales = vld1q_s8(scale);
3202 let q6scales0 = vmovl_s8(vget_low_s8(scales));
3203 let q6scales1 = vmovl_s8(vget_high_s8(scales));
3204 let q8sums0 = vld1q_s16(bsums);
3205 let q8sums1 = vld1q_s16(bsums.add(8));
3206 let prod = vaddq_s32(
3207 vaddq_s32(
3208 vmull_s16(vget_low_s16(q8sums0), vget_low_s16(q6scales0)),
3209 vmull_s16(vget_high_s16(q8sums0), vget_high_s16(q6scales0)),
3210 ),
3211 vaddq_s32(
3212 vmull_s16(vget_low_s16(q8sums1), vget_low_s16(q6scales1)),
3213 vmull_s16(vget_high_s16(q8sums1), vget_high_s16(q6scales1)),
3214 ),
3215 );
3216 let isum_mins = vaddvq_s32(prod);
3217 let mut isum = 0i32;
3218 let mut q6 = ql;
3219 let mut qhp = qh;
3220 let mut q8p = q8;
3221 let mut sc = scale;
3222 for _ in 0..2 {
3223 let qh0 = vld1q_u8(qhp);
3224 let qh1 = vld1q_u8(qhp.add(16));
3225 qhp = qhp.add(32);
3226 let q6_0 = vld1q_u8(q6);
3227 let q6_1 = vld1q_u8(q6.add(16));
3228 let q6_2 = vld1q_u8(q6.add(32));
3229 let q6_3 = vld1q_u8(q6.add(48));
3230 q6 = q6.add(64);
3231 let q8_0 = vld1q_s8(q8p);
3232 let q8_1 = vld1q_s8(q8p.add(16));
3233 let q8_2 = vld1q_s8(q8p.add(32));
3234 let q8_3 = vld1q_s8(q8p.add(48));
3235 q8p = q8p.add(64);
3236
3237 let h0 = vshlq_n_u8(vandq_u8(mone, qh0), 4);
3238 let h1 = vshlq_n_u8(vandq_u8(mone, qh1), 4);
3239 let mut shifted = vshrq_n_u8(qh0, 2);
3240 let h2 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3241 shifted = vshrq_n_u8(qh1, 2);
3242 let h3 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3243
3244 let b0 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_0, m4b), h0));
3245 let b1 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_1, m4b), h1));
3246 let b2 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_2, m4b), h2));
3247 let b3 = vreinterpretq_s8_u8(vorrq_u8(vandq_u8(q6_3, m4b), h3));
3248 let z = vdupq_n_s32(0);
3249 isum += vaddvq_s32(neon_sdot(z, b0, q8_0)) * (*sc.add(0) as i32)
3250 + vaddvq_s32(neon_sdot(z, b1, q8_1)) * (*sc.add(1) as i32)
3251 + vaddvq_s32(neon_sdot(z, b2, q8_2)) * (*sc.add(2) as i32)
3252 + vaddvq_s32(neon_sdot(z, b3, q8_3)) * (*sc.add(3) as i32);
3253 sc = sc.add(4);
3254
3255 let q8_0 = vld1q_s8(q8p);
3256 let q8_1 = vld1q_s8(q8p.add(16));
3257 let q8_2 = vld1q_s8(q8p.add(32));
3258 let q8_3 = vld1q_s8(q8p.add(48));
3259 q8p = q8p.add(64);
3260 shifted = vshrq_n_u8(qh0, 4);
3261 let h0 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3262 shifted = vshrq_n_u8(qh1, 4);
3263 let h1 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3264 shifted = vshrq_n_u8(qh0, 6);
3265 let h2 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3266 shifted = vshrq_n_u8(qh1, 6);
3267 let h3 = vshlq_n_u8(vandq_u8(mone, shifted), 4);
3268 let b0 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_0, 4), h0));
3269 let b1 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_1, 4), h1));
3270 let b2 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_2, 4), h2));
3271 let b3 = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q6_3, 4), h3));
3272 isum += vaddvq_s32(neon_sdot(z, b0, q8_0)) * (*sc.add(0) as i32)
3273 + vaddvq_s32(neon_sdot(z, b1, q8_1)) * (*sc.add(1) as i32)
3274 + vaddvq_s32(neon_sdot(z, b2, q8_2)) * (*sc.add(2) as i32)
3275 + vaddvq_s32(neon_sdot(z, b3, q8_3)) * (*sc.add(3) as i32);
3276 sc = sc.add(4);
3277 }
3278 acc += d_all * da * (isum - 32 * isum_mins) as f32;
3279 }
3280 acc
3281 }
3282
3283 #[target_feature(enable = "neon")]
3292 pub unsafe fn dot_q4_0_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3293 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
3294 let bias = vdupq_n_s16(8);
3295 let low_mask = vdupq_n_u8(0x0F);
3296
3297 let mut acc = 0f32;
3298 for (b, block) in row_bytes
3299 .as_chunks::<Q4_0_BLOCK_BYTES>()
3300 .0
3301 .iter()
3302 .enumerate()
3303 {
3304 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
3305 let base = b * Q4_0_BLOCK_ELEMS;
3306 let nibbles = vld1q_u8(block.as_ptr().add(2));
3307
3308 let lo_nibbles = vandq_u8(nibbles, low_mask); let hi_nibbles = vshrq_n_u8(nibbles, 4); let mut block_acc = vdupq_n_f32(0.0);
3312 for (group_idx, nib_u8) in [lo_nibbles, hi_nibbles].into_iter().enumerate() {
3313 let lo16 = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(nib_u8))), bias);
3314 let hi16 = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(nib_u8))), bias);
3315 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
3316 let lo32 = vmovl_s16(vget_low_s16(half16));
3317 let hi32 = vmovl_s16(vget_high_s16(half16));
3318 let f_lo = vcvtq_f32_s32(lo32);
3319 let f_hi = vcvtq_f32_s32(hi32);
3320 let elem_base = base + group_idx * 16 + half_idx * 8;
3321 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
3322 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
3323 block_acc = vfmaq_f32(block_acc, f_lo, x_lo);
3324 block_acc = vfmaq_f32(block_acc, f_hi, x_hi);
3325 }
3326 }
3327 acc += vaddvq_f32(block_acc) * scale;
3328 }
3329 acc
3330 }
3331
3332 #[inline]
3338 #[target_feature(enable = "neon")]
3339 unsafe fn widen_u8x16_to_f32_quads(
3340 v: uint8x16_t,
3341 ) -> (float32x4_t, float32x4_t, float32x4_t, float32x4_t) {
3342 let u16_lo = vmovl_u8(vget_low_u8(v)); let u16_hi = vmovl_u8(vget_high_u8(v)); (
3345 vcvtq_f32_u32(vmovl_u16(vget_low_u16(u16_lo))), vcvtq_f32_u32(vmovl_u16(vget_high_u16(u16_lo))), vcvtq_f32_u32(vmovl_u16(vget_low_u16(u16_hi))), vcvtq_f32_u32(vmovl_u16(vget_high_u16(u16_hi))), )
3350 }
3351
3352 #[inline]
3359 #[target_feature(enable = "neon")]
3360 unsafe fn fma_affine16(
3361 quads: (float32x4_t, float32x4_t, float32x4_t, float32x4_t),
3362 d: f32,
3363 min_vec: float32x4_t,
3364 x: &[f32],
3365 x_base: usize,
3366 mut acc: float32x4_t,
3367 ) -> float32x4_t {
3368 let (q0, q1, q2, q3) = quads;
3369 let mut i = 0usize;
3370 for q in [q0, q1, q2, q3] {
3371 let w = vsubq_f32(vmulq_n_f32(q, d), min_vec);
3372 let xv = vld1q_f32(x.as_ptr().add(x_base + i));
3373 acc = vfmaq_f32(acc, w, xv);
3374 i += 4;
3375 }
3376 acc
3377 }
3378
3379 #[target_feature(enable = "neon")]
3388 pub unsafe fn dot_q4_k_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3389 debug_assert_eq!(row_bytes.len() % Q4_K_BLOCK_BYTES, 0);
3390 let low_mask = vdupq_n_u8(0x0F);
3391 let mut acc = 0f32;
3392 let mut x_base = 0usize;
3393 for block in row_bytes.as_chunks::<Q4_K_BLOCK_BYTES>().0 {
3394 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
3395 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
3396 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
3397 let qs = &block[16..144];
3398
3399 let mut vec_acc = vdupq_n_f32(0.0);
3402 let mut is = 0usize;
3403 let mut q_off = 0usize;
3404 for _ in 0..4 {
3405 let (sc1, m1) = q4_k_scale_min(is, &scales);
3406 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
3407 let d1 = d * sc1 as f32;
3408 let min1_vec = vdupq_n_f32(dmin * m1 as f32);
3409 let d2 = d * sc2 as f32;
3410 let min2_vec = vdupq_n_f32(dmin * m2 as f32);
3411
3412 for g in 0..2 {
3413 let raw16 = vld1q_u8(qs.as_ptr().add(q_off + g * 16));
3414 let lo_nib = vandq_u8(raw16, low_mask);
3415 let hi_nib = vshrq_n_u8(raw16, 4);
3416 vec_acc = fma_affine16(
3417 widen_u8x16_to_f32_quads(lo_nib),
3418 d1,
3419 min1_vec,
3420 x,
3421 x_base + g * 16,
3422 vec_acc,
3423 );
3424 vec_acc = fma_affine16(
3425 widen_u8x16_to_f32_quads(hi_nib),
3426 d2,
3427 min2_vec,
3428 x,
3429 x_base + 32 + g * 16,
3430 vec_acc,
3431 );
3432 }
3433 q_off += 32;
3434 x_base += 64;
3435 is += 2;
3436 }
3437 acc += vaddvq_f32(vec_acc);
3438 }
3439 acc
3440 }
3441
3442 #[target_feature(enable = "neon")]
3452 pub unsafe fn dot_q5_k_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3453 debug_assert_eq!(row_bytes.len() % Q5_K_BLOCK_BYTES, 0);
3454 let low_mask = vdupq_n_u8(0x0F);
3455 let sixteen = vdupq_n_u8(16);
3456 let mut acc = 0f32;
3457 let mut x_base = 0usize;
3458 for block in row_bytes.as_chunks::<Q5_K_BLOCK_BYTES>().0 {
3459 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
3460 let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
3461 let scales: [u8; Q4_K_SCALE_BYTES] = block[4..16].try_into().unwrap();
3462 let qh = &block[16..48];
3463 let qs = &block[48..176];
3464
3465 let mut is = 0usize;
3466 let (mut u1, mut u2) = (1u8, 2u8);
3467 for oi in 0..4 {
3468 let (sc1, m1) = q4_k_scale_min(is, &scales);
3469 let (sc2, m2) = q4_k_scale_min(is + 1, &scales);
3470 let d1 = d * sc1 as f32;
3471 let min1_vec = vdupq_n_f32(dmin * m1 as f32);
3472 let d2 = d * sc2 as f32;
3473 let min2_vec = vdupq_n_f32(dmin * m2 as f32);
3474 let ql = &qs[oi * 32..oi * 32 + 32];
3475 let u1_vec = vdupq_n_u8(u1);
3476 let u2_vec = vdupq_n_u8(u2);
3477
3478 let mut lo_acc = vdupq_n_f32(0.0);
3479 let mut hi_acc = vdupq_n_f32(0.0);
3480 for g in 0..2 {
3481 let raw16 = vld1q_u8(ql.as_ptr().add(g * 16));
3482 let qh16 = vld1q_u8(qh.as_ptr().add(g * 16));
3483
3484 let lo_nib = vandq_u8(raw16, low_mask);
3485 let hi_nib = vshrq_n_u8(raw16, 4);
3486 let hi_bit1 = vandq_u8(vtstq_u8(qh16, u1_vec), sixteen);
3487 let hi_bit2 = vandq_u8(vtstq_u8(qh16, u2_vec), sixteen);
3488
3489 lo_acc = fma_affine16(
3490 widen_u8x16_to_f32_quads(vorrq_u8(lo_nib, hi_bit1)),
3491 d1,
3492 min1_vec,
3493 x,
3494 x_base + g * 16,
3495 lo_acc,
3496 );
3497 hi_acc = fma_affine16(
3498 widen_u8x16_to_f32_quads(vorrq_u8(hi_nib, hi_bit2)),
3499 d2,
3500 min2_vec,
3501 x,
3502 x_base + 32 + g * 16,
3503 hi_acc,
3504 );
3505 }
3506 acc += vaddvq_f32(lo_acc) + vaddvq_f32(hi_acc);
3507 x_base += 64;
3508 is += 2;
3509 u1 <<= 2;
3510 u2 <<= 2;
3511 }
3512 }
3513 acc
3514 }
3515
3516 #[inline]
3524 #[target_feature(enable = "neon")]
3525 unsafe fn widen_u8x16_centered_to_f32_quads(
3526 v: uint8x16_t,
3527 bias16: int16x8_t,
3528 ) -> (float32x4_t, float32x4_t, float32x4_t, float32x4_t) {
3529 let s16_lo = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(v))), bias16);
3530 let s16_hi = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(v))), bias16);
3531 (
3532 vcvtq_f32_s32(vmovl_s16(vget_low_s16(s16_lo))),
3533 vcvtq_f32_s32(vmovl_s16(vget_high_s16(s16_lo))),
3534 vcvtq_f32_s32(vmovl_s16(vget_low_s16(s16_hi))),
3535 vcvtq_f32_s32(vmovl_s16(vget_high_s16(s16_hi))),
3536 )
3537 }
3538
3539 #[inline]
3546 #[target_feature(enable = "neon")]
3547 unsafe fn fma_scaled16(
3548 quads: (float32x4_t, float32x4_t, float32x4_t, float32x4_t),
3549 scale: f32,
3550 x: &[f32],
3551 x_base: usize,
3552 mut acc: float32x4_t,
3553 ) -> float32x4_t {
3554 let (q0, q1, q2, q3) = quads;
3555 let mut i = 0usize;
3556 for q in [q0, q1, q2, q3] {
3557 let xv = vld1q_f32(x.as_ptr().add(x_base + i));
3558 acc = vfmaq_f32(acc, vmulq_n_f32(q, scale), xv);
3559 i += 4;
3560 }
3561 acc
3562 }
3563
3564 #[inline]
3576 #[target_feature(enable = "neon")]
3577 #[allow(clippy::too_many_arguments)]
3578 unsafe fn q6_k_group<const QH_SHIFT: i32, const HI_NIBBLE: bool>(
3579 ql: &[u8],
3580 ql_off: usize,
3581 qh: &[u8],
3582 sc: &[u8],
3583 sc_base: usize,
3584 d: f32,
3585 x: &[f32],
3586 x_base: usize,
3587 out_off: usize,
3588 low_mask: uint8x16_t,
3589 two_bit_mask: uint8x16_t,
3590 bias16: int16x8_t,
3591 ) -> f32 {
3592 let mut acc = 0f32;
3593 for sub in 0..2usize {
3594 let byte_off = sub * 16;
3595 let ql_raw = vld1q_u8(ql.as_ptr().add(ql_off + byte_off));
3596 let qh_raw = vld1q_u8(qh.as_ptr().add(byte_off));
3597
3598 let nib = if HI_NIBBLE {
3599 vshrq_n_u8::<4>(ql_raw)
3600 } else {
3601 vandq_u8(ql_raw, low_mask)
3602 };
3603 let qh_field = vandq_u8(vshrq_n_u8::<QH_SHIFT>(qh_raw), two_bit_mask);
3611 let raw6 = vorrq_u8(nib, vshlq_n_u8::<4>(qh_field));
3612
3613 let scale = d * (sc[sc_base + sub] as i8) as f32;
3614 let quads = widen_u8x16_centered_to_f32_quads(raw6, bias16);
3615 let acc_vec = fma_scaled16(
3616 quads,
3617 scale,
3618 x,
3619 x_base + out_off + sub * 16,
3620 vdupq_n_f32(0.0),
3621 );
3622 acc += vaddvq_f32(acc_vec);
3623 }
3624 acc
3625 }
3626
3627 #[inline]
3633 #[target_feature(enable = "neon")]
3634 #[allow(clippy::too_many_arguments)]
3635 unsafe fn q6_k_group_q1(
3636 ql: &[u8],
3637 qh: &[u8],
3638 sc: &[u8],
3639 d: f32,
3640 x: &[f32],
3641 x_base: usize,
3642 low_mask: uint8x16_t,
3643 two_bit_mask: uint8x16_t,
3644 bias16: int16x8_t,
3645 ) -> f32 {
3646 let mut acc = 0f32;
3647 #[allow(clippy::needless_range_loop)]
3651 for sub in 0..2usize {
3652 let byte_off = sub * 16;
3653 let ql_raw = vld1q_u8(ql.as_ptr().add(byte_off));
3654 let qh_raw = vld1q_u8(qh.as_ptr().add(byte_off));
3655
3656 let nib = vandq_u8(ql_raw, low_mask);
3657 let qh_field = vandq_u8(qh_raw, two_bit_mask);
3658 let raw6 = vorrq_u8(nib, vshlq_n_u8::<4>(qh_field));
3659
3660 let scale = d * (sc[sub] as i8) as f32;
3661 let quads = widen_u8x16_centered_to_f32_quads(raw6, bias16);
3662 let acc_vec = fma_scaled16(quads, scale, x, x_base + sub * 16, vdupq_n_f32(0.0));
3663 acc += vaddvq_f32(acc_vec);
3664 }
3665 acc
3666 }
3667
3668 #[target_feature(enable = "neon")]
3674 pub unsafe fn dot_q6_k_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3675 debug_assert_eq!(row_bytes.len() % Q6_K_BLOCK_BYTES, 0);
3676 debug_assert_eq!(
3677 row_bytes.len() / Q6_K_BLOCK_BYTES * Q6_K_BLOCK_ELEMS,
3678 x.len()
3679 );
3680 let low_mask = vdupq_n_u8(0x0F);
3681 let two_bit_mask = vdupq_n_u8(0x03);
3682 let bias16 = vdupq_n_s16(32);
3683
3684 let mut acc = 0f32;
3685 let mut x_base = 0usize;
3686 for block in row_bytes.as_chunks::<Q6_K_BLOCK_BYTES>().0 {
3687 let ql_full = &block[0..128];
3688 let qh_full = &block[128..192];
3689 let sc_full = &block[192..208];
3690 let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
3691
3692 for half in 0..2 {
3693 let ql = &ql_full[half * 64..half * 64 + 64];
3694 let qh = &qh_full[half * 32..half * 32 + 32];
3695 let sc = &sc_full[half * 8..half * 8 + 8];
3696 let half_base = x_base + half * 128;
3697
3698 acc += q6_k_group_q1(ql, qh, sc, d, x, half_base, low_mask, two_bit_mask, bias16);
3700 acc += q6_k_group::<2, false>(
3702 ql,
3703 32,
3704 qh,
3705 sc,
3706 2,
3707 d,
3708 x,
3709 half_base,
3710 32,
3711 low_mask,
3712 two_bit_mask,
3713 bias16,
3714 );
3715 acc += q6_k_group::<4, true>(
3717 ql,
3718 0,
3719 qh,
3720 sc,
3721 4,
3722 d,
3723 x,
3724 half_base,
3725 64,
3726 low_mask,
3727 two_bit_mask,
3728 bias16,
3729 );
3730 acc += q6_k_group::<6, true>(
3732 ql,
3733 32,
3734 qh,
3735 sc,
3736 6,
3737 d,
3738 x,
3739 half_base,
3740 96,
3741 low_mask,
3742 two_bit_mask,
3743 bias16,
3744 );
3745 }
3746 x_base += Q6_K_BLOCK_ELEMS;
3747 }
3748 acc
3749 }
3750
3751 #[inline]
3768 #[target_feature(enable = "neon")]
3769 unsafe fn mxfp4_nibbles_to_f32_quads(
3770 nib: uint8x16_t,
3771 ) -> (float32x4_t, float32x4_t, float32x4_t, float32x4_t) {
3772 let sign_bit = vandq_u8(nib, vdupq_n_u8(0x8));
3773 let e = vandq_u8(vshrq_n_u8(nib, 1), vdupq_n_u8(0x3));
3774 let m = vandq_u8(nib, vdupq_n_u8(0x1));
3775
3776 let pow2_table: [u8; 16] = [1, 1, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
3777 let bias_table: [u8; 16] = [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
3778 let pow2_u8 = vqtbl1q_u8(vld1q_u8(pow2_table.as_ptr()), e);
3779 let bias_u8 = vqtbl1q_u8(vld1q_u8(bias_table.as_ptr()), e);
3780
3781 let (p0, p1, p2, p3) = widen_u8x16_to_f32_quads(pow2_u8);
3782 let (b0, b1, b2, b3) = widen_u8x16_to_f32_quads(bias_u8);
3783 let (m0, m1, m2, m3) = widen_u8x16_to_f32_quads(m);
3784 let (s0, s1, s2, s3) = widen_u8x16_to_f32_quads(sign_bit);
3785
3786 let half = vdupq_n_f32(0.5);
3787 let quarter = vdupq_n_f32(0.25);
3788 let one = vdupq_n_f32(1.0);
3789
3790 let decode = |p: float32x4_t, b: float32x4_t, m: float32x4_t, s: float32x4_t| {
3791 let magnitude = vmulq_f32(p, vfmaq_f32(b, m, half)); let sign_mul = vfmsq_f32(one, s, quarter); vmulq_f32(magnitude, sign_mul)
3794 };
3795
3796 (
3797 decode(p0, b0, m0, s0),
3798 decode(p1, b1, m1, s1),
3799 decode(p2, b2, m2, s2),
3800 decode(p3, b3, m3, s3),
3801 )
3802 }
3803
3804 #[target_feature(enable = "neon")]
3814 pub unsafe fn dot_mxfp4_row_f32_neon(packed: &[u8], scales: &[u8], x: &[f32]) -> f32 {
3815 debug_assert_eq!(packed.len(), scales.len() * (MXFP4_GROUP_SIZE / 2));
3816 let low_mask = vdupq_n_u8(0x0F);
3817 let mut acc = 0f32;
3818 let mut x_base = 0usize;
3819 for (g, &e_byte) in scales.iter().enumerate() {
3820 let d = e8m0_scale(e_byte);
3821 let group = &packed[g * 16..(g + 1) * 16];
3822 let bytes = vld1q_u8(group.as_ptr());
3823 let lo_nib = vandq_u8(bytes, low_mask);
3824 let hi_nib = vshrq_n_u8(bytes, 4);
3825
3826 let mut block_acc = vdupq_n_f32(0.0);
3827 for (half_idx, nib) in [lo_nib, hi_nib].into_iter().enumerate() {
3828 let (v0, v1, v2, v3) = mxfp4_nibbles_to_f32_quads(nib);
3829 let elem_base = x_base + half_idx * 16;
3830 for (i, v) in [v0, v1, v2, v3].into_iter().enumerate() {
3831 let xv = vld1q_f32(x.as_ptr().add(elem_base + i * 4));
3832 block_acc = vfmaq_f32(block_acc, v, xv);
3833 }
3834 }
3835 acc += vaddvq_f32(block_acc) * d;
3836 x_base += MXFP4_GROUP_SIZE;
3837 }
3838 acc
3839 }
3840
3841 #[target_feature(enable = "neon")]
3845 pub unsafe fn dot_q8_1_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3846 debug_assert_eq!(row_bytes.len() % Q8_1_BLOCK_BYTES, 0);
3847 let mut acc = 0f32;
3848 for (b, block) in row_bytes
3849 .as_chunks::<Q8_1_BLOCK_BYTES>()
3850 .0
3851 .iter()
3852 .enumerate()
3853 {
3854 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
3855 let base = b * Q8_1_BLOCK_ELEMS;
3856 let qs = &block[4..36];
3857
3858 let mut block_acc = vdupq_n_f32(0.0);
3859 for g in 0..2 {
3860 let raw16 = vld1q_s8(qs.as_ptr().add(g * 16) as *const i8);
3861 let lo16 = vmovl_s8(vget_low_s8(raw16));
3862 let hi16 = vmovl_s8(vget_high_s8(raw16));
3863 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
3864 let lo32 = vmovl_s16(vget_low_s16(half16));
3865 let hi32 = vmovl_s16(vget_high_s16(half16));
3866 let f_lo = vcvtq_f32_s32(lo32);
3867 let f_hi = vcvtq_f32_s32(hi32);
3868 let elem_base = base + g * 16 + half_idx * 8;
3869 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
3870 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
3871 block_acc = vfmaq_f32(block_acc, f_lo, x_lo);
3872 block_acc = vfmaq_f32(block_acc, f_hi, x_hi);
3873 }
3874 }
3875 acc += vaddvq_f32(block_acc) * scale;
3876 }
3877 acc
3878 }
3879
3880 #[target_feature(enable = "neon")]
3886 pub unsafe fn dot_q4_1_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3887 debug_assert_eq!(row_bytes.len() % Q4_1_BLOCK_BYTES, 0);
3888 let low_mask = vdupq_n_u8(0x0F);
3889
3890 let mut acc = 0f32;
3891 for (b, block) in row_bytes
3892 .as_chunks::<Q4_1_BLOCK_BYTES>()
3893 .0
3894 .iter()
3895 .enumerate()
3896 {
3897 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
3898 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
3899 let base = b * Q4_1_BLOCK_ELEMS;
3900 let nibbles = vld1q_u8(block.as_ptr().add(4));
3901
3902 let lo_nibbles = vandq_u8(nibbles, low_mask); let hi_nibbles = vshrq_n_u8(nibbles, 4); let mut block_acc = vdupq_n_f32(0.0);
3906 for (group_idx, nib_u8) in [lo_nibbles, hi_nibbles].into_iter().enumerate() {
3907 let lo16 = vmovl_u8(vget_low_u8(nib_u8));
3908 let hi16 = vmovl_u8(vget_high_u8(nib_u8));
3909 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
3910 let lo32 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(half16)));
3911 let hi32 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(half16)));
3912 let elem_base = base + group_idx * 16 + half_idx * 8;
3913 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
3914 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
3915 let w_lo = vfmaq_n_f32(vdupq_n_f32(m), lo32, d);
3916 let w_hi = vfmaq_n_f32(vdupq_n_f32(m), hi32, d);
3917 block_acc = vfmaq_f32(block_acc, w_lo, x_lo);
3918 block_acc = vfmaq_f32(block_acc, w_hi, x_hi);
3919 }
3920 }
3921 acc += vaddvq_f32(block_acc);
3922 }
3923 acc
3924 }
3925
3926 #[target_feature(enable = "neon")]
3932 pub unsafe fn dot_q5_0_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3933 debug_assert_eq!(row_bytes.len() % Q5_0_BLOCK_BYTES, 0);
3934 let mut acc = 0f32;
3935 for (b, block) in row_bytes
3936 .as_chunks::<Q5_0_BLOCK_BYTES>()
3937 .0
3938 .iter()
3939 .enumerate()
3940 {
3941 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
3942 let qh = u32::from_le_bytes(block[2..6].try_into().unwrap());
3943 let qs = &block[6..22];
3944 let base = b * Q5_0_BLOCK_ELEMS;
3945
3946 let mut vals = [0i8; 32];
3947 for j in 0..16 {
3948 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
3949 vals[j] = (((qs[j] & 0x0F) | xh_0) as i32 - 16) as i8;
3950 vals[j + 16] = (((qs[j] >> 4) | xh_1) as i32 - 16) as i8;
3951 }
3952
3953 let mut block_acc = vdupq_n_f32(0.0);
3954 for g in 0..2 {
3955 let raw16 = vld1q_s8(vals.as_ptr().add(g * 16));
3956 let lo16 = vmovl_s8(vget_low_s8(raw16));
3957 let hi16 = vmovl_s8(vget_high_s8(raw16));
3958 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
3959 let lo32 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half16)));
3960 let hi32 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half16)));
3961 let elem_base = base + g * 16 + half_idx * 8;
3962 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
3963 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
3964 block_acc = vfmaq_f32(block_acc, lo32, x_lo);
3965 block_acc = vfmaq_f32(block_acc, hi32, x_hi);
3966 }
3967 }
3968 acc += vaddvq_f32(block_acc) * d;
3969 }
3970 acc
3971 }
3972
3973 #[target_feature(enable = "neon")]
3977 pub unsafe fn dot_q5_1_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
3978 debug_assert_eq!(row_bytes.len() % Q5_1_BLOCK_BYTES, 0);
3979 let mut acc = 0f32;
3980 for (b, block) in row_bytes
3981 .as_chunks::<Q5_1_BLOCK_BYTES>()
3982 .0
3983 .iter()
3984 .enumerate()
3985 {
3986 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
3987 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
3988 let qh = u32::from_le_bytes(block[4..8].try_into().unwrap());
3989 let qs = &block[8..24];
3990 let base = b * Q5_1_BLOCK_ELEMS;
3991
3992 let mut vals = [0u8; 32];
3993 for j in 0..16 {
3994 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
3995 vals[j] = (qs[j] & 0x0F) | xh_0;
3996 vals[j + 16] = (qs[j] >> 4) | xh_1;
3997 }
3998
3999 let mut block_acc = vdupq_n_f32(0.0);
4000 for g in 0..2 {
4001 let raw16 = vld1q_u8(vals.as_ptr().add(g * 16));
4002 let lo16 = vmovl_u8(vget_low_u8(raw16));
4003 let hi16 = vmovl_u8(vget_high_u8(raw16));
4004 for (half_idx, half16) in [lo16, hi16].into_iter().enumerate() {
4005 let lo32 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(half16)));
4006 let hi32 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(half16)));
4007 let elem_base = base + g * 16 + half_idx * 8;
4008 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
4009 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
4010 let w_lo = vfmaq_n_f32(vdupq_n_f32(m), lo32, d);
4011 let w_hi = vfmaq_n_f32(vdupq_n_f32(m), hi32, d);
4012 block_acc = vfmaq_f32(block_acc, w_lo, x_lo);
4013 block_acc = vfmaq_f32(block_acc, w_hi, x_hi);
4014 }
4015 }
4016 acc += vaddvq_f32(block_acc);
4017 }
4018 acc
4019 }
4020
4021 #[target_feature(enable = "neon")]
4030 pub unsafe fn dot_q2_k_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
4031 debug_assert_eq!(row_bytes.len() % Q2_K_BLOCK_BYTES, 0);
4032 let two_bit_mask = vdupq_n_u8(3);
4033 let mut acc = 0f32;
4034 let mut x_base = 0usize;
4035
4036 macro_rules! shr2 {
4043 (0, $v:expr) => {
4044 vandq_u8($v, two_bit_mask)
4045 };
4046 ($shift:literal, $v:expr) => {
4047 vandq_u8(vshrq_n_u8($v, $shift), two_bit_mask)
4048 };
4049 }
4050
4051 macro_rules! q2_k_sub_block {
4052 ($shift:tt, $q:expr, $scales:expr, $is:expr, $d:expr, $dmin:expr, $x:expr, $x_base:expr, $acc:expr) => {{
4053 let sc1 = $scales[$is];
4054 $is += 1;
4055 let dl1 = $d * (sc1 & 0x0F) as f32;
4056 let min1_vec = vdupq_n_f32($dmin * (sc1 >> 4) as f32);
4057 let sc2 = $scales[$is];
4058 $is += 1;
4059 let dl2 = $d * (sc2 & 0x0F) as f32;
4060 let min2_vec = vdupq_n_f32($dmin * (sc2 >> 4) as f32);
4061
4062 let lo16 = vld1q_u8($q.as_ptr());
4063 let hi16 = vld1q_u8($q.as_ptr().add(16));
4064 let lo2 = shr2!($shift, lo16);
4065 let hi2 = shr2!($shift, hi16);
4066
4067 let lo_acc = fma_affine16(
4068 widen_u8x16_to_f32_quads(lo2),
4069 dl1,
4070 min1_vec,
4071 $x,
4072 $x_base,
4073 vdupq_n_f32(0.0),
4074 );
4075 let hi_acc = fma_affine16(
4076 widen_u8x16_to_f32_quads(hi2),
4077 dl2,
4078 min2_vec,
4079 $x,
4080 $x_base + 16,
4081 vdupq_n_f32(0.0),
4082 );
4083 $acc += vaddvq_f32(lo_acc) + vaddvq_f32(hi_acc);
4084 $x_base += 32;
4085 }};
4086 }
4087
4088 for block in row_bytes.as_chunks::<Q2_K_BLOCK_BYTES>().0 {
4089 let scales: &[u8; Q2_K_SCALE_BYTES] = block[0..16].try_into().unwrap();
4090 let qs = &block[16..80];
4091 let d = f16::from_le_bytes([block[80], block[81]]).to_f32();
4092 let dmin = f16::from_le_bytes([block[82], block[83]]).to_f32();
4093
4094 let mut is = 0usize;
4095 for n in 0..2 {
4096 let q = &qs[n * 32..n * 32 + 32];
4097 q2_k_sub_block!(0, q, scales, is, d, dmin, x, x_base, acc);
4098 q2_k_sub_block!(2, q, scales, is, d, dmin, x, x_base, acc);
4099 q2_k_sub_block!(4, q, scales, is, d, dmin, x, x_base, acc);
4100 q2_k_sub_block!(6, q, scales, is, d, dmin, x, x_base, acc);
4101 }
4102 }
4103 acc
4104 }
4105
4106 #[target_feature(enable = "neon")]
4116 pub unsafe fn dot_q3_k_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
4117 debug_assert_eq!(row_bytes.len() % Q3_K_BLOCK_BYTES, 0);
4118 let two_bit_mask = vdupq_n_u8(3);
4119 let four = vdupq_n_u8(4);
4120 let mut acc = 0f32;
4121 let mut x_base = 0usize;
4122
4123 macro_rules! shr2 {
4126 (0, $v:expr) => {
4127 vandq_u8($v, two_bit_mask)
4128 };
4129 ($shift:literal, $v:expr) => {
4130 vandq_u8(vshrq_n_u8($v, $shift), two_bit_mask)
4131 };
4132 }
4133
4134 macro_rules! q3_k_sub_block {
4135 ($shift:tt, $q:expr, $hmask:expr, $m_vec:expr, $dl1:expr, $dl2:expr, $x:expr, $x_base:expr, $acc:expr) => {{
4136 let lo16 = vld1q_u8($q.as_ptr());
4137 let hi16 = vld1q_u8($q.as_ptr().add(16));
4138 let lo2 = shr2!($shift, lo16);
4139 let hi2 = shr2!($shift, hi16);
4140
4141 let hmask_lo = vld1q_u8($hmask.as_ptr());
4142 let hmask_hi = vld1q_u8($hmask.as_ptr().add(16));
4143 let bit_clear_lo = vmvnq_u8(vtstq_u8(hmask_lo, $m_vec));
4148 let bit_clear_hi = vmvnq_u8(vtstq_u8(hmask_hi, $m_vec));
4149 let bias_lo = vandq_u8(bit_clear_lo, four);
4150 let bias_hi = vandq_u8(bit_clear_hi, four);
4151
4152 let raw_lo_i16_lo = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(lo2))), {
4153 vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(bias_lo)))
4154 });
4155 let raw_lo_i16_hi =
4156 vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(lo2))), {
4157 vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(bias_lo)))
4158 });
4159 let raw_hi_i16_lo = vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(hi2))), {
4160 vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(bias_hi)))
4161 });
4162 let raw_hi_i16_hi =
4163 vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(hi2))), {
4164 vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(bias_hi)))
4165 });
4166
4167 let mut lo_acc = vdupq_n_f32(0.0);
4168 let mut hi_acc = vdupq_n_f32(0.0);
4169 for (i, half16) in [raw_lo_i16_lo, raw_lo_i16_hi].into_iter().enumerate() {
4170 let lo32 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half16)));
4171 let hi32 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half16)));
4172 let elem_base = $x_base + i * 8;
4173 let x_lo = vld1q_f32($x.as_ptr().add(elem_base));
4174 let x_hi = vld1q_f32($x.as_ptr().add(elem_base + 4));
4175 lo_acc = vfmaq_f32(lo_acc, lo32, x_lo);
4176 lo_acc = vfmaq_f32(lo_acc, hi32, x_hi);
4177 }
4178 for (i, half16) in [raw_hi_i16_lo, raw_hi_i16_hi].into_iter().enumerate() {
4179 let lo32 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half16)));
4180 let hi32 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half16)));
4181 let elem_base = $x_base + 16 + i * 8;
4182 let x_lo = vld1q_f32($x.as_ptr().add(elem_base));
4183 let x_hi = vld1q_f32($x.as_ptr().add(elem_base + 4));
4184 hi_acc = vfmaq_f32(hi_acc, lo32, x_lo);
4185 hi_acc = vfmaq_f32(hi_acc, hi32, x_hi);
4186 }
4187 $acc += vaddvq_f32(lo_acc) * $dl1 + vaddvq_f32(hi_acc) * $dl2;
4188 $x_base += 32;
4189 }};
4190 }
4191
4192 for block in row_bytes.as_chunks::<Q3_K_BLOCK_BYTES>().0 {
4193 let hmask = &block[0..32];
4194 let qs = &block[32..96];
4195 let scales_raw: &[u8; Q3_K_SCALE_BYTES] = block[96..108].try_into().unwrap();
4196 let d_all = f16::from_le_bytes([block[108], block[109]]).to_f32();
4197 let scales = q3_k_unpack_scales(scales_raw);
4198
4199 let mut is = 0usize;
4200 let mut m = 1u8;
4201 for n in 0..2 {
4202 let q = &qs[n * 32..n * 32 + 32];
4203 for shift in [0u32, 2, 4, 6] {
4204 let dl1 = d_all * (scales[is] as f32 - 32.0);
4205 let dl2 = d_all * (scales[is + 1] as f32 - 32.0);
4206 is += 2;
4207 let m_vec = vdupq_n_u8(m);
4208 match shift {
4209 0 => q3_k_sub_block!(0, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
4210 2 => q3_k_sub_block!(2, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
4211 4 => q3_k_sub_block!(4, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
4212 6 => q3_k_sub_block!(6, q, hmask, m_vec, dl1, dl2, x, x_base, acc),
4213 _ => unreachable!(),
4214 }
4215 m <<= 1;
4216 }
4217 }
4218 }
4219 acc
4220 }
4221
4222 #[target_feature(enable = "neon")]
4231 pub unsafe fn dot_iq4_nl_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
4232 debug_assert_eq!(row_bytes.len() % IQ4_NL_BLOCK_BYTES, 0);
4233 let low_mask = vdupq_n_u8(0x0F);
4234 let codebook = vld1q_s8(KVALUES_IQ4NL.as_ptr());
4235 let mut acc = 0f32;
4236 let mut x_base = 0usize;
4237 for block in row_bytes.as_chunks::<IQ4_NL_BLOCK_BYTES>().0 {
4238 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4239 let qs = &block[2..18];
4240 let bytes = vld1q_u8(qs.as_ptr());
4241 let lo_idx = vandq_u8(bytes, low_mask);
4242 let hi_idx = vshrq_n_u8(bytes, 4);
4243 let lo_vals = vqtbl1q_s8(codebook, lo_idx);
4244 let hi_vals = vqtbl1q_s8(codebook, hi_idx);
4245
4246 let mut block_acc = vdupq_n_f32(0.0);
4247 for (half_idx, vals) in [lo_vals, hi_vals].into_iter().enumerate() {
4248 let lo16 = vmovl_s8(vget_low_s8(vals));
4249 let hi16 = vmovl_s8(vget_high_s8(vals));
4250 for (i, half16) in [lo16, hi16].into_iter().enumerate() {
4251 let lo32 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half16)));
4252 let hi32 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half16)));
4253 let elem_base = x_base + half_idx * 16 + i * 8;
4254 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
4255 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
4256 block_acc = vfmaq_f32(block_acc, lo32, x_lo);
4257 block_acc = vfmaq_f32(block_acc, hi32, x_hi);
4258 }
4259 }
4260 acc += vaddvq_f32(block_acc) * d;
4261 x_base += IQ4_NL_BLOCK_ELEMS;
4262 }
4263 acc
4264 }
4265
4266 #[target_feature(enable = "neon")]
4271 pub unsafe fn dot_iq4_xs_f32_neon(row_bytes: &[u8], x: &[f32]) -> f32 {
4272 debug_assert_eq!(row_bytes.len() % IQ4_XS_BLOCK_BYTES, 0);
4273 let low_mask = vdupq_n_u8(0x0F);
4274 let codebook = vld1q_s8(KVALUES_IQ4NL.as_ptr());
4275 let mut acc = 0f32;
4276 let mut x_base = 0usize;
4277 for block in row_bytes.as_chunks::<IQ4_XS_BLOCK_BYTES>().0 {
4278 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4279 let scales_h = u16::from_le_bytes([block[2], block[3]]);
4280 let scales_l = &block[4..8];
4281 let qs = &block[8..136];
4282
4283 for ib in 0..8 {
4284 let ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0xf)
4285 | (((scales_h >> (2 * ib)) & 3) as u8) << 4;
4286 let dl = d * (ls as f32 - 32.0);
4287 let sub = &qs[ib * 16..ib * 16 + 16];
4288 let bytes = vld1q_u8(sub.as_ptr());
4289 let lo_idx = vandq_u8(bytes, low_mask);
4290 let hi_idx = vshrq_n_u8(bytes, 4);
4291 let lo_vals = vqtbl1q_s8(codebook, lo_idx);
4292 let hi_vals = vqtbl1q_s8(codebook, hi_idx);
4293
4294 let mut sub_acc = vdupq_n_f32(0.0);
4295 for (half_idx, vals) in [lo_vals, hi_vals].into_iter().enumerate() {
4296 let lo16 = vmovl_s8(vget_low_s8(vals));
4297 let hi16 = vmovl_s8(vget_high_s8(vals));
4298 for (i, half16) in [lo16, hi16].into_iter().enumerate() {
4299 let lo32 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half16)));
4300 let hi32 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half16)));
4301 let elem_base = x_base + half_idx * 16 + i * 8;
4302 let x_lo = vld1q_f32(x.as_ptr().add(elem_base));
4303 let x_hi = vld1q_f32(x.as_ptr().add(elem_base + 4));
4304 sub_acc = vfmaq_f32(sub_acc, lo32, x_lo);
4305 sub_acc = vfmaq_f32(sub_acc, hi32, x_hi);
4306 }
4307 }
4308 acc += vaddvq_f32(sub_acc) * dl;
4309 x_base += 32;
4310 }
4311 }
4312 acc
4313 }
4314}
4315
4316pub fn dot_q4_0_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4320 #[cfg(target_arch = "x86_64")]
4321 {
4322 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4323 return unsafe { simd_x86::dot_q4_0_f32_avx2(row_bytes, x) };
4324 }
4325 }
4326 #[cfg(target_arch = "aarch64")]
4327 {
4328 if std::arch::is_aarch64_feature_detected!("neon") {
4329 return unsafe { simd_aarch64::dot_q4_0_f32_neon(row_bytes, x) };
4330 }
4331 }
4332 dot_q4_0_f32_scalar(row_bytes, x)
4333}
4334
4335pub fn dot_q4_0_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4336 debug_assert_eq!(row_bytes.len() % Q4_0_BLOCK_BYTES, 0);
4337 let mut acc = 0f32;
4338 for (b, block) in row_bytes
4339 .as_chunks::<Q4_0_BLOCK_BYTES>()
4340 .0
4341 .iter()
4342 .enumerate()
4343 {
4344 let scale = f16::from_le_bytes([block[0], block[1]]).to_f32();
4345 let nibbles = &block[2..18];
4346 let base = b * Q4_0_BLOCK_ELEMS;
4347 let mut block_acc = 0f32;
4348 for i in 0..16 {
4349 let byte = nibbles[i];
4350 let lo = (byte & 0x0F) as i32 - 8;
4351 let hi = ((byte >> 4) & 0x0F) as i32 - 8;
4352 block_acc += (lo as f32) * x[base + i];
4353 block_acc += (hi as f32) * x[base + i + 16];
4354 }
4355 acc += block_acc * scale;
4356 }
4357 acc
4358}
4359
4360pub fn dequant_q4_1(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4364 if !src.len().is_multiple_of(Q4_1_BLOCK_BYTES) {
4365 return Err(QuantError::Misaligned(src.len(), Q4_1_BLOCK_BYTES));
4366 }
4367 let n_blocks = src.len() / Q4_1_BLOCK_BYTES;
4368 let mut out = vec![0f32; n_blocks * Q4_1_BLOCK_ELEMS];
4369 for (b, block) in src.as_chunks::<Q4_1_BLOCK_BYTES>().0.iter().enumerate() {
4370 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4371 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
4372 let nibbles = &block[4..20];
4373 let base = b * Q4_1_BLOCK_ELEMS;
4374 for i in 0..16 {
4375 let byte = nibbles[i];
4376 out[base + i] = (byte & 0x0F) as f32 * d + m;
4377 out[base + i + 16] = (byte >> 4) as f32 * d + m;
4378 }
4379 }
4380 Ok(out)
4381}
4382
4383pub fn dot_q4_1_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4386 #[cfg(target_arch = "x86_64")]
4387 {
4388 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4389 return unsafe { simd_x86::dot_q4_1_f32_avx2(row_bytes, x) };
4390 }
4391 }
4392 #[cfg(target_arch = "aarch64")]
4393 {
4394 if std::arch::is_aarch64_feature_detected!("neon") {
4395 return unsafe { simd_aarch64::dot_q4_1_f32_neon(row_bytes, x) };
4396 }
4397 }
4398 dot_q4_1_f32_scalar(row_bytes, x)
4399}
4400
4401pub fn dot_q4_1_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4402 debug_assert_eq!(row_bytes.len() % Q4_1_BLOCK_BYTES, 0);
4403 let mut acc = 0f32;
4404 for (b, block) in row_bytes
4405 .as_chunks::<Q4_1_BLOCK_BYTES>()
4406 .0
4407 .iter()
4408 .enumerate()
4409 {
4410 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4411 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
4412 let nibbles = &block[4..20];
4413 let base = b * Q4_1_BLOCK_ELEMS;
4414 for i in 0..16 {
4415 let byte = nibbles[i];
4416 acc += ((byte & 0x0F) as f32 * d + m) * x[base + i];
4417 acc += ((byte >> 4) as f32 * d + m) * x[base + i + 16];
4418 }
4419 }
4420 acc
4421}
4422
4423#[inline]
4429fn q5_fifth_bits(qh: u32, j: usize) -> (u8, u8) {
4430 let xh_0 = ((qh >> j) << 4) as u8 & 0x10;
4431 let xh_1 = (qh >> (j + 12)) as u8 & 0x10;
4432 (xh_0, xh_1)
4433}
4434
4435pub fn dequant_q5_0(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4439 if !src.len().is_multiple_of(Q5_0_BLOCK_BYTES) {
4440 return Err(QuantError::Misaligned(src.len(), Q5_0_BLOCK_BYTES));
4441 }
4442 let n_blocks = src.len() / Q5_0_BLOCK_BYTES;
4443 let mut out = vec![0f32; n_blocks * Q5_0_BLOCK_ELEMS];
4444 for (b, block) in src.as_chunks::<Q5_0_BLOCK_BYTES>().0.iter().enumerate() {
4445 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4446 let qh = u32::from_le_bytes(block[2..6].try_into().unwrap());
4447 let qs = &block[6..22];
4448 let base = b * Q5_0_BLOCK_ELEMS;
4449 for j in 0..16 {
4450 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
4451 let x0 = ((qs[j] & 0x0F) | xh_0) as i32 - 16;
4452 let x1 = ((qs[j] >> 4) | xh_1) as i32 - 16;
4453 out[base + j] = x0 as f32 * d;
4454 out[base + j + 16] = x1 as f32 * d;
4455 }
4456 }
4457 Ok(out)
4458}
4459
4460pub fn dot_q5_0_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4463 #[cfg(target_arch = "x86_64")]
4464 {
4465 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4466 return unsafe { simd_x86::dot_q5_0_f32_avx2(row_bytes, x) };
4467 }
4468 }
4469 #[cfg(target_arch = "aarch64")]
4470 {
4471 if std::arch::is_aarch64_feature_detected!("neon") {
4472 return unsafe { simd_aarch64::dot_q5_0_f32_neon(row_bytes, x) };
4473 }
4474 }
4475 dot_q5_0_f32_scalar(row_bytes, x)
4476}
4477
4478pub fn dot_q5_0_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4479 debug_assert_eq!(row_bytes.len() % Q5_0_BLOCK_BYTES, 0);
4480 let mut acc = 0f32;
4481 for (b, block) in row_bytes
4482 .as_chunks::<Q5_0_BLOCK_BYTES>()
4483 .0
4484 .iter()
4485 .enumerate()
4486 {
4487 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4488 let qh = u32::from_le_bytes(block[2..6].try_into().unwrap());
4489 let qs = &block[6..22];
4490 let base = b * Q5_0_BLOCK_ELEMS;
4491 for j in 0..16 {
4492 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
4493 let x0 = ((qs[j] & 0x0F) | xh_0) as i32 - 16;
4494 let x1 = ((qs[j] >> 4) | xh_1) as i32 - 16;
4495 acc += (x0 as f32 * d) * x[base + j];
4496 acc += (x1 as f32 * d) * x[base + j + 16];
4497 }
4498 }
4499 acc
4500}
4501
4502pub fn dequant_q5_1(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4506 if !src.len().is_multiple_of(Q5_1_BLOCK_BYTES) {
4507 return Err(QuantError::Misaligned(src.len(), Q5_1_BLOCK_BYTES));
4508 }
4509 let n_blocks = src.len() / Q5_1_BLOCK_BYTES;
4510 let mut out = vec![0f32; n_blocks * Q5_1_BLOCK_ELEMS];
4511 for (b, block) in src.as_chunks::<Q5_1_BLOCK_BYTES>().0.iter().enumerate() {
4512 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4513 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
4514 let qh = u32::from_le_bytes(block[4..8].try_into().unwrap());
4515 let qs = &block[8..24];
4516 let base = b * Q5_1_BLOCK_ELEMS;
4517 for j in 0..16 {
4518 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
4519 let x0 = (qs[j] & 0x0F) | xh_0;
4520 let x1 = (qs[j] >> 4) | xh_1;
4521 out[base + j] = x0 as f32 * d + m;
4522 out[base + j + 16] = x1 as f32 * d + m;
4523 }
4524 }
4525 Ok(out)
4526}
4527
4528pub fn dot_q5_1_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4531 #[cfg(target_arch = "x86_64")]
4532 {
4533 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4534 return unsafe { simd_x86::dot_q5_1_f32_avx2(row_bytes, x) };
4535 }
4536 }
4537 #[cfg(target_arch = "aarch64")]
4538 {
4539 if std::arch::is_aarch64_feature_detected!("neon") {
4540 return unsafe { simd_aarch64::dot_q5_1_f32_neon(row_bytes, x) };
4541 }
4542 }
4543 dot_q5_1_f32_scalar(row_bytes, x)
4544}
4545
4546pub fn dot_q5_1_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4547 debug_assert_eq!(row_bytes.len() % Q5_1_BLOCK_BYTES, 0);
4548 let mut acc = 0f32;
4549 for (b, block) in row_bytes
4550 .as_chunks::<Q5_1_BLOCK_BYTES>()
4551 .0
4552 .iter()
4553 .enumerate()
4554 {
4555 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4556 let m = f16::from_le_bytes([block[2], block[3]]).to_f32();
4557 let qh = u32::from_le_bytes(block[4..8].try_into().unwrap());
4558 let qs = &block[8..24];
4559 let base = b * Q5_1_BLOCK_ELEMS;
4560 for j in 0..16 {
4561 let (xh_0, xh_1) = q5_fifth_bits(qh, j);
4562 let x0 = (qs[j] & 0x0F) | xh_0;
4563 let x1 = (qs[j] >> 4) | xh_1;
4564 acc += (x0 as f32 * d + m) * x[base + j];
4565 acc += (x1 as f32 * d + m) * x[base + j + 16];
4566 }
4567 }
4568 acc
4569}
4570
4571pub fn dequant_q8_1(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4577 if !src.len().is_multiple_of(Q8_1_BLOCK_BYTES) {
4578 return Err(QuantError::Misaligned(src.len(), Q8_1_BLOCK_BYTES));
4579 }
4580 let n_blocks = src.len() / Q8_1_BLOCK_BYTES;
4581 let mut out = Vec::with_capacity(n_blocks * Q8_1_BLOCK_ELEMS);
4582 for block in src.as_chunks::<Q8_1_BLOCK_BYTES>().0 {
4583 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4584 for i in 0..Q8_1_BLOCK_ELEMS {
4585 let q = block[4 + i] as i8;
4586 out.push(q as f32 * d);
4587 }
4588 }
4589 Ok(out)
4590}
4591
4592pub fn dot_q8_1_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4598 #[cfg(target_arch = "x86_64")]
4599 {
4600 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4601 return unsafe { simd_x86::dot_q8_1_f32_avx2(row_bytes, x) };
4602 }
4603 }
4604 #[cfg(target_arch = "aarch64")]
4605 {
4606 if std::arch::is_aarch64_feature_detected!("neon") {
4607 return unsafe { simd_aarch64::dot_q8_1_f32_neon(row_bytes, x) };
4608 }
4609 }
4610 dot_q8_1_f32_scalar(row_bytes, x)
4611}
4612
4613pub fn dot_q8_1_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4614 debug_assert_eq!(row_bytes.len() % Q8_1_BLOCK_BYTES, 0);
4615 let mut acc = 0f32;
4616 for (b, block) in row_bytes
4617 .as_chunks::<Q8_1_BLOCK_BYTES>()
4618 .0
4619 .iter()
4620 .enumerate()
4621 {
4622 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4623 let base = b * Q8_1_BLOCK_ELEMS;
4624 let mut block_acc = 0f32;
4625 for i in 0..Q8_1_BLOCK_ELEMS {
4626 let q = block[4 + i] as i8;
4627 block_acc += (q as f32) * x[base + i];
4628 }
4629 acc += block_acc * d;
4630 }
4631 acc
4632}
4633
4634pub fn dequant_q2_k(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4642 if !src.len().is_multiple_of(Q2_K_BLOCK_BYTES) {
4643 return Err(QuantError::Misaligned(src.len(), Q2_K_BLOCK_BYTES));
4644 }
4645 let n_blocks = src.len() / Q2_K_BLOCK_BYTES;
4646 let mut out = Vec::with_capacity(n_blocks * Q2_K_BLOCK_ELEMS);
4647 for block in src.as_chunks::<Q2_K_BLOCK_BYTES>().0 {
4648 let scales: &[u8; Q2_K_SCALE_BYTES] = block[0..16].try_into().unwrap();
4649 let qs = &block[16..80];
4650 let d = f16::from_le_bytes([block[80], block[81]]).to_f32();
4651 let dmin = f16::from_le_bytes([block[82], block[83]]).to_f32();
4652
4653 let mut is = 0usize;
4654 for n in 0..2 {
4655 let q = &qs[n * 32..n * 32 + 32];
4656 let mut shift = 0u32;
4657 for _j in 0..4 {
4658 let sc1 = scales[is];
4659 is += 1;
4660 let (dl1, ml1) = (d * (sc1 & 0x0F) as f32, dmin * (sc1 >> 4) as f32);
4661 for &byte in &q[0..16] {
4662 let raw = (byte >> shift) & 3;
4663 out.push(dl1 * raw as f32 - ml1);
4664 }
4665
4666 let sc2 = scales[is];
4667 is += 1;
4668 let (dl2, ml2) = (d * (sc2 & 0x0F) as f32, dmin * (sc2 >> 4) as f32);
4669 for &byte in &q[16..32] {
4670 let raw = (byte >> shift) & 3;
4671 out.push(dl2 * raw as f32 - ml2);
4672 }
4673 shift += 2;
4674 }
4675 }
4676 }
4677 Ok(out)
4678}
4679
4680pub fn dot_q2_k_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4683 #[cfg(target_arch = "x86_64")]
4684 {
4685 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4686 return unsafe { simd_x86::dot_q2_k_f32_avx2(row_bytes, x) };
4687 }
4688 }
4689 #[cfg(target_arch = "aarch64")]
4690 {
4691 if std::arch::is_aarch64_feature_detected!("neon") {
4692 return unsafe { simd_aarch64::dot_q2_k_f32_neon(row_bytes, x) };
4693 }
4694 }
4695 dot_q2_k_f32_scalar(row_bytes, x)
4696}
4697
4698pub fn dot_q2_k_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4699 debug_assert_eq!(row_bytes.len() % Q2_K_BLOCK_BYTES, 0);
4700 let mut acc = 0f32;
4701 let mut x_base = 0usize;
4702 for block in row_bytes.as_chunks::<Q2_K_BLOCK_BYTES>().0 {
4703 let scales: &[u8; Q2_K_SCALE_BYTES] = block[0..16].try_into().unwrap();
4704 let qs = &block[16..80];
4705 let d = f16::from_le_bytes([block[80], block[81]]).to_f32();
4706 let dmin = f16::from_le_bytes([block[82], block[83]]).to_f32();
4707
4708 let mut is = 0usize;
4709 for n in 0..2 {
4710 let q = &qs[n * 32..n * 32 + 32];
4711 let mut shift = 0u32;
4712 for _j in 0..4 {
4713 let sc1 = scales[is];
4714 is += 1;
4715 let (dl1, ml1) = (d * (sc1 & 0x0F) as f32, dmin * (sc1 >> 4) as f32);
4716 for l in 0..16 {
4717 let raw = (q[l] >> shift) & 3;
4718 acc += (dl1 * raw as f32 - ml1) * x[x_base + l];
4719 }
4720
4721 let sc2 = scales[is];
4722 is += 1;
4723 let (dl2, ml2) = (d * (sc2 & 0x0F) as f32, dmin * (sc2 >> 4) as f32);
4724 for l in 0..16 {
4725 let raw = (q[l + 16] >> shift) & 3;
4726 acc += (dl2 * raw as f32 - ml2) * x[x_base + l + 16];
4727 }
4728 shift += 2;
4729 x_base += 32;
4730 }
4731 }
4732 }
4733 acc
4734}
4735
4736fn q3_k_unpack_scales(raw: &[u8; Q3_K_SCALE_BYTES]) -> [i8; 16] {
4746 const KMASK1: u8 = 0x03;
4747 const KMASK2: u8 = 0x0F;
4748 let mut out = [0u8; 16];
4749 for j in 0..4 {
4750 let (a0, a1, tmp) = (raw[j], raw[4 + j], raw[8 + j]);
4751 out[j] = (a0 & KMASK2) | ((tmp & KMASK1) << 4);
4756 out[4 + j] = (a1 & KMASK2) | (((tmp >> 2) & KMASK1) << 4);
4757 out[8 + j] = (a0 >> 4) | (((tmp >> 4) & KMASK1) << 4);
4758 out[12 + j] = (a1 >> 4) | (((tmp >> 6) & KMASK1) << 4);
4759 }
4760 out.map(|b| b as i8)
4764}
4765
4766pub fn dequant_q3_k(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4774 if !src.len().is_multiple_of(Q3_K_BLOCK_BYTES) {
4775 return Err(QuantError::Misaligned(src.len(), Q3_K_BLOCK_BYTES));
4776 }
4777 let n_blocks = src.len() / Q3_K_BLOCK_BYTES;
4778 let mut out = Vec::with_capacity(n_blocks * Q3_K_BLOCK_ELEMS);
4779 for block in src.as_chunks::<Q3_K_BLOCK_BYTES>().0 {
4780 let hmask = &block[0..32];
4781 let qs = &block[32..96];
4782 let scales_raw: &[u8; Q3_K_SCALE_BYTES] = block[96..108].try_into().unwrap();
4783 let d_all = f16::from_le_bytes([block[108], block[109]]).to_f32();
4784 let scales = q3_k_unpack_scales(scales_raw);
4785
4786 let mut is = 0usize;
4787 let mut m = 1u8;
4788 for n in 0..2 {
4789 let q = &qs[n * 32..n * 32 + 32];
4790 let mut shift = 0u32;
4791 for _j in 0..4 {
4792 let dl1 = d_all * (scales[is] as f32 - 32.0);
4793 is += 1;
4794 for l in 0..16 {
4795 let raw = ((q[l] >> shift) & 3) as i32;
4796 let bias = if hmask[l] & m != 0 { 0 } else { 4 };
4797 out.push(dl1 * (raw - bias) as f32);
4798 }
4799
4800 let dl2 = d_all * (scales[is] as f32 - 32.0);
4801 is += 1;
4802 for l in 0..16 {
4803 let raw = ((q[l + 16] >> shift) & 3) as i32;
4804 let bias = if hmask[l + 16] & m != 0 { 0 } else { 4 };
4805 out.push(dl2 * (raw - bias) as f32);
4806 }
4807 shift += 2;
4808 m <<= 1;
4809 }
4810 }
4811 }
4812 Ok(out)
4813}
4814
4815pub fn dot_q3_k_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4818 #[cfg(target_arch = "x86_64")]
4819 {
4820 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4821 return unsafe { simd_x86::dot_q3_k_f32_avx2(row_bytes, x) };
4822 }
4823 }
4824 #[cfg(target_arch = "aarch64")]
4825 {
4826 if std::arch::is_aarch64_feature_detected!("neon") {
4827 return unsafe { simd_aarch64::dot_q3_k_f32_neon(row_bytes, x) };
4828 }
4829 }
4830 dot_q3_k_f32_scalar(row_bytes, x)
4831}
4832
4833pub fn dot_q3_k_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4834 debug_assert_eq!(row_bytes.len() % Q3_K_BLOCK_BYTES, 0);
4835 let mut acc = 0f32;
4836 let mut x_base = 0usize;
4837 for block in row_bytes.as_chunks::<Q3_K_BLOCK_BYTES>().0 {
4838 let hmask = &block[0..32];
4839 let qs = &block[32..96];
4840 let scales_raw: &[u8; Q3_K_SCALE_BYTES] = block[96..108].try_into().unwrap();
4841 let d_all = f16::from_le_bytes([block[108], block[109]]).to_f32();
4842 let scales = q3_k_unpack_scales(scales_raw);
4843
4844 let mut is = 0usize;
4845 let mut m = 1u8;
4846 for n in 0..2 {
4847 let q = &qs[n * 32..n * 32 + 32];
4848 let mut shift = 0u32;
4849 for _j in 0..4 {
4850 let dl1 = d_all * (scales[is] as f32 - 32.0);
4851 is += 1;
4852 for l in 0..16 {
4853 let raw = ((q[l] >> shift) & 3) as i32;
4854 let bias = if hmask[l] & m != 0 { 0 } else { 4 };
4855 acc += (dl1 * (raw - bias) as f32) * x[x_base + l];
4856 }
4857
4858 let dl2 = d_all * (scales[is] as f32 - 32.0);
4859 is += 1;
4860 for l in 0..16 {
4861 let raw = ((q[l + 16] >> shift) & 3) as i32;
4862 let bias = if hmask[l + 16] & m != 0 { 0 } else { 4 };
4863 acc += (dl2 * (raw - bias) as f32) * x[x_base + l + 16];
4864 }
4865 shift += 2;
4866 m <<= 1;
4867 x_base += 32;
4868 }
4869 }
4870 }
4871 acc
4872}
4873
4874pub const IQ4_NL_BLOCK_BYTES: usize = 18;
4875pub const IQ4_NL_BLOCK_ELEMS: usize = 32;
4876pub const IQ4_XS_BLOCK_BYTES: usize = 136;
4877pub const IQ4_XS_BLOCK_ELEMS: usize = 256;
4878
4879const KVALUES_IQ4NL: [i8; 16] = [
4884 -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113,
4885];
4886
4887pub fn dequant_iq4_nl(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4888 if !src.len().is_multiple_of(IQ4_NL_BLOCK_BYTES) {
4889 return Err(QuantError::Misaligned(src.len(), IQ4_NL_BLOCK_BYTES));
4890 }
4891 let n_blocks = src.len() / IQ4_NL_BLOCK_BYTES;
4892 let mut out = Vec::with_capacity(n_blocks * IQ4_NL_BLOCK_ELEMS);
4893 for block in src.as_chunks::<IQ4_NL_BLOCK_BYTES>().0 {
4894 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4895 let qs = &block[2..18];
4896 let mut lo = [0f32; 16];
4897 let mut hi = [0f32; 16];
4898 for (j, &byte) in qs.iter().enumerate() {
4899 lo[j] = d * KVALUES_IQ4NL[(byte & 0xf) as usize] as f32;
4900 hi[j] = d * KVALUES_IQ4NL[(byte >> 4) as usize] as f32;
4901 }
4902 out.extend_from_slice(&lo);
4903 out.extend_from_slice(&hi);
4904 }
4905 Ok(out)
4906}
4907
4908pub fn dot_iq4_nl_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4911 #[cfg(target_arch = "x86_64")]
4912 {
4913 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4914 return unsafe { simd_x86::dot_iq4_nl_f32_avx2(row_bytes, x) };
4915 }
4916 }
4917 #[cfg(target_arch = "aarch64")]
4918 {
4919 if std::arch::is_aarch64_feature_detected!("neon") {
4920 return unsafe { simd_aarch64::dot_iq4_nl_f32_neon(row_bytes, x) };
4921 }
4922 }
4923 dot_iq4_nl_f32_scalar(row_bytes, x)
4924}
4925
4926pub fn dot_iq4_nl_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4927 debug_assert_eq!(row_bytes.len() % IQ4_NL_BLOCK_BYTES, 0);
4928 let mut acc = 0f32;
4929 let mut x_base = 0usize;
4930 for block in row_bytes.as_chunks::<IQ4_NL_BLOCK_BYTES>().0 {
4931 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4932 let qs = &block[2..18];
4933 for (j, &byte) in qs.iter().enumerate() {
4934 acc += (d * KVALUES_IQ4NL[(byte & 0xf) as usize] as f32) * x[x_base + j];
4935 acc += (d * KVALUES_IQ4NL[(byte >> 4) as usize] as f32) * x[x_base + 16 + j];
4936 }
4937 x_base += IQ4_NL_BLOCK_ELEMS;
4938 }
4939 acc
4940}
4941
4942pub fn dequant_iq4_xs(src: &[u8]) -> Result<Vec<f32>, QuantError> {
4943 if !src.len().is_multiple_of(IQ4_XS_BLOCK_BYTES) {
4944 return Err(QuantError::Misaligned(src.len(), IQ4_XS_BLOCK_BYTES));
4945 }
4946 let n_blocks = src.len() / IQ4_XS_BLOCK_BYTES;
4947 let mut out = Vec::with_capacity(n_blocks * IQ4_XS_BLOCK_ELEMS);
4948 for block in src.as_chunks::<IQ4_XS_BLOCK_BYTES>().0 {
4949 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4950 let scales_h = u16::from_le_bytes([block[2], block[3]]);
4951 let scales_l = &block[4..8];
4952 let qs = &block[8..136];
4953
4954 for ib in 0..8 {
4955 let ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0xf)
4956 | (((scales_h >> (2 * ib)) & 3) as u8) << 4;
4957 let dl = d * (ls as f32 - 32.0);
4958 let sub = &qs[ib * 16..ib * 16 + 16];
4959 let mut lo = [0f32; 16];
4960 let mut hi = [0f32; 16];
4961 for (j, &byte) in sub.iter().enumerate() {
4962 lo[j] = dl * KVALUES_IQ4NL[(byte & 0xf) as usize] as f32;
4963 hi[j] = dl * KVALUES_IQ4NL[(byte >> 4) as usize] as f32;
4964 }
4965 out.extend_from_slice(&lo);
4966 out.extend_from_slice(&hi);
4967 }
4968 }
4969 Ok(out)
4970}
4971
4972pub fn dot_iq4_xs_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
4975 #[cfg(target_arch = "x86_64")]
4976 {
4977 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
4978 return unsafe { simd_x86::dot_iq4_xs_f32_avx2(row_bytes, x) };
4979 }
4980 }
4981 #[cfg(target_arch = "aarch64")]
4982 {
4983 if std::arch::is_aarch64_feature_detected!("neon") {
4984 return unsafe { simd_aarch64::dot_iq4_xs_f32_neon(row_bytes, x) };
4985 }
4986 }
4987 dot_iq4_xs_f32_scalar(row_bytes, x)
4988}
4989
4990pub fn dot_iq4_xs_f32_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
4991 debug_assert_eq!(row_bytes.len() % IQ4_XS_BLOCK_BYTES, 0);
4992 let mut acc = 0f32;
4993 let mut x_base = 0usize;
4994 for block in row_bytes.as_chunks::<IQ4_XS_BLOCK_BYTES>().0 {
4995 let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
4996 let scales_h = u16::from_le_bytes([block[2], block[3]]);
4997 let scales_l = &block[4..8];
4998 let qs = &block[8..136];
4999
5000 for ib in 0..8 {
5001 let ls = ((scales_l[ib / 2] >> (4 * (ib % 2))) & 0xf)
5002 | (((scales_h >> (2 * ib)) & 3) as u8) << 4;
5003 let dl = d * (ls as f32 - 32.0);
5004 let sub = &qs[ib * 16..ib * 16 + 16];
5005 for (j, &byte) in sub.iter().enumerate() {
5006 acc += (dl * KVALUES_IQ4NL[(byte & 0xf) as usize] as f32) * x[x_base + j];
5007 acc += (dl * KVALUES_IQ4NL[(byte >> 4) as usize] as f32) * x[x_base + 16 + j];
5008 }
5009 x_base += 32;
5010 }
5011 }
5012 acc
5013}
5014
5015pub const MXFP4_GROUP_SIZE: usize = 32;
5020
5021const KVALUES_MXFP4: [f32; 16] = [
5030 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
5031];
5032
5033fn e8m0_scale(e: u8) -> f32 {
5043 if e == 0 {
5044 f32::from_bits(0x0040_0000)
5045 } else {
5046 f32::from_bits((e as u32) << 23)
5047 }
5048}
5049
5050pub fn dequant_mxfp4_row(packed: &[u8], scales: &[u8]) -> Result<Vec<f32>, QuantError> {
5063 let expected_packed_len = scales.len() * (MXFP4_GROUP_SIZE / 2);
5064 if packed.len() != expected_packed_len {
5065 return Err(QuantError::Mxfp4RowMismatch(
5066 packed.len(),
5067 expected_packed_len,
5068 ));
5069 }
5070 let mut out = Vec::with_capacity(scales.len() * MXFP4_GROUP_SIZE);
5071 for (g, &e) in scales.iter().enumerate() {
5072 let d = e8m0_scale(e);
5073 let group = &packed[g * (MXFP4_GROUP_SIZE / 2)..(g + 1) * (MXFP4_GROUP_SIZE / 2)];
5074 let mut lo = [0f32; MXFP4_GROUP_SIZE / 2];
5075 let mut hi = [0f32; MXFP4_GROUP_SIZE / 2];
5076 for (j, &byte) in group.iter().enumerate() {
5077 lo[j] = d * KVALUES_MXFP4[(byte & 0xf) as usize];
5078 hi[j] = d * KVALUES_MXFP4[(byte >> 4) as usize];
5079 }
5080 out.extend_from_slice(&lo);
5081 out.extend_from_slice(&hi);
5082 }
5083 Ok(out)
5084}
5085
5086pub fn dot_mxfp4_row_f32(packed: &[u8], scales: &[u8], x: &[f32]) -> f32 {
5094 #[cfg(target_arch = "x86_64")]
5095 {
5096 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
5097 return unsafe { simd_x86::dot_mxfp4_row_f32_avx2(packed, scales, x) };
5098 }
5099 }
5100 #[cfg(target_arch = "aarch64")]
5101 {
5102 if std::arch::is_aarch64_feature_detected!("neon") {
5103 return unsafe { simd_aarch64::dot_mxfp4_row_f32_neon(packed, scales, x) };
5104 }
5105 }
5106 dot_mxfp4_row_f32_scalar(packed, scales, x)
5107}
5108
5109pub fn dot_mxfp4_row_f32_scalar(packed: &[u8], scales: &[u8], x: &[f32]) -> f32 {
5110 debug_assert_eq!(packed.len(), scales.len() * (MXFP4_GROUP_SIZE / 2));
5111 let mut acc = 0f32;
5112 let mut x_base = 0usize;
5113 for (g, &e) in scales.iter().enumerate() {
5114 let d = e8m0_scale(e);
5115 let group = &packed[g * (MXFP4_GROUP_SIZE / 2)..(g + 1) * (MXFP4_GROUP_SIZE / 2)];
5116 for (j, &byte) in group.iter().enumerate() {
5117 acc += (d * KVALUES_MXFP4[(byte & 0xf) as usize]) * x[x_base + j];
5118 acc += (d * KVALUES_MXFP4[(byte >> 4) as usize]) * x[x_base + MXFP4_GROUP_SIZE / 2 + j];
5119 }
5120 x_base += MXFP4_GROUP_SIZE;
5121 }
5122 acc
5123}
5124
5125pub const IQ1_S_BLOCK_BYTES: usize = 50;
5154pub const IQ1_S_BLOCK_ELEMS: usize = 256;
5155pub const IQ1_M_BLOCK_BYTES: usize = 56;
5159pub const IQ1_M_BLOCK_ELEMS: usize = 256;
5160pub const IQ2_XXS_BLOCK_BYTES: usize = 66;
5163pub const IQ2_XXS_BLOCK_ELEMS: usize = 256;
5164pub const IQ2_XS_BLOCK_BYTES: usize = 74;
5167pub const IQ2_XS_BLOCK_ELEMS: usize = 256;
5168pub const IQ2_S_BLOCK_BYTES: usize = 82;
5171pub const IQ2_S_BLOCK_ELEMS: usize = 256;
5172pub const IQ3_XXS_BLOCK_BYTES: usize = 98;
5175pub const IQ3_XXS_BLOCK_ELEMS: usize = 256;
5176pub const IQ3_S_BLOCK_BYTES: usize = 110;
5180pub const IQ3_S_BLOCK_ELEMS: usize = 256;
5181
5182const IQ1S_DELTA: f32 = 0.125;
5188
5189#[inline]
5194fn iq_sign(signs: u8, j: usize) -> f32 {
5195 if signs & iq_tables::KMASK_IQ2XS[j] != 0 {
5196 -1.0
5197 } else {
5198 1.0
5199 }
5200}
5201
5202#[inline]
5203fn read_f16(bytes: &[u8]) -> f32 {
5204 f16::from_le_bytes([bytes[0], bytes[1]]).to_f32()
5205}
5206
5207#[inline]
5210fn for_each_iq1_s(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5211 let d = read_f16(block);
5212 let qs = &block[2..34];
5213 let qh = &block[34..50];
5214 let mut idx = 0usize;
5215 for ib in 0..8 {
5216 let h = u16::from_le_bytes([qh[2 * ib], qh[2 * ib + 1]]);
5217 let dl = d * (2.0 * ((h >> 12) & 7) as f32 + 1.0);
5218 let delta = if h & 0x8000 != 0 {
5219 -IQ1S_DELTA
5220 } else {
5221 IQ1S_DELTA
5222 };
5223 for l in 0..4 {
5224 let grid_index = qs[4 * ib + l] as usize | ((((h >> (3 * l)) & 7) as usize) << 8);
5225 let row = iq_tables::IQ1S_GRID[grid_index];
5226 for j in 0..8 {
5227 let v = ((row >> (8 * j)) & 0xFF) as u8 as i8;
5228 emit(idx, dl * (v as f32 + delta));
5229 idx += 1;
5230 }
5231 }
5232 }
5233}
5234
5235#[inline]
5237fn for_each_iq2_xxs(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5238 let d = read_f16(block);
5239 let qs: Vec<u16> = block[2..66]
5240 .as_chunks::<2>()
5241 .0
5242 .iter()
5243 .map(|c| u16::from_le_bytes([c[0], c[1]]))
5244 .collect();
5245 let mut idx = 0usize;
5246 for ib32 in 0..8 {
5247 let g = &qs[4 * ib32..4 * ib32 + 4];
5248 let aux32_1 = g[2] as u32 | ((g[3] as u32) << 16);
5249 let db = d * (0.5 + (aux32_1 >> 28) as f32) * 0.25;
5250 let aux8 = [
5251 (g[0] & 0xFF) as usize,
5252 (g[0] >> 8) as usize,
5253 (g[1] & 0xFF) as usize,
5254 (g[1] >> 8) as usize,
5255 ];
5256 for (l, &code) in aux8.iter().enumerate() {
5257 let row = iq_tables::IQ2XXS_GRID[code];
5258 let signs = iq_tables::KSIGNS_IQ2XS[((aux32_1 >> (7 * l)) & 127) as usize];
5259 for j in 0..8 {
5260 let mag = ((row >> (8 * j)) & 0xFF) as f32;
5261 emit(idx, db * mag * iq_sign(signs, j));
5262 idx += 1;
5263 }
5264 }
5265 }
5266}
5267
5268#[inline]
5270fn for_each_iq3_xxs(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5271 let d = read_f16(block);
5272 let qs = &block[2..66];
5273 let sas = &block[66..98];
5274 let mut idx = 0usize;
5275 for ib32 in 0..8 {
5276 let aux32 = u32::from_le_bytes([
5277 sas[4 * ib32],
5278 sas[4 * ib32 + 1],
5279 sas[4 * ib32 + 2],
5280 sas[4 * ib32 + 3],
5281 ]);
5282 let db = d * (0.5 + (aux32 >> 28) as f32) * 0.5;
5283 for l in 0..4 {
5284 let signs = iq_tables::KSIGNS_IQ2XS[((aux32 >> (7 * l)) & 127) as usize];
5285 let g1 = iq_tables::IQ3XXS_GRID[qs[8 * ib32 + 2 * l] as usize];
5286 let g2 = iq_tables::IQ3XXS_GRID[qs[8 * ib32 + 2 * l + 1] as usize];
5287 for j in 0..4 {
5288 emit(
5289 idx + j,
5290 db * ((g1 >> (8 * j)) & 0xFF) as f32 * iq_sign(signs, j),
5291 );
5292 }
5293 for j in 0..4 {
5294 emit(
5295 idx + 4 + j,
5296 db * ((g2 >> (8 * j)) & 0xFF) as f32 * iq_sign(signs, j + 4),
5297 );
5298 }
5299 idx += 8;
5300 }
5301 }
5302}
5303
5304#[inline]
5313fn for_each_iq2_xs(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5314 let d = read_f16(block);
5315 let qs = &block[2..66];
5316 let scales = &block[66..74];
5317 let mut idx = 0usize;
5318 for ib32 in 0..8 {
5319 let db = [
5320 d * (0.5 + (scales[ib32] & 0xF) as f32) * 0.25,
5321 d * (0.5 + (scales[ib32] >> 4) as f32) * 0.25,
5322 ];
5323 for l in 0..4 {
5324 let code = u16::from_le_bytes([qs[8 * ib32 + 2 * l], qs[8 * ib32 + 2 * l + 1]]);
5325 let row = iq_tables::IQ2XS_GRID[(code & 511) as usize];
5326 let signs = iq_tables::KSIGNS_IQ2XS[(code >> 9) as usize];
5327 for j in 0..8 {
5328 let mag = ((row >> (8 * j)) & 0xFF) as f32;
5329 emit(idx, db[l / 2] * mag * iq_sign(signs, j));
5330 idx += 1;
5331 }
5332 }
5333 }
5334}
5335
5336#[inline]
5348fn for_each_iq2_s(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5349 let d = read_f16(block);
5350 let qs = &block[2..34];
5351 let sign_bytes = &block[34..66];
5352 let qh = &block[66..74];
5353 let scales = &block[74..82];
5354 let mut idx = 0usize;
5355 for ib32 in 0..8 {
5356 let db = [
5357 d * (0.5 + (scales[ib32] & 0xF) as f32) * 0.25,
5358 d * (0.5 + (scales[ib32] >> 4) as f32) * 0.25,
5359 ];
5360 for l in 0..4 {
5361 let hi = ((qh[ib32] as usize) << (8 - 2 * l)) & 0x300;
5362 let row = iq_tables::IQ2S_GRID[qs[4 * ib32 + l] as usize | hi];
5363 let signs = sign_bytes[4 * ib32 + l];
5364 for j in 0..8 {
5365 let mag = ((row >> (8 * j)) & 0xFF) as f32;
5366 emit(idx, db[l / 2] * mag * iq_sign(signs, j));
5367 idx += 1;
5368 }
5369 }
5370 }
5371}
5372
5373#[inline]
5388fn for_each_iq3_s(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5389 let d = read_f16(block);
5390 let qs = &block[2..66];
5391 let qh = &block[66..74];
5392 let sign_bytes = &block[74..106];
5393 let scales = &block[106..110];
5394 let mut idx = 0usize;
5395 for ib32 in 0..8 {
5396 let nibble = if ib32 % 2 == 0 {
5397 scales[ib32 / 2] & 0xF
5398 } else {
5399 scales[ib32 / 2] >> 4
5400 };
5401 let db = d * (1.0 + 2.0 * nibble as f32);
5402 for l in 0..4 {
5403 let h = qh[ib32] as usize;
5407 let i1 = qs[8 * ib32 + 2 * l] as usize | ((h << (8 - 2 * l)) & 256);
5408 let i2 = qs[8 * ib32 + 2 * l + 1] as usize | ((h << (7 - 2 * l)) & 256);
5409 let g1 = iq_tables::IQ3S_GRID[i1];
5410 let g2 = iq_tables::IQ3S_GRID[i2];
5411 let signs = sign_bytes[4 * ib32 + l];
5412 for j in 0..4 {
5413 emit(
5414 idx + j,
5415 db * ((g1 >> (8 * j)) & 0xFF) as f32 * iq_sign(signs, j),
5416 );
5417 }
5418 for j in 0..4 {
5419 emit(
5420 idx + 4 + j,
5421 db * ((g2 >> (8 * j)) & 0xFF) as f32 * iq_sign(signs, j + 4),
5422 );
5423 }
5424 idx += 8;
5425 }
5426 }
5427}
5428
5429#[inline]
5443fn for_each_iq1_m(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5444 let qs = &block[0..32];
5445 let qh = &block[32..48];
5446 let scales = &block[48..56];
5447 let sc: [u16; 4] =
5448 std::array::from_fn(|k| u16::from_le_bytes([scales[2 * k], scales[2 * k + 1]]));
5449 let d = f16::from_bits(
5451 (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000),
5452 )
5453 .to_f32();
5454 let mut idx = 0usize;
5455 for ib in 0..8 {
5456 let shift = 6 * (ib % 2);
5457 let dl = [
5458 d * (2.0 * ((sc[ib / 2] >> shift) & 7) as f32 + 1.0),
5459 d * (2.0 * ((sc[ib / 2] >> (shift + 3)) & 7) as f32 + 1.0),
5460 ];
5461 let (h0, h1) = (qh[2 * ib] as usize, qh[2 * ib + 1] as usize);
5462 let grid_idx = [
5465 qs[4 * ib] as usize | ((h0 << 8) & 0x700),
5466 qs[4 * ib + 1] as usize | ((h0 << 4) & 0x700),
5467 qs[4 * ib + 2] as usize | ((h1 << 8) & 0x700),
5468 qs[4 * ib + 3] as usize | ((h1 << 4) & 0x700),
5469 ];
5470 let delta = [
5471 if h0 & 0x08 != 0 {
5472 -IQ1S_DELTA
5473 } else {
5474 IQ1S_DELTA
5475 },
5476 if h0 & 0x80 != 0 {
5477 -IQ1S_DELTA
5478 } else {
5479 IQ1S_DELTA
5480 },
5481 if h1 & 0x08 != 0 {
5482 -IQ1S_DELTA
5483 } else {
5484 IQ1S_DELTA
5485 },
5486 if h1 & 0x80 != 0 {
5487 -IQ1S_DELTA
5488 } else {
5489 IQ1S_DELTA
5490 },
5491 ];
5492 for l in 0..4 {
5493 let row = iq_tables::IQ1S_GRID[grid_idx[l]];
5494 for j in 0..8 {
5495 let v = ((row >> (8 * j)) & 0xFF) as u8 as i8;
5496 emit(idx, dl[l / 2] * (v as f32 + delta[l]));
5497 idx += 1;
5498 }
5499 }
5500 }
5501}
5502
5503macro_rules! iq_dequant_and_dot {
5504 ($dequant:ident, $dot_scalar:ident, $walk:ident, $bytes:ident, $elems:ident) => {
5505 pub fn $dequant(src: &[u8]) -> Result<Vec<f32>, QuantError> {
5506 if !src.len().is_multiple_of($bytes) {
5507 return Err(QuantError::Misaligned(src.len(), $bytes));
5508 }
5509 let n_blocks = src.len() / $bytes;
5510 let mut out = vec![0f32; n_blocks * $elems];
5511 for (b, block) in src.chunks_exact($bytes).enumerate() {
5512 let base = b * $elems;
5513 $walk(block, |i, v| out[base + i] = v);
5514 }
5515 Ok(out)
5516 }
5517
5518 pub fn $dot_scalar(row_bytes: &[u8], x: &[f32]) -> f32 {
5519 debug_assert_eq!(row_bytes.len() % $bytes, 0);
5520 let mut acc = 0f32;
5521 let mut x_base = 0usize;
5522 for block in row_bytes.chunks_exact($bytes) {
5523 $walk(block, |i, v| acc += v * x[x_base + i]);
5524 x_base += $elems;
5525 }
5526 acc
5527 }
5528 };
5529}
5530
5531macro_rules! iq_dispatch {
5537 ($dot:ident, $dot_scalar:ident, $avx2:ident) => {
5538 pub fn $dot(row_bytes: &[u8], x: &[f32]) -> f32 {
5539 #[cfg(target_arch = "x86_64")]
5540 {
5541 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
5542 return unsafe { simd_x86::$avx2(row_bytes, x) };
5543 }
5544 }
5545 $dot_scalar(row_bytes, x)
5546 }
5547 };
5548}
5549
5550iq_dispatch!(dot_iq1_s_f32, dot_iq1_s_f32_scalar, dot_iq1_s_f32_avx2);
5551iq_dispatch!(
5552 dot_iq2_xxs_f32,
5553 dot_iq2_xxs_f32_scalar,
5554 dot_iq2_xxs_f32_avx2
5555);
5556iq_dispatch!(
5557 dot_iq3_xxs_f32,
5558 dot_iq3_xxs_f32_scalar,
5559 dot_iq3_xxs_f32_avx2
5560);
5561
5562pub fn dot_iq2_xs_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
5570 dot_iq2_xs_f32_scalar(row_bytes, x)
5571}
5572
5573pub fn dot_iq2_s_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
5574 dot_iq2_s_f32_scalar(row_bytes, x)
5575}
5576
5577pub fn dot_iq3_s_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
5578 dot_iq3_s_f32_scalar(row_bytes, x)
5579}
5580
5581pub fn dot_iq1_m_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
5582 dot_iq1_m_f32_scalar(row_bytes, x)
5583}
5584
5585pub fn dot_mxfp4_gguf_f32(row_bytes: &[u8], x: &[f32]) -> f32 {
5589 dot_mxfp4_gguf_f32_scalar(row_bytes, x)
5590}
5591
5592iq_dequant_and_dot!(
5593 dequant_iq1_s,
5594 dot_iq1_s_f32_scalar,
5595 for_each_iq1_s,
5596 IQ1_S_BLOCK_BYTES,
5597 IQ1_S_BLOCK_ELEMS
5598);
5599iq_dequant_and_dot!(
5600 dequant_iq2_xxs,
5601 dot_iq2_xxs_f32_scalar,
5602 for_each_iq2_xxs,
5603 IQ2_XXS_BLOCK_BYTES,
5604 IQ2_XXS_BLOCK_ELEMS
5605);
5606iq_dequant_and_dot!(
5607 dequant_iq3_xxs,
5608 dot_iq3_xxs_f32_scalar,
5609 for_each_iq3_xxs,
5610 IQ3_XXS_BLOCK_BYTES,
5611 IQ3_XXS_BLOCK_ELEMS
5612);
5613iq_dequant_and_dot!(
5614 dequant_iq2_xs,
5615 dot_iq2_xs_f32_scalar,
5616 for_each_iq2_xs,
5617 IQ2_XS_BLOCK_BYTES,
5618 IQ2_XS_BLOCK_ELEMS
5619);
5620iq_dequant_and_dot!(
5621 dequant_iq2_s,
5622 dot_iq2_s_f32_scalar,
5623 for_each_iq2_s,
5624 IQ2_S_BLOCK_BYTES,
5625 IQ2_S_BLOCK_ELEMS
5626);
5627iq_dequant_and_dot!(
5628 dequant_iq3_s,
5629 dot_iq3_s_f32_scalar,
5630 for_each_iq3_s,
5631 IQ3_S_BLOCK_BYTES,
5632 IQ3_S_BLOCK_ELEMS
5633);
5634iq_dequant_and_dot!(
5635 dequant_iq1_m,
5636 dot_iq1_m_f32_scalar,
5637 for_each_iq1_m,
5638 IQ1_M_BLOCK_BYTES,
5639 IQ1_M_BLOCK_ELEMS
5640);
5641
5642pub const MXFP4_GGUF_BLOCK_BYTES: usize = 17;
5653pub const MXFP4_GGUF_BLOCK_ELEMS: usize = 32;
5654
5655#[inline]
5658fn for_each_mxfp4_gguf(block: &[u8], mut emit: impl FnMut(usize, f32)) {
5659 let d = e8m0_scale(block[0]);
5660 for (j, &byte) in block[1..17].iter().enumerate() {
5661 emit(j, d * KVALUES_MXFP4[(byte & 0x0F) as usize]);
5662 emit(j + 16, d * KVALUES_MXFP4[(byte >> 4) as usize]);
5663 }
5664}
5665
5666iq_dequant_and_dot!(
5667 dequant_mxfp4_gguf,
5668 dot_mxfp4_gguf_f32_scalar,
5669 for_each_mxfp4_gguf,
5670 MXFP4_GGUF_BLOCK_BYTES,
5671 MXFP4_GGUF_BLOCK_ELEMS
5672);
5673
5674#[cfg(test)]
5675mod tests {
5676 use super::*;
5677
5678 #[test]
5679 fn turbo4_kv_blocks_roundtrip_reasonable() {
5680 let x: Vec<f32> = (0..64).map(|i| (i as f32 * 0.17).sin() * 2.0).collect();
5681 let packed = pack_turbo4_kv_blocks(&x);
5682 assert_eq!(packed.len(), 2 * TURBO4_KV_BLOCK_BYTES);
5683 let y = unpack_turbo4_kv_blocks(&packed).unwrap();
5684 assert_eq!(y.len(), 64);
5685 let mut err = 0.0f32;
5686 for (a, b) in x.iter().zip(y.iter()) {
5687 err += (a - b).abs();
5688 }
5689 err /= x.len() as f32;
5690 assert!(err < 0.2, "mean abs err {err}");
5691 }
5692
5693 #[test]
5694 fn q8_0_roundtrip_is_within_quantization_error() {
5695 let original: Vec<f32> = (0..32).map(|i| (i as f32 - 16.0) * 0.37).collect();
5696 let packed = quantize_q8_0(&original);
5697 assert_eq!(packed.len(), Q8_0_BLOCK_BYTES);
5698 let restored = dequant_q8_0(&packed).unwrap();
5699 assert_eq!(restored.len(), 32);
5700 for (a, b) in original.iter().zip(restored.iter()) {
5701 assert!((a - b).abs() < 0.1, "a={a} b={b}");
5702 }
5703 }
5704
5705 #[test]
5706 fn quantize_activations_q8_reconstructs_within_quant_error() {
5707 let x: Vec<f32> = (0..64)
5708 .map(|i| ((i as f32) * 0.13 - 4.0).sin() * 3.0)
5709 .collect();
5710 let act = quantize_activations_q8(&x);
5711 assert_eq!(act.n_blocks(), 2);
5712 assert_eq!(act.q.len(), 64);
5713 for (b, chunk) in x.as_chunks::<32>().0.iter().enumerate() {
5714 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
5715 let tol = amax / 127.0 + 1e-6;
5716 for (i, &v) in chunk.iter().enumerate() {
5717 let recon = act.q[b * 32 + i] as f32 * act.d[b];
5718 assert!((recon - v).abs() <= tol, "b={b} i={i} v={v} recon={recon}");
5719 }
5720 }
5721 }
5722
5723 #[test]
5724 fn quantize_activations_q8_handles_all_zero_block() {
5725 let act = quantize_activations_q8(&[0f32; 32]);
5726 assert_eq!(act.d[0], 0.0);
5727 assert!(act.q.iter().all(|&q| q == 0));
5728 }
5729
5730 #[test]
5731 fn quantize_activations_q8_parallel_matches_serial() {
5732 let x: Vec<f32> = (0..512)
5733 .map(|i| ((i as f32) * 0.07 - 8.0).sin() * 2.5)
5734 .collect();
5735 let got = quantize_activations_q8(&x);
5736 let n_blocks = x.len() / Q8_0_BLOCK_ELEMS;
5737 let mut q = vec![0i8; n_blocks * Q8_0_BLOCK_ELEMS];
5738 let mut d = vec![0f32; n_blocks];
5739 for (b, chunk) in x.as_chunks::<Q8_0_BLOCK_ELEMS>().0.iter().enumerate() {
5740 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
5741 let scale = amax / 127.0;
5742 let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
5743 d[b] = scale;
5744 let base = b * Q8_0_BLOCK_ELEMS;
5745 for (i, &v) in chunk.iter().enumerate() {
5746 let qi = (v * inv).round();
5747 q[base + i] = qi.clamp(-127.0, 127.0) as i8;
5748 }
5749 }
5750 assert_eq!(got.q, q);
5751 assert_eq!(got.d, d);
5752 }
5753
5754 #[test]
5755 fn quantize_activations_q8_k_parallel_matches_serial() {
5756 let x: Vec<f32> = (0..1024)
5757 .map(|i| ((i as f32) * 0.05 - 12.0).cos() * 1.7)
5758 .collect();
5759 let got = quantize_activations_q8_k(&x);
5760 let n_blocks = x.len() / Q4_K_BLOCK_ELEMS;
5761 let mut q = vec![0i8; n_blocks * Q4_K_BLOCK_ELEMS];
5762 let mut d = vec![0f32; n_blocks];
5763 let mut bsums = vec![0i16; n_blocks * 16];
5764 for (b, chunk) in x.as_chunks::<Q4_K_BLOCK_ELEMS>().0.iter().enumerate() {
5765 let amax = chunk.iter().fold(0f32, |m, &v| m.max(v.abs()));
5766 let scale = amax / 127.0;
5767 let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
5768 d[b] = scale;
5769 let base = b * Q4_K_BLOCK_ELEMS;
5770 for (i, &v) in chunk.iter().enumerate() {
5771 let qi = (v * inv).round();
5772 q[base + i] = qi.clamp(-127.0, 127.0) as i8;
5773 }
5774 let bsum_base = b * 16;
5775 for g in 0..16 {
5776 let mut s = 0i32;
5777 let off = base + g * 16;
5778 for i in 0..16 {
5779 s += q[off + i] as i32;
5780 }
5781 bsums[bsum_base + g] = s as i16;
5782 }
5783 }
5784 assert_eq!(got.q, q);
5785 assert_eq!(got.d, d);
5786 assert_eq!(got.bsums, bsums);
5787 }
5788
5789 #[test]
5790 fn dot_q4_k_q8_matches_scalar_and_tracks_float_dot() {
5791 let n_blocks = 3;
5792 let cols = n_blocks * Q4_K_BLOCK_ELEMS;
5793 let x: Vec<f32> = (0..cols)
5794 .map(|i| ((i as f32) * 0.017 - 2.1).sin() * 1.8)
5795 .collect();
5796 let mut weights = Vec::with_capacity(n_blocks * Q4_K_BLOCK_BYTES);
5800 for b in 0..n_blocks {
5801 weights.extend_from_slice(&f16::from_f32(0.05 + b as f32 * 0.01).to_le_bytes());
5802 weights.extend_from_slice(&f16::from_f32(0.01 + b as f32 * 0.002).to_le_bytes());
5803 for i in 0..12u8 {
5805 weights.push(20 + i.wrapping_mul(3));
5806 }
5807 for i in 0..128u8 {
5808 weights.push(i.wrapping_mul(17).wrapping_add(b as u8));
5809 }
5810 }
5811 let act = quantize_activations_q8_k(&x);
5812 let dispatched = dot_q4_k_q8(&weights, &act);
5813 let scalar = dot_q4_k_q8_scalar(&weights, &act);
5814 assert_eq!(dispatched, scalar, "dispatch must match scalar");
5815 let float_dot = dot_q4_k_f32(&weights, &x);
5816 let err = (dispatched - float_dot).abs();
5817 let scale = float_dot.abs().max(1.0);
5818 assert!(
5819 err / scale < 0.05,
5820 "int-dot vs f32 relative err {err}/{scale} too large (int={dispatched} f32={float_dot})"
5821 );
5822 }
5823
5824 #[test]
5825 #[cfg(target_arch = "aarch64")]
5826 fn dot_q4_k_q8_i8mm_matches_scalar_when_available() {
5827 if !std::arch::is_aarch64_feature_detected!("i8mm") {
5828 return;
5829 }
5830 let n_blocks = 3;
5831 let cols = n_blocks * Q4_K_BLOCK_ELEMS;
5832 let x: Vec<f32> = (0..cols)
5833 .map(|i| ((i as f32) * 0.017 - 2.1).sin() * 1.8)
5834 .collect();
5835 let mut weights = Vec::with_capacity(n_blocks * Q4_K_BLOCK_BYTES);
5836 for b in 0..n_blocks {
5837 weights.extend_from_slice(&f16::from_f32(0.05 + b as f32 * 0.01).to_le_bytes());
5838 weights.extend_from_slice(&f16::from_f32(0.01 + b as f32 * 0.002).to_le_bytes());
5839 for i in 0..12u8 {
5840 weights.push(20 + i.wrapping_mul(3));
5841 }
5842 for i in 0..128u8 {
5843 weights.push(i.wrapping_mul(17).wrapping_add(b as u8));
5844 }
5845 }
5846 let act = quantize_activations_q8_k(&x);
5847 let scalar = dot_q4_k_q8_scalar(&weights, &act);
5848 let i8mm = unsafe { simd_aarch64::dot_q4_k_q8_neon_i8mm(&weights, &act) };
5849 assert_eq!(i8mm, scalar, "i8mm must match scalar");
5850 let dispatched = dot_q4_k_q8(&weights, &act);
5851 assert_eq!(
5852 dispatched, scalar,
5853 "dispatch must match scalar on i8mm host"
5854 );
5855 }
5856
5857 #[test]
5858 fn dot_q5_k_q8_matches_scalar_and_tracks_float_dot() {
5859 let x: Vec<f32> = (0..Q5_K_BLOCK_ELEMS)
5860 .map(|i| ((i as f32) * 0.013 - 1.7).sin() * 1.5)
5861 .collect();
5862 let act = quantize_activations_q8_k(&x);
5863 let dispatched = dot_q5_k_q8(&Q5_K_TEST_BLOCK, &act);
5864 let scalar = dot_q5_k_q8_scalar(&Q5_K_TEST_BLOCK, &act);
5865 assert_eq!(dispatched, scalar, "dispatch must match scalar");
5866 let float_dot = dot_q5_k_f32(&Q5_K_TEST_BLOCK, &x);
5867 let err = (dispatched - float_dot).abs();
5868 let scale = float_dot.abs().max(1.0);
5869 assert!(
5870 err / scale < 0.05,
5871 "Q5_K int-dot vs f32 relative err {err}/{scale} (int={dispatched} f32={float_dot})"
5872 );
5873 }
5874
5875 #[test]
5876 fn gemm_q5_k_q8_row_matches_per_act_dots() {
5877 let acts: Vec<_> = (0..Q5_K_GEMM_NC)
5878 .map(|j| {
5879 let x: Vec<f32> = (0..Q5_K_BLOCK_ELEMS)
5880 .map(|i| ((i as f32) * 0.013 - 1.7 + j as f32).sin() * 1.5)
5881 .collect();
5882 quantize_activations_q8_k(&x)
5883 })
5884 .collect();
5885 let mut out = vec![0f32; acts.len()];
5886 gemm_q5_k_q8_row(&Q5_K_TEST_BLOCK, &acts, &mut out);
5887 for (j, act) in acts.iter().enumerate() {
5888 let want = dot_q5_k_q8(&Q5_K_TEST_BLOCK, act);
5889 let err = (out[j] - want).abs();
5890 assert!(
5891 err < 1e-4,
5892 "act {j}: gemm {got} vs dot {want}",
5893 got = out[j]
5894 );
5895 }
5896 }
5897
5898 #[test]
5899 fn gemm_q6_k_q8_row_matches_per_act_dots() {
5900 let acts: Vec<_> = (0..Q6_K_GEMM_NC)
5901 .map(|j| {
5902 let x: Vec<f32> = (0..Q6_K_BLOCK_ELEMS)
5903 .map(|i| ((i as f32) * 0.011 - 0.9 + j as f32).cos() * 1.9)
5904 .collect();
5905 quantize_activations_q8_k(&x)
5906 })
5907 .collect();
5908 let mut out = vec![0f32; acts.len()];
5909 gemm_q6_k_q8_row(&Q6_K_TEST_BLOCK, &acts, &mut out);
5910 for (j, act) in acts.iter().enumerate() {
5911 let want = dot_q6_k_q8(&Q6_K_TEST_BLOCK, act);
5912 let err = (out[j] - want).abs();
5913 assert!(
5914 err < 1e-3,
5915 "act {j}: gemm {got} vs dot {want}",
5916 got = out[j]
5917 );
5918 }
5919 }
5920
5921 #[test]
5922 fn dot_q6_k_q8_matches_scalar_and_tracks_float_dot() {
5923 let x: Vec<f32> = (0..Q6_K_BLOCK_ELEMS)
5924 .map(|i| ((i as f32) * 0.011 - 0.9).cos() * 1.9)
5925 .collect();
5926 let act = quantize_activations_q8_k(&x);
5927 let dispatched = dot_q6_k_q8(&Q6_K_TEST_BLOCK, &act);
5928 let scalar = dot_q6_k_q8_scalar(&Q6_K_TEST_BLOCK, &act);
5929 assert_eq!(dispatched, scalar, "dispatch must match scalar");
5930 let float_dot = dot_q6_k_f32(&Q6_K_TEST_BLOCK, &x);
5931 let err = (dispatched - float_dot).abs();
5932 let scale = float_dot.abs().max(1.0);
5933 assert!(
5934 err / scale < 0.05,
5935 "Q6_K int-dot vs f32 relative err {err}/{scale} (int={dispatched} f32={float_dot})"
5936 );
5937 }
5938
5939 #[test]
5940 fn dot_q8_0_q8_dispatch_matches_scalar_and_float_dot() {
5941 let n_blocks = 5;
5944 let cols = n_blocks * Q8_0_BLOCK_ELEMS;
5945 let x: Vec<f32> = (0..cols)
5946 .map(|i| ((i as f32) * 0.019 - 1.3).cos() * 2.7)
5947 .collect();
5948
5949 let mut weights = Vec::with_capacity(n_blocks * Q8_0_BLOCK_BYTES);
5950 for b in 0..n_blocks {
5951 weights.extend_from_slice(&f16::from_f32(0.021 + b as f32 * 0.004).to_le_bytes());
5952 for i in 0..Q8_0_BLOCK_ELEMS {
5953 weights.push(((i as i32 * 7 + b as i32 * 3) % 255 - 127) as i8 as u8);
5954 }
5955 }
5956
5957 let act = quantize_activations_q8(&x);
5958 let dispatched = dot_q8_0_q8(&weights, &act);
5959 let scalar = dot_q8_0_q8_scalar(&weights, &act);
5960 assert_eq!(
5961 dispatched.to_bits(),
5962 scalar.to_bits(),
5963 "SIMD int dot must match scalar int dot bit-for-bit"
5964 );
5965
5966 let float_dot = dot_q8_0_f32(&weights, &x);
5967 let rel = (dispatched - float_dot).abs() / float_dot.abs().max(1e-6);
5970 assert!(
5971 rel < 0.02,
5972 "int dot {dispatched} vs float {float_dot} rel={rel}"
5973 );
5974 }
5975
5976 #[test]
5977 fn dot_q4_0_q8_dispatch_matches_scalar_and_float_dot() {
5978 let n_blocks = 5;
5979 let cols = n_blocks * Q4_0_BLOCK_ELEMS;
5980 let x: Vec<f32> = (0..cols)
5981 .map(|i| ((i as f32) * 0.019 - 1.3).cos() * 2.7)
5982 .collect();
5983
5984 let mut weights = Vec::with_capacity(n_blocks * Q4_0_BLOCK_BYTES);
5985 for b in 0..n_blocks {
5986 weights.extend_from_slice(&f16::from_f32(0.021 + b as f32 * 0.004).to_le_bytes());
5987 for i in 0..16 {
5988 weights.push(((i as u32 * 13 + b as u32 * 7) % 256) as u8);
5989 }
5990 }
5991
5992 let act = quantize_activations_q8(&x);
5993 let dispatched = dot_q4_0_q8(&weights, &act);
5994 let scalar = dot_q4_0_q8_scalar(&weights, &act);
5995 assert_eq!(
5996 dispatched.to_bits(),
5997 scalar.to_bits(),
5998 "SIMD Q4_0 int dot must match scalar bit-for-bit"
5999 );
6000
6001 let float_dot = dot_q4_0_f32(&weights, &x);
6002 let rel = (dispatched - float_dot).abs() / float_dot.abs().max(1e-6);
6003 assert!(
6004 rel < 0.03,
6005 "Q4_0 int dot {dispatched} vs float {float_dot} rel={rel}"
6006 );
6007 }
6008
6009 #[test]
6010 fn q4_0_zero_nibble_maps_to_negative_bias() {
6011 let mut block = Vec::new();
6013 block.extend_from_slice(&f16::from_f32(1.0).to_le_bytes());
6014 block.extend_from_slice(&[0u8; 16]); let out = dequant_q4_0(&block).unwrap();
6016 assert_eq!(out.len(), 32);
6017 assert!(out.iter().all(|&v| v == -8.0));
6018 }
6019
6020 #[test]
6021 fn rejects_misaligned_buffers() {
6022 let bad = vec![0u8; 5];
6023 assert!(dequant_q8_0(&bad).is_err());
6024 assert!(dequant_q4_0(&bad).is_err());
6025 }
6026
6027 #[test]
6028 fn q4_1_affine_nibble_maps_to_scale_plus_min() {
6029 let mut block = Vec::new();
6032 block.extend_from_slice(&f16::from_f32(2.0).to_le_bytes());
6033 block.extend_from_slice(&f16::from_f32(5.0).to_le_bytes());
6034 block.extend_from_slice(&[0x11u8; 16]); let out = dequant_q4_1(&block).unwrap();
6036 assert_eq!(out.len(), 32);
6037 assert!(out.iter().all(|&v| (v - 7.0).abs() < 1e-6));
6038 }
6039
6040 #[test]
6041 fn q5_0_fifth_bit_extends_range_past_a_plain_nibble() {
6042 let mut block = Vec::new();
6051 block.extend_from_slice(&f16::from_f32(1.0).to_le_bytes());
6052 let mut qh = [0u8; 4];
6053 qh[0] |= 1 << 0; block.extend_from_slice(&qh);
6055 let mut qs = [0u8; 16];
6056 qs[0] = 0x0F; block.extend_from_slice(&qs);
6058 let out = dequant_q5_0(&block).unwrap();
6059 assert_eq!(out.len(), 32);
6060 assert_eq!(out[0], 15.0);
6062 assert_eq!(out[1], -16.0);
6064 }
6065
6066 #[test]
6067 fn q5_1_fifth_bit_without_bias_subtraction() {
6068 let mut block = Vec::new();
6069 block.extend_from_slice(&f16::from_f32(1.0).to_le_bytes());
6070 block.extend_from_slice(&f16::from_f32(0.0).to_le_bytes());
6071 let mut qh = [0u8; 4];
6072 qh[0] |= 1 << 0;
6073 block.extend_from_slice(&qh);
6074 let mut qs = [0u8; 16];
6075 qs[0] = 0x0F;
6076 block.extend_from_slice(&qs);
6077 let out = dequant_q5_1(&block).unwrap();
6078 assert_eq!(out.len(), 32);
6079 assert_eq!(out[0], 31.0);
6081 assert_eq!(out[1], 0.0);
6082 }
6083
6084 #[test]
6085 fn q8_1_matches_q8_0_math_ignoring_the_extra_sum_field() {
6086 let mut block = Vec::new();
6087 block.extend_from_slice(&f16::from_f32(0.5).to_le_bytes());
6088 block.extend_from_slice(&f16::from_f32(999.0).to_le_bytes()); let qs: Vec<i8> = (0..32).map(|i| i - 16).collect();
6090 block.extend_from_slice(&i8_to_u8_bytes(&qs));
6091 let out = dequant_q8_1(&block).unwrap();
6092 assert_eq!(out.len(), 32);
6093 for (i, &v) in out.iter().enumerate() {
6094 assert_eq!(v, (i as f32 - 16.0) * 0.5);
6095 }
6096 }
6097
6098 fn i8_to_u8_bytes(src: &[i8]) -> Vec<u8> {
6102 src.iter().map(|&b| b as u8).collect()
6103 }
6104
6105 #[test]
6106 fn legacy_formats_fused_dot_matches_dequant_then_dot() {
6107 let x: Vec<f32> = (0..32).map(|i| ((i as f32) * 0.07).sin()).collect();
6108
6109 let mut q4_1 = Vec::new();
6110 q4_1.extend_from_slice(&f16::from_f32(0.3).to_le_bytes());
6111 q4_1.extend_from_slice(&f16::from_f32(-1.2).to_le_bytes());
6112 q4_1.extend_from_slice(
6113 &(0..16)
6114 .map(|i| (i as u8) | ((15 - i as u8) << 4))
6115 .collect::<Vec<u8>>(),
6116 );
6117 let expected: f32 = dequant_q4_1(&q4_1)
6118 .unwrap()
6119 .iter()
6120 .zip(x.iter())
6121 .map(|(a, b)| a * b)
6122 .sum();
6123 let fused = dot_q4_1_f32(&q4_1, &x);
6124 assert!(
6125 (fused - expected).abs() < 1e-3,
6126 "Q4_1: fused={fused} expected={expected}"
6127 );
6128
6129 let mut q5_0 = Vec::new();
6130 q5_0.extend_from_slice(&f16::from_f32(0.4).to_le_bytes());
6131 q5_0.extend_from_slice(&[0xA5, 0x3C, 0x00, 0xFF]);
6132 q5_0.extend_from_slice(
6133 &(0..16)
6134 .map(|i| (i as u8) | ((15 - i as u8) << 4))
6135 .collect::<Vec<u8>>(),
6136 );
6137 let expected: f32 = dequant_q5_0(&q5_0)
6138 .unwrap()
6139 .iter()
6140 .zip(x.iter())
6141 .map(|(a, b)| a * b)
6142 .sum();
6143 let fused = dot_q5_0_f32(&q5_0, &x);
6144 assert!(
6145 (fused - expected).abs() < 1e-3,
6146 "Q5_0: fused={fused} expected={expected}"
6147 );
6148
6149 let mut q5_1 = Vec::new();
6150 q5_1.extend_from_slice(&f16::from_f32(0.2).to_le_bytes());
6151 q5_1.extend_from_slice(&f16::from_f32(0.9).to_le_bytes());
6152 q5_1.extend_from_slice(&[0x12, 0x34, 0x56, 0x78]);
6153 q5_1.extend_from_slice(
6154 &(0..16)
6155 .map(|i| (i as u8) | ((15 - i as u8) << 4))
6156 .collect::<Vec<u8>>(),
6157 );
6158 let expected: f32 = dequant_q5_1(&q5_1)
6159 .unwrap()
6160 .iter()
6161 .zip(x.iter())
6162 .map(|(a, b)| a * b)
6163 .sum();
6164 let fused = dot_q5_1_f32(&q5_1, &x);
6165 assert!(
6166 (fused - expected).abs() < 1e-3,
6167 "Q5_1: fused={fused} expected={expected}"
6168 );
6169
6170 let mut q8_1 = Vec::new();
6171 q8_1.extend_from_slice(&f16::from_f32(0.6).to_le_bytes());
6172 q8_1.extend_from_slice(&f16::from_f32(0.0).to_le_bytes());
6173 let qs: Vec<i8> = (0..32).map(|i| ((i * 7) % 61) as i8 - 30).collect();
6174 q8_1.extend_from_slice(&i8_to_u8_bytes(&qs));
6175 let expected: f32 = dequant_q8_1(&q8_1)
6176 .unwrap()
6177 .iter()
6178 .zip(x.iter())
6179 .map(|(a, b)| a * b)
6180 .sum();
6181 let fused = dot_q8_1_f32(&q8_1, &x);
6182 assert!(
6183 (fused - expected).abs() < 1e-3,
6184 "Q8_1: fused={fused} expected={expected}"
6185 );
6186 }
6187
6188 #[test]
6189 fn legacy_formats_reject_misaligned_buffers() {
6190 let bad = vec![0u8; 5];
6191 assert!(dequant_q4_1(&bad).is_err());
6192 assert!(dequant_q5_0(&bad).is_err());
6193 assert!(dequant_q5_1(&bad).is_err());
6194 assert!(dequant_q8_1(&bad).is_err());
6195 }
6196
6197 #[test]
6198 fn bf16_widening_is_exact_for_round_values() {
6199 for v in [0.0f32, 1.0, -1.0, 2.5, -0.5, 100.0, -100.0] {
6203 let bf16_bits = (v.to_bits() >> 16) as u16;
6204 let bytes = bf16_bits.to_le_bytes();
6205 let restored = dequant_bf16(&bytes).unwrap();
6206 assert_eq!(restored, vec![v], "bf16 round-trip mismatch for {v}");
6207 }
6208 }
6209
6210 #[test]
6211 fn bf16_widening_matches_hand_computed_bits() {
6212 let bytes = 0x3F80u16.to_le_bytes();
6215 let out = dequant_bf16(&bytes).unwrap();
6216 assert_eq!(out, vec![1.0f32]);
6217 assert_eq!(out[0].to_bits(), 0x3F800000);
6218 }
6219
6220 #[test]
6221 fn bf16_rejects_odd_length_buffers() {
6222 let bad = vec![0u8; 3];
6223 assert!(dequant_bf16(&bad).is_err());
6224 }
6225
6226 #[test]
6227 fn f16_widening_is_exact_and_covers_the_special_values() {
6228 let subnormal = f16::from_bits(0x0001); let cases: Vec<f16> = [0.0f32, -0.0, 1.0, -1.0, 2.5, -0.5, 65504.0, -65504.0]
6233 .iter()
6234 .map(|&v| f16::from_f32(v))
6235 .chain(std::iter::once(subnormal))
6236 .collect();
6237 let bytes: Vec<u8> = cases.iter().flat_map(|h| h.to_le_bytes()).collect();
6238 let out = dequant_f16(&bytes).unwrap();
6239 assert_eq!(out.len(), cases.len());
6240 for (got, want) in out.iter().zip(cases.iter()) {
6241 assert_eq!(got.to_bits(), want.to_f32().to_bits());
6242 }
6243 assert_eq!(out[8], 2f32.powi(-24));
6244
6245 let inf = f16::INFINITY.to_le_bytes();
6247 assert!(dequant_f16(&inf).unwrap()[0].is_infinite());
6248 }
6249
6250 #[test]
6251 fn f16_rejects_odd_length_buffers() {
6252 let bad = vec![0u8; 5];
6253 assert!(dequant_f16(&bad).is_err());
6254 }
6255
6256 #[test]
6257 fn fused_q8_0_dot_matches_dequant_then_dot() {
6258 let original: Vec<f32> = (0..32).map(|i| (i as f32 - 16.0) * 0.37).collect();
6259 let packed = quantize_q8_0(&original);
6260 let x: Vec<f32> = (0..32).map(|i| (i as f32) * 0.01 - 0.16).collect();
6261
6262 let dequanted = dequant_q8_0(&packed).unwrap();
6263 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6264
6265 let fused = dot_q8_0_f32(&packed, &x);
6266 assert!(
6267 (fused - expected).abs() < 1e-3,
6268 "fused={fused} expected={expected}"
6269 );
6270 }
6271
6272 #[test]
6273 fn dispatched_dot_matches_scalar_reference_across_many_blocks() {
6274 let n_blocks = 5;
6279 let original: Vec<f32> = (0..n_blocks * 32)
6280 .map(|i| ((i as f32) - (n_blocks * 16) as f32) * 0.29)
6281 .collect();
6282 let packed = quantize_q8_0(&original);
6283 let x: Vec<f32> = (0..n_blocks * 32)
6284 .map(|i| ((i as f32) * 0.013).sin())
6285 .collect();
6286
6287 let dispatched = dot_q8_0_f32(&packed, &x);
6288 let scalar = dot_q8_0_f32_scalar(&packed, &x);
6289 assert!(
6290 (dispatched - scalar).abs() < 1e-2,
6291 "dispatched={dispatched} scalar={scalar} (should match regardless of which SIMD path the host CPU takes)"
6292 );
6293 }
6294
6295 #[cfg(target_arch = "x86_64")]
6296 #[test]
6297 fn avx2_kernel_matches_scalar_directly_when_available() {
6298 if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
6299 eprintln!("skipping: host CPU lacks AVX2/FMA");
6300 return;
6301 }
6302 let n_blocks = 8;
6303 let original: Vec<f32> = (0..n_blocks * 32)
6304 .map(|i| ((i % 37) as f32 - 18.0) * 0.11)
6305 .collect();
6306 let packed = quantize_q8_0(&original);
6307 let x: Vec<f32> = (0..n_blocks * 32)
6308 .map(|i| ((i as f32) * 0.07).cos())
6309 .collect();
6310
6311 let simd = unsafe { simd_x86::dot_q8_0_f32_avx2(&packed, &x) };
6312 let scalar = dot_q8_0_f32_scalar(&packed, &x);
6313 assert!(
6314 (simd - scalar).abs() < 1e-2,
6315 "AVX2 kernel diverged from scalar: simd={simd} scalar={scalar}"
6316 );
6317 }
6318
6319 #[cfg(target_arch = "x86_64")]
6320 #[test]
6321 fn avx2_q4_0_kernel_matches_scalar_directly_when_available() {
6322 if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
6323 eprintln!("skipping: host CPU lacks AVX2/FMA");
6324 return;
6325 }
6326 let n_blocks = 6;
6331 let mut packed = Vec::new();
6332 for b in 0..n_blocks {
6333 packed.extend_from_slice(&half::f16::from_f32(0.05 + b as f32 * 0.01).to_le_bytes());
6334 for i in 0..16u8 {
6335 let lo = (i + b as u8) % 16;
6336 let hi = (15 - i + b as u8) % 16;
6337 packed.push(lo | (hi << 4));
6338 }
6339 }
6340 let x: Vec<f32> = (0..n_blocks * 32)
6341 .map(|i| ((i as f32) * 0.09).sin())
6342 .collect();
6343
6344 let simd = unsafe { simd_x86::dot_q4_0_f32_avx2(&packed, &x) };
6345 let scalar = dot_q4_0_f32_scalar(&packed, &x);
6346 assert!(
6347 (simd - scalar).abs() < 1e-2,
6348 "AVX2 Q4_0 kernel diverged from scalar: simd={simd} scalar={scalar}"
6349 );
6350 }
6351
6352 #[cfg(target_arch = "aarch64")]
6353 #[test]
6354 fn neon_kernel_matches_scalar_directly_when_available() {
6355 if !std::arch::is_aarch64_feature_detected!("neon") {
6356 eprintln!("skipping: host CPU lacks NEON (unexpected on real aarch64 hardware)");
6357 return;
6358 }
6359 let n_blocks = 8;
6360 let original: Vec<f32> = (0..n_blocks * 32)
6361 .map(|i| ((i % 37) as f32 - 18.0) * 0.11)
6362 .collect();
6363 let packed = quantize_q8_0(&original);
6364 let x: Vec<f32> = (0..n_blocks * 32)
6365 .map(|i| ((i as f32) * 0.07).cos())
6366 .collect();
6367
6368 let simd = unsafe { simd_aarch64::dot_q8_0_f32_neon(&packed, &x) };
6369 let scalar = dot_q8_0_f32_scalar(&packed, &x);
6370 assert!(
6371 (simd - scalar).abs() < 1e-2,
6372 "NEON kernel diverged from scalar: simd={simd} scalar={scalar}"
6373 );
6374 }
6375
6376 #[cfg(target_arch = "aarch64")]
6377 #[test]
6378 fn neon_q4_0_kernel_matches_scalar_directly_when_available() {
6379 if !std::arch::is_aarch64_feature_detected!("neon") {
6380 eprintln!("skipping: host CPU lacks NEON (unexpected on real aarch64 hardware)");
6381 return;
6382 }
6383 let n_blocks = 6;
6388 let mut packed = Vec::new();
6389 for b in 0..n_blocks {
6390 packed.extend_from_slice(&half::f16::from_f32(0.05 + b as f32 * 0.01).to_le_bytes());
6391 for i in 0..16u8 {
6392 let lo = (i + b as u8) % 16;
6393 let hi = (15 - i + b as u8) % 16;
6394 packed.push(lo | (hi << 4));
6395 }
6396 }
6397 let x: Vec<f32> = (0..n_blocks * 32)
6398 .map(|i| ((i as f32) * 0.09).sin())
6399 .collect();
6400
6401 let simd = unsafe { simd_aarch64::dot_q4_0_f32_neon(&packed, &x) };
6402 let scalar = dot_q4_0_f32_scalar(&packed, &x);
6403 assert!(
6404 (simd - scalar).abs() < 1e-2,
6405 "NEON Q4_0 kernel diverged from scalar: simd={simd} scalar={scalar}"
6406 );
6407 }
6408
6409 #[test]
6410 fn dispatched_q4_0_matches_scalar_reference() {
6411 let n_blocks = 4;
6412 let mut packed = Vec::new();
6413 for b in 0..n_blocks {
6414 packed.extend_from_slice(&half::f16::from_f32(0.2).to_le_bytes());
6415 for i in 0..16u8 {
6416 packed.push((i % 16) | (((15 - i + b as u8) % 16) << 4));
6417 }
6418 }
6419 let x: Vec<f32> = (0..n_blocks * 32)
6420 .map(|i| (i as f32) * 0.02 - 1.0)
6421 .collect();
6422
6423 let dispatched = dot_q4_0_f32(&packed, &x);
6424 let scalar = dot_q4_0_f32_scalar(&packed, &x);
6425 assert!(
6426 (dispatched - scalar).abs() < 1e-2,
6427 "dispatched={dispatched} scalar={scalar}"
6428 );
6429 }
6430
6431 #[test]
6432 fn fused_q4_0_dot_matches_dequant_then_dot() {
6433 let mut block = Vec::new();
6434 block.extend_from_slice(&f16::from_f32(1.0).to_le_bytes());
6435 block.extend_from_slice(&[0x12u8; 16]); let x: Vec<f32> = (0..32).map(|i| (i as f32) * 0.1).collect();
6437
6438 let dequanted = dequant_q4_0(&block).unwrap();
6439 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6440 let fused = dot_q4_0_f32(&block, &x);
6441 assert!(
6442 (fused - expected).abs() < 1e-3,
6443 "fused={fused} expected={expected}"
6444 );
6445 }
6446
6447 const Q4_K_TEST_BLOCK: [u8; 144] = [
6453 0x66, 0x2a, 0x66, 0x2a, 0x02, 0x02, 0x02, 0x02, 0x4f, 0x4b, 0x10, 0x12, 0x42, 0xe4, 0xc1,
6454 0xb2, 0x64, 0xa8, 0x70, 0x2d, 0x6a, 0xa6, 0x76, 0x79, 0xa6, 0xf7, 0x5a, 0xda, 0x37, 0x87,
6455 0x38, 0xd5, 0xf9, 0xfa, 0xc2, 0x98, 0x33, 0x94, 0x48, 0x59, 0x46, 0x73, 0xb2, 0x3b, 0x28,
6456 0x18, 0x2e, 0x02, 0xe4, 0x5d, 0x86, 0xa9, 0x93, 0x39, 0x51, 0x75, 0x5f, 0xb6, 0xac, 0x0a,
6457 0x17, 0x35, 0x8d, 0xf7, 0x97, 0x7a, 0x95, 0xf5, 0x51, 0xc9, 0xdd, 0xb8, 0xdf, 0x7a, 0x69,
6458 0xdb, 0xcb, 0xfe, 0xa6, 0xf0, 0x69, 0xf6, 0xf2, 0xc6, 0xad, 0xb4, 0x68, 0x9f, 0xad, 0x7f,
6459 0xd6, 0x40, 0x8f, 0x14, 0xca, 0xdb, 0xa9, 0x7d, 0x89, 0xb6, 0xad, 0x96, 0xa9, 0x69, 0x96,
6460 0xaa, 0x98, 0x79, 0x06, 0x9a, 0x86, 0x74, 0xff, 0xde, 0x8e, 0xf0, 0xf0, 0x3f, 0xcd, 0xdd,
6461 0x7d, 0x7f, 0x0c, 0x3d, 0x0e, 0x7f, 0x88, 0x8f, 0xf7, 0x95, 0x83, 0x13, 0x11, 0x85, 0x55,
6462 0x0c, 0x5c, 0x7b, 0x9e, 0x51, 0x48, 0x69, 0x67, 0x1e,
6463 ];
6464 const Q4_K_GOLDEN: [f32; 256] = [
6465 -0.349915, 0.0499878, -0.749817, 0.549866, 0.249939, -0.149963, -0.149963, 0.149963,
6466 -0.149963, -0.0499878, 0.249939, 0.249939, -0.0499878, -0.0499878, 0.0499878, -0.249939,
6467 0.149963, 0.249939, -0.549866, 0.0499878, -0.44989, -0.349915, 0.0499878, 0.149963,
6468 -0.149963, -0.44989, -0.549866, 0.349915, 0.0499878, 0.0499878, 0.649841, -0.549866,
6469 0.0499878, 0.44989, 0.149963, -0.349915, 0.0499878, 0.44989, 0.149963, 0.149963, 0.44989,
6470 0.949768, -0.0499878, 0.749817, -0.249939, 0.249939, -0.249939, 0.749817, 0.949768,
6471 0.949768, 0.649841, 0.349915, -0.249939, 0.349915, -0.149963, -0.0499878, -0.149963,
6472 0.149963, 0.549866, -0.249939, -0.349915, -0.44989, -0.349915, -0.549866, -0.399902,
6473 0.499878, -0.199951, 0.0999756, -0.499878, 0.0999756, -0.699829, -0.299927, 0.699829,
6474 -0.199951, 0.399902, 0.199951, -0.0999756, -0.299927, 0.499878, -0.0999756, -0.0999756,
6475 0.199951, -0.299927, -0.299927, -0.699829, 0.0999756, 0.499878, 0.0, 0.699829, 0.199951,
6476 0.0999756, 0.299927, 0.299927, 0.599854, -0.199951, -0.799805, 0.499878, -0.399902,
6477 -0.0999756, 0.0999756, 0.0, -0.599854, -0.399902, -0.199951, -0.399902, 0.199951,
6478 0.0999756, -0.89978, -0.799805, -0.599854, -0.0999756, 0.599854, 0.0, -0.199951, 0.0,
6479 0.599854, -0.399902, 0.299927, 0.399902, 0.199951, 0.399902, -0.199951, -0.299927,
6480 0.399902, 0.299927, 0.599854, 0.0999756, 0.599854, -0.0999756, -0.399902, -0.799805,
6481 -0.399902, 0.299927, -0.599854, -0.199951, 0.499878, 0.299927, 0.499878, -0.399902,
6482 -0.999756, 0.499878, -0.599854, 0.0, 0.0999756, -0.0999756, 0.299927, -0.0999756,
6483 -0.399902, 0.299927, -0.399902, -0.0999756, -0.0999756, -0.399902, 0.0, -0.199951,
6484 -0.0999756, -0.399902, 0.0, -0.399902, -0.599854, -0.299927, 1.49963, 1.49963, 0.89978,
6485 0.499878, 0.699829, -0.299927, 0.299927, 0.499878, -0.0999756, 1.09973, -0.699829,
6486 0.0999756, -1.29968, 0.89978, 1.09973, 0.499878, -0.0999756, 0.0999756, 0.699829, 0.499878,
6487 0.299927, 0.499878, -0.299927, 0.299927, 0.499878, 0.299927, -0.0999756, -1.49963,
6488 0.299927, 0.0999756, -0.0999756, 0.149963, 0.0999756, 0.0999756, -0.599854, -0.599854,
6489 0.149963, 0.0499878, 0.0499878, 0.0499878, 0.149963, 0.0, 0.0499878, 0.0999756, 0.149963,
6490 -0.199951, 0.149963, -0.249939, -0.349915, -0.44989, -0.44989, -0.549866, -0.349915,
6491 -0.349915, 0.0, 0.0, -0.0499878, 0.0999756, -0.549866, -0.199951, -0.149963, -0.249939,
6492 0.0999756, 0.949768, 0.749817, 0.249939, 0.949768, 0.949768, -0.249939, 0.649841, 0.749817,
6493 0.149963, 0.149963, -0.549866, -0.249939, -0.549866, 0.149963, 0.249939, 0.249939,
6494 0.949768, 0.349915, 0.249939, -0.44989, -0.44989, 0.249939, -0.0499878, -0.549866,
6495 -0.0499878, 0.149963, 0.349915, -0.0499878, -0.149963, 0.0499878, 0.0499878, -0.44989,
6496 ];
6497
6498 #[rustfmt::skip]
6500 const Q5_K_TEST_BLOCK: [u8; 176] = [
6501 0x66, 0x2a, 0x66, 0x2a, 0x01, 0x01, 0x01, 0x01, 0x4f, 0x4b, 0x10, 0x12, 0x41, 0xe2, 0xc1,
6502 0xb1, 0x72, 0x2f, 0x20, 0x07, 0x31, 0x0c, 0x38, 0xb3, 0x9c, 0xb8, 0xad, 0x2f, 0x9a, 0xea,
6503 0x17, 0xd0, 0xee, 0x93, 0x9e, 0x3e, 0x74, 0xbb, 0x28, 0x18, 0x39, 0x25, 0xb6, 0x09, 0x18,
6504 0x29, 0x1c, 0x1d, 0x29, 0x41, 0x40, 0x0a, 0x74, 0x7d, 0xfd, 0x21, 0xdd, 0x6d, 0x45, 0x73,
6505 0x0e, 0x1e, 0xc0, 0x4a, 0xfc, 0xf3, 0x8e, 0x24, 0x6b, 0x34, 0x7d, 0xbe, 0x94, 0xde, 0x59,
6506 0x7a, 0x35, 0x30, 0x36, 0x0a, 0xf9, 0x4a, 0x9b, 0xa2, 0x26, 0x21, 0xa2, 0xfa, 0xdf, 0x4b,
6507 0x29, 0x64, 0x6f, 0xbb, 0xca, 0x0f, 0x3c, 0xda, 0x20, 0xf4, 0x93, 0x86, 0xab, 0x6e, 0xb9,
6508 0xe5, 0xd5, 0xa0, 0x82, 0xd6, 0x41, 0xff, 0x12, 0xbc, 0x34, 0xbb, 0xab, 0xb8, 0x20, 0x2f,
6509 0xbb, 0x5f, 0x0c, 0x10, 0xcf, 0x49, 0xc5, 0x86, 0x5c, 0xdf, 0xff, 0x78, 0x44, 0x26, 0x3b,
6510 0xc2, 0x23, 0x3d, 0x2b, 0xe9, 0x00, 0x12, 0xf8, 0xea, 0xe2, 0x9e, 0x5e, 0x50, 0x20, 0x9f,
6511 0x9d, 0x8d, 0x7d, 0x7f, 0xcc, 0x1d, 0x0e, 0x13, 0xf8, 0xc2, 0xf1, 0x3d, 0x08, 0x2f, 0x23,
6512 0x13, 0xac, 0x0d, 0xa7, 0xe7, 0x20, 0xa3, 0x90, 0xb7, 0xc8, 0x28,
6513 ];
6514 const Q5_K_GOLDEN: [f32; 256] = [
6515 -0.299927, 0.0999756, -0.749817, 0.549866, 0.249939, -0.0999756, -0.0999756, 0.0999756,
6516 -0.0999756, -0.0999756, 0.299927, 0.199951, -0.0499878, -0.0499878, 0.0499878, -0.249939,
6517 -0.149963, 0.199951, -0.0499878, -0.549866, -0.199951, 0.249939, -0.0999756, -0.0499878,
6518 0.249939, 0.749817, -0.299927, 0.549866, -0.499878, 0.0499878, -0.44989, 0.549866,
6519 0.349915, 0.44989, -0.349915, 0.249939, -0.199951, -0.199951, 0.199951, 0.349915,
6520 0.0999756, -0.249939, -0.349915, 0.599854, 0.249939, 0.299927, 0.849792, -0.349915,
6521 0.999756, 0.999756, 0.649841, 0.349915, -0.249939, 0.399902, -0.199951, 0.0, -0.0999756,
6522 0.0999756, 0.499878, -0.199951, -0.399902, -0.399902, -0.399902, -0.549866, -0.349915,
6523 0.499878, -0.249939, 0.0999756, -0.499878, 0.0499878, -0.699829, -0.299927, 0.749817,
6524 -0.249939, 0.44989, 0.199951, -0.0499878, -0.249939, 0.499878, -0.0499878, 0.599854,
6525 -0.299927, 0.0, 0.199951, 0.149963, -0.499878, -0.249939, -0.0999756, -0.349915, 0.249939,
6526 0.249939, -0.799805, -0.699829, -0.499878, 0.0499878, 0.749817, -0.149963, 0.0999756,
6527 -0.44989, -0.399902, -0.799805, 0.0, 0.399902, -0.149963, 0.549866, 0.0999756, 0.0,
6528 0.199951, 0.199951, 0.44989, -0.299927, -0.89978, 0.0499878, -0.249939, 0.0, 0.649841,
6529 -0.44989, 0.299927, 0.399902, 0.199951, 0.44989, -0.199951, -0.249939, 0.399902, 0.299927,
6530 0.549866, 0.0999756, 0.649841, -0.0999756, -0.399902, -0.799805, -0.44989, 0.349915,
6531 -0.599854, -0.199951, 0.549866, 0.349915, 0.549866, -0.399902, -0.999756, 0.549866,
6532 -0.549866, 0.0499878, 0.0999756, -0.399902, 0.549866, 0.549866, 0.199951, 0.0, 0.0999756,
6533 -0.44989, -0.0999756, -0.0499878, -0.349915, 0.349915, -0.549866, -0.199951, -0.89978,
6534 0.199951, 0.299927, 0.199951, 1.19971, 0.399902, -0.399902, 1.09973, -0.399902, 0.299927,
6535 0.299927, -0.399902, 0.599854, 0.0999756, 0.199951, -0.299927, 0.499878, -0.299927,
6536 -0.699829, 0.599854, -0.199951, 0.0, 0.799805, 0.499878, 0.299927, 0.399902, -0.299927,
6537 0.299927, 0.399902, 0.299927, -0.0999756, -1.49963, 0.199951, 0.0, -0.0999756, 0.299927,
6538 0.0999756, 0.0999756, -0.599854, -0.599854, 0.149963, 0.0499878, 0.0499878, 0.0499878,
6539 0.149963, 0.0, 0.0499878, 0.0999756, 0.349915, -0.199951, 0.299927, 0.249939, 0.0499878,
6540 -0.199951, 0.149963, 0.349915, -0.44989, 0.0, 0.0499878, -0.249939, -0.249939, -0.599854,
6541 -0.44989, -0.599854, -0.249939, -0.199951, -0.199951, 0.149963, -0.0999756, -0.299927,
6542 -0.299927, -0.44989, -0.0999756, -0.0999756, 0.649841, 0.599854, 0.599854, 0.849792,
6543 -0.499878, 0.249939, 0.299927, 0.199951, 0.849792, 0.999756, 0.399902, 0.249939, -0.44989,
6544 -0.44989, 0.299927, -0.0499878, -0.549866, -0.0499878, 0.149963, 0.349915, -0.0499878,
6545 -0.0999756, 0.0, 0.0499878, -0.44989,
6546 ];
6547
6548 #[test]
6549 fn q5_k_dequant_matches_independent_python_reference() {
6550 let got = dequant_q5_k(&Q5_K_TEST_BLOCK).unwrap();
6551 assert_eq!(got.len(), Q5_K_GOLDEN.len());
6552 for (i, (a, b)) in got.iter().zip(Q5_K_GOLDEN.iter()).enumerate() {
6553 assert!(
6554 (a - b).abs() < 1e-3,
6555 "Q5_K element {i}: rust={a} python={b}"
6556 );
6557 }
6558 }
6559
6560 #[test]
6561 fn q5_k_fused_dot_matches_dequant_then_dot() {
6562 let dequanted = dequant_q5_k(&Q5_K_TEST_BLOCK).unwrap();
6563 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).sin()).collect();
6564 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6565 let fused = dot_q5_k_f32(&Q5_K_TEST_BLOCK, &x);
6566 assert!(
6567 (fused - expected).abs() < 1e-2,
6568 "fused={fused} expected={expected}"
6569 );
6570 }
6571
6572 #[test]
6573 fn q5_k_rejects_misaligned_buffers() {
6574 let bad = vec![0u8; 5];
6575 assert!(dequant_q5_k(&bad).is_err());
6576 }
6577
6578 const Q6_K_TEST_BLOCK: [u8; 210] = [
6579 0xe0, 0xa5, 0x40, 0x5c, 0x8d, 0x3a, 0x0a, 0x26, 0xfb, 0x4b, 0x6e, 0x9a, 0xdf, 0x3e, 0xa3,
6580 0xc4, 0xf8, 0x2b, 0x1d, 0x95, 0x76, 0x7d, 0x3b, 0xcd, 0xfd, 0xef, 0xc2, 0x0b, 0x07, 0x63,
6581 0x29, 0xfb, 0x81, 0x57, 0xbe, 0xbe, 0x06, 0xf7, 0x3a, 0x92, 0xc4, 0x43, 0xff, 0xad, 0xac,
6582 0x7e, 0x0f, 0x00, 0x2a, 0x4f, 0xf0, 0xf8, 0xa9, 0xfa, 0x3c, 0x90, 0x6d, 0x73, 0x2d, 0x5a,
6583 0xe6, 0xc6, 0x46, 0xf2, 0x0d, 0x55, 0x4c, 0x25, 0x38, 0x71, 0x2b, 0x35, 0x38, 0x82, 0x16,
6584 0x37, 0x5f, 0x32, 0x61, 0x02, 0xdd, 0x2f, 0x6f, 0x7b, 0x1f, 0xb4, 0x1a, 0x1b, 0x3e, 0x4f,
6585 0x11, 0xa3, 0x17, 0x40, 0x5a, 0x5f, 0x76, 0xcd, 0x19, 0x27, 0x9b, 0xc7, 0xc8, 0xf7, 0xf7,
6586 0xee, 0xf4, 0x86, 0xd9, 0xfd, 0xa7, 0xfe, 0x9e, 0xac, 0x70, 0x53, 0x5b, 0x76, 0xfb, 0x39,
6587 0xf8, 0x4b, 0x98, 0xfe, 0xd0, 0x06, 0x21, 0x4c, 0x4d, 0xbe, 0x10, 0x2b, 0x06, 0x65, 0xc9,
6588 0x5e, 0xf9, 0x95, 0x72, 0xae, 0x99, 0xd9, 0x7e, 0x15, 0xbd, 0x5e, 0x6d, 0xe8, 0x25, 0x8a,
6589 0xd5, 0x99, 0xc6, 0x6b, 0x69, 0xc7, 0x84, 0xc6, 0xa4, 0xf7, 0xb9, 0x6d, 0x68, 0x45, 0x0e,
6590 0x65, 0x69, 0xeb, 0xe6, 0xeb, 0xe9, 0x28, 0xa6, 0xb9, 0x96, 0xf2, 0xe8, 0xa7, 0x9b, 0x6e,
6591 0x79, 0x8a, 0x68, 0x65, 0x59, 0x98, 0x8b, 0x44, 0x41, 0x98, 0x9a, 0x56, 0x01, 0x01, 0x01,
6592 0x02, 0x01, 0x01, 0x01, 0x01, 0x02, 0x01, 0x02, 0x02, 0x01, 0x01, 0x01, 0x02, 0x1f, 0x25,
6593 ];
6594 const Q6_K_GOLDEN: [f32; 256] = [
6595 -0.320068, 0.100021, -0.640137, 0.56012, 0.260056, -0.120026, -0.120026, 0.120026,
6596 -0.100021, -0.100021, 0.28006, 0.200043, -0.0200043, -0.0400085, 0.0600128, -0.240051,
6597 -0.160034, 0.220047, -0.0600128, -0.540115, -0.200043, 0.260056, -0.100021, -0.0600128,
6598 0.260056, 0.620132, -0.28006, 0.540115, -0.500107, 0.0600128, -0.460098, 0.540115,
6599 0.340073, 0.460098, -0.360077, 0.28006, -0.200043, -0.180038, 0.200043, 0.360077,
6600 0.0800171, -0.260056, -0.340073, 0.580124, 0.240051, 0.28006, 0.620132, -0.320068, 1.04022,
6601 1.24026, 0.640137, 0.320068, -0.28006, 0.400085, -0.160034, 0.0, -0.120026, 0.120026,
6602 0.520111, -0.240051, -0.400085, -0.400085, -0.400085, -0.56012, -0.360077, 0.520111,
6603 -0.240051, 0.100021, -0.480103, 0.0600128, -0.640137, -0.28006, 0.620132, -0.240051,
6604 0.440094, 0.180038, -0.0600128, -0.260056, 0.520111, -0.0800171, 0.620132, -0.28006,
6605 0.0200043, 0.180038, 0.14003, -0.500107, -0.260056, -0.0800171, -0.340073, 0.28006,
6606 0.240051, -0.640137, -0.640137, -0.520111, 0.0400085, 0.620132, -0.160034, 0.100021,
6607 -0.42009, -0.42009, -0.640137, -0.0200043, 0.380081, -0.14003, 0.56012, 0.0800171,
6608 -0.0200043, 0.200043, 0.200043, 0.460098, -0.320068, -0.640137, 0.0400085, -0.240051,
6609 -0.0200043, 0.620132, -0.440094, 0.300064, 0.380081, 0.180038, 0.440094, -0.180038,
6610 -0.28006, 0.42009, 0.28006, 0.56012, 0.0800171, 0.620132, -0.120026, -0.440094, -0.800171,
6611 -0.440094, 0.320068, -0.600128, -0.200043, 0.840179, 0.320068, 0.720154, -0.400085,
6612 -1.00021, 0.600128, -0.56012, 0.0400085, 0.0800171, -0.380081, 0.620132, 0.620132,
6613 0.220047, -0.0200043, 0.0800171, -0.440094, -0.100021, -0.0400085, -0.340073, 0.340073,
6614 -0.580124, -0.180038, -0.640137, 0.200043, 0.300064, 0.240051, 1.16025, 0.360077,
6615 -0.360077, 1.08023, -0.360077, 0.320068, 0.28006, -0.360077, 0.56012, 0.160034, 0.240051,
6616 -0.28006, 0.520111, -0.360077, -0.720154, 0.56012, -0.160034, 0.0, 0.760162, 0.440094,
6617 0.240051, 0.440094, -0.28006, 0.320068, 0.440094, 0.320068, -0.0800171, -1.28027, 0.240051,
6618 0.0400085, -0.160034, 0.320068, 0.100021, 0.0800171, -0.600128, -0.580124, 0.14003,
6619 0.0400085, 0.0600128, 0.0600128, 0.160034, 0.0200043, 0.0600128, 0.100021, 0.380081,
6620 -0.200043, 0.320068, 0.260056, 0.0400085, -0.200043, 0.14003, 0.340073, -0.42009,
6621 0.0200043, 0.0200043, -0.260056, -0.240051, -0.620132, -0.440094, -0.620132, -0.240051,
6622 -0.220047, -0.220047, 0.14003, -0.0800171, -0.300064, -0.28006, -0.460098, -0.0800171,
6623 -0.0800171, 0.620132, 0.620132, 0.600128, 0.620132, -0.480103, 0.260056, 0.300064,
6624 0.200043, 0.620132, 1.00021, 0.400085, 0.28006, -0.440094, -0.440094, 0.28006, -0.0400085,
6625 -0.520111, -0.0400085, 0.160034, 0.360077, -0.0400085, -0.120026, 0.0, 0.0800171,
6626 -0.480103,
6627 ];
6628
6629 const Q6_K_SIGNED_SCALES_TEST_BLOCK: [u8; 210] = [
6636 0xe0, 0xa5, 0x40, 0x5c, 0x8d, 0x3a, 0x0a, 0x26, 0xfb, 0x4b, 0x6e, 0x9a, 0xdf, 0x3e, 0xa3,
6637 0xc4, 0x18, 0xe5, 0xf3, 0x7b, 0x9a, 0x93, 0xd5, 0x43, 0x13, 0x20, 0x4e, 0xf5, 0xf9, 0xad,
6638 0xe7, 0x05, 0x81, 0x57, 0xbe, 0xbe, 0x06, 0xf7, 0x3a, 0x92, 0xc4, 0x43, 0xff, 0xad, 0xac,
6639 0x7e, 0x0f, 0x00, 0xe6, 0xc0, 0x10, 0x08, 0x67, 0x16, 0xd4, 0x70, 0xa3, 0x9d, 0xe3, 0xb6,
6640 0x2a, 0x4a, 0xca, 0x0e, 0x0d, 0x55, 0x4c, 0x25, 0x38, 0x71, 0x2b, 0x35, 0x38, 0x82, 0x16,
6641 0x37, 0x5f, 0x32, 0x61, 0x02, 0x33, 0xe1, 0xa1, 0x95, 0xf1, 0x5c, 0xf6, 0xf5, 0xd2, 0xc1,
6642 0xff, 0x6d, 0xf9, 0xcf, 0xb6, 0xb1, 0x76, 0xcd, 0x19, 0x27, 0x9b, 0xc7, 0xc8, 0xf7, 0xf7,
6643 0xee, 0xf4, 0x86, 0xd9, 0xfd, 0xa7, 0xfe, 0x72, 0x64, 0x90, 0xbd, 0xb5, 0x9a, 0x15, 0xd7,
6644 0x18, 0xc5, 0x78, 0x12, 0x3f, 0x0a, 0xef, 0xc4, 0x4d, 0xbe, 0x10, 0x2b, 0x06, 0x65, 0xc9,
6645 0x5e, 0xf9, 0x95, 0x72, 0xae, 0x99, 0xd9, 0x7e, 0x15, 0x42, 0xa1, 0x96, 0x17, 0xda, 0x75,
6646 0x2a, 0x6a, 0x39, 0x94, 0x96, 0x38, 0x7b, 0x39, 0x5b, 0x08, 0xb9, 0x6d, 0x68, 0x45, 0x0e,
6647 0x65, 0x69, 0xeb, 0xe6, 0xeb, 0xe9, 0x28, 0xa6, 0xb9, 0x96, 0xf2, 0x17, 0x58, 0x68, 0x91,
6648 0x86, 0x75, 0x97, 0x9a, 0xa6, 0x67, 0x74, 0xbb, 0xbe, 0xa7, 0x65, 0xa9, 0x01, 0xff, 0x01,
6649 0xfe, 0x01, 0xff, 0x01, 0xff, 0x02, 0xff, 0x02, 0xfe, 0x01, 0xff, 0x01, 0xfe, 0x1f, 0x25,
6650 ];
6651 const Q6_K_SIGNED_SCALES_GOLDEN: [f32; 256] = [
6652 -0.320068, 0.100021, -0.640137, 0.56012, 0.260056, -0.120026, -0.120026, 0.120026,
6653 -0.100021, -0.100021, 0.28006, 0.200043, -0.0200043, -0.0400085, 0.0600128, -0.240051,
6654 -0.160034, 0.220047, -0.0600128, -0.540115, -0.200043, 0.260056, -0.100021, -0.0600128,
6655 0.260056, 0.640137, -0.28006, 0.540115, -0.500107, 0.0600128, -0.460098, 0.540115,
6656 0.340073, 0.460098, -0.360077, 0.28006, -0.200043, -0.180038, 0.200043, 0.360077,
6657 0.0800171, -0.260056, -0.340073, 0.580124, 0.240051, 0.28006, 0.620132, -0.320068, 1.04022,
6658 1.28027, 0.640137, 0.320068, -0.28006, 0.400085, -0.160034, -0.0, -0.120026, 0.120026,
6659 0.520111, -0.240051, -0.400085, -0.400085, -0.400085, -0.56012, -0.360077, 0.520111,
6660 -0.240051, 0.100021, -0.480103, 0.0600128, -0.640137, -0.28006, 0.620132, -0.240051,
6661 0.440094, 0.180038, -0.0600128, -0.260056, 0.520111, -0.0800171, 0.620132, -0.28006,
6662 0.0200043, 0.180038, 0.14003, -0.500107, -0.260056, -0.0800171, -0.340073, 0.28006,
6663 0.240051, -0.620132, -0.620132, -0.520111, 0.0400085, 0.640137, -0.160034, 0.100021,
6664 -0.42009, -0.42009, -0.640137, -0.0200043, 0.380081, -0.14003, 0.56012, 0.0800171,
6665 -0.0200043, 0.200043, 0.200043, 0.460098, -0.320068, -0.640137, 0.0400085, -0.240051,
6666 -0.0200043, 0.640137, -0.440094, 0.300064, 0.380081, 0.180038, 0.440094, -0.180038,
6667 -0.28006, 0.42009, 0.28006, 0.56012, 0.0800171, 0.640137, -0.120026, -0.440094, -0.800171,
6668 -0.440094, 0.320068, -0.600128, -0.200043, 0.840179, 0.320068, 0.720154, -0.400085,
6669 -1.00021, 0.600128, -0.56012, 0.0400085, 0.0800171, -0.380081, 0.620132, 0.620132,
6670 0.220047, -0.0200043, 0.0800171, -0.440094, -0.100021, -0.0400085, -0.340073, 0.340073,
6671 -0.580124, -0.180038, -0.620132, 0.200043, 0.300064, 0.240051, 1.16025, 0.360077,
6672 -0.360077, 1.08023, -0.360077, 0.320068, 0.28006, -0.360077, 0.56012, 0.160034, 0.240051,
6673 -0.28006, 0.520111, -0.360077, -0.720154, 0.56012, -0.160034, -0.0, 0.760162, 0.440094,
6674 0.240051, 0.440094, -0.28006, 0.320068, 0.440094, 0.320068, -0.0800171, -1.24026, 0.240051,
6675 0.0400085, -0.160034, 0.320068, 0.100021, 0.0800171, -0.600128, -0.580124, 0.14003,
6676 0.0400085, 0.0600128, 0.0600128, 0.160034, 0.0200043, 0.0600128, 0.100021, 0.380081,
6677 -0.200043, 0.320068, 0.260056, 0.0400085, -0.200043, 0.14003, 0.340073, -0.42009,
6678 0.0200043, 0.0200043, -0.260056, -0.240051, -0.620132, -0.440094, -0.620132, -0.240051,
6679 -0.220047, -0.220047, 0.14003, -0.0800171, -0.300064, -0.28006, -0.460098, -0.0800171,
6680 -0.0800171, 0.620132, 0.620132, 0.600128, 0.620132, -0.480103, 0.260056, 0.300064,
6681 0.200043, 0.620132, 1.00021, 0.400085, 0.28006, -0.440094, -0.440094, 0.28006, -0.0400085,
6682 -0.520111, -0.0400085, 0.160034, 0.360077, -0.0400085, -0.120026, -0.0, 0.0800171,
6683 -0.480103,
6684 ];
6685
6686 #[test]
6687 fn q4_k_dequant_matches_independent_python_reference() {
6688 let got = dequant_q4_k(&Q4_K_TEST_BLOCK).unwrap();
6689 assert_eq!(got.len(), Q4_K_GOLDEN.len());
6690 for (i, (a, b)) in got.iter().zip(Q4_K_GOLDEN.iter()).enumerate() {
6691 assert!(
6692 (a - b).abs() < 1e-3,
6693 "Q4_K element {i}: rust={a} python={b}"
6694 );
6695 }
6696 }
6697
6698 #[test]
6699 fn q4_k_fused_dot_matches_dequant_then_dot() {
6700 let dequanted = dequant_q4_k(&Q4_K_TEST_BLOCK).unwrap();
6701 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.017).sin()).collect();
6702 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6703 let fused = dot_q4_k_f32(&Q4_K_TEST_BLOCK, &x);
6704 assert!(
6705 (fused - expected).abs() < 1e-2,
6706 "fused={fused} expected={expected}"
6707 );
6708 }
6709
6710 #[test]
6711 fn q6_k_dequant_matches_independent_python_reference() {
6712 let got = dequant_q6_k(&Q6_K_TEST_BLOCK).unwrap();
6713 assert_eq!(got.len(), Q6_K_GOLDEN.len());
6714 for (i, (a, b)) in got.iter().zip(Q6_K_GOLDEN.iter()).enumerate() {
6715 assert!(
6716 (a - b).abs() < 1e-3,
6717 "Q6_K element {i}: rust={a} python={b}"
6718 );
6719 }
6720 }
6721
6722 #[test]
6723 fn q6_k_fused_dot_matches_dequant_then_dot() {
6724 let dequanted = dequant_q6_k(&Q6_K_TEST_BLOCK).unwrap();
6725 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.021).cos()).collect();
6726 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6727 let fused = dot_q6_k_f32(&Q6_K_TEST_BLOCK, &x);
6728 assert!(
6729 (fused - expected).abs() < 1e-2,
6730 "fused={fused} expected={expected}"
6731 );
6732 }
6733
6734 const IQ1_S_TEST_BLOCK: [u8; 50] = [
6741 0x0a, 0x2f, 0xfa, 0x06, 0x1e, 0x37, 0x6f, 0xe3, 0x62, 0xd0, 0xb6, 0xa4, 0x25, 0xae, 0x76,
6742 0x14, 0x72, 0x5b, 0xfa, 0x05, 0xd1, 0xf1, 0x2a, 0x4c, 0xad, 0x29, 0xae, 0xf4, 0xcf, 0x0c,
6743 0x96, 0x51, 0x58, 0x03, 0x6d, 0xd3, 0x10, 0x92, 0x70, 0xff, 0x61, 0x58, 0xc8, 0x30, 0x25,
6744 0x64, 0x49, 0x85, 0xc0, 0x24,
6745 ];
6746 const IQ1_S_GOLDEN: [f32; 256] = [
6747 1.05861, 1.05861, 1.05861, -0.15123, -0.15123, -0.15123, -1.36107, 1.05861, -1.36107,
6748 -1.36107, 1.05861, -0.15123, -1.36107, -0.15123, 1.05861, -0.15123, -0.15123, -0.15123,
6749 -1.36107, -0.15123, -0.15123, -0.15123, 1.05861, -0.15123, -1.36107, -0.15123, -1.36107,
6750 -0.15123, -0.15123, -0.15123, -0.15123, -1.36107, -0.371201, 0.288712, -0.0412445,
6751 0.288712, -0.0412445, -0.0412445, -0.371201, -0.371201, 0.288712, -0.0412445, -0.0412445,
6752 -0.0412445, -0.0412445, -0.0412445, -0.371201, -0.0412445, -0.0412445, -0.371201,
6753 -0.0412445, -0.0412445, -0.0412445, -0.0412445, -0.371201, -0.371201, 0.288712, 0.288712,
6754 0.288712, 0.288712, -0.371201, -0.371201, 0.288712, -0.371201, 1.44356, 1.44356, 1.44356,
6755 -1.856, 1.44356, 1.44356, -1.856, -1.856, -1.856, 1.44356, -0.206223, -1.856, -1.856,
6756 -0.206223, -0.206223, 1.44356, 1.44356, -0.206223, -0.206223, -0.206223, -0.206223,
6757 -0.206223, 1.44356, -0.206223, -0.206223, 1.44356, -1.856, 1.44356, -0.206223, -0.206223,
6758 1.44356, 1.44356, 0.15123, 1.36107, 1.36107, 1.36107, 1.36107, 0.15123, 0.15123, -1.05861,
6759 0.15123, 0.15123, -1.05861, 1.36107, 0.15123, 0.15123, 0.15123, 0.15123, 1.36107, 1.36107,
6760 -1.05861, 1.36107, 1.36107, 0.15123, 0.15123, -1.05861, 0.15123, 1.36107, -1.05861,
6761 -1.05861, -1.05861, 1.36107, 0.15123, 0.15123, 0.866135, 0.0962372, -0.67366, 0.0962372,
6762 0.866135, -0.67366, 0.0962372, -0.67366, 0.866135, 0.0962372, 0.0962372, 0.866135,
6763 0.866135, -0.67366, 0.0962372, -0.67366, -0.67366, 0.866135, 0.0962372, 0.0962372,
6764 -0.67366, 0.0962372, 0.0962372, 0.0962372, 0.866135, -0.67366, 0.0962372, 0.866135,
6765 0.0962372, -0.67366, 0.0962372, -0.67366, 1.60854, 0.178726, 1.60854, 0.178726, 0.178726,
6766 0.178726, 1.60854, 0.178726, 1.60854, 0.178726, -1.25108, 1.60854, 1.60854, 0.178726,
6767 0.178726, 0.178726, 0.178726, 0.178726, 1.60854, 1.60854, 0.178726, 1.60854, -1.25108,
6768 -1.25108, -1.25108, -1.25108, 0.178726, -1.25108, 1.60854, 0.178726, 1.60854, -1.25108,
6769 0.0962372, -0.123734, -0.0137482, -0.0137482, 0.0962372, 0.0962372, -0.0137482, -0.123734,
6770 -0.123734, 0.0962372, -0.123734, 0.0962372, 0.0962372, -0.123734, 0.0962372, -0.123734,
6771 -0.0137482, 0.0962372, -0.0137482, -0.0137482, 0.0962372, -0.123734, -0.123734, 0.0962372,
6772 -0.123734, -0.0137482, -0.0137482, 0.0962372, -0.123734, -0.0137482, 0.0962372, -0.123734,
6773 0.618668, 0.618668, -0.481186, 0.618668, -0.481186, 0.618668, -0.481186, -0.481186,
6774 -0.481186, 0.0687408, 0.618668, 0.0687408, -0.481186, 0.0687408, -0.481186, -0.481186,
6775 0.0687408, 0.0687408, 0.618668, 0.618668, 0.618668, 0.618668, -0.481186, 0.0687408,
6776 0.618668, 0.0687408, -0.481186, 0.0687408, -0.481186, 0.0687408, 0.618668, -0.481186,
6777 ];
6778
6779 const IQ2_XXS_TEST_BLOCK: [u8; 66] = [
6780 0x29, 0x30, 0xd9, 0x33, 0x95, 0x4c, 0x08, 0x1e, 0xad, 0x79, 0x49, 0xf2, 0x8d, 0x5f, 0x93,
6781 0xea, 0x78, 0x18, 0x98, 0xb9, 0x94, 0x14, 0xad, 0xce, 0xca, 0x1d, 0xab, 0x81, 0x53, 0x4a,
6782 0x68, 0xd0, 0x59, 0x96, 0x36, 0x5d, 0xbe, 0x20, 0xc4, 0xff, 0xe4, 0x2c, 0xcd, 0x2f, 0x4f,
6783 0x4f, 0x67, 0x53, 0xc6, 0xd5, 0xa2, 0xfb, 0xc7, 0xf3, 0xe2, 0x6b, 0xf1, 0x99, 0x23, 0x1e,
6784 0x2d, 0x5e, 0x8c, 0x78, 0xc2, 0x31,
6785 ];
6786 const IQ2_XXS_GOLDEN: [f32; 256] = [
6787 1.95007, 1.95007, 1.95007, -6.09398, 6.09398, 1.95007, 1.95007, -10.4816, 1.95007, 1.95007,
6788 -1.95007, -10.4816, -6.09398, -6.09398, 1.95007, 1.95007, 6.09398, 6.09398, -1.95007,
6789 10.4816, -6.09398, -1.95007, 1.95007, -6.09398, -1.95007, 1.95007, -1.95007, -6.09398,
6790 1.95007, 1.95007, -6.09398, 1.95007, -0.390015, -1.2188, 0.390015, 0.390015, -0.390015,
6791 0.390015, 1.2188, -0.390015, -0.390015, 0.390015, -0.390015, 0.390015, -0.390015, 1.2188,
6792 -1.2188, 2.09633, -0.390015, -0.390015, 2.09633, 1.2188, 0.390015, -0.390015, -0.390015,
6793 1.2188, -0.390015, -2.09633, 1.2188, 0.390015, 1.2188, 1.2188, -1.2188, -0.390015,
6794 -0.390015, 2.09633, -0.390015, -1.2188, 2.09633, -0.390015, 0.390015, 1.2188, -0.390015,
6795 0.390015, -1.2188, -2.09633, -0.390015, 1.2188, 1.2188, 1.2188, -0.390015, -0.390015,
6796 0.390015, -2.09633, 1.2188, -0.390015, 0.390015, 1.2188, 2.09633, -0.390015, -2.09633,
6797 -2.09633, 0.390015, -0.390015, -0.390015, -0.390015, 13.2767, 2.47009, 2.47009, -13.2767,
6798 7.71904, -13.2767, -2.47009, -7.71904, 2.47009, 2.47009, 13.2767, 2.47009, 2.47009,
6799 -13.2767, 13.2767, -2.47009, -2.47009, -2.47009, -13.2767, 2.47009, 7.71904, -2.47009,
6800 -7.71904, -2.47009, 2.47009, -2.47009, 7.71904, 2.47009, -2.47009, -2.47009, 7.71904,
6801 -2.47009, 0.650024, 0.650024, -2.03133, 0.650024, 3.49388, 2.03133, -0.650024, 0.650024,
6802 -2.03133, -3.49388, -0.650024, -2.03133, -0.650024, -2.03133, -2.03133, -0.650024,
6803 -0.650024, -0.650024, 0.650024, 0.650024, -0.650024, 3.49388, 2.03133, -2.03133, -2.03133,
6804 -0.650024, -0.650024, 0.650024, 0.650024, -2.03133, -0.650024, -0.650024, -10.9692,
6805 -10.9692, -3.51013, 3.51013, 3.51013, -10.9692, -18.867, -10.9692, 3.51013, -18.867,
6806 -3.51013, 3.51013, 10.9692, -10.9692, 3.51013, -3.51013, -3.51013, 3.51013, 10.9692,
6807 -18.867, -3.51013, 3.51013, 10.9692, -3.51013, 3.51013, -3.51013, -10.9692, -18.867,
6808 3.51013, -3.51013, 10.9692, 3.51013, 2.47009, -2.47009, 2.47009, 2.47009, 13.2767,
6809 -7.71904, -2.47009, -7.71904, -7.71904, -13.2767, -2.47009, 2.47009, -7.71904, 2.47009,
6810 -13.2767, -13.2767, -2.47009, 13.2767, -13.2767, 7.71904, 2.47009, 2.47009, -13.2767,
6811 -7.71904, -13.2767, -2.47009, -13.2767, -2.47009, 2.47009, 7.71904, -7.71904, -13.2767,
6812 2.84386, 0.910034, -4.89143, -0.910034, 0.910034, 2.84386, 0.910034, 0.910034, -4.89143,
6813 0.910034, 4.89143, 0.910034, -4.89143, -0.910034, -0.910034, 0.910034, -0.910034, 0.910034,
6814 0.910034, -0.910034, 2.84386, 2.84386, 0.910034, 0.910034, 0.910034, -0.910034, -0.910034,
6815 -4.89143, 0.910034, 2.84386, 2.84386, -0.910034,
6816 ];
6817
6818 const IQ3_XXS_TEST_BLOCK: [u8; 98] = [
6819 0x71, 0x31, 0x16, 0x0a, 0x79, 0x04, 0x5d, 0x87, 0xae, 0x2a, 0x4a, 0x43, 0xfd, 0x02, 0xba,
6820 0x6c, 0x10, 0x42, 0x80, 0xe5, 0x1d, 0x08, 0x22, 0xcb, 0x21, 0x54, 0xf9, 0xaa, 0x8e, 0xc2,
6821 0xf2, 0x34, 0x66, 0x1e, 0x2a, 0xef, 0x19, 0xae, 0x48, 0x47, 0x29, 0xa0, 0x72, 0xd1, 0x31,
6822 0xc0, 0x65, 0x49, 0xde, 0x79, 0x32, 0xe6, 0x4d, 0xb6, 0x55, 0x3f, 0x4d, 0xf1, 0x18, 0xbb,
6823 0x18, 0x59, 0x4c, 0x31, 0xa3, 0xb2, 0x34, 0xdd, 0xf6, 0x4a, 0x91, 0x51, 0x3f, 0x3e, 0x40,
6824 0x69, 0xad, 0xbf, 0x1a, 0xd0, 0x05, 0xfb, 0xbe, 0x8b, 0x0b, 0xdd, 0xdf, 0x7d, 0x94, 0x74,
6825 0x92, 0x3e, 0xff, 0x04, 0x2a, 0xc4, 0xea, 0xc9,
6826 ];
6827 const IQ3_XXS_GOLDEN: [f32; 256] = [
6828 1.5304, 23.7211, -4.59119, 1.5304, -10.7128, -23.7211, 1.5304, -1.5304, 7.65198, -7.65198,
6829 7.65198, -7.65198, -10.7128, -4.59119, 1.5304, 1.5304, -4.59119, -23.7211, 16.8344,
6830 -4.59119, -13.7736, 23.7211, -16.8344, -7.65198, -13.7736, -1.5304, -1.5304, 13.7736,
6831 -23.7211, 10.7128, -13.7736, -1.5304, -3.57092, 1.19031, 5.95154, 3.57092, -18.4498,
6832 10.7128, 1.19031, 3.57092, -5.95154, -1.19031, 13.0934, 18.4498, 10.7128, -1.19031,
6833 1.19031, -1.19031, -18.4498, 1.19031, -8.33215, -10.7128, -13.0934, -1.19031, -3.57092,
6834 5.95154, -3.57092, 5.95154, 3.57092, 1.19031, -10.7128, -8.33215, -1.19031, 3.57092,
6835 3.91101, 60.6207, 27.3771, 19.5551, 35.1991, 35.1991, -35.1991, -50.8431, 11.733, -27.3771,
6836 19.5551, 3.91101, -11.733, 27.3771, -3.91101, -3.91101, -43.0211, 60.6207, -19.5551,
6837 3.91101, -50.8431, -19.5551, 11.733, 43.0211, -60.6207, 43.0211, -19.5551, -3.91101,
6838 -11.733, -27.3771, -27.3771, 11.733, 5.27136, -68.5277, 36.8995, -81.7061, -68.5277,
6839 36.8995, 68.5277, -36.8995, 26.3568, 15.8141, 5.27136, 36.8995, 57.985, -5.27136, 81.7061,
6840 -47.4423, -5.27136, -47.4423, -15.8141, 81.7061, -47.4423, 68.5277, 68.5277, 5.27136,
6841 26.3568, 26.3568, 5.27136, -26.3568, -36.8995, 36.8995, -26.3568, -5.27136, 71.1634,
6842 -32.1383, -41.3207, -4.59119, -22.9559, -32.1383, 4.59119, -71.1634, -41.3207, -4.59119,
6843 -22.9559, 4.59119, -41.3207, 4.59119, 4.59119, 41.3207, 4.59119, -22.9559, -13.7736,
6844 -13.7736, 13.7736, -13.7736, 13.7736, 13.7736, 32.1383, 13.7736, 41.3207, -4.59119,
6845 13.7736, -13.7736, -32.1383, -32.1383, -39.5352, -33.1586, -7.65198, -12.7533, -17.8546,
6846 28.0573, -17.8546, 28.0573, -12.7533, -17.8546, 28.0573, -2.55066, -17.8546, -22.9559,
6847 -28.0573, 22.9559, -2.55066, 12.7533, 2.55066, 12.7533, -12.7533, 12.7533, -7.65198,
6848 -7.65198, 22.9559, 33.1586, -2.55066, 33.1586, 12.7533, -12.7533, 12.7533, 12.7533,
6849 0.85022, -1.87048, 1.87048, 0.170044, -1.87048, 0.170044, 1.87048, 2.21057, -0.85022,
6850 0.510132, -0.85022, -0.510132, -2.63568, -1.19031, -0.85022, 1.5304, 2.21057, 1.5304,
6851 -1.19031, -0.510132, -1.19031, -0.85022, -0.170044, -0.510132, -0.85022, -0.510132,
6852 -0.85022, 0.510132, 2.21057, -0.85022, 0.510132, 2.63568, 21.2555, -4.2511, 21.2555,
6853 -4.2511, 46.7621, -38.2599, 29.7577, -38.2599, 21.2555, 4.2511, 21.2555, -4.2511, 4.2511,
6854 46.7621, 38.2599, -12.7533, -4.2511, -12.7533, 21.2555, -12.7533, 21.2555, -29.7577,
6855 46.7621, 4.2511, -65.892, -38.2599, -38.2599, -29.7577, 29.7577, 46.7621, -4.2511,
6856 -38.2599,
6857 ];
6858
6859 #[test]
6860 fn iq1_s_dequant_matches_independent_python_reference() {
6861 let got = dequant_iq1_s(&IQ1_S_TEST_BLOCK).unwrap();
6862 assert_eq!(got.len(), IQ1_S_GOLDEN.len());
6863 for (i, (a, b)) in got.iter().zip(IQ1_S_GOLDEN.iter()).enumerate() {
6864 assert!(
6865 (a - b).abs() < 1e-3,
6866 "IQ1_S element {i}: rust={a} python={b}"
6867 );
6868 }
6869 }
6870
6871 #[test]
6872 fn iq2_xxs_dequant_matches_independent_python_reference() {
6873 let got = dequant_iq2_xxs(&IQ2_XXS_TEST_BLOCK).unwrap();
6874 assert_eq!(got.len(), IQ2_XXS_GOLDEN.len());
6875 for (i, (a, b)) in got.iter().zip(IQ2_XXS_GOLDEN.iter()).enumerate() {
6876 assert!(
6877 (a - b).abs() < 1e-3,
6878 "IQ2_XXS element {i}: rust={a} python={b}"
6879 );
6880 }
6881 }
6882
6883 #[test]
6884 fn iq3_xxs_dequant_matches_independent_python_reference() {
6885 let got = dequant_iq3_xxs(&IQ3_XXS_TEST_BLOCK).unwrap();
6886 assert_eq!(got.len(), IQ3_XXS_GOLDEN.len());
6887 for (i, (a, b)) in got.iter().zip(IQ3_XXS_GOLDEN.iter()).enumerate() {
6888 assert!(
6889 (a - b).abs() < 1e-3,
6890 "IQ3_XXS element {i}: rust={a} python={b}"
6891 );
6892 }
6893 }
6894
6895 #[test]
6896 fn iq_lowbit_fused_dots_match_dequant_then_dot() {
6897 type DequantFn = fn(&[u8]) -> Result<Vec<f32>, QuantError>;
6898 type DotFn = fn(&[u8], &[f32]) -> f32;
6899 let x: Vec<f32> = (0..1024).map(|i| ((i as f32) * 0.027).sin()).collect();
6900 let cases: [(&[u8], usize, DequantFn, DotFn); 3] = [
6901 (&IQ1_S_TEST_BLOCK, 4, dequant_iq1_s, dot_iq1_s_f32),
6902 (&IQ2_XXS_TEST_BLOCK, 4, dequant_iq2_xxs, dot_iq2_xxs_f32),
6903 (&IQ3_XXS_TEST_BLOCK, 4, dequant_iq3_xxs, dot_iq3_xxs_f32),
6904 ];
6905 for (block, n, dequant, dot) in cases {
6906 let packed = repeat_block(block, n);
6907 let dequanted = dequant(&packed).unwrap();
6908 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
6909 let fused = dot(&packed, &x[..dequanted.len()]);
6910 assert!(
6911 (fused - expected).abs() < 1e-2,
6912 "fused={fused} expected={expected}"
6913 );
6914 }
6915 }
6916
6917 #[cfg(target_arch = "x86_64")]
6922 #[test]
6923 fn avx2_iq_kernels_match_scalar_directly_on_random_blocks() {
6924 if !(is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma")) {
6925 eprintln!("skipping: host CPU lacks AVX2+FMA");
6926 return;
6927 }
6928 type ScalarFn = fn(&[u8], &[f32]) -> f32;
6929 type Avx2Fn = unsafe fn(&[u8], &[f32]) -> f32;
6930 let cases: [(&str, usize, ScalarFn, Avx2Fn); 3] = [
6931 (
6932 "iq1_s",
6933 IQ1_S_BLOCK_BYTES,
6934 dot_iq1_s_f32_scalar,
6935 simd_x86::dot_iq1_s_f32_avx2,
6936 ),
6937 (
6938 "iq2_xxs",
6939 IQ2_XXS_BLOCK_BYTES,
6940 dot_iq2_xxs_f32_scalar,
6941 simd_x86::dot_iq2_xxs_f32_avx2,
6942 ),
6943 (
6944 "iq3_xxs",
6945 IQ3_XXS_BLOCK_BYTES,
6946 dot_iq3_xxs_f32_scalar,
6947 simd_x86::dot_iq3_xxs_f32_avx2,
6948 ),
6949 ];
6950 for (name, block_bytes, scalar, avx2) in cases {
6951 for trial in 0..16u32 {
6952 let n_blocks = 3;
6953 let mut bytes =
6954 pseudo_random_bytes(trial.wrapping_mul(97) + 5, n_blocks * block_bytes);
6955 for b in 0..n_blocks {
6956 let d = half::f16::from_f32(0.05 + 0.01 * trial as f32).to_le_bytes();
6958 bytes[b * block_bytes] = d[0];
6959 bytes[b * block_bytes + 1] = d[1];
6960 }
6961 let x: Vec<f32> = (0..n_blocks * 256)
6962 .map(|i| ((i as f32) * 0.017 + trial as f32).sin())
6963 .collect();
6964 let s = scalar(&bytes, &x);
6965 let v = unsafe { avx2(&bytes, &x) };
6966 let tol = 2e-3_f32.max(s.abs() * 1e-3);
6973 assert!(
6974 (s - v).abs() < tol,
6975 "{name} trial {trial}: scalar={s} avx2={v}"
6976 );
6977 }
6978 }
6979 }
6980
6981 #[cfg(target_arch = "x86_64")]
6986 fn pseudo_random_bytes(seed: u32, len: usize) -> Vec<u8> {
6987 let mut state = seed.wrapping_mul(2654435761).wrapping_add(1);
6988 (0..len)
6989 .map(|_| {
6990 state = state.wrapping_mul(1664525).wrapping_add(1013904223);
6991 (state >> 16) as u8
6992 })
6993 .collect()
6994 }
6995
6996 #[test]
6997 fn iq_lowbit_dequant_rejects_misaligned_buffers() {
6998 let bad = vec![0u8; 7];
6999 assert!(dequant_iq1_s(&bad).is_err());
7000 assert!(dequant_iq2_xxs(&bad).is_err());
7001 assert!(dequant_iq3_xxs(&bad).is_err());
7002 assert!(dequant_iq2_xs(&bad).is_err());
7003 assert!(dequant_iq2_s(&bad).is_err());
7004 assert!(dequant_iq3_s(&bad).is_err());
7005 assert!(dequant_iq1_m(&bad).is_err());
7006 }
7007
7008 #[test]
7022 fn iq_tier_dequant_matches_real_ggml_exactly() {
7023 type DequantFn = fn(&[u8]) -> Result<Vec<f32>, QuantError>;
7024 let cases: [(&str, &[u8], &[f32], DequantFn); 4] = [
7025 (
7026 "IQ2_XS",
7027 &iq_tier_goldens::IQ2_XS_TEST_BLOCKS,
7028 &iq_tier_goldens::IQ2_XS_GOLDEN,
7029 dequant_iq2_xs,
7030 ),
7031 (
7032 "IQ2_S",
7033 &iq_tier_goldens::IQ2_S_TEST_BLOCKS,
7034 &iq_tier_goldens::IQ2_S_GOLDEN,
7035 dequant_iq2_s,
7036 ),
7037 (
7038 "IQ3_S",
7039 &iq_tier_goldens::IQ3_S_TEST_BLOCKS,
7040 &iq_tier_goldens::IQ3_S_GOLDEN,
7041 dequant_iq3_s,
7042 ),
7043 (
7044 "IQ1_M",
7045 &iq_tier_goldens::IQ1_M_TEST_BLOCKS,
7046 &iq_tier_goldens::IQ1_M_GOLDEN,
7047 dequant_iq1_m,
7048 ),
7049 ];
7050 for (name, blocks, golden, dequant) in cases {
7051 let got = dequant(blocks).unwrap();
7052 assert_eq!(got.len(), golden.len(), "{name}: element count");
7053 for (i, (a, b)) in got.iter().zip(golden.iter()).enumerate() {
7054 assert_eq!(
7055 a.to_bits(),
7056 b.to_bits(),
7057 "{name} element {i} (block {}, offset {}): rust={a} ggml={b}",
7058 i / 256,
7059 i % 256
7060 );
7061 }
7062 }
7063 }
7064
7065 #[test]
7075 fn iq_tier_all_ones_block_reaches_the_maximum_grid_index() {
7076 let d = f16::from_le_bytes([
7080 iq_tier_goldens::IQ2_XS_TEST_BLOCKS[0],
7081 iq_tier_goldens::IQ2_XS_TEST_BLOCKS[1],
7082 ])
7083 .to_f32();
7084 let mag = (iq_tables::IQ2XS_GRID[511] & 0xFF) as f32;
7085 assert_eq!(
7086 iq_tier_goldens::IQ2_XS_GOLDEN[0],
7087 -(d * (0.5 + 15.0) * 0.25) * mag
7088 );
7089
7090 let d = f16::from_le_bytes([
7093 iq_tier_goldens::IQ2_S_TEST_BLOCKS[0],
7094 iq_tier_goldens::IQ2_S_TEST_BLOCKS[1],
7095 ])
7096 .to_f32();
7097 let mag = (iq_tables::IQ2S_GRID[1023] & 0xFF) as f32;
7098 assert_eq!(
7099 iq_tier_goldens::IQ2_S_GOLDEN[0],
7100 -(d * (0.5 + 15.0) * 0.25) * mag
7101 );
7102
7103 let d = f16::from_le_bytes([
7106 iq_tier_goldens::IQ3_S_TEST_BLOCKS[0],
7107 iq_tier_goldens::IQ3_S_TEST_BLOCKS[1],
7108 ])
7109 .to_f32();
7110 let mag = (iq_tables::IQ3S_GRID[511] & 0xFF) as f32;
7111 assert_eq!(iq_tier_goldens::IQ3_S_GOLDEN[0], -(d * 31.0) * mag);
7112
7113 let sc: [u16; 4] = std::array::from_fn(|k| {
7119 u16::from_le_bytes([
7120 iq_tier_goldens::IQ1_M_TEST_BLOCKS[48 + 2 * k],
7121 iq_tier_goldens::IQ1_M_TEST_BLOCKS[48 + 2 * k + 1],
7122 ])
7123 });
7124 let d = f16::from_bits(
7125 (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000),
7126 )
7127 .to_f32();
7128 let v = (iq_tables::IQ1S_GRID[2047] & 0xFF) as u8 as i8;
7129 assert_eq!(
7130 iq_tier_goldens::IQ1_M_GOLDEN[0],
7131 d * 15.0 * (v as f32 - IQ1S_DELTA)
7132 );
7133 }
7134
7135 #[test]
7141 fn iq_tier_fused_dots_match_dequant_then_dot() {
7142 type DequantFn = fn(&[u8]) -> Result<Vec<f32>, QuantError>;
7143 type DotFn = fn(&[u8], &[f32]) -> f32;
7144 let x: Vec<f32> = (0..1024).map(|i| ((i as f32) * 0.031).cos()).collect();
7145 let cases: [(&str, &[u8], DequantFn, DotFn); 4] = [
7146 (
7147 "IQ2_XS",
7148 &iq_tier_goldens::IQ2_XS_TEST_BLOCKS,
7149 dequant_iq2_xs,
7150 dot_iq2_xs_f32,
7151 ),
7152 (
7153 "IQ2_S",
7154 &iq_tier_goldens::IQ2_S_TEST_BLOCKS,
7155 dequant_iq2_s,
7156 dot_iq2_s_f32,
7157 ),
7158 (
7159 "IQ3_S",
7160 &iq_tier_goldens::IQ3_S_TEST_BLOCKS,
7161 dequant_iq3_s,
7162 dot_iq3_s_f32,
7163 ),
7164 (
7165 "IQ1_M",
7166 &iq_tier_goldens::IQ1_M_TEST_BLOCKS,
7167 dequant_iq1_m,
7168 dot_iq1_m_f32,
7169 ),
7170 ];
7171 for (name, blocks, dequant, dot) in cases {
7172 let dequanted = dequant(blocks).unwrap();
7173 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7174 let fused = dot(blocks, &x[..dequanted.len()]);
7175 assert!(
7176 (fused - expected).abs() <= expected.abs() * 1e-5 + 1e-3,
7177 "{name}: fused={fused} expected={expected}"
7178 );
7179 }
7180 }
7181
7182 const MXFP4_GGUF_TEST_BLOCKS: [u8; 68] = [
7188 0x79, 0xb4, 0x8d, 0xe2, 0x62, 0x5d, 0xbb, 0x9d, 0x54, 0xe6, 0xdb, 0x94, 0x59, 0x7d, 0x28,
7189 0xf9, 0x79, 0x7a, 0xfc, 0xc1, 0xfa, 0x1e, 0x53, 0x5b, 0x0e, 0xc2, 0x5a, 0x2f, 0x0c, 0x82,
7190 0x4d, 0xcb, 0x11, 0x28, 0x7b, 0x7c, 0xb6, 0x45, 0xe0, 0xb0, 0x52, 0x40, 0x51, 0xec, 0x30,
7191 0x1a, 0xd2, 0x17, 0xf3, 0xbb, 0xfc, 0x7c, 0x8f, 0xf0, 0x67, 0x83, 0x88, 0x9d, 0x79, 0xdb,
7192 0xf4, 0x45, 0x29, 0x78, 0xe6, 0xf4, 0x99, 0xea,
7193 ];
7194 const MXFP4_GGUF_GOLDEN: [f32; 128] = [
7195 0.03125, -0.046875, 0.015625, 0.015625, -0.046875, -0.0234375, -0.046875, 0.03125, 0.0625,
7196 -0.0234375, 0.03125, -0.0078125, -0.046875, 0.0, -0.0078125, -0.0078125, -0.0234375, 0.0,
7197 -0.0625, 0.0625, 0.046875, -0.0234375, -0.0078125, 0.046875, -0.0625, -0.046875,
7198 -0.0078125, 0.046875, 0.09375, 0.015625, -0.09375, 0.09375, -0.0625, 0.015625, -0.03125,
7199 -0.125, 0.046875, -0.046875, -0.125, 0.03125, -0.03125, -0.1875, -0.0625, 0.03125,
7200 -0.09375, -0.046875, 0.015625, 0.0, -0.1875, -0.0625, -0.1875, 0.015625, 0.09375, 0.09375,
7201 0.0, -0.0625, 0.09375, 0.03125, 0.0, 0.0, 0.0625, -0.0625, 0.015625, 0.03125, -0.125, 0.25,
7202 0.1875, 0.0, 0.0, 0.0625, 0.0, 0.03125, -0.125, 0.0, -0.0625, 0.0625, 0.375, 0.09375,
7203 -0.09375, -0.125, 0.375, -0.09375, 0.125, -0.25, -0.09375, 0.1875, 0.125, 0.1875, -0.25,
7204 0.09375, 0.03125, -0.1875, 0.03125, -0.375, -0.09375, -0.375, -0.75, 0.0, 0.75, 0.1875,
7205 0.0, -0.375, -0.0625, -0.1875, 0.25, 0.375, -0.0625, 0.0, 0.5, 0.25, -0.0625, -0.125, 0.0,
7206 -0.75, 0.5, 0.0, 0.0, -0.0625, 0.75, -0.375, -0.75, 0.25, 0.125, 0.75, -0.5, -0.75,
7207 -0.0625, -0.5,
7208 ];
7209
7210 #[test]
7211 fn mxfp4_gguf_dequant_matches_independent_python_reference() {
7212 let got = dequant_mxfp4_gguf(&MXFP4_GGUF_TEST_BLOCKS).unwrap();
7213 assert_eq!(got.len(), MXFP4_GGUF_GOLDEN.len());
7214 for (i, (a, b)) in got.iter().zip(MXFP4_GGUF_GOLDEN.iter()).enumerate() {
7215 assert!(
7216 (a - b).abs() < 1e-3,
7217 "MXFP4-GGUF element {i}: rust={a} python={b}"
7218 );
7219 }
7220 }
7221
7222 #[test]
7223 fn mxfp4_gguf_fused_dot_matches_dequant_then_dot() {
7224 let dequanted = dequant_mxfp4_gguf(&MXFP4_GGUF_TEST_BLOCKS).unwrap();
7225 let x: Vec<f32> = (0..128).map(|i| ((i as f32) * 0.031).cos()).collect();
7226 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7227 let fused = dot_mxfp4_gguf_f32(&MXFP4_GGUF_TEST_BLOCKS, &x);
7228 assert!(
7229 (fused - expected).abs() < 1e-2,
7230 "fused={fused} expected={expected}"
7231 );
7232 }
7233
7234 #[test]
7239 fn mxfp4_gguf_block_form_agrees_with_two_buffer_form() {
7240 let mut packed = Vec::new();
7241 let mut scales = Vec::new();
7242 for block in MXFP4_GGUF_TEST_BLOCKS
7243 .as_chunks::<MXFP4_GGUF_BLOCK_BYTES>()
7244 .0
7245 {
7246 scales.push(block[0]);
7247 packed.extend_from_slice(&block[1..17]);
7248 }
7249 let x: Vec<f32> = (0..128).map(|i| ((i as f32) * 0.019).sin()).collect();
7250 let a = dot_mxfp4_gguf_f32(&MXFP4_GGUF_TEST_BLOCKS, &x);
7251 let b = dot_mxfp4_row_f32(&packed, &scales, &x);
7252 assert!((a - b).abs() < 1e-4, "block={a} two-buffer={b}");
7253 }
7254
7255 const Q6_K_SIGNED_TEST_BLOCK: [u8; 210] = [
7264 0x10, 0x5b, 0x5f, 0x45, 0x4a, 0xa0, 0x3f, 0x10, 0xf2, 0x7f, 0xdd, 0xf5, 0x25, 0x03, 0xc3,
7265 0x12, 0x74, 0xe1, 0x4e, 0x42, 0xf1, 0x04, 0xe1, 0xad, 0xc6, 0x55, 0x59, 0x4b, 0x5a, 0xfc,
7266 0xf5, 0x3f, 0xc5, 0x0b, 0xac, 0x7b, 0x4c, 0xd4, 0x19, 0xa6, 0x27, 0xdd, 0xf4, 0x7d, 0x9c,
7267 0xfc, 0x03, 0xd2, 0x5f, 0xe3, 0xff, 0x9c, 0xa6, 0x74, 0xa0, 0xe1, 0xbe, 0xf0, 0x26, 0xdb,
7268 0x4b, 0x23, 0xa0, 0xbc, 0xb1, 0x94, 0xd7, 0x7e, 0xcf, 0xf7, 0x97, 0xb4, 0xac, 0x1f, 0xb1,
7269 0x9f, 0xb7, 0xbe, 0xa3, 0xb5, 0xd2, 0xd4, 0x6d, 0x9c, 0x3d, 0xf3, 0x5f, 0x0e, 0x64, 0xbf,
7270 0x54, 0x40, 0xc8, 0xef, 0x9d, 0xc3, 0xf3, 0x4c, 0xb0, 0xf8, 0x54, 0xcf, 0xf3, 0x12, 0xcc,
7271 0x2f, 0x0c, 0xee, 0xab, 0x5d, 0x8d, 0x0b, 0x19, 0xb2, 0x99, 0xbd, 0x4a, 0xec, 0x04, 0xb3,
7272 0xf6, 0xc1, 0xb9, 0xf8, 0x1d, 0xfe, 0x51, 0xea, 0x99, 0xe5, 0x75, 0x5b, 0x98, 0x28, 0x05,
7273 0x18, 0x8a, 0x9f, 0xda, 0xb7, 0xb6, 0xe5, 0x5b, 0x3a, 0x52, 0x49, 0xcc, 0x72, 0xff, 0x61,
7274 0x91, 0x95, 0xa2, 0xa1, 0x5d, 0xd5, 0xc4, 0x7d, 0xb1, 0x0b, 0xda, 0xa9, 0xa2, 0x97, 0x1e,
7275 0x7e, 0xe9, 0xa2, 0xd6, 0xdd, 0x0e, 0x94, 0x21, 0xa4, 0x67, 0x92, 0xad, 0x46, 0xab, 0xe1,
7276 0xe2, 0x3b, 0x21, 0x69, 0x2a, 0x1e, 0xd3, 0xea, 0xa4, 0xdf, 0xa6, 0xd2, 0xff, 0x01, 0xfe,
7277 0xff, 0x01, 0xff, 0x01, 0x01, 0x02, 0xff, 0xff, 0x01, 0xfe, 0x02, 0x01, 0xff, 0x1f, 0x25,
7278 ];
7279 const Q6_K_SIGNED_GOLDEN: [f32; 256] = [
7280 0.320068, 0.100021, 0.0200043, -0.42009, 0.440094, 0.640137, 0.0200043, 0.640137,
7281 -0.0400085, -0.620132, -0.260056, -0.42009, -0.100021, 0.260056, -0.380081, -0.0400085,
7282 0.0800171, -0.300064, -0.360077, 0.0400085, 0.340073, -0.240051, -0.300064, -0.0600128,
7283 0.120026, -0.220047, -0.14003, -0.100021, -0.440094, -0.0800171, -0.220047, 0.620132,
7284 -0.200043, 0.200043, 0.160034, -0.440094, -0.480103, -0.160034, 0.28006, -0.240051,
7285 -0.28006, -1.16025, -0.160034, 0.120026, 0.160034, 0.160034, -0.120026, -0.0800171,
7286 0.340073, -0.0600128, -0.620132, 0.400085, -0.440094, 0.56012, 0.640137, 0.300064,
7287 0.360077, 0.640137, -0.440094, 0.100021, 0.100021, -0.380081, 0.640137, -0.240051,
7288 -0.300064, 0.100021, 0.42009, -0.240051, -0.240051, 0.200043, -0.580124, -0.300064,
7289 -0.340073, -0.180038, -0.0600128, 0.620132, 0.360077, 0.0, -0.0800171, 0.340073, 0.180038,
7290 0.360077, 0.56012, -0.400085, -0.620132, -0.0, 0.0400085, 0.120026, -0.240051, -0.100021,
7291 0.220047, 0.240051, 0.540115, -0.620132, -0.620132, 0.580124, 0.240051, 0.320068,
7292 -0.120026, -0.180038, 0.0800171, -0.380081, -0.620132, -0.440094, 0.0400085, 0.260056,
7293 0.620132, 0.14003, 0.180038, 0.620132, -0.320068, -0.380081, -0.220047, -0.0400085,
7294 0.620132, -0.14003, 0.520111, -0.180038, 0.200043, 0.28006, 0.220047, 0.300064, -0.28006,
7295 0.580124, 0.400085, -0.28006, 0.200043, -0.42009, 0.0400085, -0.480103, 0.28006, 1.20026,
7296 0.600128, 0.28006, -0.360077, 0.160034, 0.480103, -0.0400085, 0.0400085, -0.680145,
7297 -0.360077, -0.720154, 0.760162, 0.200043, 0.28006, -0.0800171, -0.580124, 0.0800171,
7298 -0.260056, -0.380081, 0.0200043, 0.0400085, -0.0800171, -0.300064, -0.400085, -0.0,
7299 0.480103, -0.620132, -0.260056, -0.0600128, -0.0600128, -0.240051, 0.640137, 0.160034,
7300 -0.400085, -0.620132, -0.0600128, 0.600128, 0.0800171, -0.620132, -0.56012, 0.0400085,
7301 0.42009, 0.0600128, 0.0600128, 0.42009, 0.500107, -0.28006, 0.180038, -0.380081, -0.440094,
7302 0.240051, -0.56012, 0.0600128, 0.120026, 0.340073, -0.460098, 0.160034, -0.0600128,
7303 0.600128, -0.300064, -0.440094, 0.200043, -0.360077, -0.520111, 0.360077, 0.160034,
7304 -1.24026, -0.360077, -0.440094, 0.240051, 0.600128, 0.840179, 0.28006, -0.440094,
7305 -0.440094, -0.400085, 0.200043, 0.520111, -0.760162, 0.240051, 0.360077, 0.120026, 1.24026,
7306 0.200043, 0.0, 0.240051, -0.200043, -0.440094, 0.160034, 0.480103, -0.0800171, 0.360077,
7307 -0.160034, 0.620132, 0.0800171, 0.220047, 0.300064, -0.540115, -0.0800171, 0.620132,
7308 0.0200043, 0.56012, 0.360077, -0.640137, 0.28006, -0.440094, 0.100021, -0.160034, 0.0,
7309 -0.0200043, 0.100021, -0.180038, -0.540115, -0.400085, 0.360077, 0.640137, 0.100021,
7310 0.340073, 0.400085, -0.540115, -0.620132, -0.0200043, -0.620132, -0.100021, -0.600128,
7311 ];
7312
7313 #[test]
7314 fn q6_k_signed_scale_dequant_matches_independent_python_reference() {
7315 let got = dequant_q6_k(&Q6_K_SIGNED_TEST_BLOCK).unwrap();
7316 assert_eq!(got.len(), Q6_K_SIGNED_GOLDEN.len());
7317 for (i, (a, b)) in got.iter().zip(Q6_K_SIGNED_GOLDEN.iter()).enumerate() {
7318 assert!(
7319 (a - b).abs() < 1e-3,
7320 "Q6_K signed-scale element {i}: rust={a} python={b}"
7321 );
7322 }
7323 }
7324
7325 #[test]
7326 fn q6_k_signed_scale_fused_dot_matches_dequant_then_dot() {
7327 let dequanted = dequant_q6_k(&Q6_K_SIGNED_TEST_BLOCK).unwrap();
7328 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).sin()).collect();
7329 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7330 let fused = dot_q6_k_f32(&Q6_K_SIGNED_TEST_BLOCK, &x);
7331 assert!(
7332 (fused - expected).abs() < 1e-2,
7333 "fused={fused} expected={expected}"
7334 );
7335 }
7336
7337 #[test]
7338 fn dispatched_q6_k_matches_scalar_on_signed_scales() {
7339 let n_blocks = 4;
7344 let packed = repeat_block(&Q6_K_SIGNED_TEST_BLOCK, n_blocks);
7345 let x: Vec<f32> = (0..256 * n_blocks)
7346 .map(|i| ((i as f32) * 0.019).sin())
7347 .collect();
7348 let dispatched = dot_q6_k_f32(&packed, &x);
7349 let scalar = dot_q6_k_f32_scalar(&packed, &x);
7350 assert!(
7351 (dispatched - scalar).abs() < 1e-1,
7352 "dispatched={dispatched} scalar={scalar}"
7353 );
7354 }
7355
7356 #[test]
7357 fn q6_k_dequant_matches_python_reference_with_negative_scales() {
7358 let got = dequant_q6_k(&Q6_K_SIGNED_SCALES_TEST_BLOCK).unwrap();
7363 assert_eq!(got.len(), Q6_K_SIGNED_SCALES_GOLDEN.len());
7364 for (i, (a, b)) in got.iter().zip(Q6_K_SIGNED_SCALES_GOLDEN.iter()).enumerate() {
7365 assert!(
7366 (a - b).abs() < 1e-3,
7367 "Q6_K signed-scale element {i}: rust={a} python={b}"
7368 );
7369 }
7370 }
7371
7372 #[test]
7373 fn q6_k_fused_dot_matches_dequant_then_dot_with_negative_scales() {
7374 let dequanted = dequant_q6_k(&Q6_K_SIGNED_SCALES_TEST_BLOCK).unwrap();
7375 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).sin()).collect();
7376 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7377 let fused = dot_q6_k_f32(&Q6_K_SIGNED_SCALES_TEST_BLOCK, &x);
7378 assert!(
7379 (fused - expected).abs() < 1e-2,
7380 "fused={fused} expected={expected}"
7381 );
7382 }
7383
7384 #[test]
7385 fn q6_k_scalar_dot_matches_python_reference_with_negative_scales() {
7386 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).sin()).collect();
7391 let expected: f32 = Q6_K_SIGNED_SCALES_GOLDEN
7392 .iter()
7393 .zip(x.iter())
7394 .map(|(a, b)| a * b)
7395 .sum();
7396 let scalar = dot_q6_k_f32_scalar(&Q6_K_SIGNED_SCALES_TEST_BLOCK, &x);
7397 assert!(
7398 (scalar - expected).abs() < 1e-2,
7399 "scalar={scalar} expected={expected}"
7400 );
7401 }
7402
7403 #[test]
7404 fn q4_k_and_q6_k_reject_misaligned_buffers() {
7405 let bad = vec![0u8; 5];
7406 assert!(dequant_q4_k(&bad).is_err());
7407 assert!(dequant_q6_k(&bad).is_err());
7408 }
7409
7410 const Q2_K_TEST_BLOCK: [u8; 84] = [
7417 0x92, 0x32, 0xc9, 0x0e, 0x0f, 0xf8, 0x10, 0xf0, 0xd1, 0x82, 0xca, 0x81, 0x7f, 0x11, 0xdb,
7418 0xff, 0x78, 0xf8, 0xab, 0xc5, 0x60, 0x0c, 0xc0, 0xbc, 0xa6, 0x52, 0x56, 0x1b, 0xc0, 0x36,
7419 0x6b, 0x6e, 0xbb, 0x53, 0x32, 0x90, 0x0a, 0x41, 0x67, 0x97, 0x48, 0x76, 0x86, 0x23, 0xd5,
7420 0x8e, 0x9e, 0x02, 0xc1, 0x1b, 0xea, 0x9c, 0xb7, 0x55, 0xc3, 0x1b, 0xf4, 0x59, 0xc6, 0xef,
7421 0x11, 0x61, 0xbc, 0x54, 0xd7, 0x8a, 0x6d, 0xed, 0x9e, 0xe7, 0x48, 0x69, 0x8e, 0x3a, 0x30,
7422 0x6c, 0xd8, 0xdc, 0x85, 0xc1, 0xec, 0x35, 0x14, 0x32,
7423 ];
7424 const Q2_K_GOLDEN: [f32; 256] = [
7425 -1.70947, -1.70947, 0.51123, -0.969238, -1.70947, -1.70947, -1.70947, -1.70947, -0.229004,
7426 -0.229004, -0.229004, 0.51123, -1.70947, -0.229004, 0.51123, -0.229004, 1.65088, 1.65088,
7427 0.910645, -0.569824, 0.910645, 0.17041, 1.65088, 1.65088, -0.569824, 0.910645, 0.910645,
7428 1.65088, 0.17041, 0.910645, 0.910645, 0.910645, 4.38281, 4.38281, 4.38281, 1.05176,
7429 -2.2793, 7.71387, -2.2793, 7.71387, 1.05176, -2.2793, 1.05176, 4.38281, -2.2793, 1.05176,
7430 4.38281, 7.71387, 10.3633, 0.0, 0.0, 0.0, 10.3633, 0.0, 5.18164, 5.18164, 10.3633, 5.18164,
7431 5.18164, 0.0, 5.18164, 15.5449, 15.5449, 0.0, 16.6553, 16.6553, 11.1035, 0.0, 11.1035, 0.0,
7432 0.0, 16.6553, 11.1035, 5.55176, 5.55176, 5.55176, 0.0, 16.6553, 11.1035, 11.1035, 6.03369,
7433 0.111816, 6.03369, 0.111816, -2.84912, -2.84912, 3.07275, 0.111816, -2.84912, 6.03369,
7434 -2.84912, 3.07275, 0.111816, -2.84912, 0.111816, -2.84912, -0.189941, -0.189941, -0.189941,
7435 -0.189941, -0.189941, -0.189941, -0.189941, -0.189941, -0.189941, -0.189941, -0.189941,
7436 -0.189941, -0.189941, -0.189941, -0.189941, -0.189941, -2.84912, -2.84912, -2.84912,
7437 -2.84912, -2.84912, -2.84912, -2.84912, -2.84912, -2.84912, -2.84912, -2.84912, -2.84912,
7438 -2.84912, -2.84912, -2.84912, -2.84912, -2.09912, -1.35889, -1.729, -2.46924, -1.35889,
7439 -2.09912, -1.35889, -1.35889, -2.46924, -2.09912, -1.729, -1.35889, -2.09912, -2.09912,
7440 -2.46924, -2.46924, 0.701172, -0.0390625, -0.779297, -0.779297, -0.0390625, 0.701172,
7441 -1.51953, -0.779297, -0.0390625, -0.0390625, -1.51953, -1.51953, -1.51953, -1.51953,
7442 -0.779297, -0.779297, -2.2793, 5.12305, 5.12305, 8.82422, 1.42188, 1.42188, -2.2793,
7443 5.12305, 1.42188, 5.12305, 1.42188, 8.82422, -2.2793, -2.2793, 8.82422, 1.42188, -1.14941,
7444 -0.779297, -0.40918, -0.40918, -0.40918, -1.14941, -0.779297, -0.779297, -0.40918,
7445 -0.779297, -1.51953, -0.40918, -0.779297, -0.40918, -1.14941, -1.51953, -1.32959, 4.22217,
7446 9.77393, 4.22217, 15.3257, 4.22217, -1.32959, 4.22217, 15.3257, 4.22217, -1.32959, 9.77393,
7447 4.22217, 9.77393, 15.3257, 4.22217, 0.180176, -0.189941, 0.550293, 0.550293, 0.180176,
7448 0.550293, -0.189941, 0.550293, -0.189941, 0.92041, 0.92041, 0.550293, 0.180176, 0.180176,
7449 -0.189941, -0.189941, 9.74463, -2.46924, 9.74463, 5.67334, 5.67334, 1.60205, 9.74463,
7450 -2.46924, 9.74463, 1.60205, 9.74463, 9.74463, -2.46924, 1.60205, 5.67334, 1.60205, 13.8062,
7451 8.25439, 2.70264, 13.8062, 8.25439, 13.8062, 2.70264, 2.70264, 8.25439, -2.84912, -2.84912,
7452 2.70264, 13.8062, 13.8062, 8.25439, 13.8062,
7453 ];
7454
7455 const Q3_K_TEST_BLOCK: [u8; 110] = [
7456 0x56, 0xf2, 0xb4, 0x2b, 0xd5, 0x6f, 0x51, 0x71, 0x3c, 0x0a, 0xb9, 0x1d, 0xd0, 0xb9, 0x3b,
7457 0xb3, 0x0f, 0xff, 0x8c, 0xb2, 0x83, 0x3a, 0x3d, 0x24, 0xb1, 0x12, 0x56, 0xe3, 0x23, 0x54,
7458 0xf2, 0xfa, 0x7f, 0xdf, 0x31, 0xe1, 0x18, 0x26, 0x6e, 0xcd, 0x5b, 0x38, 0xee, 0xbd, 0x9f,
7459 0x8c, 0x57, 0x47, 0x0b, 0x11, 0xcb, 0xfb, 0xb4, 0x83, 0xa0, 0x4e, 0x0b, 0xd4, 0xa7, 0x85,
7460 0xe0, 0x60, 0xf3, 0xb3, 0xe3, 0x95, 0x43, 0xc6, 0x05, 0x05, 0x77, 0x53, 0xed, 0x23, 0xcc,
7461 0x6a, 0x0e, 0x89, 0xa1, 0x79, 0x85, 0xf6, 0x6e, 0x5a, 0x23, 0x63, 0xbe, 0x53, 0xfa, 0xa2,
7462 0x2b, 0xe9, 0xcd, 0xce, 0xf8, 0x3d, 0x6f, 0xd0, 0x42, 0x6e, 0x3b, 0x7f, 0x23, 0x26, 0xd3,
7463 0xb9, 0x18, 0xbf, 0xa4, 0x34,
7464 ];
7465 const Q3_K_GOLDEN: [f32; 256] = [
7466 -8.99121, -8.99121, -26.9736, 8.99121, 0.0, 17.9824, 17.9824, 8.99121, -8.99121, -35.9648,
7467 17.9824, 8.99121, -8.99121, 0.0, 26.9736, 26.9736, -13.9219, -4.64062, 4.64062, 4.64062,
7468 -0.0, 4.64062, -0.0, 9.28125, -13.9219, 18.5625, 4.64062, -4.64062, -0.0, 18.5625, 4.64062,
7469 4.64062, -26.1035, -26.1035, 34.8047, -0.0, 17.4023, -8.70117, 8.70117, 8.70117, 17.4023,
7470 -17.4023, 8.70117, 8.70117, 8.70117, 8.70117, -8.70117, -8.70117, 17.4023, 0.0, -17.4023,
7471 17.4023, 8.70117, 0.0, -34.8047, -8.70117, -17.4023, 8.70117, 8.70117, 8.70117, 0.0,
7472 -34.8047, 0.0, 0.0, -18.2725, 18.2725, -18.2725, 12.1816, -6.09082, -12.1816, 12.1816,
7473 24.3633, -6.09082, 6.09082, 12.1816, -18.2725, 18.2725, 24.3633, 18.2725, 24.3633, 0.0,
7474 4.35059, 0.0, -4.35059, -4.35059, -17.4023, 8.70117, 0.0, -17.4023, -13.0518, 8.70117,
7475 -17.4023, -8.70117, 8.70117, -4.35059, -4.35059, -2.61035, -0.870117, -3.48047, 2.61035,
7476 -3.48047, 0.0, -2.61035, -0.870117, 0.870117, 0.0, 2.61035, 1.74023, -1.74023, 1.74023,
7477 0.870117, -2.61035, 0.0, 0.0, 19.1426, -6.38086, -12.7617, 12.7617, 12.7617, -19.1426,
7478 -25.5234, -6.38086, -12.7617, -12.7617, -6.38086, -19.1426, -6.38086, 12.7617, -8.70117,
7479 -2.90039, -8.70117, 5.80078, -2.90039, 8.70117, -8.70117, -8.70117, -2.90039, 2.90039,
7480 -0.0, -5.80078, -5.80078, -2.90039, -2.90039, -2.90039, -25.2334, 16.8223, -16.8223,
7481 16.8223, -8.41113, 25.2334, 16.8223, -8.41113, 16.8223, 16.8223, 25.2334, -25.2334,
7482 -25.2334, 16.8223, 0.0, 8.41113, 13.9219, -3.48047, -0.0, -3.48047, 10.4414, -3.48047,
7483 10.4414, -0.0, -10.4414, 13.9219, -10.4414, 6.96094, 3.48047, -6.96094, -0.0, -6.96094,
7484 -19.1426, 6.38086, -6.38086, 12.7617, -25.5234, 0.0, 19.1426, 0.0, 12.7617, -25.5234,
7485 -12.7617, 12.7617, 19.1426, -6.38086, 12.7617, 19.1426, 11.0215, 5.51074, -22.043, -22.043,
7486 0.0, 0.0, 16.5322, 5.51074, -11.0215, -11.0215, -22.043, -11.0215, 0.0, -22.043, -11.0215,
7487 -5.51074, -8.12109, 6.09082, -4.06055, -6.09082, -4.06055, -4.06055, -2.03027, -6.09082,
7488 -2.03027, -4.06055, 4.06055, 4.06055, -8.12109, 0.0, 6.09082, 6.09082, 8.70117, -17.4023,
7489 -8.70117, 8.70117, -0.0, 34.8047, 26.1035, 26.1035, 8.70117, 34.8047, -26.1035, 26.1035,
7490 -0.0, -17.4023, 17.4023, -8.70117, -1.16016, 1.74023, 0.580078, 0.580078, 0.0, -1.74023,
7491 -1.16016, -1.74023, 1.74023, -1.16016, -2.32031, 1.74023, -0.580078, -0.580078, 1.74023,
7492 0.0,
7493 ];
7494
7495 #[test]
7496 fn q2_k_dequant_matches_independent_python_reference() {
7497 let got = dequant_q2_k(&Q2_K_TEST_BLOCK).unwrap();
7498 assert_eq!(got.len(), Q2_K_GOLDEN.len());
7499 for (i, (a, b)) in got.iter().zip(Q2_K_GOLDEN.iter()).enumerate() {
7500 assert!(
7501 (a - b).abs() < 1e-3,
7502 "Q2_K element {i}: rust={a} python={b}"
7503 );
7504 }
7505 }
7506
7507 #[test]
7508 fn q2_k_fused_dot_matches_dequant_then_dot() {
7509 let dequanted = dequant_q2_k(&Q2_K_TEST_BLOCK).unwrap();
7510 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.019).sin()).collect();
7511 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7512 let fused = dot_q2_k_f32(&Q2_K_TEST_BLOCK, &x);
7513 assert!(
7514 (fused - expected).abs() < 1e-1,
7515 "fused={fused} expected={expected}"
7516 );
7517 }
7518
7519 #[test]
7520 fn q3_k_dequant_matches_independent_python_reference() {
7521 let got = dequant_q3_k(&Q3_K_TEST_BLOCK).unwrap();
7522 assert_eq!(got.len(), Q3_K_GOLDEN.len());
7523 for (i, (a, b)) in got.iter().zip(Q3_K_GOLDEN.iter()).enumerate() {
7524 assert!(
7525 (a - b).abs() < 1e-3,
7526 "Q3_K element {i}: rust={a} python={b}"
7527 );
7528 }
7529 }
7530
7531 #[test]
7532 fn q3_k_fused_dot_matches_dequant_then_dot() {
7533 let dequanted = dequant_q3_k(&Q3_K_TEST_BLOCK).unwrap();
7534 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).cos()).collect();
7535 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7536 let fused = dot_q3_k_f32(&Q3_K_TEST_BLOCK, &x);
7537 assert!(
7538 (fused - expected).abs() < 1e-1,
7539 "fused={fused} expected={expected}"
7540 );
7541 }
7542
7543 #[test]
7544 fn q2_k_and_q3_k_reject_misaligned_buffers() {
7545 let bad = vec![0u8; 5];
7546 assert!(dequant_q2_k(&bad).is_err());
7547 assert!(dequant_q3_k(&bad).is_err());
7548 }
7549
7550 const IQ4_NL_TEST_BLOCK: [u8; 18] = [
7556 0xf6, 0x34, 0x3c, 0x7f, 0x90, 0x6a, 0xdc, 0x0f, 0x77, 0xfc, 0xb9, 0x1c, 0xdf, 0x74, 0xe0,
7557 0x40, 0x5d, 0xf3,
7558 ];
7559 const IQ4_NL_GOLDEN: [f32; 32] = [
7560 16.4331, 35.0366, -39.3774, 7.75146, 16.4331, 35.0366, -3.10059, 16.4331, 4.03076, 16.4331,
7561 35.0366, -15.1929, -39.3774, -39.3774, 21.394, -20.1538, -20.1538, -3.10059, 4.03076,
7562 -6.82129, 21.394, -39.3774, -3.10059, 35.0366, 11.7822, -32.2461, 21.394, -3.10059,
7563 27.5952, -15.1929, -10.8521, 35.0366,
7564 ];
7565
7566 const IQ4_XS_TEST_BLOCK: [u8; 136] = [
7567 0x5c, 0x33, 0xb4, 0x39, 0xd1, 0x64, 0x97, 0x82, 0xcb, 0xbd, 0x88, 0x95, 0xf3, 0x60, 0x2a,
7568 0xb5, 0xe7, 0x24, 0xd3, 0xee, 0xfe, 0x71, 0x13, 0xbe, 0x70, 0x84, 0x48, 0x79, 0x7b, 0x3e,
7569 0xf0, 0x55, 0xdc, 0xb2, 0xb2, 0xde, 0x32, 0xa1, 0x5b, 0x02, 0x01, 0xdc, 0x2a, 0xbb, 0xf7,
7570 0x0b, 0x8a, 0x88, 0xdd, 0x0b, 0x02, 0x7e, 0x5e, 0x76, 0x87, 0x30, 0x1e, 0x1c, 0xcf, 0x48,
7571 0xd7, 0x61, 0xf3, 0x51, 0x52, 0x17, 0x98, 0x0a, 0x87, 0xcf, 0x02, 0x91, 0xc8, 0xee, 0xc0,
7572 0x91, 0x69, 0x2a, 0x4f, 0x64, 0x68, 0xa7, 0xb2, 0xe6, 0x98, 0x21, 0x81, 0x75, 0x53, 0x2a,
7573 0x8d, 0x12, 0xae, 0xe0, 0xea, 0x0c, 0x75, 0xff, 0x22, 0x5e, 0x25, 0x19, 0xda, 0x2e, 0x51,
7574 0x4e, 0x81, 0xdc, 0x0e, 0x78, 0x86, 0xd7, 0x58, 0xb5, 0xb7, 0xf6, 0x45, 0xa9, 0x0a, 0x83,
7575 0xfd, 0x2a, 0x12, 0x7d, 0xf0, 0x12, 0x97, 0xe2, 0xfe, 0xf4, 0xd0, 0xa2, 0x11, 0x14, 0x78,
7576 0xdb,
7577 ];
7578 const IQ4_XS_GOLDEN: [f32; 256] = [
7579 -270.917, -491.928, -7.12939, 249.529, 463.411, 905.433, -178.235, 249.529, 71.2939,
7580 349.34, 463.411, -634.516, -634.516, 741.457, 463.411, -634.516, -377.858, -270.917,
7581 -7.12939, -92.6821, -805.622, 156.847, 591.74, -270.917, -634.516, 591.74, -491.928,
7582 -634.516, -805.622, 71.2939, 741.457, -270.917, 87.6226, 33.8071, -0.689941, -8.96924,
7583 -26.2178, -61.4048, 87.6226, 24.1479, -36.5669, 57.2651, 57.2651, -61.4048, 57.2651,
7584 71.7539, -26.2178, 57.2651, 6.89941, -0.689941, 33.8071, 6.89941, 6.89941, 44.8462,
7585 -77.9634, 24.1479, -47.606, -26.2178, -26.2178, -47.606, 44.8462, -17.2485, 24.1479,
7586 87.6226, -478.359, 243.779, 114.99, 174.785, -45.9961, 174.785, 114.99, 4.59961, 317.373,
7587 174.785, -381.768, 409.365, 409.365, -101.191, -45.9961, -584.15, -584.15, 317.373,
7588 -381.768, 174.785, 519.756, -584.15, 4.59961, 4.59961, 317.373, -584.15, -584.15, -45.9961,
7589 -160.986, -45.9961, 4.59961, -298.975, 122.81, 73.1338, 155.927, 1.37988, -13.7988,
7590 -143.508, -89.6924, -143.508, -114.53, -13.7988, 1.37988, 34.4971, -13.7988, 155.927,
7591 -114.53, -143.508, -143.508, -143.508, 73.1338, -67.6143, 95.2119, -30.3574, 155.927,
7592 -48.2959, -48.2959, -143.508, 17.9385, -175.245, 1.37988, 73.1338, -175.245, 17.9385,
7593 -2.06982, -184.214, 262.868, 215.262, -26.9077, -51.7456, -233.89, 101.421, -2.06982,
7594 20.6982, 171.795, 45.5361, -2.06982, 215.262, 215.262, 72.4438, -109.701, -184.214,
7595 -109.701, -26.9077, 45.5361, 171.795, 101.421, 45.5361, 45.5361, -51.7456, -78.6533,
7596 -184.214, -26.9077, 171.795, -2.06982, 20.6982, -134.539, 51.7456, 142.818, -171.795,
7597 184.214, -262.868, 51.7456, 109.701, -72.4438, 233.89, -171.795, 184.214, -72.4438,
7598 26.9077, 51.7456, 184.214, -72.4438, -171.795, 2.06982, -215.262, 51.7456, 184.214,
7599 184.214, -262.868, -20.6982, 233.89, -171.795, -72.4438, -171.795, -215.262, 142.818,
7600 -171.795, -430.523, 368.429, -430.523, 219.401, 368.429, 4.13965, -91.0723, -41.3965,
7601 4.13965, -144.888, -41.3965, -91.0723, -144.888, 53.8154, 103.491, -269.077, -144.888,
7602 -202.843, 4.13965, 285.636, -525.735, -41.3965, 4.13965, 285.636, -144.888, 157.307,
7603 157.307, 467.78, -202.843, 103.491, -525.735, 4.13965, -380.848, -137.988, 458.121,
7604 -380.848, 700.98, 458.121, 55.1953, 458.121, -491.238, 270.457, 700.98, 458.121, 574.031,
7605 270.457, -5.51953, -209.742, -623.707, 458.121, 574.031, 55.1953, -623.707, 574.031,
7606 -71.7539, -491.238, -623.707, -623.707, -380.848, -137.988, 574.031, 574.031, 55.1953,
7607 -380.848,
7608 ];
7609
7610 #[test]
7611 fn iq4_nl_dequant_matches_independent_python_reference() {
7612 let got = dequant_iq4_nl(&IQ4_NL_TEST_BLOCK).unwrap();
7613 assert_eq!(got.len(), IQ4_NL_GOLDEN.len());
7614 for (i, (a, b)) in got.iter().zip(IQ4_NL_GOLDEN.iter()).enumerate() {
7615 assert!(
7616 (a - b).abs() < 1e-2,
7617 "IQ4_NL element {i}: rust={a} python={b}"
7618 );
7619 }
7620 }
7621
7622 #[test]
7623 fn iq4_nl_fused_dot_matches_dequant_then_dot() {
7624 let dequanted = dequant_iq4_nl(&IQ4_NL_TEST_BLOCK).unwrap();
7625 let x: Vec<f32> = (0..32).map(|i| ((i as f32) * 0.019).sin()).collect();
7626 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7627 let fused = dot_iq4_nl_f32(&IQ4_NL_TEST_BLOCK, &x);
7628 assert!(
7629 (fused - expected).abs() < 1e-1,
7630 "fused={fused} expected={expected}"
7631 );
7632 }
7633
7634 #[test]
7635 fn iq4_xs_dequant_matches_independent_python_reference() {
7636 let got = dequant_iq4_xs(&IQ4_XS_TEST_BLOCK).unwrap();
7637 assert_eq!(got.len(), IQ4_XS_GOLDEN.len());
7638 for (i, (a, b)) in got.iter().zip(IQ4_XS_GOLDEN.iter()).enumerate() {
7639 assert!(
7640 (a - b).abs() < 1e-1,
7641 "IQ4_XS element {i}: rust={a} python={b}"
7642 );
7643 }
7644 }
7645
7646 #[test]
7647 fn iq4_xs_fused_dot_matches_dequant_then_dot() {
7648 let dequanted = dequant_iq4_xs(&IQ4_XS_TEST_BLOCK).unwrap();
7649 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.023).cos()).collect();
7650 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7651 let fused = dot_iq4_xs_f32(&IQ4_XS_TEST_BLOCK, &x);
7652 assert!(
7653 (fused - expected).abs() < 1e-1,
7654 "fused={fused} expected={expected}"
7655 );
7656 }
7657
7658 #[test]
7659 fn iq4_nl_and_iq4_xs_reject_misaligned_buffers() {
7660 let bad = vec![0u8; 5];
7661 assert!(dequant_iq4_nl(&bad).is_err());
7662 assert!(dequant_iq4_xs(&bad).is_err());
7663 }
7664
7665 const MXFP4_TEST_PACKED: [u8; 32] = [
7670 0xaa, 0xf9, 0x12, 0xda, 0x04, 0xac, 0xce, 0x2d, 0xbf, 0x4c, 0xc3, 0x06, 0x67, 0x59, 0xd1,
7671 0xa3, 0xea, 0xf1, 0x8f, 0x5d, 0xe5, 0xe6, 0x9e, 0x77, 0x73, 0x9c, 0x6f, 0x14, 0x5f, 0x1f,
7672 0xd9, 0x5e,
7673 ];
7674 const MXFP4_TEST_SCALES: [u8; 2] = [0x00, 0x82];
7675 const MXFP4_GOLDEN: [f32; 64] = [
7676 -5.87747e-39,
7677 -2.93874e-39,
7678 5.87747e-39,
7679 -5.87747e-39,
7680 1.17549e-38,
7681 -1.17549e-38,
7682 -2.35099e-38,
7683 -1.76324e-38,
7684 -3.52648e-38,
7685 -1.17549e-38,
7686 8.81621e-39,
7687 2.35099e-38,
7688 3.52648e-38,
7689 -2.93874e-39,
7690 2.93874e-39,
7691 8.81621e-39,
7692 -5.87747e-39,
7693 -3.52648e-38,
7694 2.93874e-39,
7695 -1.76324e-38,
7696 0.0,
7697 -5.87747e-39,
7698 -1.17549e-38,
7699 5.87747e-39,
7700 -8.81621e-39,
7701 1.17549e-38,
7702 -1.17549e-38,
7703 0.0,
7704 2.35099e-38,
7705 1.76324e-38,
7706 -1.76324e-38,
7707 -5.87747e-39,
7708 -8.0,
7709 4.0,
7710 -48.0,
7711 -24.0,
7712 24.0,
7713 32.0,
7714 -32.0,
7715 48.0,
7716 12.0,
7717 -16.0,
7718 -48.0,
7719 16.0,
7720 -48.0,
7721 -48.0,
7722 -4.0,
7723 -32.0,
7724 -32.0,
7725 -48.0,
7726 -0.0,
7727 24.0,
7728 -32.0,
7729 -32.0,
7730 -4.0,
7731 48.0,
7732 48.0,
7733 -4.0,
7734 32.0,
7735 4.0,
7736 24.0,
7737 4.0,
7738 -24.0,
7739 24.0,
7740 ];
7741
7742 #[test]
7743 fn mxfp4_dequant_matches_independent_python_reference() {
7744 let got = dequant_mxfp4_row(&MXFP4_TEST_PACKED, &MXFP4_TEST_SCALES).unwrap();
7745 assert_eq!(got.len(), MXFP4_GOLDEN.len());
7746 for (i, (a, b)) in got.iter().zip(MXFP4_GOLDEN.iter()).enumerate() {
7747 let tol = 1e-38f32.max(b.abs() * 1e-3);
7748 assert!(
7749 (a - b).abs() < tol,
7750 "MXFP4 element {i}: rust={a} python={b}"
7751 );
7752 }
7753 }
7754
7755 #[test]
7756 fn mxfp4_fused_dot_matches_dequant_then_dot() {
7757 let dequanted = dequant_mxfp4_row(&MXFP4_TEST_PACKED, &MXFP4_TEST_SCALES).unwrap();
7758 let x: Vec<f32> = (0..64).map(|i| ((i as f32) * 0.037).sin()).collect();
7759 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
7760 let fused = dot_mxfp4_row_f32(&MXFP4_TEST_PACKED, &MXFP4_TEST_SCALES, &x);
7761 assert!(
7762 (fused - expected).abs() < 1e-3,
7763 "fused={fused} expected={expected}"
7764 );
7765 }
7766
7767 #[test]
7768 fn mxfp4_scale_byte_zero_and_max_match_the_e8m0_formula() {
7769 assert!((e8m0_scale(0) - 2f32.powi(-127)).abs() < 1e-45);
7772 assert_eq!(e8m0_scale(127), 1.0);
7773 assert_eq!(e8m0_scale(128), 2.0);
7774 }
7775
7776 #[test]
7777 fn mxfp4_simd_dispatch_matches_scalar_across_every_possible_packed_byte_value() {
7778 let packed: Vec<u8> = (0..=255u8).collect();
7789 let n_groups = packed.len() / (MXFP4_GROUP_SIZE / 2);
7790 let scales: Vec<u8> = (0..n_groups).map(|i| ((i * 17 + 3) % 180) as u8).collect();
7796 let x: Vec<f32> = (0..n_groups * MXFP4_GROUP_SIZE)
7797 .map(|i| ((i as f32) * 0.013).cos())
7798 .collect();
7799
7800 let scalar = dot_mxfp4_row_f32_scalar(&packed, &scales, &x);
7801 let dispatched = dot_mxfp4_row_f32(&packed, &scales, &x);
7802 assert!(
7803 (scalar - dispatched).abs() < scalar.abs() * 1e-3 + 1e-3,
7804 "scalar={scalar} dispatched (SIMD)={dispatched}"
7805 );
7806
7807 #[cfg(target_arch = "aarch64")]
7808 {
7809 let neon = unsafe { simd_aarch64::dot_mxfp4_row_f32_neon(&packed, &scales, &x) };
7810 assert!(
7811 (scalar - neon).abs() < scalar.abs() * 1e-3 + 1e-3,
7812 "scalar={scalar} neon={neon}"
7813 );
7814 }
7815 }
7816
7817 #[test]
7818 fn mxfp4_rejects_a_packed_scales_length_mismatch() {
7819 let bad_packed = vec![0u8; 15]; let scales = [0u8; 1];
7821 assert!(matches!(
7822 dequant_mxfp4_row(&bad_packed, &scales),
7823 Err(QuantError::Mxfp4RowMismatch(15, 16))
7824 ));
7825 }
7826
7827 fn repeat_block(block: &[u8], n: usize) -> Vec<u8> {
7830 block
7831 .iter()
7832 .copied()
7833 .cycle()
7834 .take(block.len() * n)
7835 .collect()
7836 }
7837
7838 #[test]
7839 fn dispatched_q4_k_matches_scalar_reference_across_many_blocks() {
7840 let n_blocks = 4;
7841 let packed = repeat_block(&Q4_K_TEST_BLOCK, n_blocks);
7842 let x: Vec<f32> = (0..256 * n_blocks)
7843 .map(|i| ((i as f32) * 0.013).sin())
7844 .collect();
7845 let dispatched = dot_q4_k_f32(&packed, &x);
7846 let scalar = dot_q4_k_f32_scalar(&packed, &x);
7847 assert!(
7848 (dispatched - scalar).abs() < 1e-1,
7849 "dispatched={dispatched} scalar={scalar}"
7850 );
7851 }
7852
7853 #[test]
7854 fn dispatched_q5_k_matches_scalar_reference_across_many_blocks() {
7855 let n_blocks = 4;
7856 let packed = repeat_block(&Q5_K_TEST_BLOCK, n_blocks);
7857 let x: Vec<f32> = (0..256 * n_blocks)
7858 .map(|i| ((i as f32) * 0.011).cos())
7859 .collect();
7860 let dispatched = dot_q5_k_f32(&packed, &x);
7861 let scalar = dot_q5_k_f32_scalar(&packed, &x);
7862 assert!(
7863 (dispatched - scalar).abs() < 1e-1,
7864 "dispatched={dispatched} scalar={scalar}"
7865 );
7866 }
7867
7868 #[test]
7869 fn dispatched_q6_k_matches_scalar_reference_across_many_blocks() {
7870 let n_blocks = 4;
7871 let packed = repeat_block(&Q6_K_TEST_BLOCK, n_blocks);
7872 let x: Vec<f32> = (0..256 * n_blocks)
7873 .map(|i| ((i as f32) * 0.019).sin())
7874 .collect();
7875 let dispatched = dot_q6_k_f32(&packed, &x);
7876 let scalar = dot_q6_k_f32_scalar(&packed, &x);
7877 assert!(
7878 (dispatched - scalar).abs() < 1e-1,
7879 "dispatched={dispatched} scalar={scalar}"
7880 );
7881 }
7882
7883 #[test]
7884 fn dispatched_q6_k_matches_scalar_reference_with_negative_scales() {
7885 let n_blocks = 4;
7890 let packed = repeat_block(&Q6_K_SIGNED_SCALES_TEST_BLOCK, n_blocks);
7891 let x: Vec<f32> = (0..256 * n_blocks)
7892 .map(|i| ((i as f32) * 0.019).sin())
7893 .collect();
7894 let dispatched = dot_q6_k_f32(&packed, &x);
7895 let scalar = dot_q6_k_f32_scalar(&packed, &x);
7896 assert!(
7897 (dispatched - scalar).abs() < 1e-1,
7898 "dispatched={dispatched} scalar={scalar}"
7899 );
7900 }
7901
7902 #[cfg(target_arch = "aarch64")]
7903 #[test]
7904 fn neon_q4_k_kernel_matches_scalar_directly_when_available() {
7905 if !std::arch::is_aarch64_feature_detected!("neon") {
7906 eprintln!("skipping: host CPU lacks NEON");
7907 return;
7908 }
7909 let n_blocks = 4;
7910 let packed = repeat_block(&Q4_K_TEST_BLOCK, n_blocks);
7911 let x: Vec<f32> = (0..256 * n_blocks)
7912 .map(|i| ((i as f32) * 0.037).cos())
7913 .collect();
7914 let simd = unsafe { simd_aarch64::dot_q4_k_f32_neon(&packed, &x) };
7915 let scalar = dot_q4_k_f32_scalar(&packed, &x);
7916 assert!(
7917 (simd - scalar).abs() < 1e-1,
7918 "NEON Q4_K kernel diverged from scalar: simd={simd} scalar={scalar}"
7919 );
7920 }
7921
7922 #[cfg(target_arch = "aarch64")]
7923 #[test]
7924 fn neon_q5_k_q8_kernel_matches_scalar_directly_when_available() {
7925 if !std::arch::is_aarch64_feature_detected!("neon") {
7926 eprintln!("skipping: host CPU lacks NEON");
7927 return;
7928 }
7929 let n_blocks = 4;
7930 let packed = repeat_block(&Q5_K_TEST_BLOCK, n_blocks);
7931 let x: Vec<f32> = (0..256 * n_blocks)
7932 .map(|i| ((i as f32) * 0.029).sin())
7933 .collect();
7934 let act = quantize_activations_q8_k(&x);
7935 let dispatched = dot_q5_k_q8(&packed, &act);
7936 let scalar = dot_q5_k_q8_scalar(&packed, &act);
7937 assert_eq!(
7938 dispatched,
7939 scalar,
7940 "Q5_K×Q8_K dispatch must match scalar (dotprod={})",
7941 std::arch::is_aarch64_feature_detected!("dotprod")
7942 );
7943 if std::arch::is_aarch64_feature_detected!("dotprod") {
7944 let sdot = unsafe { simd_aarch64::dot_q5_k_q8_neon_sdot(&packed, &act) };
7945 assert_eq!(sdot, scalar, "NEON SDOT Q5_K×Q8_K diverged from scalar");
7946 }
7947 if std::arch::is_aarch64_feature_detected!("neon") {
7948 let neon = unsafe { simd_aarch64::dot_q5_k_q8_neon(&packed, &act) };
7949 assert_eq!(neon, scalar, "NEON widen Q5_K×Q8_K diverged from scalar");
7950 }
7951 }
7952
7953 #[cfg(target_arch = "aarch64")]
7954 #[test]
7955 fn neon_q5_k_kernel_matches_scalar_directly_when_available() {
7956 if !std::arch::is_aarch64_feature_detected!("neon") {
7957 eprintln!("skipping: host CPU lacks NEON");
7958 return;
7959 }
7960 let n_blocks = 4;
7961 let packed = repeat_block(&Q5_K_TEST_BLOCK, n_blocks);
7962 let x: Vec<f32> = (0..256 * n_blocks)
7963 .map(|i| ((i as f32) * 0.029).sin())
7964 .collect();
7965 let simd = unsafe { simd_aarch64::dot_q5_k_f32_neon(&packed, &x) };
7966 let scalar = dot_q5_k_f32_scalar(&packed, &x);
7967 assert!(
7968 (simd - scalar).abs() < 1e-1,
7969 "NEON Q5_K kernel diverged from scalar: simd={simd} scalar={scalar}"
7970 );
7971 }
7972
7973 #[cfg(target_arch = "aarch64")]
7974 #[test]
7975 fn neon_q6_k_kernel_matches_scalar_directly_when_available() {
7976 if !std::arch::is_aarch64_feature_detected!("neon") {
7977 eprintln!("skipping: host CPU lacks NEON");
7978 return;
7979 }
7980 let n_blocks = 4;
7981 let packed = repeat_block(&Q6_K_TEST_BLOCK, n_blocks);
7982 let x: Vec<f32> = (0..256 * n_blocks)
7983 .map(|i| ((i as f32) * 0.041).cos())
7984 .collect();
7985 let simd = unsafe { simd_aarch64::dot_q6_k_f32_neon(&packed, &x) };
7986 let scalar = dot_q6_k_f32_scalar(&packed, &x);
7987 assert!(
7988 (simd - scalar).abs() < 1e-1,
7989 "NEON Q6_K kernel diverged from scalar: simd={simd} scalar={scalar}"
7990 );
7991 }
7992
7993 #[cfg(target_arch = "aarch64")]
7994 #[test]
7995 fn neon_q6_k_kernel_matches_scalar_directly_on_negative_scales() {
7996 if !std::arch::is_aarch64_feature_detected!("neon") {
7997 eprintln!("skipping: host CPU lacks NEON");
7998 return;
7999 }
8000 let n_blocks = 4;
8001 let packed = repeat_block(&Q6_K_SIGNED_SCALES_TEST_BLOCK, n_blocks);
8002 let x: Vec<f32> = (0..256 * n_blocks)
8003 .map(|i| ((i as f32) * 0.041).cos())
8004 .collect();
8005 let simd = unsafe { simd_aarch64::dot_q6_k_f32_neon(&packed, &x) };
8006 let scalar = dot_q6_k_f32_scalar(&packed, &x);
8007 assert!(
8008 (simd - scalar).abs() < 1e-1,
8009 "NEON Q6_K kernel diverged from scalar on negative scales: simd={simd} scalar={scalar}"
8010 );
8011 }
8012
8013 #[test]
8014 fn q4_k_scalar_matches_independent_python_reference_via_dispatch_entrypoint() {
8015 let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.017).sin()).collect();
8022 let dequanted = dequant_q4_k(&Q4_K_TEST_BLOCK).unwrap();
8023 let expected: f32 = dequanted.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
8024 let dispatched = dot_q4_k_f32(&Q4_K_TEST_BLOCK, &x);
8025 assert!((dispatched - expected).abs() < 1e-2);
8026 }
8027
8028 fn q4_1_test_block() -> Vec<u8> {
8031 let mut b = Vec::new();
8032 b.extend_from_slice(&f16::from_f32(0.3).to_le_bytes());
8033 b.extend_from_slice(&f16::from_f32(-1.2).to_le_bytes());
8034 b.extend_from_slice(
8035 &(0..16)
8036 .map(|i| (i as u8) | ((15 - i as u8) << 4))
8037 .collect::<Vec<u8>>(),
8038 );
8039 b
8040 }
8041
8042 fn q5_0_test_block() -> Vec<u8> {
8043 let mut b = Vec::new();
8044 b.extend_from_slice(&f16::from_f32(0.4).to_le_bytes());
8045 b.extend_from_slice(&[0xA5, 0x3C, 0x00, 0xFF]);
8046 b.extend_from_slice(
8047 &(0..16)
8048 .map(|i| (i as u8) | ((15 - i as u8) << 4))
8049 .collect::<Vec<u8>>(),
8050 );
8051 b
8052 }
8053
8054 fn q5_1_test_block() -> Vec<u8> {
8055 let mut b = Vec::new();
8056 b.extend_from_slice(&f16::from_f32(0.2).to_le_bytes());
8057 b.extend_from_slice(&f16::from_f32(0.9).to_le_bytes());
8058 b.extend_from_slice(&[0x12, 0x34, 0x56, 0x78]);
8059 b.extend_from_slice(
8060 &(0..16)
8061 .map(|i| (i as u8) | ((15 - i as u8) << 4))
8062 .collect::<Vec<u8>>(),
8063 );
8064 b
8065 }
8066
8067 fn q8_1_test_block() -> Vec<u8> {
8068 let mut b = Vec::new();
8069 b.extend_from_slice(&f16::from_f32(0.6).to_le_bytes());
8070 b.extend_from_slice(&f16::from_f32(0.0).to_le_bytes());
8071 let qs: Vec<i8> = (0..32).map(|i| ((i * 7) % 61) as i8 - 30).collect();
8072 b.extend_from_slice(&i8_to_u8_bytes(&qs));
8073 b
8074 }
8075
8076 #[test]
8077 fn dispatched_matches_scalar_for_the_8_newly_simd_formats_across_many_blocks() {
8078 let n_blocks = 4;
8079
8080 let q4_1 = repeat_block(&q4_1_test_block(), n_blocks);
8081 let x32 = |seed: f32| -> Vec<f32> {
8082 (0..32 * n_blocks)
8083 .map(|i| ((i as f32) * seed).sin())
8084 .collect()
8085 };
8086 let x = x32(0.031);
8087 assert!((dot_q4_1_f32(&q4_1, &x) - dot_q4_1_f32_scalar(&q4_1, &x)).abs() < 1e-1);
8088
8089 let q5_0 = repeat_block(&q5_0_test_block(), n_blocks);
8090 let x = x32(0.037);
8091 assert!((dot_q5_0_f32(&q5_0, &x) - dot_q5_0_f32_scalar(&q5_0, &x)).abs() < 1e-1);
8092
8093 let q5_1 = repeat_block(&q5_1_test_block(), n_blocks);
8094 let x = x32(0.041);
8095 assert!((dot_q5_1_f32(&q5_1, &x) - dot_q5_1_f32_scalar(&q5_1, &x)).abs() < 1e-1);
8096
8097 let q8_1 = repeat_block(&q8_1_test_block(), n_blocks);
8098 let x = x32(0.043);
8099 assert!((dot_q8_1_f32(&q8_1, &x) - dot_q8_1_f32_scalar(&q8_1, &x)).abs() < 1e-1);
8100
8101 let q2_k = repeat_block(&Q2_K_TEST_BLOCK, n_blocks);
8102 let x256 = |seed: f32| -> Vec<f32> {
8103 (0..256 * n_blocks)
8104 .map(|i| ((i as f32) * seed).cos())
8105 .collect()
8106 };
8107 let x = x256(0.013);
8108 assert!((dot_q2_k_f32(&q2_k, &x) - dot_q2_k_f32_scalar(&q2_k, &x)).abs() < 1e-1);
8109
8110 let q3_k = repeat_block(&Q3_K_TEST_BLOCK, n_blocks);
8111 let x = x256(0.017);
8112 assert!((dot_q3_k_f32(&q3_k, &x) - dot_q3_k_f32_scalar(&q3_k, &x)).abs() < 1e-1);
8113
8114 let iq4_nl = repeat_block(&IQ4_NL_TEST_BLOCK, n_blocks);
8115 let x = x32(0.019);
8116 assert!((dot_iq4_nl_f32(&iq4_nl, &x) - dot_iq4_nl_f32_scalar(&iq4_nl, &x)).abs() < 1e-1);
8117
8118 let iq4_xs = repeat_block(&IQ4_XS_TEST_BLOCK, n_blocks);
8119 let x = x256(0.023);
8120 assert!((dot_iq4_xs_f32(&iq4_xs, &x) - dot_iq4_xs_f32_scalar(&iq4_xs, &x)).abs() < 1e-1);
8121 }
8122
8123 #[cfg(target_arch = "aarch64")]
8124 #[test]
8125 fn neon_kernels_match_scalar_directly_for_the_8_newly_simd_formats() {
8126 if !std::arch::is_aarch64_feature_detected!("neon") {
8127 eprintln!("skipping: host CPU lacks NEON");
8128 return;
8129 }
8130 let n_blocks = 4;
8131 let x32 = |seed: f32| -> Vec<f32> {
8132 (0..32 * n_blocks)
8133 .map(|i| ((i as f32) * seed).sin())
8134 .collect()
8135 };
8136 let x256 = |seed: f32| -> Vec<f32> {
8137 (0..256 * n_blocks)
8138 .map(|i| ((i as f32) * seed).cos())
8139 .collect()
8140 };
8141
8142 let q4_1 = repeat_block(&q4_1_test_block(), n_blocks);
8143 let x = x32(0.031);
8144 let simd = unsafe { simd_aarch64::dot_q4_1_f32_neon(&q4_1, &x) };
8145 assert!((simd - dot_q4_1_f32_scalar(&q4_1, &x)).abs() < 1e-1);
8146
8147 let q5_0 = repeat_block(&q5_0_test_block(), n_blocks);
8148 let x = x32(0.037);
8149 let simd = unsafe { simd_aarch64::dot_q5_0_f32_neon(&q5_0, &x) };
8150 assert!((simd - dot_q5_0_f32_scalar(&q5_0, &x)).abs() < 1e-1);
8151
8152 let q5_1 = repeat_block(&q5_1_test_block(), n_blocks);
8153 let x = x32(0.041);
8154 let simd = unsafe { simd_aarch64::dot_q5_1_f32_neon(&q5_1, &x) };
8155 assert!((simd - dot_q5_1_f32_scalar(&q5_1, &x)).abs() < 1e-1);
8156
8157 let q8_1 = repeat_block(&q8_1_test_block(), n_blocks);
8158 let x = x32(0.043);
8159 let simd = unsafe { simd_aarch64::dot_q8_1_f32_neon(&q8_1, &x) };
8160 assert!((simd - dot_q8_1_f32_scalar(&q8_1, &x)).abs() < 1e-1);
8161
8162 let q2_k = repeat_block(&Q2_K_TEST_BLOCK, n_blocks);
8163 let x = x256(0.013);
8164 let simd = unsafe { simd_aarch64::dot_q2_k_f32_neon(&q2_k, &x) };
8165 assert!((simd - dot_q2_k_f32_scalar(&q2_k, &x)).abs() < 1e-1);
8166
8167 let q3_k = repeat_block(&Q3_K_TEST_BLOCK, n_blocks);
8168 let x = x256(0.017);
8169 let simd = unsafe { simd_aarch64::dot_q3_k_f32_neon(&q3_k, &x) };
8170 assert!((simd - dot_q3_k_f32_scalar(&q3_k, &x)).abs() < 1e-1);
8171
8172 let iq4_nl = repeat_block(&IQ4_NL_TEST_BLOCK, n_blocks);
8173 let x = x32(0.019);
8174 let simd = unsafe { simd_aarch64::dot_iq4_nl_f32_neon(&iq4_nl, &x) };
8175 assert!((simd - dot_iq4_nl_f32_scalar(&iq4_nl, &x)).abs() < 1e-1);
8176
8177 let iq4_xs = repeat_block(&IQ4_XS_TEST_BLOCK, n_blocks);
8178 let x = x256(0.023);
8179 let simd = unsafe { simd_aarch64::dot_iq4_xs_f32_neon(&iq4_xs, &x) };
8180 assert!((simd - dot_iq4_xs_f32_scalar(&iq4_xs, &x)).abs() < 1e-1);
8181 }
8182
8183 #[cfg(target_arch = "x86_64")]
8184 #[test]
8185 fn avx2_kernels_match_scalar_directly_for_the_8_newly_simd_formats() {
8186 if !(is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma")) {
8187 eprintln!("skipping: host CPU lacks AVX2+FMA");
8188 return;
8189 }
8190 let n_blocks = 4;
8191 let x32 = |seed: f32| -> Vec<f32> {
8192 (0..32 * n_blocks)
8193 .map(|i| ((i as f32) * seed).sin())
8194 .collect()
8195 };
8196 let x256 = |seed: f32| -> Vec<f32> {
8197 (0..256 * n_blocks)
8198 .map(|i| ((i as f32) * seed).cos())
8199 .collect()
8200 };
8201
8202 let q4_1 = repeat_block(&q4_1_test_block(), n_blocks);
8203 let x = x32(0.031);
8204 let simd = unsafe { simd_x86::dot_q4_1_f32_avx2(&q4_1, &x) };
8205 assert!((simd - dot_q4_1_f32_scalar(&q4_1, &x)).abs() < 1e-1);
8206
8207 let q5_0 = repeat_block(&q5_0_test_block(), n_blocks);
8208 let x = x32(0.037);
8209 let simd = unsafe { simd_x86::dot_q5_0_f32_avx2(&q5_0, &x) };
8210 assert!((simd - dot_q5_0_f32_scalar(&q5_0, &x)).abs() < 1e-1);
8211
8212 let q5_1 = repeat_block(&q5_1_test_block(), n_blocks);
8213 let x = x32(0.041);
8214 let simd = unsafe { simd_x86::dot_q5_1_f32_avx2(&q5_1, &x) };
8215 assert!((simd - dot_q5_1_f32_scalar(&q5_1, &x)).abs() < 1e-1);
8216
8217 let q8_1 = repeat_block(&q8_1_test_block(), n_blocks);
8218 let x = x32(0.043);
8219 let simd = unsafe { simd_x86::dot_q8_1_f32_avx2(&q8_1, &x) };
8220 assert!((simd - dot_q8_1_f32_scalar(&q8_1, &x)).abs() < 1e-1);
8221
8222 let q2_k = repeat_block(&Q2_K_TEST_BLOCK, n_blocks);
8223 let x = x256(0.013);
8224 let simd = unsafe { simd_x86::dot_q2_k_f32_avx2(&q2_k, &x) };
8225 assert!((simd - dot_q2_k_f32_scalar(&q2_k, &x)).abs() < 1e-1);
8226
8227 let q3_k = repeat_block(&Q3_K_TEST_BLOCK, n_blocks);
8228 let x = x256(0.017);
8229 let simd = unsafe { simd_x86::dot_q3_k_f32_avx2(&q3_k, &x) };
8230 assert!((simd - dot_q3_k_f32_scalar(&q3_k, &x)).abs() < 1e-1);
8231
8232 let iq4_nl = repeat_block(&IQ4_NL_TEST_BLOCK, n_blocks);
8233 let x = x32(0.019);
8234 let simd = unsafe { simd_x86::dot_iq4_nl_f32_avx2(&iq4_nl, &x) };
8235 assert!((simd - dot_iq4_nl_f32_scalar(&iq4_nl, &x)).abs() < 1e-1);
8236
8237 let iq4_xs = repeat_block(&IQ4_XS_TEST_BLOCK, n_blocks);
8238 let x = x256(0.023);
8239 let simd = unsafe { simd_x86::dot_iq4_xs_f32_avx2(&iq4_xs, &x) };
8240 assert!((simd - dot_iq4_xs_f32_scalar(&iq4_xs, &x)).abs() < 1e-1);
8241 }
8242}