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