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