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