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