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