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