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