1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::cell::UnsafeCell;
21use std::sync::Arc;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24mod mxfp4;
26
27#[cfg(target_arch = "x86_64")]
30mod moe_fast;
31
32pub static MATMAT_GATE: AtomicU64 = AtomicU64::new(0);
35pub static MATMAT_GPU_HIT: AtomicU64 = AtomicU64::new(0);
36
37pub enum QTensor {
38 F32 {
39 data: Vec<f32>,
40 rows: usize,
41 cols: usize,
42 },
43 Mapped {
44 model: Arc<CmfModel>,
45 idx: usize,
47 dtype: TensorDtype,
48 rows: usize,
49 cols: usize,
50 row_scale: Vec<f32>,
52 col_field: Vec<f32>,
54 vbit_offsets: Vec<usize>,
58 repack: Vec<u8>,
66 },
67 Blr(Box<crate::blr::BlrTensor>),
73}
74
75fn repack_enabled() -> bool {
82 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
83 *ON.get_or_init(|| {
84 std::env::var("CMF_REPACK")
85 .map(|v| v == "1")
86 .unwrap_or(cfg!(target_os = "android"))
87 })
88}
89
90fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
94 #[cfg(target_arch = "aarch64")]
95 let arch_ok = sdot_enabled();
96 #[cfg(not(target_arch = "aarch64"))]
97 let arch_ok = false;
98 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
99 return Vec::new();
100 }
101 q8_repack_layout(bytes, rows, cols)
102}
103
104fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
107 let groups = rows / 4;
108 let mut rep = vec![0u8; groups * 4 * cols];
109 for g in 0..groups {
110 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
111 for c in 0..cols / 16 {
112 for lane in 0..4 {
113 let src = (g * 4 + lane) * cols + c * 16;
114 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
115 .copy_from_slice(&bytes[src..src + 16]);
116 }
117 }
118 }
119 rep
120}
121
122fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
125 let ng = cols / GROUP_SIZE;
126 let bits = &bytes[..rows];
127 let mut offsets = Vec::with_capacity(rows + 1);
128 let mut off = rows + rows * ng * 2;
129 for r in 0..rows {
130 offsets.push(off);
131 off += (cols * bits[r] as usize).div_ceil(8);
132 }
133 offsets.push(off);
134 offsets
135}
136
137fn blocked_enabled() -> bool {
143 use std::sync::atomic::Ordering::Relaxed;
144 match BLOCKED_OVERRIDE.load(Relaxed) {
145 1 => false,
146 2 => true,
147 _ => {
148 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
149 *ON.get_or_init(|| {
150 std::env::var("CMF_X86_BLOCKED")
151 .map(|v| v != "0")
152 .unwrap_or(true)
153 })
154 }
155 }
156}
157
158static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
159
160pub fn set_blocked_override(on: Option<bool>) {
167 let v = match on {
168 None => 0,
169 Some(false) => 1,
170 Some(true) => 2,
171 };
172 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
173}
174
175fn gpu_lmhead_enabled() -> bool {
176 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
177 *ON.get_or_init(|| {
178 std::env::var("CMF_GPU_LMHEAD")
179 .map(|v| v != "0")
180 .unwrap_or(true)
181 })
182}
183
184fn gpu_split_frac() -> f32 {
185 if FULL_GPU_Q8.get() {
188 return 1.0;
189 }
190 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
191 *F.get_or_init(|| {
192 std::env::var("CMF_GPU_SPLIT")
193 .ok()
194 .and_then(|v| v.parse::<f32>().ok())
195 .unwrap_or(0.5)
196 .clamp(0.0, 1.0)
197 })
198}
199
200impl QTensor {
201 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
202 debug_assert_eq!(data.len(), rows * cols);
203 Self::F32 { data, rows, cols }
204 }
205
206 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
209 let idx = model
212 .tensor_index(name)
213 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
214 let entry = &model.tensors[idx];
215 if entry.shape.len() != 2 {
216 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
217 }
218 let (rows, cols) = (entry.shape[0], entry.shape[1]);
219 let bytes = model.entry_bytes(entry);
220
221 match entry.dtype {
222 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
223 let n = rows * cols;
224 let scales_off = n;
225 let row_scale: Vec<f32> = (0..rows)
226 .map(|o| {
227 f16_to_f32(u16::from_le_bytes([
228 bytes[scales_off + o * 2],
229 bytes[scales_off + o * 2 + 1],
230 ]))
231 })
232 .collect();
233 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
234 let col_off = n + rows * 2;
235 (0..cols)
236 .map(|i| {
237 f16_to_f32(u16::from_le_bytes([
238 bytes[col_off + i * 2],
239 bytes[col_off + i * 2 + 1],
240 ]))
241 })
242 .collect()
243 } else {
244 Vec::new()
245 };
246 Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale,
253 col_field,
254 vbit_offsets: Vec::new(),
255 repack: q8_repack(bytes, rows, cols),
256 })
257 }
258 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
260 model: model.clone(),
261 idx,
262 dtype: entry.dtype,
263 rows,
264 cols,
265 row_scale: Vec::new(),
266 col_field: Vec::new(),
267 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
268 repack: Vec::new(),
269 }),
270 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
274 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
275 let offsets: Vec<usize> = (0..=rows)
276 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
277 .collect();
278 Ok(Self::Mapped {
279 model: model.clone(),
280 idx,
281 dtype: entry.dtype,
282 rows,
283 cols,
284 row_scale: Vec::new(),
285 col_field: Vec::new(),
286 vbit_offsets: offsets,
287 repack: Vec::new(),
288 })
289 }
290 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
296 model: model.clone(),
297 idx,
298 dtype: entry.dtype,
299 rows,
300 cols,
301 row_scale: Vec::new(),
302 col_field: Vec::new(),
303 vbit_offsets: Vec::new(),
304 repack: Vec::new(),
305 }),
306 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
309 model: model.clone(),
310 idx,
311 dtype: entry.dtype,
312 rows,
313 cols,
314 row_scale: Vec::new(),
315 col_field: Vec::new(),
316 vbit_offsets: Vec::new(),
317 repack: Vec::new(),
318 }),
319 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
321 model: model.clone(),
322 idx,
323 dtype: entry.dtype,
324 rows,
325 cols,
326 row_scale: Vec::new(),
327 col_field: Vec::new(),
328 vbit_offsets: Vec::new(),
329 repack: Vec::new(),
330 }),
331 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332 model: model.clone(),
333 idx,
334 dtype: entry.dtype,
335 rows,
336 cols,
337 row_scale: Vec::new(),
338 col_field: Vec::new(),
339 vbit_offsets: Vec::new(),
340 repack: Vec::new(),
341 }),
342 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
345 model: model.clone(),
346 idx,
347 dtype: entry.dtype,
348 rows,
349 cols,
350 row_scale: Vec::new(),
351 col_field: Vec::new(),
352 vbit_offsets: Vec::new(),
353 repack: Vec::new(),
354 }),
355 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
356 model: model.clone(),
357 idx,
358 dtype: entry.dtype,
359 rows,
360 cols,
361 row_scale: Vec::new(),
362 col_field: Vec::new(),
363 vbit_offsets: Vec::new(),
364 repack: Vec::new(),
365 }),
366 TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
369 model: model.clone(),
370 idx,
371 dtype: entry.dtype,
372 rows,
373 cols,
374 row_scale: Vec::new(),
375 col_field: Vec::new(),
376 vbit_offsets: Vec::new(),
377 repack: Vec::new(),
378 }),
379 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
383 model: model.clone(),
384 idx,
385 dtype: entry.dtype,
386 rows,
387 cols,
388 row_scale: Vec::new(),
389 col_field: Vec::new(),
390 vbit_offsets: Vec::new(),
391 repack: Vec::new(),
392 }),
393 TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
396 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
397 model, idx,
398 )?)))
399 }
400 _ => {
402 let mut data = vec![0.0f32; rows * cols];
403 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
404 Ok(Self::from_f32(data, rows, cols))
405 }
406 }
407 }
408
409 pub fn is_blr(&self) -> bool {
412 matches!(self, Self::Blr(_))
413 }
414
415 pub fn is_mxfp4(&self) -> bool {
418 matches!(
419 self,
420 Self::Mapped {
421 dtype: TensorDtype::Mxfp4,
422 ..
423 }
424 )
425 }
426
427 fn per_tensor_only(&self) -> bool {
430 self.is_blr() || self.is_mxfp4()
431 }
432
433 pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
439 match self {
440 Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
441 _ => None,
442 }
443 }
444
445 pub(crate) fn graph_weight_or_blr(
449 &self,
450 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
451 self.graph_weight()
452 .or_else(|| self.graph_weight_descriptor())
453 .or_else(|| {
454 self.graph_blr()
455 .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
456 })
457 }
458
459 pub(crate) fn is_q1(&self) -> bool {
463 matches!(
464 self,
465 Self::Mapped {
466 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
467 ..
468 }
469 )
470 }
471
472 fn bin_w(&self) -> Option<BinW<'_>> {
475 match self {
476 Self::Mapped {
477 dtype: TensorDtype::Q1,
478 ..
479 } => Some(BinW::Q1(self.quant_bytes())),
480 Self::Mapped {
481 dtype: TensorDtype::Q1G,
482 rows,
483 cols,
484 ..
485 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
486 _ => None,
487 }
488 }
489
490 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
493 match self {
494 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
495 _ => None,
496 }
497 }
498
499 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
506 if self.has_prism_contract() {
507 return None;
508 }
509 match self {
510 #[cfg(target_os = "macos")]
511 Self::Mapped {
512 dtype: TensorDtype::Q1T,
513 ..
514 } if !crate::gpu::metal_q1t_enabled() => None,
515 Self::Mapped {
516 idx,
517 dtype:
518 TensorDtype::Q1
519 | TensorDtype::Q1G
520 | TensorDtype::Q1T
521 | TensorDtype::Q4Block
522 | TensorDtype::Q4Tiled
523 | TensorDtype::Q4TiledP
527 | TensorDtype::Q8Row
528 | TensorDtype::Q8_2f,
529 rows,
530 cols,
531 ..
532 } => Some((*idx, *rows, *cols)),
533 _ => None,
534 }
535 }
536
537 #[cfg(target_os = "macos")]
544 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
545 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
546 let name = &model.tensors[idx].name;
547 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
548 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
549 if forward && affine {
550 let e = model.tensors.get(idx)?;
551 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
552 }
553 }
554 self.q1_parts()
555 }
556
557 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
563 if self.has_prism_contract() {
564 return None;
565 }
566 match self {
567 Self::Mapped {
568 idx,
569 dtype: TensorDtype::Q4Tiled,
570 rows,
571 cols,
572 ..
573 } => Some((*idx, *rows, *cols)),
574 _ => None,
575 }
576 }
577
578 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
582 if self.has_prism_contract() {
583 return None;
584 }
585 match self {
586 Self::Mapped {
587 idx,
588 dtype: TensorDtype::Q4TiledP,
589 rows,
590 cols,
591 ..
592 } => Some((*idx, *rows, *cols)),
593 _ => None,
594 }
595 }
596
597 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
602 if self.has_prism_contract() {
603 return None;
604 }
605 match self {
606 Self::Mapped {
607 idx,
608 dtype: TensorDtype::Q8Row,
609 rows,
610 cols,
611 row_scale,
612 col_field,
613 ..
614 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
615 _ => None,
616 }
617 }
618
619 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
622 if self.has_prism_contract() {
623 return None;
624 }
625 match self {
626 Self::Mapped {
627 idx,
628 dtype: TensorDtype::Q8_2f,
629 rows,
630 cols,
631 row_scale,
632 col_field,
633 ..
634 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
635 _ => None,
636 }
637 }
638
639 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
643 match self {
644 Self::Mapped { dtype, .. } => Some(*dtype),
645 _ => None,
646 }
647 }
648
649 pub fn model_idx(&self) -> Option<usize> {
654 match self {
655 Self::Mapped { idx, .. } => Some(*idx),
656 _ => None,
657 }
658 }
659
660 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
665 match self {
666 Self::Mapped { model, .. } => Some(model.clone()),
667 _ => None,
668 }
669 }
670
671 pub fn row_spans(&self, r: usize, out: &mut Vec<(usize, usize)>) -> bool {
677 let Self::Mapped { dtype, .. } = self else {
678 return false;
679 };
680 let Some(spans) = row_span_offsets(*dtype, self.rows(), self.cols(), r) else {
681 return false;
682 };
683 let base = self.quant_bytes().as_ptr() as usize;
684 out.extend(spans.iter().filter(|s| s.1 > 0).map(|&(o, n)| (base + o, n)));
685 true
686 }
687
688 pub(crate) fn has_prism_contract(&self) -> bool {
693 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
694 }
695
696 pub fn rows(&self) -> usize {
697 match self {
698 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
699 Self::Blr(t) => t.rows,
700 }
701 }
702
703 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
706 if self.has_prism_contract() {
707 return None;
708 }
709 match self {
710 Self::Mapped {
711 model,
712 idx,
713 dtype: TensorDtype::Q4Tiled,
714 ..
715 } => Some((model, *idx)),
716 _ => None,
717 }
718 }
719
720 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
723 if self.has_prism_contract() {
724 return None;
725 }
726 match self {
727 Self::Mapped {
728 model,
729 idx,
730 dtype: TensorDtype::Q4TiledP,
731 ..
732 } => Some((model, *idx)),
733 _ => None,
734 }
735 }
736
737 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
745 if self.has_prism_contract() {
746 return None;
747 }
748 match self {
749 Self::Mapped {
750 model,
751 idx,
752 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
753 ..
754 } => Some((model, *idx)),
755 _ => None,
756 }
757 }
758
759 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
762 if self.has_prism_contract() {
763 return None;
764 }
765 match self {
766 Self::Mapped {
767 model,
768 idx,
769 dtype: TensorDtype::Q8_2f,
770 ..
771 } => Some((model, *idx)),
772 _ => None,
773 }
774 }
775
776 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
779 if self.has_prism_contract() {
780 return None;
781 }
782 match self {
783 Self::Mapped {
784 model,
785 idx,
786 dtype: TensorDtype::Q2TiledP,
787 ..
788 } => Some((model, *idx)),
789 _ => None,
790 }
791 }
792
793 pub fn cols(&self) -> usize {
794 match self {
795 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
796 Self::Blr(t) => t.cols,
797 }
798 }
799
800 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
803 if self.has_prism_contract() {
804 return None;
805 }
806 match self {
807 Self::Mapped {
808 model,
809 idx,
810 dtype: TensorDtype::Q1,
811 ..
812 } => Some((model, *idx)),
813 _ => None,
814 }
815 }
816
817 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
828 if self.has_prism_contract() {
829 return None;
830 }
831 self.graph_weight_descriptor()
832 }
833
834 pub(crate) fn graph_weight_descriptor(
839 &self,
840 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
841 match self {
842 Self::Mapped {
843 model,
844 idx,
845 dtype: TensorDtype::Q8Row,
846 row_scale,
847 ..
848 } => Some((model, *idx, 0, row_scale.as_slice())),
849 Self::Mapped {
850 model,
851 idx,
852 dtype: TensorDtype::Q1,
853 ..
854 } => Some((model, *idx, 1, &[])),
855 Self::Mapped {
858 model,
859 idx,
860 dtype: TensorDtype::Q1G,
861 ..
862 } => Some((model, *idx, 10, &[])),
863 Self::Mapped {
868 model,
869 idx,
870 dtype: TensorDtype::Q4Tiled,
871 ..
872 } => Some((model, *idx, 5, &[])),
873 Self::Mapped {
877 model,
878 idx,
879 dtype: TensorDtype::Q4TiledP,
880 ..
881 } => Some((model, *idx, 6, &[])),
882 Self::Mapped {
883 model,
884 idx,
885 dtype: TensorDtype::Q4Block,
886 ..
887 } => Some((model, *idx, 2, &[])),
888 Self::Mapped {
893 model,
894 idx,
895 dtype: TensorDtype::Q8_2f,
896 ..
897 } => Some((model, *idx, 7, &[])),
898 Self::Mapped {
899 model,
900 idx,
901 dtype: TensorDtype::Q1T,
902 ..
903 } => Some((model, *idx, 3, &[])),
904 Self::Mapped {
908 model,
909 idx,
910 dtype: TensorDtype::Q2TiledP,
911 ..
912 } => Some((model, *idx, 9, &[])),
913 _ => None,
914 }
915 }
916
917 pub fn q8_2f_codes(&self) -> Option<(&[u8], &[f32], &[f32])> {
921 match self {
922 Self::Mapped { dtype: TensorDtype::Q8_2f, row_scale, col_field, .. } => {
923 Some((self.quant_bytes(), row_scale, col_field))
924 }
925 _ => None,
926 }
927 }
928
929 pub fn as_f32(&self) -> Option<&[f32]> {
930 match self {
931 Self::F32 { data, .. } => Some(data),
932 Self::Mapped { .. } | Self::Blr(_) => None,
933 }
934 }
935
936 pub fn mapped_bytes(&self) -> Option<&[u8]> {
938 match self {
939 Self::Mapped { .. } => Some(self.quant_bytes()),
940 _ => None,
941 }
942 }
943
944 fn quant_bytes(&self) -> &[u8] {
945 match self {
946 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
947 Self::Blr(t) => t.bytes(),
948 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
949 }
950 }
951
952 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
954 let cols = self.cols();
955 debug_assert_eq!(dst.len(), cols);
956 match self {
957 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
958 Self::Blr(t) => t.row_f32(r, dst),
959 Self::Mapped {
960 model,
961 idx,
962 dtype,
963 row_scale,
964 col_field,
965 vbit_offsets,
966 ..
967 } => {
968 if *dtype == TensorDtype::Q4Tiled {
969 let bytes = self.quant_bytes();
970 let gpr = cols / GROUP_SIZE;
971 for gi in 0..gpr {
972 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
973 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
974 for (k, &b) in tile[2..].iter().enumerate() {
975 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
976 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
977 }
978 }
979 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
980 crate::prism::inverse_embedding(model, dst);
981 }
982 return;
983 }
984 if *dtype == TensorDtype::Q4TiledP {
985 let bytes = self.quant_bytes();
986 let gpr = cols / GROUP_SIZE;
987 let v = Q4tpView::new(bytes, self.rows(), cols);
988 let mut sc = vec![0f32; gpr];
989 v.scales_into(r, gpr, &mut sc);
990 for gi in 0..gpr {
991 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
992 let s = sc[gi];
993 for (k, &b) in tile.iter().enumerate() {
994 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
995 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
996 }
997 }
998 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
999 crate::prism::inverse_embedding(model, dst);
1000 }
1001 return;
1002 }
1003 if *dtype == TensorDtype::Q2TiledP {
1004 let bytes = self.quant_bytes();
1005 let gpr = cols / GROUP_SIZE;
1006 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
1007 let mut sc = vec![0f32; gpr];
1008 v.scales_into(r, gpr, &mut sc);
1009 for gi in 0..gpr {
1010 let ch =
1011 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
1012 let s = sc[gi];
1013 for (k, &b) in ch.iter().enumerate() {
1014 for j in 0..4 {
1015 let center = if crate::prism::is_affine_target(
1016 model,
1017 &model.tensors[*idx].name,
1018 ) {
1019 1.0
1020 } else {
1021 1.5
1022 };
1023 dst[gi * GROUP_SIZE + k * 4 + j] =
1024 (((b >> (2 * j)) & 3) as f32 - center) * s;
1025 }
1026 }
1027 }
1028 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1029 crate::prism::inverse_embedding(model, dst);
1030 }
1031 return;
1032 }
1033 if *dtype == TensorDtype::Q4Block {
1034 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
1035 let gpr = cols / GROUP_SIZE;
1036 for gi in 0..gpr {
1037 let g = r * gpr + gi;
1038 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1039 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1040 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
1041 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
1042 }
1043 }
1044 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1045 crate::prism::inverse_embedding(model, dst);
1046 }
1047 return;
1048 }
1049 if *dtype == TensorDtype::Mxfp4 {
1050 mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1051 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1052 crate::prism::inverse_embedding(model, dst);
1053 }
1054 return;
1055 }
1056 if *dtype == TensorDtype::Q1 {
1057 let bytes = self.quant_bytes();
1058 let gpr = cols / GROUP_SIZE;
1059 for gi in 0..gpr {
1060 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1061 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1062 for (j, &b) in tile[2..].iter().enumerate() {
1063 for k in 0..8 {
1064 dst[gi * GROUP_SIZE + j * 8 + k] =
1065 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1066 }
1067 }
1068 }
1069 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1070 crate::prism::inverse_embedding(model, dst);
1071 }
1072 return;
1073 }
1074 if *dtype == TensorDtype::Q1G {
1075 let bytes = self.quant_bytes();
1078 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1079 let gpr = cols / GROUP_SIZE;
1080 for gi in 0..gpr {
1081 let t = r * gpr + gi;
1082 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1083 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1084 for k in 0..8 {
1085 dst[gi * GROUP_SIZE + j * 8 + k] =
1086 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1087 }
1088 }
1089 }
1090 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1091 crate::prism::inverse_embedding(model, dst);
1092 }
1093 return;
1094 }
1095 if *dtype == TensorDtype::Q1T {
1096 let bytes = self.quant_bytes();
1097 let gpr = cols / GROUP_SIZE;
1098 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1099 for gi in 0..gpr {
1100 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1101 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1102 bytes[off],
1103 bytes[off + 1],
1104 ]));
1105 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1106 for k in 0..GROUP_SIZE {
1107 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1108 {
1109 1 => s,
1110 2 => -s,
1111 _ => 0.0,
1112 };
1113 }
1114 }
1115 let rows = self.rows();
1117 let entries = base_len + (rows + 1) * 4;
1118 if entries <= bytes.len() {
1119 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1120 let r0 = u32::from_le_bytes([
1121 ptrs[r * 4],
1122 ptrs[r * 4 + 1],
1123 ptrs[r * 4 + 2],
1124 ptrs[r * 4 + 3],
1125 ]) as usize;
1126 let r1 = u32::from_le_bytes([
1127 ptrs[(r + 1) * 4],
1128 ptrs[(r + 1) * 4 + 1],
1129 ptrs[(r + 1) * 4 + 2],
1130 ptrs[(r + 1) * 4 + 3],
1131 ]) as usize;
1132 let off = entries + r0 * 4;
1133 for i in 0..r1 - r0 {
1134 let item = &bytes[off + i * 4..off + i * 4 + 4];
1135 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1136 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1137 item[2], item[3],
1138 ]));
1139 if c < cols {
1140 dst[c] = v;
1141 }
1142 }
1143 }
1144 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1145 crate::prism::inverse_embedding(model, dst);
1146 }
1147 return;
1148 }
1149 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1150 let bytes = self.quant_bytes();
1151 let rows = self.rows();
1152 let ng = cols / GROUP_SIZE;
1153 let bits = &bytes[..rows];
1154 let sc_off = rows;
1155 let off = vbit_offsets[r];
1158 let b = bits[r] as usize;
1159 let l = ((1usize << (b - 1)) - 1) as f32;
1160 let data = &bytes[off..];
1161 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1162 for (i, d) in dst.iter_mut().enumerate() {
1163 while nbits < b {
1164 acc = (acc << 8) | data[byte_idx] as u64;
1165 byte_idx += 1;
1166 nbits += 8;
1167 }
1168 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1169 nbits -= b;
1170 let so = (r * ng + i / GROUP_SIZE) * 2;
1171 let sv = f16_to_f32(u16::from_le_bytes([
1172 bytes[sc_off + so],
1173 bytes[sc_off + so + 1],
1174 ]));
1175 *d = (u - l) * sv;
1176 }
1177 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1178 crate::prism::inverse_embedding(model, dst);
1179 }
1180 return;
1181 }
1182 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1183 let s = row_scale[r];
1184 match dtype {
1185 TensorDtype::Q8Row => {
1186 for (d, &b) in dst.iter_mut().zip(q) {
1187 *d = (b as i8) as f32 * s;
1188 }
1189 }
1190 TensorDtype::Q8_2f => {
1191 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1192 *d = (b as i8) as f32 * s * col_field[i];
1193 }
1194 }
1195 _ => unreachable!(),
1196 }
1197 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1198 crate::prism::inverse_embedding(model, dst);
1199 }
1200 }
1201 }
1202 }
1203
1204 pub fn sparse_col_ok(&self) -> bool {
1209 match self {
1210 Self::F32 { .. } => true,
1211 Self::Blr(_) => false,
1212 Self::Mapped { dtype, .. } => {
1213 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1214 }
1215 }
1216 }
1217
1218 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1222 let inter = self.cols();
1223 let hidden = self.rows();
1224 debug_assert_eq!(out.len(), hidden);
1225 match self {
1226 Self::F32 { data, .. } => {
1227 for (k, o) in out.iter_mut().enumerate() {
1228 *o += w * data[k * inter + c];
1229 }
1230 }
1231 Self::Blr(t) => t.add_col_scaled(c, w, out),
1232 Self::Mapped {
1233 dtype,
1234 row_scale,
1235 col_field,
1236 ..
1237 } => {
1238 let q = self.quant_bytes();
1239 let colf = if *dtype == TensorDtype::Q8_2f {
1240 col_field[c]
1241 } else {
1242 1.0
1243 };
1244 let wc = w * colf;
1245 for (k, o) in out.iter_mut().enumerate() {
1246 let b = q[k * inter + c] as i8 as f32;
1247 *o += wc * b * row_scale[k];
1248 }
1249 }
1250 }
1251 }
1252
1253 #[inline]
1262 pub fn prefetch_row(&self, r: usize) {
1263 let Self::Mapped { dtype, .. } = self else {
1264 return;
1265 };
1266 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1267 return;
1268 }
1269 let cols = self.cols();
1270 let q = self.quant_bytes();
1271 let (a, b) = (r * cols, (r + 1) * cols);
1272 if b > q.len() {
1273 return;
1274 }
1275 let mut j = a;
1276 while j < b {
1277 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1278 j += 512;
1279 }
1280 }
1281
1282 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1291 let cols = self.cols();
1292 debug_assert_eq!(out.len(), cols);
1293 match self {
1294 Self::F32 { data, .. } => {
1295 let row = &data[r * cols..(r + 1) * cols];
1296 for (o, v) in out.iter_mut().zip(row) {
1297 *o += w * v;
1298 }
1299 }
1300 Self::Blr(t) => {
1301 let tmp = &mut scratch[..cols];
1302 t.row_f32(r, tmp);
1303 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1304 *o += w * v;
1305 }
1306 }
1307 Self::Mapped {
1308 dtype,
1309 row_scale,
1310 col_field,
1311 ..
1312 } => match dtype {
1313 TensorDtype::Q8Row => {
1314 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1315 let ws = w * row_scale[r];
1316 let row: &[i8] =
1317 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1318 axpy_i8_f32(out, row, ws);
1319 }
1320 TensorDtype::Q8_2f => {
1321 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1322 let ws = w * row_scale[r];
1323 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1324 *o += ws * c * (*b as i8 as f32);
1325 }
1326 }
1327 _ => {
1328 self.row_f32(r, scratch);
1329 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1330 *o += w * v;
1331 }
1332 }
1333 },
1334 }
1335 }
1336
1337 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1341 let cols = self.cols();
1342 match self {
1343 Self::F32 { data, .. } => {
1344 let row = &data[r * cols..(r + 1) * cols];
1345 row.iter().zip(x).map(|(w, v)| w * v).sum()
1346 }
1347 Self::Blr(t) => {
1348 let tmp = &mut scratch[..cols];
1349 t.row_f32(r, tmp);
1350 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1351 }
1352 Self::Mapped {
1353 model,
1354 idx,
1355 dtype,
1356 row_scale,
1357 col_field,
1358 ..
1359 } => {
1360 let prism_forward =
1361 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1362 if prism_forward {
1363 let transformed = crate::prism::forward(model, &x[..cols]);
1364 let gpr = cols / GROUP_SIZE;
1365 match dtype {
1366 TensorDtype::Q2TiledP => {
1367 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1368 let mut sc = vec![0f32; gpr];
1369 v.scales_into(r, gpr, &mut sc);
1370 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1371 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1372 }
1373 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1374 }
1375 _ => {
1376 self.row_f32(r, scratch);
1377 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1378 }
1379 }
1380 }
1381 match dtype {
1382 TensorDtype::Q8Row => {
1383 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1384 dot_i8_f32(q, x) * row_scale[r]
1385 }
1386 TensorDtype::Q8_2f => {
1387 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1388 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1389 }
1390 _ => {
1391 self.row_f32(r, scratch);
1392 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1393 }
1394 }
1395 }
1396 }
1397 }
1398
1399 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1402 if self.exact_cpu_default() {
1403 return float_activations_scope(|| self.matvec(x, out, pool));
1404 }
1405 match self {
1406 Self::F32 { data, .. } => {
1410 if !crate::f32_backend::matvec(data, x, out) {
1411 matvec_rows(pool, data, x, out);
1412 }
1413 }
1414 Self::Blr(t) => {
1415 if crate::gpu::enabled_here()
1416 && crate::blr::gpu_wanted()
1417 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1418 {
1419 return;
1420 }
1421 t.matvec(x, out, pool)
1422 }
1423 Self::Mapped {
1424 model,
1425 idx,
1426 dtype,
1427 rows,
1428 cols,
1429 row_scale,
1430 col_field,
1431 vbit_offsets,
1432 repack,
1433 } => {
1434 let _ = (model, idx);
1435 assert!(
1445 out.len() >= *rows && x.len() >= *cols,
1446 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1447 out.len(),
1448 x.len(),
1449 );
1450 let prism_forward =
1451 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1452 if *dtype == TensorDtype::Q2TiledP
1453 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1454 {
1455 use std::sync::atomic::{AtomicUsize, Ordering};
1456 static N: AtomicUsize = AtomicUsize::new(0);
1457 let n = N.fetch_add(1, Ordering::Relaxed);
1458 if n < 128 {
1459 eprintln!(
1460 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1461 model.tensors[*idx].name,
1462 prism_forward,
1463 rows,
1464 cols,
1465 crate::gpu::enabled_here(),
1466 crate::gpu::q2tp_gpu_opt_in(),
1467 crate::gpu::cur_layer(),
1468 );
1469 }
1470 }
1471 if prism_forward {
1476 let transformed = crate::prism::forward(model, &x[..*cols]);
1477 match dtype {
1478 TensorDtype::Q4Block => {
1479 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1480 }
1481 TensorDtype::Q4Tiled => {
1482 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1483 }
1484 TensorDtype::Q4TiledP => {
1485 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1486 }
1487 TensorDtype::Q2TiledP => {
1488 let affine =
1489 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1490 if *rows * *cols >= 8_388_608
1491 && crate::gpu::enabled_here()
1492 && crate::gpu::q2tp_gpu_opt_in()
1493 {
1494 let gpu_ok = if affine {
1495 crate::gpu::q2tp_affine_matvec(
1496 model,
1497 *idx,
1498 &transformed,
1499 *rows,
1500 *cols,
1501 out,
1502 )
1503 } else {
1504 crate::gpu::q2tp_matvec(
1505 model,
1506 *idx,
1507 &transformed,
1508 *rows,
1509 *cols,
1510 out,
1511 )
1512 };
1513 if gpu_ok {
1514 return;
1515 }
1516 }
1517 if affine {
1518 q2tp_affine_matvec(
1519 self.quant_bytes(),
1520 &transformed,
1521 *rows,
1522 *cols,
1523 out,
1524 pool,
1525 )
1526 } else {
1527 q2tp_matvec(
1528 self.quant_bytes(),
1529 &transformed,
1530 *rows,
1531 *cols,
1532 out,
1533 pool,
1534 )
1535 }
1536 }
1537 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1538 self.bin_w().expect("binary dtype"),
1539 &transformed,
1540 *rows,
1541 *cols,
1542 out,
1543 pool,
1544 ),
1545 TensorDtype::Q1T => {
1546 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1547 }
1548 TensorDtype::Mxfp4 => {
1549 mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1550 }
1551 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1552 self.quant_bytes(),
1553 vbit_offsets,
1554 &transformed,
1555 *rows,
1556 *cols,
1557 out,
1558 pool,
1559 ),
1560 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1561 self.quant_bytes(),
1562 repack,
1563 row_scale,
1564 &transformed,
1565 col_field,
1566 *dtype,
1567 *rows,
1568 *cols,
1569 out,
1570 pool,
1571 ),
1572 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1573 }
1574 return;
1575 }
1576 if *dtype == TensorDtype::Q4Block {
1577 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1581 let t0 = std::time::Instant::now();
1582 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1583 crate::gpu::ProbeArm::Gpu => {
1584 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1585 crate::gpu::probe_record(
1586 crate::gpu::OpClass::Matvec,
1587 true,
1588 t0.elapsed(),
1589 );
1590 return;
1591 }
1592 }
1593 crate::gpu::ProbeArm::CpuTimed => {
1594 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1595 crate::gpu::probe_record(
1596 crate::gpu::OpClass::Matvec,
1597 false,
1598 t0.elapsed(),
1599 );
1600 return;
1601 }
1602 crate::gpu::ProbeArm::Cpu => {}
1603 }
1604 }
1605 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1606 return;
1607 }
1608 if *dtype == TensorDtype::Q4Tiled {
1609 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1614 let t0 = std::time::Instant::now();
1615 let cls = crate::gpu::matvec_class(*rows, *cols);
1616 match crate::gpu::probe_arm(cls) {
1617 crate::gpu::ProbeArm::Gpu => {
1618 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1619 crate::gpu::probe_record(cls, true, t0.elapsed());
1620 return;
1621 }
1622 }
1623 crate::gpu::ProbeArm::CpuTimed => {
1624 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1625 crate::gpu::probe_record(cls, false, t0.elapsed());
1626 return;
1627 }
1628 crate::gpu::ProbeArm::Cpu => {}
1629 }
1630 }
1631 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1632 return;
1633 }
1634 if *dtype == TensorDtype::Q4TiledP {
1635 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1641 let t0 = std::time::Instant::now();
1642 let cls = crate::gpu::matvec_class(*rows, *cols);
1643 match crate::gpu::probe_arm(cls) {
1644 crate::gpu::ProbeArm::Gpu => {
1645 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1646 crate::gpu::probe_record(cls, true, t0.elapsed());
1647 return;
1648 }
1649 }
1650 crate::gpu::ProbeArm::CpuTimed => {
1651 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1652 crate::gpu::probe_record(cls, false, t0.elapsed());
1653 return;
1654 }
1655 crate::gpu::ProbeArm::Cpu => {}
1656 }
1657 }
1658 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1659 return;
1660 }
1661 if *dtype == TensorDtype::Q2TiledP {
1662 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1663 return;
1664 }
1665 if *dtype == TensorDtype::Mxfp4 {
1666 self.mxfp4_matvec(x, out, pool);
1667 return;
1668 }
1669 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1670 let w = self.bin_w().expect("binary dtype");
1675 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1676 let t0 = std::time::Instant::now();
1677 let arm = if crate::gpu::q1_force() {
1678 crate::gpu::ProbeArm::Gpu
1679 } else {
1680 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1681 };
1682 match arm {
1683 crate::gpu::ProbeArm::Gpu => {
1684 let ok = if *dtype == TensorDtype::Q1G {
1685 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1686 } else {
1687 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1688 };
1689 if ok {
1690 crate::gpu::probe_record(
1691 crate::gpu::OpClass::Matvec,
1692 true,
1693 t0.elapsed(),
1694 );
1695 return;
1696 }
1697 }
1698 crate::gpu::ProbeArm::CpuTimed => {
1699 bin_matvec(w, x, *rows, *cols, out, pool);
1700 crate::gpu::probe_record(
1701 crate::gpu::OpClass::Matvec,
1702 false,
1703 t0.elapsed(),
1704 );
1705 return;
1706 }
1707 crate::gpu::ProbeArm::Cpu => {}
1708 }
1709 }
1710 bin_matvec(w, x, *rows, *cols, out, pool);
1711 return;
1712 }
1713 if *dtype == TensorDtype::Q1T {
1714 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1718 let t0 = std::time::Instant::now();
1719 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1720 crate::gpu::ProbeArm::Gpu => {
1721 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1722 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1723 crate::gpu::probe_record(
1724 crate::gpu::OpClass::Matvec,
1725 true,
1726 t0.elapsed(),
1727 );
1728 return;
1729 }
1730 }
1731 crate::gpu::ProbeArm::CpuTimed => {
1732 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1733 crate::gpu::probe_record(
1734 crate::gpu::OpClass::Matvec,
1735 false,
1736 t0.elapsed(),
1737 );
1738 return;
1739 }
1740 crate::gpu::ProbeArm::Cpu => {}
1741 }
1742 }
1743 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1744 return;
1745 }
1746 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1747 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1748 return;
1749 }
1750 let xs = prescale(x, col_field, *dtype);
1751 if *rows >= crate::gpu::min_rows()
1756 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1757 && gpu_lmhead_enabled()
1758 && crate::gpu::enabled_here()
1759 {
1760 let t0 = std::time::Instant::now();
1763 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1764 crate::gpu::ProbeArm::Gpu => {}
1765 crate::gpu::ProbeArm::CpuTimed => {
1766 qmatvec(
1767 self.quant_bytes(),
1768 repack,
1769 row_scale,
1770 x,
1771 col_field,
1772 *dtype,
1773 *rows,
1774 *cols,
1775 out,
1776 pool,
1777 );
1778 crate::gpu::probe_record(
1779 crate::gpu::OpClass::Matvec,
1780 false,
1781 t0.elapsed(),
1782 );
1783 return;
1784 }
1785 crate::gpu::ProbeArm::Cpu => {
1786 qmatvec(
1787 self.quant_bytes(),
1788 repack,
1789 row_scale,
1790 x,
1791 col_field,
1792 *dtype,
1793 *rows,
1794 *cols,
1795 out,
1796 pool,
1797 );
1798 return;
1799 }
1800 }
1801 let frac = gpu_split_frac();
1802 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1803 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1804 let bytes = self.quant_bytes();
1805 let ok = std::thread::scope(|sc| {
1806 let g = sc.spawn(|| {
1807 crate::gpu::q8_matvec_range(
1808 model,
1809 *idx,
1810 cpu_rows,
1811 &row_scale[cpu_rows..],
1812 &xs,
1813 *rows - cpu_rows,
1814 *cols,
1815 out_gpu,
1816 )
1817 });
1818 if cpu_rows > 0 {
1819 let rep_cpu = if repack.is_empty() {
1822 &[][..]
1823 } else {
1824 &repack[..(cpu_rows / 4) * 4 * *cols]
1825 };
1826 qmatvec(
1827 &bytes[..cpu_rows * *cols],
1828 rep_cpu,
1829 &row_scale[..cpu_rows],
1830 x,
1831 col_field,
1832 *dtype,
1833 cpu_rows,
1834 *cols,
1835 out_cpu,
1836 pool,
1837 );
1838 }
1839 g.join().unwrap_or(false)
1840 });
1841 if ok {
1842 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1843 return;
1844 }
1845 qmatvec(
1848 &bytes[cpu_rows * *cols..(*rows) * *cols],
1849 &[],
1850 &row_scale[cpu_rows..],
1851 x,
1852 col_field,
1853 *dtype,
1854 *rows - cpu_rows,
1855 *cols,
1856 out_gpu,
1857 pool,
1858 );
1859 return;
1860 }
1861 qmatvec(
1862 self.quant_bytes(),
1863 repack,
1864 row_scale,
1865 x,
1866 col_field,
1867 *dtype,
1868 *rows,
1869 *cols,
1870 out,
1871 pool,
1872 );
1873 }
1874 }
1875 }
1876
1877 pub fn matvec2(
1879 &self,
1880 x1: &[f32],
1881 x2: &[f32],
1882 o1: &mut [f32],
1883 o2: &mut [f32],
1884 pool: Option<&Pool>,
1885 ) {
1886 if self.exact_cpu_default() {
1887 return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1888 }
1889 match self {
1890 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1891 Self::Blr(_) => {
1892 self.matvec(x1, o1, pool);
1893 self.matvec(x2, o2, pool);
1894 }
1895 Self::Mapped {
1896 model,
1897 idx,
1898 dtype,
1899 rows,
1900 cols,
1901 row_scale,
1902 col_field,
1903 vbit_offsets,
1904 ..
1905 } => {
1906 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1907 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1908 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1909 match dtype {
1910 TensorDtype::Q4Block => {
1911 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1912 }
1913 TensorDtype::Q4Tiled => {
1914 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1915 }
1916 TensorDtype::Q4TiledP => {
1917 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1918 }
1919 TensorDtype::Q2TiledP => {
1920 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1921 q2tp_affine_matvec2(
1922 self.quant_bytes(),
1923 &tx1,
1924 &tx2,
1925 *rows,
1926 *cols,
1927 o1,
1928 o2,
1929 pool,
1930 )
1931 } else {
1932 q2tp_matvec2(
1933 self.quant_bytes(),
1934 &tx1,
1935 &tx2,
1936 *rows,
1937 *cols,
1938 o1,
1939 o2,
1940 pool,
1941 )
1942 }
1943 }
1944 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1945 self.bin_w().expect("binary dtype"),
1946 &tx1,
1947 &tx2,
1948 *rows,
1949 *cols,
1950 o1,
1951 o2,
1952 pool,
1953 ),
1954 TensorDtype::Q1T => {
1955 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1956 }
1957 TensorDtype::Mxfp4 => mxfp4::matvec2(
1958 self.quant_bytes(),
1959 &tx1,
1960 &tx2,
1961 *rows,
1962 *cols,
1963 o1,
1964 o2,
1965 pool,
1966 ),
1967 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1968 self.quant_bytes(),
1969 vbit_offsets,
1970 &tx1,
1971 &tx2,
1972 *rows,
1973 *cols,
1974 o1,
1975 o2,
1976 pool,
1977 ),
1978 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1979 self.quant_bytes(),
1980 row_scale,
1981 &tx1,
1982 &tx2,
1983 col_field,
1984 *dtype,
1985 *rows,
1986 *cols,
1987 o1,
1988 o2,
1989 pool,
1990 ),
1991 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1992 }
1993 return;
1994 }
1995 if *dtype == TensorDtype::Q4Block {
1996 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1997 return;
1998 }
1999 if *dtype == TensorDtype::Q4Tiled {
2000 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2001 return;
2002 }
2003 if *dtype == TensorDtype::Q4TiledP {
2004 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2005 return;
2006 }
2007 if *dtype == TensorDtype::Q2TiledP {
2008 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2009 return;
2010 }
2011 if *dtype == TensorDtype::Mxfp4 {
2012 mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2013 return;
2014 }
2015 if let Some(w) = self.bin_w() {
2016 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
2017 return;
2018 }
2019 if *dtype == TensorDtype::Q1T {
2020 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2026 return;
2027 }
2028 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2029 vbitmatvec2(
2030 self.quant_bytes(),
2031 vbit_offsets,
2032 x1,
2033 x2,
2034 *rows,
2035 *cols,
2036 o1,
2037 o2,
2038 pool,
2039 );
2040 return;
2041 }
2042 qmatvec2(
2043 self.quant_bytes(),
2044 row_scale,
2045 x1,
2046 x2,
2047 col_field,
2048 *dtype,
2049 *rows,
2050 *cols,
2051 o1,
2052 o2,
2053 pool,
2054 );
2055 }
2056 }
2057 }
2058}
2059
2060impl QTensor {
2061 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2069 if self.has_prism_contract() {
2070 return None;
2071 }
2072 match self {
2073 Self::Mapped {
2074 model, idx, dtype, ..
2075 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2076 _ => None,
2077 }
2078 }
2079
2080 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2081 if self.exact_cpu_default() {
2082 return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2083 }
2084 let cols = self.cols();
2085 let rows = self.rows();
2086 debug_assert_eq!(xs_all.len(), b * cols);
2087 debug_assert_eq!(out.len(), b * rows);
2088 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2089 if crate::gptq_capture::capturing() {
2093 if let Self::Mapped { model, idx, .. } = self {
2094 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2095 }
2096 if let Self::Blr(t) = self {
2097 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2098 }
2099 }
2100 match self {
2101 Self::Blr(t) => {
2102 if crate::gpu::enabled_here()
2103 && crate::blr::gpu_wanted()
2104 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2105 {
2106 return;
2107 }
2108 t.matmat(xs_all, b, out, pool)
2109 }
2110 Self::F32 { data, .. } => {
2111 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2112 return;
2113 }
2114 let out_addr = SendMut(out.as_mut_ptr());
2115 let kernel = |start: usize, end: usize, r0: usize, r1: usize| {
2122 let mut bi = start;
2123 while bi < end {
2124 let n = (end - bi).min(4);
2125 for o in r0..r1 {
2126 let row = &data[o * cols..(o + 1) * cols];
2127 if n == 4 {
2128 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2129 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2130 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2131 for j in 0..cols {
2132 let w = row[j];
2133 a0 += w * x0[j];
2134 a1 += w * x1[j];
2135 a2 += w * x2[j];
2136 a3 += w * x3[j];
2137 }
2138 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2139 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2140 }
2141 } else {
2142 for k in 0..n {
2143 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2144 let mut acc = 0f32;
2145 for j in 0..cols {
2146 acc += row[j] * x[j];
2147 }
2148 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2149 }
2150 }
2151 }
2152 bi += n;
2153 }
2154 };
2155 match pool {
2156 Some(pool) if b < 256 && rows >= 256 => {
2162 pool.run_rows(rows, &|r0, r1| kernel(0, b, r0, r1))
2163 }
2164 _ => dispatch_rows(pool, b, &|s, e| kernel(s, e, 0, rows)),
2165 }
2166 }
2167 Self::Mapped {
2168 model,
2169 idx,
2170 dtype,
2171 row_scale,
2172 col_field,
2173 vbit_offsets,
2174 ..
2175 } => {
2176 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2177 let mut transformed = Vec::with_capacity(xs_all.len());
2178 for bi in 0..b {
2179 transformed.extend_from_slice(&crate::prism::forward(
2180 model,
2181 &xs_all[bi * cols..(bi + 1) * cols],
2182 ));
2183 }
2184 match dtype {
2185 TensorDtype::Q4Block => {
2186 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2187 }
2188 TensorDtype::Q4Tiled => {
2189 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2190 }
2191 TensorDtype::Q4TiledP => {
2192 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2193 }
2194 TensorDtype::Q2TiledP => {
2195 let affine =
2196 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2197 let gpu_batch_ok = if affine {
2203 b >= 2
2204 } else {
2205 b >= 32 && b * rows * cols >= 128_000_000
2206 };
2207 if gpu_batch_ok
2208 && cols % 32 == 0
2209 && crate::gpu::enabled_here()
2210 && crate::gpu::q2tp_gpu_opt_in()
2211 {
2212 let gpu_ok = if affine {
2213 crate::gpu::q2tp_affine_matmat(
2214 model,
2215 *idx,
2216 &transformed,
2217 b,
2218 rows,
2219 cols,
2220 out,
2221 )
2222 } else {
2223 crate::gpu::q2tp_matmat(
2224 model,
2225 *idx,
2226 &transformed,
2227 b,
2228 rows,
2229 cols,
2230 out,
2231 )
2232 };
2233 if gpu_ok {
2234 return;
2235 }
2236 }
2237 if affine
2241 && b == 1
2242 && cols % 32 == 0
2243 && crate::gpu::enabled_here()
2244 && crate::gpu::q2tp_gpu_opt_in()
2245 && crate::gpu::q2tp_affine_matvec(
2246 model,
2247 *idx,
2248 &transformed[..cols],
2249 rows,
2250 cols,
2251 &mut out[..rows],
2252 )
2253 {
2254 return;
2255 }
2256 if affine {
2257 q2tp_affine_matmat(
2258 self.quant_bytes(),
2259 &transformed,
2260 b,
2261 rows,
2262 cols,
2263 out,
2264 pool,
2265 )
2266 } else {
2267 q2tp_matmat(
2268 self.quant_bytes(),
2269 &transformed,
2270 b,
2271 rows,
2272 cols,
2273 out,
2274 pool,
2275 )
2276 }
2277 }
2278 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2279 self.bin_w().expect("binary dtype"),
2280 &transformed,
2281 b,
2282 rows,
2283 cols,
2284 out,
2285 pool,
2286 ),
2287 TensorDtype::Q1T => {
2288 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2289 }
2290 TensorDtype::Mxfp4 => {
2291 mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2292 }
2293 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2294 self.quant_bytes(),
2295 vbit_offsets,
2296 &transformed,
2297 b,
2298 rows,
2299 cols,
2300 out,
2301 pool,
2302 ),
2303 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2304 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2305 .map(|bi| {
2306 prescale(
2307 &transformed[bi * cols..(bi + 1) * cols],
2308 col_field,
2309 *dtype,
2310 )
2311 })
2312 .collect();
2313 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2314 }
2315 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2316 }
2317 return;
2318 }
2319 if *dtype == TensorDtype::Q4Block {
2320 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2321 return;
2322 }
2323 if *dtype == TensorDtype::Mxfp4 {
2324 self.mxfp4_matmat(xs_all, b, out, pool);
2325 return;
2326 }
2327 if *dtype == TensorDtype::Q4TiledP {
2328 if crate::gpu::cuda_active()
2349 && b >= 2
2350 && rows * cols >= (1 << 20)
2351 && cols % 32 == 0
2352 && !row_exact()
2353 && crate::gpu::enabled_here()
2354 {
2355 if let Self::Mapped { model, idx, .. } = self {
2356 if crate::gpu::q4tp_matmat(model, *idx, xs_all, b, rows, cols, out) {
2357 return;
2358 }
2359 }
2360 }
2361 if b >= 32
2362 && b * rows * cols >= 128_000_000
2363 && cols % 32 == 0
2364 && !row_exact()
2365 && !crate::gpu::mm_killed()
2366 && crate::gpu::enabled_here()
2367 {
2368 let class = if b >= 128 {
2369 crate::gpu::OpClass::MatmatWide
2370 } else {
2371 crate::gpu::OpClass::Matmat
2372 };
2373 if let Self::Mapped { model, idx, .. } = self {
2374 if crate::mm_ab::on() {
2386 let mut g = vec![0f32; b * rows];
2387 let t = std::time::Instant::now();
2388 let took = crate::gpu::q4tp_matmat(
2389 model, *idx, xs_all, b, rows, cols, &mut g,
2390 );
2391 let dg = t.elapsed();
2392 let t = std::time::Instant::now();
2393 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2394 let dc = t.elapsed();
2395 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2396 return;
2397 }
2398 let t0 = std::time::Instant::now();
2399 let resident = crate::gpu::weight_is_resident(model, *idx);
2403 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2404 crate::gpu::ProbeArm::Gpu => {
2405 if crate::gpu::q4tp_matmat(
2406 model, *idx, xs_all, b, rows, cols, out,
2407 ) {
2408 let el = t0.elapsed();
2409 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2419 let budget = std::time::Duration::from_secs_f64(
2420 flops / 1.5e12 * 8.0 + 0.020,
2421 );
2422 crate::gpu::mm_budget_check(
2423 "q4tp matmat",
2424 el,
2425 budget,
2426 crate::gpu::probe_was_cold() || !resident,
2427 );
2428 crate::gpu::probe_record(class, true, el);
2429 return;
2430 }
2431 }
2432 crate::gpu::ProbeArm::CpuTimed => {
2433 q4tp_matmat(
2434 self.quant_bytes(),
2435 xs_all,
2436 b,
2437 rows,
2438 cols,
2439 out,
2440 pool,
2441 );
2442 crate::gpu::probe_record(class, false, t0.elapsed());
2443 return;
2444 }
2445 crate::gpu::ProbeArm::Cpu => {}
2446 }
2447 }
2448 }
2449 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2450 return;
2451 }
2452 if *dtype == TensorDtype::Q2TiledP {
2453 if b >= 32
2459 && b * rows * cols >= 128_000_000
2460 && cols % 32 == 0
2461 && !crate::gpu::mm_killed()
2462 && crate::gpu::enabled_here()
2463 {
2464 let class = if b >= 128 {
2465 crate::gpu::OpClass::MatmatWide
2466 } else {
2467 crate::gpu::OpClass::Matmat
2468 };
2469 if let Self::Mapped { model, idx, .. } = self {
2470 let t0 = std::time::Instant::now();
2471 match crate::gpu::probe_arm(class) {
2472 crate::gpu::ProbeArm::Gpu => {
2473 if crate::gpu::q2tp_matmat(
2474 model, *idx, xs_all, b, rows, cols, out,
2475 ) {
2476 crate::gpu::probe_record(class, true, t0.elapsed());
2477 return;
2478 }
2479 }
2480 crate::gpu::ProbeArm::CpuTimed => {
2481 q2tp_matmat(
2482 self.quant_bytes(),
2483 xs_all,
2484 b,
2485 rows,
2486 cols,
2487 out,
2488 pool,
2489 );
2490 crate::gpu::probe_record(class, false, t0.elapsed());
2491 return;
2492 }
2493 crate::gpu::ProbeArm::Cpu => {}
2494 }
2495 }
2496 }
2497 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2502 return;
2503 }
2504 if *dtype == TensorDtype::Q4Tiled {
2505 if b >= 32
2517 && b * rows * cols >= 128_000_000
2518 && cols % 32 == 0
2519 && !crate::gpu::mm_killed()
2520 && crate::gpu::enabled_here()
2521 {
2522 let class = if b >= 128 {
2523 crate::gpu::OpClass::MatmatWide
2524 } else {
2525 crate::gpu::OpClass::Matmat
2526 };
2527 if let Self::Mapped { model, idx, .. } = self {
2528 let t0 = std::time::Instant::now();
2529 match crate::gpu::probe_arm(class) {
2530 crate::gpu::ProbeArm::Gpu => {
2531 if crate::gpu::q4t_matmat(
2532 model, *idx, xs_all, b, rows, cols, out,
2533 ) {
2534 let el = t0.elapsed();
2535 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2545 let budget = std::time::Duration::from_secs_f64(
2546 flops / 1.5e12 * 8.0 + 0.020,
2547 );
2548 crate::gpu::mm_budget_check(
2549 "q4t matmat",
2550 el,
2551 budget,
2552 crate::gpu::probe_was_cold(),
2553 );
2554 crate::gpu::probe_record(class, true, el);
2555 return;
2556 }
2557 }
2558 crate::gpu::ProbeArm::CpuTimed => {
2559 q4t_matmat(
2560 self.quant_bytes(),
2561 xs_all,
2562 b,
2563 rows,
2564 cols,
2565 out,
2566 pool,
2567 );
2568 crate::gpu::probe_record(class, false, t0.elapsed());
2569 return;
2570 }
2571 crate::gpu::ProbeArm::Cpu => {}
2572 }
2573 }
2574 }
2575 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2576 return;
2577 }
2578 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2579 let w = self.bin_w().expect("binary dtype");
2583 if b >= 32
2584 && b * rows * cols >= 128_000_000
2585 && cols % 64 == 0
2586 && crate::gpu::enabled_here()
2587 {
2588 if let Self::Mapped { model, idx, .. } = self {
2589 let t0 = std::time::Instant::now();
2590 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2591 crate::gpu::ProbeArm::Gpu => {
2592 let ok = if *dtype == TensorDtype::Q1G {
2593 crate::gpu::q1g_matmat(
2594 model, *idx, xs_all, b, rows, cols, out,
2595 )
2596 } else {
2597 crate::gpu::q1_matmat(
2598 model, *idx, xs_all, b, rows, cols, out,
2599 )
2600 };
2601 if ok {
2602 crate::gpu::probe_record(
2603 crate::gpu::OpClass::Matmat,
2604 true,
2605 t0.elapsed(),
2606 );
2607 return;
2608 }
2609 }
2610 crate::gpu::ProbeArm::CpuTimed => {
2611 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2612 crate::gpu::probe_record(
2613 crate::gpu::OpClass::Matmat,
2614 false,
2615 t0.elapsed(),
2616 );
2617 return;
2618 }
2619 crate::gpu::ProbeArm::Cpu => {}
2620 }
2621 }
2622 }
2623 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2624 return;
2625 }
2626 if *dtype == TensorDtype::Q1T {
2627 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2630 if let Self::Mapped { model, idx, .. } = self {
2631 let t0 = std::time::Instant::now();
2632 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2633 crate::gpu::ProbeArm::Gpu => {
2634 if crate::gpu::q1t_matmat(
2635 model, *idx, xs_all, b, rows, cols, out,
2636 ) {
2637 crate::gpu::probe_record(
2638 crate::gpu::OpClass::Matmat,
2639 true,
2640 t0.elapsed(),
2641 );
2642 return;
2643 }
2644 }
2645 crate::gpu::ProbeArm::CpuTimed => {
2646 q1t_matmat(
2647 self.quant_bytes(),
2648 xs_all,
2649 b,
2650 rows,
2651 cols,
2652 out,
2653 pool,
2654 );
2655 crate::gpu::probe_record(
2656 crate::gpu::OpClass::Matmat,
2657 false,
2658 t0.elapsed(),
2659 );
2660 return;
2661 }
2662 crate::gpu::ProbeArm::Cpu => {}
2663 }
2664 }
2665 }
2666 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2667 return;
2668 }
2669 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2670 vbitmatmat(
2671 self.quant_bytes(),
2672 vbit_offsets,
2673 xs_all,
2674 b,
2675 rows,
2676 cols,
2677 out,
2678 pool,
2679 );
2680 return;
2681 }
2682 let _pmp = crate::cpuprof::spans::span("mm.prescale");
2683 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2684 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2685 .collect();
2686 drop(_pmp);
2687 if row_exact()
2693 && (1..=4).contains(&b)
2694 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2695 && crate::gpu::enabled_here()
2696 && crate::gpu::wgpu_active()
2697 {
2698 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2699 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2700 return;
2701 }
2702 }
2703 if b >= 8
2715 && b * rows * cols >= 128_000_000
2716 && crate::gpu::enabled_here()
2717 && std::env::var("CMF_MATMAT_GPU").as_deref() != Ok("0")
2718 {
2719 if let Self::Mapped { model, idx, .. } = self {
2720 MATMAT_GATE.fetch_add(1, Ordering::Relaxed);
2721 let t0 = std::time::Instant::now();
2722 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2723 crate::gpu::ProbeArm::Gpu
2724 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2725 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2726 {
2727 let q = self.quant_bytes();
2731 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2732 return;
2733 }
2734 crate::gpu::ProbeArm::Gpu => {
2735 if *dtype == TensorDtype::Q8_2f
2743 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2744 && crate::gpu::q8_matmat_2f(
2745 model, *idx, row_scale, col_field, xs_all, b, rows,
2746 cols, out,
2747 )
2748 {
2749 MATMAT_GPU_HIT.fetch_add(1, Ordering::Relaxed);
2750 crate::gpu::probe_record(
2751 crate::gpu::OpClass::Matmat,
2752 true,
2753 t0.elapsed(),
2754 );
2755 return;
2756 }
2757 let flat: Vec<f32> =
2758 pre.iter().flat_map(|v| v.iter().copied()).collect();
2759 if crate::gpu::q8_matmat(
2760 model, *idx, row_scale, &flat, b, rows, cols, out,
2761 ) {
2762 crate::gpu::probe_record(
2763 crate::gpu::OpClass::Matmat,
2764 true,
2765 t0.elapsed(),
2766 );
2767 return;
2768 }
2769 }
2770 crate::gpu::ProbeArm::CpuTimed => {
2771 let q = self.quant_bytes();
2772 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2773 crate::gpu::probe_record(
2774 crate::gpu::OpClass::Matmat,
2775 false,
2776 t0.elapsed(),
2777 );
2778 return;
2779 }
2780 crate::gpu::ProbeArm::Cpu => {}
2781 }
2782 }
2783 }
2784 let q = self.quant_bytes();
2785 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2786 }
2787 }
2788 }
2789}
2790
2791impl QTensor {
2792 fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2796 let (rows, cols) = (self.rows(), self.cols());
2797 if let Self::Mapped { model, idx, .. } = self {
2798 if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2799 let t0 = std::time::Instant::now();
2800 let cls = crate::gpu::matvec_class(rows, cols);
2801 match crate::gpu::probe_arm(cls) {
2802 crate::gpu::ProbeArm::Gpu => {
2803 if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2804 crate::gpu::probe_record(cls, true, t0.elapsed());
2805 return;
2806 }
2807 }
2808 crate::gpu::ProbeArm::CpuTimed => {
2809 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2810 crate::gpu::probe_record(cls, false, t0.elapsed());
2811 return;
2812 }
2813 crate::gpu::ProbeArm::Cpu => {}
2814 }
2815 }
2816 }
2817 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2818 }
2819
2820 fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2825 let (rows, cols) = (self.rows(), self.cols());
2826 if let Self::Mapped { model, idx, .. } = self {
2827 if b >= 2
2828 && b * rows * cols >= 33_554_432
2829 && !row_exact()
2830 && !crate::gpu::mm_killed()
2831 && crate::gpu::enabled_here()
2832 {
2833 let class = if b >= 128 {
2834 crate::gpu::OpClass::MatmatWide
2835 } else {
2836 crate::gpu::OpClass::Matmat
2837 };
2838 let t0 = std::time::Instant::now();
2839 match crate::gpu::probe_arm(class) {
2840 crate::gpu::ProbeArm::Gpu => {
2841 if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2842 crate::gpu::probe_record(class, true, t0.elapsed());
2843 return;
2844 }
2845 }
2846 crate::gpu::ProbeArm::CpuTimed => {
2847 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2848 crate::gpu::probe_record(class, false, t0.elapsed());
2849 return;
2850 }
2851 crate::gpu::ProbeArm::Cpu => {}
2852 }
2853 }
2854 }
2855 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2856 }
2857}
2858
2859impl QTensor {
2860 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2870 let (rows, cols) = (self.rows(), self.cols());
2871 let Self::Mapped {
2872 model,
2873 idx,
2874 dtype,
2875 row_scale,
2876 col_field,
2877 ..
2878 } = self
2879 else {
2880 return false;
2881 };
2882 if crate::prism::has_contract(model) {
2883 return false;
2884 }
2885 match *dtype {
2886 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2887 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2891 if *dtype == TensorDtype::Q8_2f
2894 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2895 && crate::gpu::q8_matmat_2f(
2896 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2897 )
2898 {
2899 return true;
2900 }
2901 let flat: Vec<f32> = (0..b)
2902 .flat_map(|bi| {
2903 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2904 })
2905 .collect();
2906 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2907 }
2908 _ => false,
2909 }
2910 }
2911
2912 pub fn matvec_many<const N: usize>(
2919 ts: [&QTensor; N],
2920 x: &[f32],
2921 mut outs: [&mut [f32]; N],
2922 pool: Option<&Pool>,
2923 ) {
2924 if Self::any_exact_cpu_default(&ts) {
2925 return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2926 }
2927 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2928 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2929 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2934 t.matvec(x, o, pool);
2935 }
2936 return;
2937 }
2938 let uniform_q8 = ts.iter().all(|t| {
2939 matches!(
2940 t,
2941 Self::Mapped {
2942 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2943 ..
2944 }
2945 )
2946 });
2947 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2948 if uniform_f32 && crate::f32_backend::active() {
2949 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2950 t.matvec(x, o, pool);
2951 }
2952 return;
2953 }
2954 let uniform_q4 = ts.iter().all(|t| {
2955 matches!(
2956 t,
2957 Self::Mapped {
2958 dtype: TensorDtype::Q4Block,
2959 ..
2960 }
2961 )
2962 });
2963 let uniform_vbit = ts.iter().all(|t| {
2964 matches!(
2965 t,
2966 Self::Mapped {
2967 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2968 ..
2969 }
2970 )
2971 });
2972 let uniform_q1 =
2975 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2976 let uniform_q1t = ts.iter().all(|t| {
2977 matches!(
2978 t,
2979 Self::Mapped {
2980 dtype: TensorDtype::Q1T,
2981 ..
2982 }
2983 )
2984 });
2985 let uniform_q4tp = ts.iter().all(|t| {
2990 matches!(
2991 t,
2992 Self::Mapped {
2993 dtype: TensorDtype::Q4TiledP,
2994 ..
2995 }
2996 )
2997 }) && ts
2998 .iter()
2999 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
3000 let Some(pool) = pool else {
3001 for (t, o) in ts.iter().zip(outs.iter_mut()) {
3002 t.matvec(x, o, None);
3003 }
3004 return;
3005 };
3006 if total_rows < 256
3007 || !(uniform_q8
3008 || uniform_f32
3009 || uniform_q4
3010 || uniform_vbit
3011 || uniform_q1
3012 || uniform_q1t
3013 || uniform_q4tp)
3014 {
3015 for (t, o) in ts.iter().zip(outs.iter_mut()) {
3016 t.matvec(x, o, Some(pool));
3017 }
3018 return;
3019 }
3020
3021 if uniform_q4tp {
3022 let cols = ts[0].cols();
3028 let gpr = cols / GROUP_SIZE;
3029 let views: [Q4tpView; N] =
3030 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
3031 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
3032 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3033 let locate = |flat: usize| -> (usize, usize) {
3035 let mut acc = 0;
3036 for (i, &r) in rows_of.iter().enumerate() {
3037 if flat < acc + r {
3038 return (i, flat - acc);
3039 }
3040 acc += r;
3041 }
3042 (rows_of.len() - 1, 0)
3043 };
3044 let (views, outs_addr) = (&views, &outs_addr);
3045 if a8w8_enabled() {
3046 let act = split_act(x);
3047 let act = &act;
3048 #[cfg(target_arch = "x86_64")]
3049 let lanes = q4tp_lanes();
3050 let run = |start: usize, end: usize| {
3051 #[cfg(target_arch = "x86_64")]
3052 if lanes > 1 && act.xg8.len() == gpr {
3053 with_krow(gpr * lanes, |sc| {
3056 unsafe {
3059 match lanes {
3060 2 => q4tp_many_lanes::<2>(
3061 &views[..],
3062 &locate,
3063 start,
3064 end,
3065 gpr,
3066 &act.xq,
3067 &act.xg8,
3068 act.sx,
3069 &act.outliers,
3070 sc,
3071 &outs_addr[..],
3072 ),
3073 8 => q4tp_many_lanes::<8>(
3074 &views[..],
3075 &locate,
3076 start,
3077 end,
3078 gpr,
3079 &act.xq,
3080 &act.xg8,
3081 act.sx,
3082 &act.outliers,
3083 sc,
3084 &outs_addr[..],
3085 ),
3086 _ => q4tp_many_lanes::<4>(
3087 &views[..],
3088 &locate,
3089 start,
3090 end,
3091 gpr,
3092 &act.xq,
3093 &act.xg8,
3094 act.sx,
3095 &act.outliers,
3096 sc,
3097 &outs_addr[..],
3098 ),
3099 }
3100 }
3101 });
3102 return;
3103 }
3104 crate::cpuprof::count_plain_rows((end - start) as u64);
3105 with_krow(gpr, |sc| {
3106 for flat in start..end {
3107 let (t, r) = locate(flat);
3108 let v = &views[t];
3109 v.scales_into(r, gpr, sc);
3110 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
3111 for &(j, xv) in &act.outliers {
3112 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3113 acc += w * s * xv;
3114 }
3115 unsafe { *outs_addr[t].at(r) = acc };
3117 }
3118 });
3119 };
3120 pool.run_rows(total_rows, &run);
3121 } else {
3122 let run = |start: usize, end: usize| {
3123 with_krow(gpr, |sc| {
3124 for flat in start..end {
3125 let (t, r) = locate(flat);
3126 let v = &views[t];
3127 v.scales_into(r, gpr, sc);
3128 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
3130 }
3131 });
3132 };
3133 pool.run_rows(total_rows, &run);
3134 }
3135 return;
3136 }
3137
3138 if uniform_q1 {
3139 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3142 if a8w8_enabled() {
3143 let act = split_act(x);
3144 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
3145 let (act, gsum) = (&act, &gsum);
3146 let closures: [_; N] = std::array::from_fn(|i| {
3147 let (w, gpr, out) = (
3148 ts[i].bin_w().expect("uniform_q1 checked"),
3149 ts[i].cols() / GROUP_SIZE,
3150 outs_addr[i],
3151 );
3152 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3153 });
3154 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3155 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3156 pool.run_many(&parts);
3157 } else {
3158 let closures: [_; N] = std::array::from_fn(|i| {
3159 let (w, gpr, out) = (
3160 ts[i].bin_w().expect("uniform_q1 checked"),
3161 ts[i].cols() / GROUP_SIZE,
3162 outs_addr[i],
3163 );
3164 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3165 });
3166 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3167 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3168 pool.run_many(&parts);
3169 }
3170 return;
3171 }
3172
3173 if uniform_q1t {
3174 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3178 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3179 if a8w8_enabled() {
3180 let act = split_act(x);
3181 let act = &act;
3182 let x_ref = x;
3183 let closures: [_; N] = std::array::from_fn(|i| {
3184 let bytes = ts[i].quant_bytes();
3185 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3186 let gpr = cols / GROUP_SIZE;
3187 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3188 let out = outs_addr[i];
3189 move |s: usize, e: usize| {
3190 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3191 }
3192 });
3193 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3194 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3195 pool.run_many(&parts);
3196 } else {
3197 let x_ref = x;
3198 let closures: [_; N] = std::array::from_fn(|i| {
3199 let bytes = ts[i].quant_bytes();
3200 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3201 let gpr = cols / GROUP_SIZE;
3202 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3203 let out = outs_addr[i];
3204 move |s: usize, e: usize| {
3205 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3206 }
3207 });
3208 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3209 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3210 pool.run_many(&parts);
3211 }
3212 return;
3213 }
3214
3215 if uniform_q4 || uniform_vbit {
3216 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3217 if a8w8_enabled() {
3219 let act = split_act(x);
3220 let act = &act;
3221 if uniform_q4 {
3222 let closures: [_; N] = std::array::from_fn(|i| {
3223 let (packed, scales) =
3224 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3225 let (gpr, cols, out) =
3226 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3227 move |s: usize, e: usize| {
3228 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3229 }
3230 });
3231 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3232 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3233 pool.run_many(&parts);
3234 } else {
3235 let closures: [_; N] = std::array::from_fn(|i| {
3236 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3237 unreachable!()
3238 };
3239 let (bytes, rows, cols, out) = (
3240 ts[i].quant_bytes(),
3241 ts[i].rows(),
3242 ts[i].cols(),
3243 outs_addr[i],
3244 );
3245 move |s: usize, e: usize| {
3246 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3247 }
3248 });
3249 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3250 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3251 pool.run_many(&parts);
3252 }
3253 return;
3254 }
3255 if uniform_q4 {
3256 let closures: [_; N] = std::array::from_fn(|i| {
3257 let (packed, scales) =
3258 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3259 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3260 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3261 });
3262 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3263 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3264 pool.run_many(&parts);
3265 } else {
3266 let closures: [_; N] = std::array::from_fn(|i| {
3267 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3268 unreachable!()
3269 };
3270 let (bytes, rows, cols, out) = (
3271 ts[i].quant_bytes(),
3272 ts[i].rows(),
3273 ts[i].cols(),
3274 outs_addr[i],
3275 );
3276 move |s: usize, e: usize| {
3277 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3278 }
3279 });
3280 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3281 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3282 pool.run_many(&parts);
3283 }
3284 return;
3285 }
3286
3287 if uniform_f32 {
3288 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3289 let closures: [_; N] = std::array::from_fn(|i| {
3290 let Self::F32 { data, cols, .. } = ts[i] else {
3291 unreachable!()
3292 };
3293 let out = outs_addr[i];
3294 move |start: usize, end: usize| {
3295 for o in start..end {
3296 let row = &data[o * cols..(o + 1) * cols];
3297 let mut sum = 0.0f32;
3298 for j in 0..*cols {
3299 sum += row[j] * x[j];
3300 }
3301 unsafe { *out.at(o) = sum };
3303 }
3304 }
3305 });
3306 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3307 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3308 pool.run_many(&parts);
3309 return;
3310 }
3311
3312 struct Ctx<'a> {
3315 bytes: &'a [u8],
3316 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3317 rep: &'a [u8],
3318 row_scale: &'a [f32],
3319 cols: usize,
3320 xs: std::borrow::Cow<'a, [f32]>,
3321 }
3322 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3323 let Self::Mapped {
3324 dtype,
3325 cols,
3326 row_scale,
3327 col_field,
3328 repack,
3329 ..
3330 } = ts[i]
3331 else {
3332 unreachable!()
3333 };
3334 Ctx {
3335 bytes: ts[i].quant_bytes(),
3336 rep: repack,
3337 row_scale,
3338 cols: *cols,
3339 xs: prescale(x, col_field, *dtype),
3340 }
3341 });
3342 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3343 #[cfg(target_arch = "aarch64")]
3344 if sdot_enabled() {
3345 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3346 let closures: [_; N] = std::array::from_fn(|i| {
3347 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3348 move |start: usize, end: usize| {
3349 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3350 }
3351 });
3352 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3353 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3354 pool.run_many(&parts);
3355 return;
3356 }
3357 #[cfg(target_arch = "x86_64")]
3358 if avx2_a8w8_enabled() {
3359 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3360 let closures: [_; N] = std::array::from_fn(|i| {
3361 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3362 move |start: usize, end: usize| {
3363 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3364 }
3365 });
3366 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3367 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3368 pool.run_many(&parts);
3369 return;
3370 }
3371 let closures: [_; N] = std::array::from_fn(|i| {
3372 let (c, out) = (&ctxs[i], outs_addr[i]);
3373 move |start: usize, end: usize| {
3374 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3375 }
3376 });
3377 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3378 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3379 pool.run_many(&parts);
3380 }
3381}
3382
3383impl QTensor {
3384 #[allow(clippy::needless_range_loop)]
3389 pub fn matvec2_many<const N: usize>(
3390 ts: [&QTensor; N],
3391 x1: &[f32],
3392 x2: &[f32],
3393 mut o1s: [&mut [f32]; N],
3394 mut o2s: [&mut [f32]; N],
3395 pool: Option<&Pool>,
3396 ) {
3397 if Self::any_exact_cpu_default(&ts) {
3398 return float_activations_scope(|| {
3399 Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3400 });
3401 }
3402 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3403 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3404 for i in 0..N {
3405 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3406 }
3407 return;
3408 }
3409 let uniform_q8 = ts.iter().all(|t| {
3410 matches!(
3411 t,
3412 Self::Mapped {
3413 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3414 ..
3415 }
3416 )
3417 });
3418 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3419 let uniform_q4 = ts.iter().all(|t| {
3420 matches!(
3421 t,
3422 Self::Mapped {
3423 dtype: TensorDtype::Q4Block,
3424 ..
3425 }
3426 )
3427 });
3428 let uniform_vbit = ts.iter().all(|t| {
3429 matches!(
3430 t,
3431 Self::Mapped {
3432 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3433 ..
3434 }
3435 )
3436 });
3437 let fusable = pool.is_some()
3438 && total_rows >= 256
3439 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3440 if !fusable {
3441 for i in 0..N {
3442 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3443 }
3444 return;
3445 }
3446 let pool = pool.unwrap();
3447
3448 if uniform_q4 || uniform_vbit {
3449 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3450 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3451 if a8w8_enabled() {
3453 let a1 = split_act(x1);
3454 let a2 = split_act(x2);
3455 let (a1, a2) = (&a1, &a2);
3456 if uniform_q4 {
3457 let closures: [_; N] = std::array::from_fn(|i| {
3458 let (packed, scales) =
3459 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3460 let (gpr, cols, o1, o2) =
3461 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3462 move |s: usize, e: usize| {
3463 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3464 }
3465 });
3466 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3467 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3468 pool.run_many(&parts);
3469 } else {
3470 let closures: [_; N] = std::array::from_fn(|i| {
3471 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3472 unreachable!()
3473 };
3474 let (bytes, rows, cols, o1, o2) = (
3475 ts[i].quant_bytes(),
3476 ts[i].rows(),
3477 ts[i].cols(),
3478 p1[i],
3479 p2[i],
3480 );
3481 move |s: usize, e: usize| {
3482 vbit_range2_a8w8(
3483 bytes,
3484 vbit_offsets,
3485 x1,
3486 x2,
3487 a1,
3488 a2,
3489 rows,
3490 cols,
3491 o1,
3492 o2,
3493 s,
3494 e,
3495 )
3496 }
3497 });
3498 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3499 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3500 pool.run_many(&parts);
3501 }
3502 return;
3503 }
3504 if uniform_q4 {
3505 let closures: [_; N] = std::array::from_fn(|i| {
3506 let (packed, scales) =
3507 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3508 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3509 move |s: usize, e: usize| {
3510 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3511 }
3512 });
3513 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3514 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3515 pool.run_many(&parts);
3516 } else {
3517 let closures: [_; N] = std::array::from_fn(|i| {
3518 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3519 unreachable!()
3520 };
3521 let (bytes, rows, cols, o1, o2) = (
3522 ts[i].quant_bytes(),
3523 ts[i].rows(),
3524 ts[i].cols(),
3525 p1[i],
3526 p2[i],
3527 );
3528 move |s: usize, e: usize| {
3529 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3530 }
3531 });
3532 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3533 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3534 pool.run_many(&parts);
3535 }
3536 return;
3537 }
3538
3539 if uniform_f32 {
3540 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3541 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3542 let closures: [_; N] = std::array::from_fn(|i| {
3543 let Self::F32 { data, cols, .. } = ts[i] else {
3544 unreachable!()
3545 };
3546 let (o1, o2) = (p1[i], p2[i]);
3547 move |start: usize, end: usize| {
3548 for o in start..end {
3549 let row = &data[o * cols..(o + 1) * cols];
3550 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3551 for j in 0..*cols {
3552 s1 += row[j] * x1[j];
3553 s2 += row[j] * x2[j];
3554 }
3555 unsafe {
3557 *o1.at(o) = s1;
3558 *o2.at(o) = s2;
3559 }
3560 }
3561 }
3562 });
3563 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3564 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3565 pool.run_many(&parts);
3566 return;
3567 }
3568
3569 struct Ctx<'a> {
3570 bytes: &'a [u8],
3571 row_scale: &'a [f32],
3572 cols: usize,
3573 xs1: std::borrow::Cow<'a, [f32]>,
3574 xs2: std::borrow::Cow<'a, [f32]>,
3575 }
3576 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3577 let Self::Mapped {
3578 dtype,
3579 cols,
3580 row_scale,
3581 col_field,
3582 ..
3583 } = ts[i]
3584 else {
3585 unreachable!()
3586 };
3587 Ctx {
3588 bytes: ts[i].quant_bytes(),
3589 row_scale,
3590 cols: *cols,
3591 xs1: prescale(x1, col_field, *dtype),
3592 xs2: prescale(x2, col_field, *dtype),
3593 }
3594 });
3595 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3596 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3597 #[cfg(target_arch = "aarch64")]
3598 if sdot_enabled() {
3599 let acts: [(SplitAct, SplitAct); N] =
3600 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3601 let closures: [_; N] = std::array::from_fn(|i| {
3602 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3603 move |start: usize, end: usize| {
3604 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3605 }
3606 });
3607 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3608 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3609 pool.run_many(&parts);
3610 return;
3611 }
3612 #[cfg(target_arch = "x86_64")]
3613 if avx2_a8w8_enabled() {
3614 let acts: [(SplitAct, SplitAct); N] =
3615 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3616 let closures: [_; N] = std::array::from_fn(|i| {
3617 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3618 move |start: usize, end: usize| {
3619 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3620 }
3621 });
3622 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3623 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3624 pool.run_many(&parts);
3625 return;
3626 }
3627 let closures: [_; N] = std::array::from_fn(|i| {
3628 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3629 move |start: usize, end: usize| {
3630 q8_range2_f32(
3631 c.bytes,
3632 c.row_scale,
3633 &c.xs1,
3634 &c.xs2,
3635 c.cols,
3636 o1,
3637 o2,
3638 start,
3639 end,
3640 )
3641 }
3642 });
3643 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3644 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3645 pool.run_many(&parts);
3646 }
3647
3648 pub fn matvec_silu_mul(
3653 gate: &QTensor,
3654 up: &QTensor,
3655 x: &[f32],
3656 out: &mut [f32],
3657 pool: Option<&Pool>,
3658 ) -> bool {
3659 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3660 }
3661
3662 pub fn matvec_silu_mul_limited(
3668 gate: &QTensor,
3669 up: &QTensor,
3670 x: &[f32],
3671 out: &mut [f32],
3672 limit: f32,
3673 pool: Option<&Pool>,
3674 ) -> bool {
3675 if gate.exact_cpu_default() || up.exact_cpu_default() {
3676 return float_activations_scope(|| {
3677 Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3678 });
3679 }
3680 if gate.has_prism_contract()
3681 || up.has_prism_contract()
3682 || gate.per_tensor_only()
3683 || up.per_tensor_only()
3684 {
3685 return false;
3689 }
3690 let inter = gate.rows();
3691 debug_assert_eq!(up.rows(), inter);
3692 debug_assert_eq!(out.len(), inter);
3693 debug_assert_eq!(gate.cols(), up.cols());
3694 if !a8w8_enabled() {
3695 return false;
3696 }
3697 let act = split_act(x);
3698 let act = &act;
3699 let x_ref = x;
3700 let out_addr = SendMut(out.as_mut_ptr());
3701
3702 match (gate, up) {
3703 (
3705 Self::Mapped {
3706 dtype: TensorDtype::Q4Block,
3707 ..
3708 },
3709 Self::Mapped {
3710 dtype: TensorDtype::Q4Block,
3711 ..
3712 },
3713 ) => {
3714 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3715 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3716 let gpr = gate.cols() / GROUP_SIZE;
3717 let cols = gate.cols();
3718 let run = move |start: usize, end: usize| {
3719 for r in start..end {
3720 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3721 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3722 for &(j, xv) in &act.outliers {
3723 let flat = r * cols + j;
3724 let gb = gp[flat / 2];
3725 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3726 let gsc = f16_to_f32(u16::from_le_bytes([
3727 gs[(flat / GROUP_SIZE) * 2],
3728 gs[(flat / GROUP_SIZE) * 2 + 1],
3729 ]));
3730 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3731 let ub = up_p[flat / 2];
3732 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3733 let usc = f16_to_f32(u16::from_le_bytes([
3734 up_s[(flat / GROUP_SIZE) * 2],
3735 up_s[(flat / GROUP_SIZE) * 2 + 1],
3736 ]));
3737 uv += ((un as i32 - 8) as f32) * usc * xv;
3738 }
3739 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3741 }
3742 };
3743 dispatch_rows(pool, inter, &run);
3744 true
3745 }
3746 (
3750 Self::Mapped {
3751 dtype: TensorDtype::Q4Tiled,
3752 ..
3753 },
3754 Self::Mapped {
3755 dtype: TensorDtype::Q4Tiled,
3756 ..
3757 },
3758 ) => {
3759 let g_bytes = gate.quant_bytes();
3760 let u_bytes = up.quant_bytes();
3761 let gpr = gate.cols() / GROUP_SIZE;
3762 let run = move |start: usize, end: usize| {
3763 for r in start..end {
3764 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3765 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3766 for &(j, xv) in &act.outliers {
3767 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3768 gv += w * s * xv;
3769 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3770 uv += w * s * xv;
3771 }
3772 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3774 }
3775 };
3776 dispatch_rows(pool, inter, &run);
3777 true
3778 }
3779 (
3782 Self::Mapped {
3783 dtype: TensorDtype::Q4TiledP,
3784 ..
3785 },
3786 Self::Mapped {
3787 dtype: TensorDtype::Q4TiledP,
3788 ..
3789 },
3790 ) => {
3791 let cols = gate.cols();
3792 let gpr = cols / GROUP_SIZE;
3793 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3794 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3795 #[cfg(target_arch = "x86_64")]
3796 let lanes = q4tp_lanes();
3797 let run = |start: usize, end: usize| {
3798 #[cfg(target_arch = "x86_64")]
3799 if lanes > 1 && act.xg8.len() == gpr {
3800 with_krows(gpr * lanes, |gsc, usc| {
3803 unsafe {
3805 match lanes {
3806 2 => q4tp_gateup_lanes::<2>(
3807 &gv_view, &uv_view, start, end, gpr, &act.xq,
3808 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3809 out_addr,
3810 ),
3811 8 => q4tp_gateup_lanes::<8>(
3812 &gv_view, &uv_view, start, end, gpr, &act.xq,
3813 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3814 out_addr,
3815 ),
3816 _ => q4tp_gateup_lanes::<4>(
3817 &gv_view, &uv_view, start, end, gpr, &act.xq,
3818 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3819 out_addr,
3820 ),
3821 }
3822 }
3823 });
3824 return;
3825 }
3826 crate::cpuprof::count_plain_rows((end - start) as u64);
3827 with_krows(gpr, |gsc, usc| {
3828 for r in start..end {
3829 gv_view.scales_into(r, gpr, gsc);
3830 uv_view.scales_into(r, gpr, usc);
3831 let mut gv =
3832 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3833 let mut uv =
3834 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3835 for &(j, xv) in &act.outliers {
3836 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3837 gv += w * s * xv;
3838 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3839 uv += w * s * xv;
3840 }
3841 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3843 }
3844 });
3845 };
3846 dispatch_rows(pool, inter, &run);
3847 true
3848 }
3849 (
3856 Self::Mapped {
3857 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3858 ..
3859 },
3860 Self::Mapped {
3861 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3862 ..
3863 },
3864 ) => {
3865 let gw = gate.bin_w().expect("binary dtype matched");
3866 let uw = up.bin_w().expect("binary dtype matched");
3867 let gpr = gate.cols() / GROUP_SIZE;
3868 let gsum = q1_group_sums(&act.xq, gpr);
3869 let gsum = &gsum;
3870 let run = move |start: usize, end: usize| {
3871 for r in start..end {
3872 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3873 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3874 for &(j, xv) in &act.outliers {
3875 let (w, s) = gw.outlier(r, gpr, j);
3876 gv += w * s * xv;
3877 let (w, s) = uw.outlier(r, gpr, j);
3878 uv += w * s * xv;
3879 }
3880 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3882 }
3883 };
3884 dispatch_rows(pool, inter, &run);
3885 true
3886 }
3887 (
3891 Self::Mapped {
3892 dtype: TensorDtype::Q2TiledP,
3893 ..
3894 },
3895 Self::Mapped {
3896 dtype: TensorDtype::Q2TiledP,
3897 ..
3898 },
3899 ) => {
3900 let cols = gate.cols();
3901 let gpr = cols / GROUP_SIZE;
3902 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3903 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3904 let gsum = q1_group_sums(&act.xq, gpr);
3905 let gsum = &gsum;
3906 let run = move |start: usize, end: usize| {
3907 with_krows(gpr, |gsc, usc| {
3908 for r in start..end {
3909 gv_view.scales_into(r, gpr, gsc);
3910 uv_view.scales_into(r, gpr, usc);
3911 let mut gv =
3912 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3913 let mut uv =
3914 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3915 for &(j, xv) in &act.outliers {
3916 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3917 gv += w * s * xv;
3918 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3919 uv += w * s * xv;
3920 }
3921 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3923 }
3924 });
3925 };
3926 dispatch_rows(pool, inter, &run);
3927 true
3928 }
3929 (
3934 Self::Mapped {
3935 dtype: TensorDtype::Q8Row,
3936 row_scale: g_rs,
3937 ..
3938 },
3939 Self::Mapped {
3940 dtype: TensorDtype::Q8Row,
3941 row_scale: u_rs,
3942 ..
3943 },
3944 ) => {
3945 let g_bytes = gate.quant_bytes();
3946 let u_bytes = up.quant_bytes();
3947 let cols = gate.cols();
3948 let run = move |start: usize, end: usize| {
3949 for r in start..end {
3950 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3951 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3952 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3954 }
3955 };
3956 dispatch_rows(pool, inter, &run);
3957 true
3958 }
3959 (
3961 Self::Mapped {
3962 dtype: TensorDtype::Q1T,
3963 ..
3964 },
3965 Self::Mapped {
3966 dtype: TensorDtype::Q1T,
3967 ..
3968 },
3969 ) => {
3970 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3971 let g_bytes = gate.quant_bytes();
3972 let u_bytes = up.quant_bytes();
3973 let gpr = gate.cols() / GROUP_SIZE;
3974 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3975 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3976 let run = move |start: usize, end: usize| {
3977 for r in start..end {
3978 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3979 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3980 for &(j, xv) in &act.outliers {
3981 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3982 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3983 }
3984 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3985 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3986 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3988 }
3989 };
3990 dispatch_rows(pool, inter, &run);
3991 true
3992 }
3993 _ => false,
3994 }
3995 }
3996
3997 pub fn moe_gate_up_many(
4012 pairs: &[(&QTensor, &QTensor)],
4013 x: &[f32],
4014 outs: &mut [Vec<f32>],
4015 pool: Option<&Pool>,
4016 ) -> bool {
4017 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
4018 }
4019
4020 pub fn moe_gate_up_many_limited(
4024 pairs: &[(&QTensor, &QTensor)],
4025 x: &[f32],
4026 outs: &mut [Vec<f32>],
4027 limit: f32,
4028 pool: Option<&Pool>,
4029 ) -> bool {
4030 if Self::any_exact_cpu_default_pairs(pairs) {
4031 return float_activations_scope(|| {
4032 Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
4033 });
4034 }
4035 if pairs.is_empty() || pairs.len() != outs.len() {
4036 return false;
4037 }
4038 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4039 return false;
4040 }
4041 if !a8w8_enabled() {
4042 if limit > 0.0 {
4043 return false;
4046 }
4047 let groups = vec![vec![0]; pairs.len()];
4048 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
4049 }
4050 let inter = pairs[0].0.rows();
4051 let cols = pairs[0].0.cols();
4052 if cols % GROUP_SIZE != 0 {
4053 return false;
4054 }
4055 let gpr = cols / GROUP_SIZE;
4056 let q2 = matches!(
4059 pairs[0].0,
4060 Self::Mapped {
4061 dtype: TensorDtype::Q2TiledP,
4062 ..
4063 }
4064 );
4065 let want = if q2 {
4066 TensorDtype::Q2TiledP
4067 } else {
4068 TensorDtype::Q4TiledP
4069 };
4070 let mut views = Vec::with_capacity(pairs.len() * 2);
4071 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
4072 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
4073 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
4074 if !both
4075 || g.rows() != inter
4076 || u.rows() != inter
4077 || g.cols() != cols
4078 || u.cols() != cols
4079 || o.len() != inter
4080 {
4081 return false;
4082 }
4083 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
4084 views.push(mk(g.quant_bytes(), inter, cols));
4085 views.push(mk(u.quant_bytes(), inter, cols));
4086 }
4087 let act = split_act(x);
4088 let gsum = if q2 {
4089 q1_group_sums(&act.xq, gpr)
4090 } else {
4091 Vec::new()
4092 };
4093 let (act, gsum) = (&act, &gsum);
4094 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4095 let (views, ptrs) = (&views, &ptrs);
4096 let run = |start: usize, end: usize| {
4097 with_krows(gpr, |gsc, usc| {
4098 for flat in start..end {
4099 let (e, r) = (flat / inter, flat % inter);
4100 let gv_view = &views[e * 2];
4101 let uv_view = &views[e * 2 + 1];
4102 gv_view.scales_into(r, gpr, gsc);
4103 uv_view.scales_into(r, gpr, usc);
4104 let (mut gv, mut uv) = if q2 {
4105 (
4106 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
4107 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
4108 )
4109 } else {
4110 (
4111 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
4112 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
4113 )
4114 };
4115 for &(j, xv) in &act.outliers {
4116 let (og, ou) = if q2 {
4117 (
4118 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
4119 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
4120 )
4121 } else {
4122 (
4123 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
4124 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
4125 )
4126 };
4127 gv += og.0 * og.1 * xv;
4128 uv += ou.0 * ou.1 * xv;
4129 }
4130 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
4132 }
4133 });
4134 };
4135 dispatch_rows(pool, pairs.len() * inter, &run);
4136 true
4137 }
4138
4139 pub fn moe_down_many(
4148 downs: &[&QTensor],
4149 gs: &[Vec<f32>],
4150 weights: &[f32],
4151 out: &mut [f32],
4152 pool: Option<&Pool>,
4153 ) -> bool {
4154 if Self::any_exact_cpu_default(downs) {
4155 return float_activations_scope(|| {
4156 Self::moe_down_many(downs, gs, weights, out, pool)
4157 });
4158 }
4159 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
4160 return false;
4161 }
4162 if downs.iter().any(|d| d.per_tensor_only()) {
4163 return false;
4164 }
4165 if !a8w8_enabled() {
4166 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
4167 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
4168 return false;
4169 }
4170 out.fill(0.0);
4171 for (row, &w) in terms.iter().zip(weights) {
4172 for (o, &v) in out.iter_mut().zip(row) {
4173 *o += w * v;
4174 }
4175 }
4176 return true;
4177 }
4178 let rows = out.len();
4179 let cols = downs[0].cols();
4180 if cols % GROUP_SIZE != 0 {
4181 return false;
4182 }
4183 let gpr = cols / GROUP_SIZE;
4184 let mut views = Vec::with_capacity(downs.len());
4185 for (d, g) in downs.iter().zip(gs.iter()) {
4186 if !matches!(
4187 d,
4188 Self::Mapped {
4189 dtype: TensorDtype::Q4TiledP,
4190 ..
4191 }
4192 ) || d.rows() != rows
4193 || d.cols() != cols
4194 || g.len() != cols
4195 {
4196 return false;
4197 }
4198 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4199 }
4200 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4202 let out_addr = SendMut(out.as_mut_ptr());
4210 let (views, acts, weights) = (&views, &acts, &weights);
4211 let run = |start: usize, end: usize| {
4212 with_krow(gpr, |sc| {
4213 for r in start..end {
4214 let mut acc = 0f32;
4215 for (e, v) in views.iter().enumerate() {
4216 v.scales_into(r, gpr, sc);
4217 let a = &acts[e];
4218 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4219 for &(j, xv) in &a.outliers {
4220 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4221 d += w * s * xv;
4222 }
4223 acc += weights[e] * d;
4224 }
4225 unsafe { *out_addr.at(r) = acc };
4227 }
4228 });
4229 };
4230 dispatch_rows(pool, rows, &run);
4231 true
4232 }
4233
4234 pub fn moe_gate_up_rows(
4246 pairs: &[(&QTensor, &QTensor)],
4247 groups: &[Vec<usize>],
4248 xs: &[f32],
4249 outs: &mut [Vec<f32>],
4250 pool: Option<&Pool>,
4251 ) -> bool {
4252 if Self::any_exact_cpu_default_pairs(pairs) {
4253 return float_activations_scope(|| {
4254 Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4255 });
4256 }
4257 if pairs.is_empty() || pairs.len() != groups.len() {
4258 return false;
4259 }
4260 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4261 return false;
4262 }
4263 let inter = pairs[0].0.rows();
4264 let cols = pairs[0].0.cols();
4265 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4266 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4267 return false;
4268 }
4269 let b = xs.len() / cols;
4270 let gpr = cols / GROUP_SIZE;
4271 let mut views = Vec::with_capacity(pairs.len() * 2);
4272 for (g, u) in pairs {
4273 let q4tp = |t: &QTensor| {
4274 matches!(
4275 t,
4276 Self::Mapped {
4277 dtype: TensorDtype::Q4TiledP,
4278 ..
4279 }
4280 )
4281 };
4282 if g.has_prism_contract()
4283 || u.has_prism_contract()
4284 || !q4tp(g)
4285 || !q4tp(u)
4286 || g.rows() != inter
4287 || u.rows() != inter
4288 || g.cols() != cols
4289 || u.cols() != cols
4290 {
4291 return false;
4292 }
4293 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4294 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4295 }
4296 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4297 return false;
4298 }
4299 let quantized = a8w8_enabled();
4300 let acts: Vec<SplitAct> = if quantized {
4301 (0..b)
4302 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4303 .collect()
4304 } else {
4305 Vec::new()
4306 };
4307 let mut offs = Vec::with_capacity(groups.len());
4308 let mut o = 0usize;
4309 for g in groups {
4310 offs.push(o);
4311 o += g.len();
4312 }
4313 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4314 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4315 let run = |start: usize, end: usize| {
4316 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4317 for flat in start..end {
4318 let (e, r) = (flat / inter, flat % inter);
4319 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4320 gv_view.scales_into(r, gpr, &mut gsc);
4321 uv_view.scales_into(r, gpr, &mut usc);
4322 for (k, &t) in groups[e].iter().enumerate() {
4323 if !quantized {
4324 let x = &xs[t * cols..(t + 1) * cols];
4325 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4326 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4327 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4328 continue;
4329 }
4330 let act = &acts[t];
4331 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4332 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4333 for &(j, xv) in &act.outliers {
4334 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4335 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4336 gv += og.0 * og.1 * xv;
4337 uv += ou.0 * ou.1 * xv;
4338 }
4339 let silu_g = gv / (1.0 + (-gv).exp());
4340 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4343 }
4344 }
4345 };
4346 dispatch_rows(pool, pairs.len() * inter, &run);
4347 true
4348 }
4349
4350 pub fn moe_down_rows(
4357 downs: &[&QTensor],
4358 group_lens: &[usize],
4359 gs: &[Vec<f32>],
4360 outs: &mut [Vec<f32>],
4361 pool: Option<&Pool>,
4362 ) -> bool {
4363 if Self::any_exact_cpu_default(downs) {
4364 return float_activations_scope(|| {
4365 Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4366 });
4367 }
4368 if downs.is_empty() || downs.len() != group_lens.len() {
4369 return false;
4370 }
4371 if downs.iter().any(|d| d.per_tensor_only()) {
4372 return false;
4373 }
4374 let rows = downs[0].rows();
4375 let cols = downs[0].cols();
4376 let n_pairs: usize = group_lens.iter().sum();
4377 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4378 return false;
4379 }
4380 let gpr = cols / GROUP_SIZE;
4381 let mut views = Vec::with_capacity(downs.len());
4382 for d in downs {
4383 if d.has_prism_contract()
4384 || !matches!(
4385 d,
4386 Self::Mapped {
4387 dtype: TensorDtype::Q4TiledP,
4388 ..
4389 }
4390 )
4391 || d.rows() != rows
4392 || d.cols() != cols
4393 {
4394 return false;
4395 }
4396 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4397 }
4398 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4399 return false;
4400 }
4401 let quantized = a8w8_enabled();
4402 let acts: Vec<SplitAct> = if quantized {
4403 gs.iter().map(|g| split_act(g)).collect()
4404 } else {
4405 Vec::new()
4406 };
4407 let mut offs = Vec::with_capacity(group_lens.len());
4408 let mut o = 0usize;
4409 for &l in group_lens {
4410 offs.push(o);
4411 o += l;
4412 }
4413 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4414 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4415 let run = |start: usize, end: usize| {
4416 let mut sc = vec![0f32; gpr];
4417 for flat in start..end {
4418 let (e, r) = (flat / rows, flat % rows);
4419 let v = &views[e];
4420 v.scales_into(r, gpr, &mut sc);
4421 for k in 0..group_lens[e] {
4422 if !quantized {
4423 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4424 unsafe { *ptrs[offs[e] + k].at(r) = d };
4425 continue;
4426 }
4427 let a = &acts[offs[e] + k];
4428 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4429 for &(j, xv) in &a.outliers {
4430 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4431 d += w * s * xv;
4432 }
4433 unsafe { *ptrs[offs[e] + k].at(r) = d };
4435 }
4436 }
4437 };
4438 dispatch_rows(pool, downs.len() * rows, &run);
4439 true
4440 }
4441
4442 pub fn moe_expert_rows_fast(
4456 experts: &[(&QTensor, &QTensor, &QTensor)],
4457 groups: &[Vec<usize>],
4458 xs: &[f32],
4459 out: &mut [f32],
4460 pool: Option<&Pool>,
4461 ) -> bool {
4462 #[cfg(target_arch = "x86_64")]
4463 {
4464 moe_fast::run(experts, groups, xs, out, pool)
4465 }
4466 #[cfg(not(target_arch = "x86_64"))]
4467 {
4468 let _ = (experts, groups, xs, out, pool);
4469 false
4470 }
4471 }
4472
4473 pub fn moe_expert_rows_fast_prof(on: bool) -> bool {
4475 #[cfg(target_arch = "x86_64")]
4476 {
4477 moe_fast::PROF_ON.swap(on, Ordering::Relaxed)
4478 }
4479 #[cfg(not(target_arch = "x86_64"))]
4480 {
4481 let _ = on;
4482 false
4483 }
4484 }
4485
4486 pub fn moe_expert_rows_fast_prof_take() -> Option<String> {
4489 #[cfg(target_arch = "x86_64")]
4490 {
4491 let p = moe_fast::prof_take();
4492 let (calls, parts, experts) = (p[0], p[1], p[2]);
4493 if calls == 0 || parts == 0 || experts == 0 {
4494 return None;
4495 }
4496 let us = |ns: u64| ns as f64 / 1e3 / experts as f64;
4498 let call = |ns: u64| ns as f64 / 1e3 / calls as f64;
4499 Some(format!(
4500 "{calls} calls of {:.1} threads, {experts} experts, {} pairs, outliers {:.1} a token row and {:.1} a gate row; \
4501 a call {:.0} us of setup ({:.0} the tokens' split) and {:.0} us of wall; thread time an expert, us: \
4502 entry {:.0}, gate/up {:.0}, gate rows {:.0}, down {:.0}, waiting {:.0}",
4503 parts as f64 / calls as f64,
4504 p[3],
4505 p[5] as f64 / p[4].max(1) as f64,
4506 p[6] as f64 / p[3].max(1) as f64,
4507 call(p[7]),
4508 call(p[8]),
4509 call(p[9]),
4510 us(p[10]),
4511 us(p[11]),
4512 us(p[12]),
4513 us(p[13]),
4514 us(p[14]),
4515 ))
4516 }
4517 #[cfg(not(target_arch = "x86_64"))]
4518 None
4519 }
4520
4521 pub fn moe_expert_rows_fast_tune(knob: usize, value: usize) -> usize {
4526 #[cfg(target_arch = "x86_64")]
4527 {
4528 moe_fast::set_tune(knob, value)
4529 }
4530 #[cfg(not(target_arch = "x86_64"))]
4531 {
4532 let _ = (knob, value);
4533 0
4534 }
4535 }
4536
4537 pub fn moe_expert_rows_fast_available() -> bool {
4540 #[cfg(target_arch = "x86_64")]
4541 {
4542 avx512vnni_enabled() && a8w8_enabled()
4543 }
4544 #[cfg(not(target_arch = "x86_64"))]
4545 {
4546 false
4547 }
4548 }
4549}
4550
4551#[cfg(target_os = "macos")]
4556mod accel_blas {
4557 #[link(name = "Accelerate", kind = "framework")]
4558 unsafe extern "C" {
4559 pub fn cblas_sgemm(
4560 order: i32,
4561 trans_a: i32,
4562 trans_b: i32,
4563 m: i32,
4564 n: i32,
4565 k: i32,
4566 alpha: f32,
4567 a: *const f32,
4568 lda: i32,
4569 b: *const f32,
4570 ldb: i32,
4571 beta: f32,
4572 c: *mut f32,
4573 ldc: i32,
4574 );
4575 }
4576}
4577
4578#[cfg(target_os = "macos")]
4579pub(crate) fn accel_gemm_enabled() -> bool {
4580 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4581 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4582}
4583
4584#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4587pub(crate) fn accel_gemm_enabled() -> bool {
4588 true
4589}
4590
4591#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4595pub(crate) fn accel_gemm_enabled() -> bool {
4596 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4597 *ON.get_or_init(|| {
4598 std::is_x86_feature_detected!("avx2")
4599 && std::is_x86_feature_detected!("fma")
4600 && std::env::var("CMF_CPU_ATTEND").map(|v| v != "0").unwrap_or(true)
4601 })
4602}
4603
4604#[cfg(not(any(target_os = "macos", target_arch = "aarch64", target_arch = "x86_64")))]
4606#[allow(dead_code)]
4607pub(crate) fn accel_gemm_enabled() -> bool {
4608 false
4609}
4610
4611#[cfg(target_arch = "aarch64")]
4618#[allow(clippy::too_many_arguments)]
4619pub(crate) fn neon_gemm_rm(
4620 m: usize,
4621 n: usize,
4622 k: usize,
4623 alpha: f32,
4624 a: &[f32],
4625 lda: usize,
4626 b_mat: &[f32],
4627 ldb: usize,
4628 b_rows_are_n: bool,
4629 c: &mut [f32],
4630 ldc: usize,
4631) {
4632 debug_assert!(a.len() >= (m - 1) * lda + k);
4633 debug_assert!(c.len() >= (m - 1) * ldc + n);
4634 unsafe {
4636 use core::arch::aarch64::*;
4637 let mut i = 0usize;
4638 while i < m {
4639 let mi = (m - i).min(4);
4640 let mut j = 0usize;
4641 while j < n {
4642 let nj = (n - j).min(8);
4643 if mi == 4 && nj == 8 {
4644 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4645 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4646 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4647 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4648 for p in 0..k {
4649 let (b0, b1) = if b_rows_are_n {
4650 let base = b_mat.as_ptr().add(j * ldb + p);
4653 let g = |o: usize| *base.add(o * ldb);
4654 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4655 } else {
4656 let base = b_mat.as_ptr().add(p * ldb + j);
4657 (
4658 [*base, *base.add(1), *base.add(2), *base.add(3)],
4659 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4660 )
4661 };
4662 let bv0 = vld1q_f32(b0.as_ptr());
4663 let bv1 = vld1q_f32(b1.as_ptr());
4664 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4665 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4666 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4667 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4668 c0a = vfmaq_f32(c0a, a0, bv0);
4669 c0b = vfmaq_f32(c0b, a0, bv1);
4670 c1a = vfmaq_f32(c1a, a1, bv0);
4671 c1b = vfmaq_f32(c1b, a1, bv1);
4672 c2a = vfmaq_f32(c2a, a2, bv0);
4673 c2b = vfmaq_f32(c2b, a2, bv1);
4674 c3a = vfmaq_f32(c3a, a3, bv0);
4675 c3b = vfmaq_f32(c3b, a3, bv1);
4676 }
4677 let al = vdupq_n_f32(alpha);
4678 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4679 .iter()
4680 .enumerate()
4681 {
4682 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4683 vst1q_f32(dst, vmulq_f32(*ca, al));
4684 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4685 }
4686 } else {
4687 for r in 0..mi {
4688 for q in 0..nj {
4689 let mut acc = 0f32;
4690 for p in 0..k {
4691 let bv = if b_rows_are_n {
4692 b_mat[(j + q) * ldb + p]
4693 } else {
4694 b_mat[p * ldb + j + q]
4695 };
4696 acc += a[(i + r) * lda + p] * bv;
4697 }
4698 c[(i + r) * ldc + j + q] = acc * alpha;
4699 }
4700 }
4701 }
4702 j += nj;
4703 }
4704 i += mi;
4705 }
4706 }
4707}
4708
4709#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4711#[allow(clippy::too_many_arguments)]
4712pub(crate) fn sgemm_rm(
4713 m: usize,
4714 n: usize,
4715 k: usize,
4716 alpha: f32,
4717 a: &[f32],
4718 lda: usize,
4719 b_mat: &[f32],
4720 ldb: usize,
4721 b_rows_are_n: bool,
4722 c: &mut [f32],
4723 ldc: usize,
4724) {
4725 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4726}
4727
4728#[cfg(target_arch = "aarch64")]
4731#[allow(clippy::too_many_arguments)]
4732pub(crate) fn micro_gemm_rm(
4733 m: usize,
4734 n: usize,
4735 k: usize,
4736 alpha: f32,
4737 a: &[f32],
4738 lda: usize,
4739 b_mat: &[f32],
4740 ldb: usize,
4741 b_rows_are_n: bool,
4742 c: &mut [f32],
4743 ldc: usize,
4744) {
4745 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4746}
4747
4748#[cfg(target_arch = "x86_64")]
4750#[allow(clippy::too_many_arguments)]
4751pub(crate) fn micro_gemm_rm(
4752 m: usize,
4753 n: usize,
4754 k: usize,
4755 alpha: f32,
4756 a: &[f32],
4757 lda: usize,
4758 b_mat: &[f32],
4759 ldb: usize,
4760 b_rows_are_n: bool,
4761 c: &mut [f32],
4762 ldc: usize,
4763) {
4764 x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4765}
4766
4767#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4769#[allow(clippy::too_many_arguments)]
4770pub(crate) fn sgemm_rm(
4771 m: usize,
4772 n: usize,
4773 k: usize,
4774 alpha: f32,
4775 a: &[f32],
4776 lda: usize,
4777 b_mat: &[f32],
4778 ldb: usize,
4779 b_rows_are_n: bool,
4780 c: &mut [f32],
4781 ldc: usize,
4782) {
4783 x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4784}
4785
4786#[cfg(target_arch = "x86_64")]
4799#[allow(clippy::too_many_arguments)]
4800pub(crate) fn x86_gemm_rm(
4801 m: usize,
4802 n: usize,
4803 k: usize,
4804 alpha: f32,
4805 a: &[f32],
4806 lda: usize,
4807 b_mat: &[f32],
4808 ldb: usize,
4809 b_rows_are_n: bool,
4810 c: &mut [f32],
4811 ldc: usize,
4812) {
4813 if m == 0 || n == 0 {
4814 return;
4815 }
4816 debug_assert!(a.len() >= (m - 1) * lda + k);
4817 debug_assert!(c.len() >= (m - 1) * ldc + n);
4818 if b_rows_are_n {
4819 debug_assert!(k == 0 || b_mat.len() >= (n - 1) * ldb + k);
4820 let mut bt = vec![0f32; k * n];
4821 for j in 0..n {
4822 for (p, &v) in b_mat[j * ldb..j * ldb + k].iter().enumerate() {
4823 bt[p * n + j] = v;
4824 }
4825 }
4826 return x86_gemm_rm(m, n, k, alpha, a, lda, &bt, n, false, c, ldc);
4827 }
4828 if k == 0 {
4829 for r in 0..m {
4830 c[r * ldc..r * ldc + n].fill(0.0);
4831 }
4832 return;
4833 }
4834 debug_assert!(b_mat.len() >= (k - 1) * ldb + n);
4835 if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
4836 unsafe { x86_gemm_avx2(m, n, k, alpha, a, lda, b_mat, ldb, c, ldc) }
4839 } else {
4840 for r in 0..m {
4841 for j in 0..n {
4842 let mut acc = 0f32;
4843 for p in 0..k {
4844 acc += a[r * lda + p] * b_mat[p * ldb + j];
4845 }
4846 c[r * ldc + j] = acc * alpha;
4847 }
4848 }
4849 }
4850}
4851
4852#[cfg(target_arch = "x86_64")]
4861#[target_feature(enable = "avx2,fma")]
4862#[allow(clippy::too_many_arguments)]
4863unsafe fn x86_gemm_avx2(
4864 m: usize,
4865 n: usize,
4866 k: usize,
4867 alpha: f32,
4868 a: &[f32],
4869 lda: usize,
4870 b: &[f32],
4871 ldb: usize,
4872 c: &mut [f32],
4873 ldc: usize,
4874) {
4875 use core::arch::x86_64::*;
4876 const NC: usize = 128;
4877 const KC: usize = 256;
4878 unsafe {
4879 let ap = a.as_ptr();
4880 let bp = b.as_ptr();
4881 let cp = c.as_mut_ptr();
4882 let al = _mm256_set1_ps(alpha);
4883 let mut jc = 0usize;
4884 while jc < n {
4885 let jend = (jc + NC).min(n);
4886 let mut pc = 0usize;
4887 while pc < k {
4888 let pend = (pc + KC).min(k);
4889 let (first, last) = (pc == 0, pend == k);
4890 let mut i = 0usize;
4891 while i < m {
4892 let mi = (m - i).min(4);
4893 let mut j = jc;
4894 while j < jend {
4895 let nj = (jend - j).min(16);
4896 if mi == 4 && nj == 16 {
4897 let a0 = ap.add(i * lda);
4898 let a1 = ap.add((i + 1) * lda);
4899 let a2 = ap.add((i + 2) * lda);
4900 let a3 = ap.add((i + 3) * lda);
4901 let c0 = cp.add(i * ldc + j);
4902 let c1 = cp.add((i + 1) * ldc + j);
4903 let c2 = cp.add((i + 2) * ldc + j);
4904 let c3 = cp.add((i + 3) * ldc + j);
4905 let (mut c00, mut c01, mut c10, mut c11, mut c20, mut c21, mut c30, mut c31) =
4906 if first {
4907 let z = _mm256_setzero_ps();
4908 (z, z, z, z, z, z, z, z)
4909 } else {
4910 (
4911 _mm256_loadu_ps(c0),
4912 _mm256_loadu_ps(c0.add(8)),
4913 _mm256_loadu_ps(c1),
4914 _mm256_loadu_ps(c1.add(8)),
4915 _mm256_loadu_ps(c2),
4916 _mm256_loadu_ps(c2.add(8)),
4917 _mm256_loadu_ps(c3),
4918 _mm256_loadu_ps(c3.add(8)),
4919 )
4920 };
4921 for p in pc..pend {
4922 let bb = bp.add(p * ldb + j);
4923 let b0 = _mm256_loadu_ps(bb);
4924 let b1 = _mm256_loadu_ps(bb.add(8));
4925 let x0 = _mm256_broadcast_ss(&*a0.add(p));
4926 c00 = _mm256_fmadd_ps(x0, b0, c00);
4927 c01 = _mm256_fmadd_ps(x0, b1, c01);
4928 let x1 = _mm256_broadcast_ss(&*a1.add(p));
4929 c10 = _mm256_fmadd_ps(x1, b0, c10);
4930 c11 = _mm256_fmadd_ps(x1, b1, c11);
4931 let x2 = _mm256_broadcast_ss(&*a2.add(p));
4932 c20 = _mm256_fmadd_ps(x2, b0, c20);
4933 c21 = _mm256_fmadd_ps(x2, b1, c21);
4934 let x3 = _mm256_broadcast_ss(&*a3.add(p));
4935 c30 = _mm256_fmadd_ps(x3, b0, c30);
4936 c31 = _mm256_fmadd_ps(x3, b1, c31);
4937 }
4938 if last {
4939 c00 = _mm256_mul_ps(c00, al);
4940 c01 = _mm256_mul_ps(c01, al);
4941 c10 = _mm256_mul_ps(c10, al);
4942 c11 = _mm256_mul_ps(c11, al);
4943 c20 = _mm256_mul_ps(c20, al);
4944 c21 = _mm256_mul_ps(c21, al);
4945 c30 = _mm256_mul_ps(c30, al);
4946 c31 = _mm256_mul_ps(c31, al);
4947 }
4948 _mm256_storeu_ps(c0, c00);
4949 _mm256_storeu_ps(c0.add(8), c01);
4950 _mm256_storeu_ps(c1, c10);
4951 _mm256_storeu_ps(c1.add(8), c11);
4952 _mm256_storeu_ps(c2, c20);
4953 _mm256_storeu_ps(c2.add(8), c21);
4954 _mm256_storeu_ps(c3, c30);
4955 _mm256_storeu_ps(c3.add(8), c31);
4956 } else {
4957 for r in 0..mi {
4960 let ar = ap.add((i + r) * lda);
4961 let cr = cp.add((i + r) * ldc + j);
4962 let mut q = 0usize;
4963 while q + 8 <= nj {
4964 let mut acc = if first {
4965 _mm256_setzero_ps()
4966 } else {
4967 _mm256_loadu_ps(cr.add(q))
4968 };
4969 for p in pc..pend {
4970 acc = _mm256_fmadd_ps(
4971 _mm256_broadcast_ss(&*ar.add(p)),
4972 _mm256_loadu_ps(bp.add(p * ldb + j + q)),
4973 acc,
4974 );
4975 }
4976 if last {
4977 acc = _mm256_mul_ps(acc, al);
4978 }
4979 _mm256_storeu_ps(cr.add(q), acc);
4980 q += 8;
4981 }
4982 while q < nj {
4983 let mut acc = if first { 0.0 } else { *cr.add(q) };
4984 for p in pc..pend {
4985 acc += *ar.add(p) * *bp.add(p * ldb + j + q);
4986 }
4987 if last {
4988 acc *= alpha;
4989 }
4990 *cr.add(q) = acc;
4991 q += 1;
4992 }
4993 }
4994 }
4995 j += nj;
4996 }
4997 i += mi;
4998 }
4999 pc = pend;
5000 }
5001 jc = jend;
5002 }
5003 }
5004}
5005
5006#[allow(clippy::too_many_arguments)]
5010pub fn sgemm_public(
5011 m: usize,
5012 n: usize,
5013 k: usize,
5014 alpha: f32,
5015 a: &[f32],
5016 lda: usize,
5017 b_mat: &[f32],
5018 ldb: usize,
5019 b_rows_are_n: bool,
5020 c: &mut [f32],
5021 ldc: usize,
5022) {
5023 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
5024 {
5025 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5026 }
5027 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
5032 {
5033 if m == 0 || n == 0 {
5034 return;
5035 }
5036 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
5037 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
5038 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
5039 unsafe {
5041 matrixmultiply::sgemm(
5042 m,
5043 k,
5044 n,
5045 alpha,
5046 a.as_ptr(),
5047 lda as isize,
5048 1,
5049 b_mat.as_ptr(),
5050 rsb,
5051 csb,
5052 0.0,
5053 c.as_mut_ptr(),
5054 ldc as isize,
5055 1,
5056 );
5057 }
5058 }
5059}
5060
5061#[cfg(target_os = "macos")]
5064#[allow(clippy::too_many_arguments)]
5065pub(crate) fn sgemm_rm(
5066 m: usize,
5067 n: usize,
5068 k: usize,
5069 alpha: f32,
5070 a: &[f32],
5071 lda: usize,
5072 b_mat: &[f32],
5073 ldb: usize,
5074 b_rows_are_n: bool,
5075 c: &mut [f32],
5076 ldc: usize,
5077) {
5078 debug_assert!(a.len() >= (m - 1) * lda + k);
5079 debug_assert!(c.len() >= (m - 1) * ldc + n);
5080 #[cfg(target_arch = "aarch64")]
5085 if std::env::var("CMF_FORCE_NEON_GEMM")
5086 .map(|v| v == "1")
5087 .unwrap_or(false)
5088 {
5089 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5090 }
5091 unsafe {
5092 accel_blas::cblas_sgemm(
5093 101, 111, if b_rows_are_n { 112 } else { 111 },
5096 m as i32,
5097 n as i32,
5098 k as i32,
5099 alpha,
5100 a.as_ptr(),
5101 lda as i32,
5102 b_mat.as_ptr(),
5103 ldb as i32,
5104 0.0,
5105 c.as_mut_ptr(),
5106 ldc as i32,
5107 );
5108 }
5109}
5110
5111#[cfg(target_os = "macos")]
5118fn qmatmat_accel(
5119 q: &[u8],
5120 row_scale: &[f32],
5121 pre: &[std::borrow::Cow<'_, [f32]>],
5122 rows: usize,
5123 cols: usize,
5124 out: &mut [f32],
5125 pool: Option<&Pool>,
5126) {
5127 const TR: usize = 2048;
5132 let b = pre.len();
5133 thread_local! {
5134 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5135 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5136 }
5137 XPANEL.with(|xp| {
5138 WTILE.with(|wt| {
5139 let mut xpanel = xp.borrow_mut();
5140 xpanel.clear();
5141 for x in pre {
5142 xpanel.extend_from_slice(x);
5143 }
5144 let mut wtile = wt.borrow_mut();
5145 wtile.resize(TR * cols, 0.0);
5146 let mut r0 = 0usize;
5147 while r0 < rows {
5148 let tr = TR.min(rows - r0);
5149 let wt_addr = SendMut(wtile.as_mut_ptr());
5151 let run = |start: usize, end: usize| {
5152 for r in start..end {
5153 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
5154 let s = row_scale[r0 + r];
5155 let dst =
5157 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5158 for (d, &v) in dst.iter_mut().zip(row) {
5159 *d = (v as i8) as f32 * s;
5160 }
5161 }
5162 };
5163 dispatch_rows(pool, tr, &run);
5164 unsafe {
5166 accel_blas::cblas_sgemm(
5167 101, 111, 112, b as i32,
5171 tr as i32,
5172 cols as i32,
5173 1.0,
5174 xpanel.as_ptr(),
5175 cols as i32,
5176 wtile.as_ptr(),
5177 cols as i32,
5178 0.0,
5179 out.as_mut_ptr().add(r0),
5180 rows as i32,
5181 );
5182 }
5183 r0 += tr;
5184 }
5185 })
5186 });
5187}
5188
5189fn qmatmat(
5190 q: &[u8],
5191 row_scale: &[f32],
5192 pre: &[std::borrow::Cow<'_, [f32]>],
5193 rows: usize,
5194 cols: usize,
5195 out: &mut [f32],
5196 pool: Option<&Pool>,
5197) {
5198 let b = pre.len();
5199 debug_assert_eq!(out.len(), b * rows);
5200 let _p = crate::cpuprof::spans::span("mm.qmatmat");
5201 #[cfg(target_os = "macos")]
5206 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5207 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
5208 return;
5209 }
5210 #[cfg(target_arch = "aarch64")]
5211 if sdot_enabled() {
5212 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5213 let out_addr = SendMut(out.as_mut_ptr());
5214 let blocked_ok = blocked_enabled();
5217 let use_i8mm = i8mm_enabled();
5218 if blocked_ok {
5219 let run = |start: usize, end: usize| {
5220 let mut o = start;
5221 while o < end {
5222 if o + 2 <= end {
5223 let r0 = &q[o * cols..(o + 1) * cols];
5224 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5225 let mut bi = 0usize;
5226 while bi + 4 <= acts.len() {
5227 let xs = [
5228 acts[bi].xq.as_slice(),
5229 acts[bi + 1].xq.as_slice(),
5230 acts[bi + 2].xq.as_slice(),
5231 acts[bi + 3].xq.as_slice(),
5232 ];
5233 let d = if use_i8mm {
5234 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
5235 } else {
5236 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
5237 };
5238 for (r, row) in [r0, r1].into_iter().enumerate() {
5239 for k in 0..4 {
5240 let act = &acts[bi + k];
5241 let mut v = d[r][k] as f32 * act.sx;
5242 for &(j, xv) in &act.outliers {
5243 v += (row[j] as i8) as f32 * xv;
5244 }
5245 unsafe {
5246 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5247 };
5248 }
5249 }
5250 bi += 4;
5251 }
5252 while bi < acts.len() {
5253 for (r, row) in [r0, r1].into_iter().enumerate() {
5254 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
5255 unsafe { *out_addr.at(bi * rows + o + r) = v };
5256 }
5257 bi += 1;
5258 }
5259 o += 2;
5260 } else {
5261 let row = &q[o * cols..(o + 1) * cols];
5262 for (bi, act) in acts.iter().enumerate() {
5263 let v = row_dot_sdot(row, act) * row_scale[o];
5264 unsafe { *out_addr.at(bi * rows + o) = v };
5265 }
5266 o += 1;
5267 }
5268 }
5269 };
5270 dispatch_rows(pool, rows, &run);
5271 return;
5272 }
5273 let run = |start: usize, end: usize| {
5274 for o in start..end {
5275 let row = &q[o * cols..(o + 1) * cols];
5276 for (bi, act) in acts.iter().enumerate() {
5277 let v = row_dot_sdot(row, act) * row_scale[o];
5278 unsafe { *out_addr.at(bi * rows + o) = v };
5279 }
5280 }
5281 };
5282 dispatch_rows(pool, rows, &run);
5283 return;
5284 }
5285 #[cfg(target_arch = "x86_64")]
5290 if avx2_a8w8_enabled() {
5291 let acts = {
5292 let _psa = crate::cpuprof::spans::span("mm.split_act");
5293 let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5294 drop(_psa);
5295 a
5296 };
5297 let out_addr = SendMut(out.as_mut_ptr());
5298 let blocked_ok = blocked_enabled();
5301 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
5302 let run = |start: usize, end: usize| {
5303 let mut o = start;
5304 while o < end {
5305 if o + 2 <= end {
5306 let r0 = &q[o * cols..(o + 1) * cols];
5307 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5308 let mut bi = 0usize;
5309 while bi + 4 <= acts.len() {
5310 let xs = [
5311 acts[bi].xq.as_slice(),
5312 acts[bi + 1].xq.as_slice(),
5313 acts[bi + 2].xq.as_slice(),
5314 acts[bi + 3].xq.as_slice(),
5315 ];
5316 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
5317 for (r, row) in [r0, r1].into_iter().enumerate() {
5318 for k in 0..4 {
5319 let act = &acts[bi + k];
5320 let mut v = d[r][k] as f32 * act.sx;
5321 for &(j, xv) in &act.outliers {
5322 v += (row[j] as i8) as f32 * xv;
5323 }
5324 unsafe {
5325 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5326 };
5327 }
5328 }
5329 bi += 4;
5330 }
5331 while bi < acts.len() {
5332 for (r, row) in [r0, r1].into_iter().enumerate() {
5333 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
5334 unsafe { *out_addr.at(bi * rows + o + r) = v };
5335 }
5336 bi += 1;
5337 }
5338 o += 2;
5339 } else {
5340 let row = &q[o * cols..(o + 1) * cols];
5341 for (bi, act) in acts.iter().enumerate() {
5342 let v = row_dot_avx2(row, act) * row_scale[o];
5343 unsafe { *out_addr.at(bi * rows + o) = v };
5344 }
5345 o += 1;
5346 }
5347 }
5348 };
5349 dispatch_rows(pool, rows, &run);
5350 return;
5351 }
5352 let run = |start: usize, end: usize| {
5353 for o in start..end {
5354 let row = &q[o * cols..(o + 1) * cols];
5355 for (bi, act) in acts.iter().enumerate() {
5356 let v = row_dot_avx2(row, act) * row_scale[o];
5357 unsafe { *out_addr.at(bi * rows + o) = v };
5358 }
5359 }
5360 };
5361 dispatch_rows(pool, rows, &run);
5362 return;
5363 }
5364 let out_addr = SendMut(out.as_mut_ptr());
5365 let run = |start: usize, end: usize| {
5366 for o in start..end {
5367 let row = &q[o * cols..(o + 1) * cols];
5368 for (bi, x) in pre.iter().enumerate() {
5369 let mut acc = 0f32;
5370 for j in 0..cols {
5371 acc += (row[j] as i8) as f32 * x[j];
5372 }
5373 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
5374 }
5375 }
5376 };
5377 dispatch_rows(pool, rows, &run);
5378}
5379
5380fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
5383 match pool {
5384 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
5385 _ => run(0, rows),
5386 }
5387}
5388
5389fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
5391 let groups = rows * cols / GROUP_SIZE;
5392 bytes.split_at(groups * 16)
5393}
5394
5395#[inline]
5400fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
5401 #[cfg(target_arch = "aarch64")]
5402 unsafe {
5403 return vbit_fill4_neon(data, buf);
5404 }
5405 #[cfg(target_arch = "x86_64")]
5406 if avx2_enabled() {
5407 return unsafe { vbit_fill4_avx2(data, buf) };
5408 }
5409 #[allow(unreachable_code)]
5410 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5411 let u = unpack8::<4>(&data[blk * 4..]);
5412 for k in 0..8 {
5413 chunk[k] = (u[k] - 7) as i8 as u8;
5414 }
5415 }
5416}
5417
5418#[cfg(target_arch = "aarch64")]
5419#[target_feature(enable = "neon")]
5420unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
5421 unsafe {
5424 use core::arch::aarch64::*;
5425 let n = buf.len();
5426 let mask = vdupq_n_u8(0x0F);
5427 let seven = vdupq_n_s8(7);
5428 let mut g = 0usize;
5429 while g * 32 + 32 <= n {
5430 let b = vld1q_u8(data.as_ptr().add(g * 16));
5431 let hi = vshrq_n_u8::<4>(b);
5432 let lo = vandq_u8(b, mask);
5433 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
5434 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
5435 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
5436 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
5437 g += 1;
5438 }
5439 }
5440}
5441
5442#[cfg(target_arch = "x86_64")]
5443#[target_feature(enable = "avx2")]
5444unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
5445 unsafe {
5447 use core::arch::x86_64::*;
5448 let n = buf.len();
5449 let mask = _mm_set1_epi8(0x0F);
5450 let seven = _mm256_set1_epi8(7);
5451 let mut g = 0usize;
5452 while g * 32 + 32 <= n {
5453 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
5454 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
5455 let lo = _mm_and_si128(b, mask);
5456 let z = _mm256_sub_epi8(
5457 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
5458 seven,
5459 );
5460 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
5461 g += 1;
5462 }
5463 }
5464}
5465
5466#[inline(always)]
5471fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
5472 let mut acc = 0u64;
5473 for i in 0..B {
5474 acc = (acc << 8) | data[i] as u64;
5475 }
5476 let mask = (1u64 << B) - 1;
5477 let mut out = [0i32; 8];
5478 for (k, o) in out.iter_mut().enumerate() {
5479 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5480 }
5481 out
5482}
5483
5484#[allow(clippy::too_many_arguments)]
5490fn vbitmatvec(
5491 bytes: &[u8],
5492 offsets: &[usize],
5493 x: &[f32],
5494 rows: usize,
5495 cols: usize,
5496 out: &mut [f32],
5497 pool: Option<&Pool>,
5498) {
5499 debug_assert_eq!(out.len(), rows);
5500 debug_assert_eq!(offsets.len(), rows + 1);
5501
5502 if a8w8_enabled() {
5506 let act = split_act(x);
5507 let out_addr = SendMut(out.as_mut_ptr());
5508 let run = move |start: usize, end: usize| {
5509 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5510 };
5511 dispatch_rows(pool, rows, &run);
5512 return;
5513 }
5514
5515 let out_addr = SendMut(out.as_mut_ptr());
5516 let run = move |start: usize, end: usize| {
5517 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5518 };
5519 dispatch_rows(pool, rows, &run);
5520}
5521
5522#[allow(clippy::too_many_arguments)]
5526fn vbit_range_a8w8(
5527 bytes: &[u8],
5528 offsets: &[usize],
5529 x: &[f32],
5530 act: &SplitAct,
5531 rows: usize,
5532 cols: usize,
5533 out: SendMut,
5534 start: usize,
5535 end: usize,
5536) {
5537 let ng = cols / GROUP_SIZE;
5538 let bits = &bytes[..rows];
5539 let sc_off = rows;
5540 let row_dot = |r: usize| -> f32 {
5541 let b = bits[r] as usize;
5542 let l = (1i32 << (b - 1)) - 1;
5543 let mask = (1u64 << b) - 1;
5544 let data = &bytes[offsets[r]..offsets[r + 1]];
5545 if b == 8 {
5546 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5548 let mut dot = 0f32;
5549 for g in 0..ng {
5550 let so = (r * ng + g) * 2;
5551 let sgf = f16_to_f32(u16::from_le_bytes([
5552 bytes[sc_off + so],
5553 bytes[sc_off + so + 1],
5554 ]));
5555 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5556 let mut gd = 0f32;
5557 for &xv in xg.iter() {
5558 if nbits < 8 {
5559 acc = (acc << 8) | data[idx] as u64;
5560 idx += 1;
5561 nbits += 8;
5562 }
5563 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5564 nbits -= 8;
5565 gd += (u - l) as f32 * xv;
5566 }
5567 dot += gd * sgf;
5568 }
5569 return dot;
5570 }
5571 thread_local! {
5575 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5576 const { std::cell::RefCell::new(Vec::new()) };
5577 }
5578 #[inline(always)]
5579 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5580 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5581 let u = unpack8::<B>(&data[blk * B..]);
5582 for k in 0..8 {
5583 chunk[k] = (u[k] - l) as i8 as u8;
5584 }
5585 }
5586 }
5587 let _ = mask;
5588 VBIT_SCRATCH.with(|scratch| {
5589 let mut buf = scratch.borrow_mut();
5590 buf.resize(cols, 0);
5591 match b {
5592 3 => fill::<3>(data, l, &mut buf),
5593 4 => vbit_fill4(data, &mut buf),
5594 5 => fill::<5>(data, l, &mut buf),
5595 6 => fill::<6>(data, l, &mut buf),
5596 _ => unreachable!(),
5597 }
5598 let mut dot = 0f32;
5599 for g in 0..ng {
5600 let so = (r * ng + g) * 2;
5601 let s = f16_to_f32(u16::from_le_bytes([
5602 bytes[sc_off + so],
5603 bytes[sc_off + so + 1],
5604 ]));
5605 let d = dot_i8_i8(
5606 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5607 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5608 ) as f32
5609 * act.sx;
5610 dot += d * s;
5611 }
5612 for &(j, xv) in &act.outliers {
5613 let so = (r * ng + j / GROUP_SIZE) * 2;
5614 let s = f16_to_f32(u16::from_le_bytes([
5615 bytes[sc_off + so],
5616 bytes[sc_off + so + 1],
5617 ]));
5618 dot += (buf[j] as i8) as f32 * s * xv;
5620 }
5621 dot
5622 })
5623 };
5624 for r in start..end {
5625 unsafe { *out.at(r) = row_dot(r) };
5627 }
5628}
5629
5630#[allow(clippy::too_many_arguments)]
5632fn vbit_range_f32(
5633 bytes: &[u8],
5634 offsets: &[usize],
5635 x: &[f32],
5636 rows: usize,
5637 cols: usize,
5638 out: SendMut,
5639 start: usize,
5640 end: usize,
5641) {
5642 let ng = cols / GROUP_SIZE;
5643 let bits = &bytes[..rows];
5644 let sc_off = rows;
5645 #[inline(always)]
5649 fn dot_row<const B: usize>(
5650 data: &[u8],
5651 bytes: &[u8],
5652 sc_off: usize,
5653 r: usize,
5654 ng: usize,
5655 x: &[f32],
5656 ) -> f32 {
5657 let l = ((1i32 << (B - 1)) - 1) as f32;
5658 let gbytes = GROUP_SIZE * B / 8;
5659 let mut dot = 0f32;
5660 for g in 0..ng {
5661 let so = (r * ng + g) * 2;
5662 let s = f16_to_f32(u16::from_le_bytes([
5663 bytes[sc_off + so],
5664 bytes[sc_off + so + 1],
5665 ]));
5666 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5667 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5668 let mut gd = 0f32;
5669 for blk in 0..GROUP_SIZE / 8 {
5670 let u = unpack8::<B>(&gd0[blk * B..]);
5671 let xb = &xg[blk * 8..blk * 8 + 8];
5672 for k in 0..8 {
5673 gd += (u[k] as f32 - l) * xb[k];
5674 }
5675 }
5676 dot += gd * s;
5677 }
5678 dot
5679 }
5680 for r in start..end {
5681 let data = &bytes[offsets[r]..offsets[r + 1]];
5682 let v = match bits[r] {
5683 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5684 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5685 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5686 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5687 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5688 b => unreachable!("vbit bit-width {b} (validated at load)"),
5689 };
5690 unsafe { *out.at(r) = v };
5692 }
5693}
5694
5695#[allow(clippy::too_many_arguments)]
5700fn vbitmatvec2(
5701 bytes: &[u8],
5702 offsets: &[usize],
5703 x1: &[f32],
5704 x2: &[f32],
5705 rows: usize,
5706 cols: usize,
5707 o1: &mut [f32],
5708 o2: &mut [f32],
5709 pool: Option<&Pool>,
5710) {
5711 debug_assert_eq!(o1.len(), rows);
5712 debug_assert_eq!(o2.len(), rows);
5713
5714 if a8w8_enabled() {
5715 let a1 = split_act(x1);
5716 let a2 = split_act(x2);
5717 let p1 = SendMut(o1.as_mut_ptr());
5718 let p2 = SendMut(o2.as_mut_ptr());
5719 let run = move |start: usize, end: usize| {
5720 vbit_range2_a8w8(
5721 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5722 )
5723 };
5724 dispatch_rows(pool, rows, &run);
5725 return;
5726 }
5727
5728 let p1 = SendMut(o1.as_mut_ptr());
5729 let p2 = SendMut(o2.as_mut_ptr());
5730 let run = move |start: usize, end: usize| {
5731 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5732 };
5733 dispatch_rows(pool, rows, &run);
5734}
5735
5736#[allow(clippy::too_many_arguments)]
5740fn vbit_range2_a8w8(
5741 bytes: &[u8],
5742 offsets: &[usize],
5743 x1: &[f32],
5744 x2: &[f32],
5745 a1: &SplitAct,
5746 a2: &SplitAct,
5747 rows: usize,
5748 cols: usize,
5749 p1: SendMut,
5750 p2: SendMut,
5751 start: usize,
5752 end: usize,
5753) {
5754 let ng = cols / GROUP_SIZE;
5755 let bits = &bytes[..rows];
5756 let sc_off = rows;
5757 let row_dots = |r: usize| -> (f32, f32) {
5758 let b = bits[r] as usize;
5759 let l = (1i32 << (b - 1)) - 1;
5760 let data = &bytes[offsets[r]..offsets[r + 1]];
5761 if b == 8 {
5762 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5765 let (mut d1, mut d2) = (0f32, 0f32);
5766 for g in 0..ng {
5767 let so = (r * ng + g) * 2;
5768 let sgf = f16_to_f32(u16::from_le_bytes([
5769 bytes[sc_off + so],
5770 bytes[sc_off + so + 1],
5771 ]));
5772 let (mut g1, mut g2) = (0f32, 0f32);
5773 for k in 0..GROUP_SIZE {
5774 if nbits < 8 {
5775 acc = (acc << 8) | data[idx] as u64;
5776 idx += 1;
5777 nbits += 8;
5778 }
5779 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5780 nbits -= 8;
5781 let w = (u - l) as f32;
5782 g1 += w * x1[g * GROUP_SIZE + k];
5783 g2 += w * x2[g * GROUP_SIZE + k];
5784 }
5785 d1 += g1 * sgf;
5786 d2 += g2 * sgf;
5787 }
5788 return (d1, d2);
5789 }
5790 thread_local! {
5791 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5792 const { std::cell::RefCell::new(Vec::new()) };
5793 }
5794 #[inline(always)]
5795 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5796 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5797 let u = unpack8::<B>(&data[blk * B..]);
5798 for k in 0..8 {
5799 chunk[k] = (u[k] - l) as i8 as u8;
5800 }
5801 }
5802 }
5803 VBIT_SCRATCH2.with(|scratch| {
5804 let mut buf = scratch.borrow_mut();
5805 buf.resize(cols, 0);
5806 match b {
5807 3 => fill::<3>(data, l, &mut buf),
5808 4 => vbit_fill4(data, &mut buf),
5809 5 => fill::<5>(data, l, &mut buf),
5810 6 => fill::<6>(data, l, &mut buf),
5811 _ => unreachable!(),
5812 }
5813 let (mut d1, mut d2) = (0f32, 0f32);
5814 for g in 0..ng {
5815 let so = (r * ng + g) * 2;
5816 let s = f16_to_f32(u16::from_le_bytes([
5817 bytes[sc_off + so],
5818 bytes[sc_off + so + 1],
5819 ]));
5820 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5821 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5822 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5823 d1 += v1 * s;
5824 d2 += v2 * s;
5825 }
5826 for &(j, xv) in &a1.outliers {
5827 let so = (r * ng + j / GROUP_SIZE) * 2;
5828 let s = f16_to_f32(u16::from_le_bytes([
5829 bytes[sc_off + so],
5830 bytes[sc_off + so + 1],
5831 ]));
5832 d1 += (buf[j] as i8) as f32 * s * xv;
5833 }
5834 for &(j, xv) in &a2.outliers {
5835 let so = (r * ng + j / GROUP_SIZE) * 2;
5836 let s = f16_to_f32(u16::from_le_bytes([
5837 bytes[sc_off + so],
5838 bytes[sc_off + so + 1],
5839 ]));
5840 d2 += (buf[j] as i8) as f32 * s * xv;
5841 }
5842 (d1, d2)
5843 })
5844 };
5845 for r in start..end {
5846 let (v1, v2) = row_dots(r);
5847 unsafe {
5849 *p1.at(r) = v1;
5850 *p2.at(r) = v2;
5851 }
5852 }
5853}
5854
5855#[allow(clippy::too_many_arguments)]
5859fn vbit_range2_f32(
5860 bytes: &[u8],
5861 offsets: &[usize],
5862 x1: &[f32],
5863 x2: &[f32],
5864 rows: usize,
5865 cols: usize,
5866 p1: SendMut,
5867 p2: SendMut,
5868 start: usize,
5869 end: usize,
5870) {
5871 let ng = cols / GROUP_SIZE;
5872 let bits = &bytes[..rows];
5873 let sc_off = rows;
5874 #[inline(always)]
5875 #[allow(clippy::too_many_arguments)]
5876 fn dot_row2<const B: usize>(
5877 data: &[u8],
5878 bytes: &[u8],
5879 sc_off: usize,
5880 r: usize,
5881 ng: usize,
5882 x1: &[f32],
5883 x2: &[f32],
5884 ) -> (f32, f32) {
5885 let l = ((1i32 << (B - 1)) - 1) as f32;
5886 let gbytes = GROUP_SIZE * B / 8;
5887 let (mut d1, mut d2) = (0f32, 0f32);
5888 for g in 0..ng {
5889 let so = (r * ng + g) * 2;
5890 let s = f16_to_f32(u16::from_le_bytes([
5891 bytes[sc_off + so],
5892 bytes[sc_off + so + 1],
5893 ]));
5894 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5895 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5896 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5897 let (mut g1, mut g2) = (0f32, 0f32);
5898 for blk in 0..GROUP_SIZE / 8 {
5899 let u = unpack8::<B>(&gd0[blk * B..]);
5900 for k in 0..8 {
5901 let w = u[k] as f32 - l;
5902 g1 += w * x1g[blk * 8 + k];
5903 g2 += w * x2g[blk * 8 + k];
5904 }
5905 }
5906 d1 += g1 * s;
5907 d2 += g2 * s;
5908 }
5909 (d1, d2)
5910 }
5911 for r in start..end {
5912 let data = &bytes[offsets[r]..offsets[r + 1]];
5913 let (v1, v2) = match bits[r] {
5914 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5915 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5916 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5917 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5918 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5919 b => unreachable!("vbit bit-width {b} (validated at load)"),
5920 };
5921 unsafe {
5923 *p1.at(r) = v1;
5924 *p2.at(r) = v2;
5925 }
5926 }
5927}
5928
5929#[inline]
5936#[allow(unreachable_code)]
5937fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5938 #[cfg(target_arch = "aarch64")]
5939 unsafe {
5940 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5941 }
5942 #[cfg(target_arch = "x86_64")]
5943 unsafe {
5944 if vnni_tiles_enabled() {
5945 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5946 }
5947 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5948 }
5949 let mut acc = 0f32;
5950 for gi in 0..gpr {
5951 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5952 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5953 let mut d = 0i32;
5954 for (k, &b) in tile[2..].iter().enumerate() {
5955 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5956 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5957 }
5958 acc += d as f32 * s;
5959 }
5960 acc
5961}
5962
5963#[cfg(target_arch = "aarch64")]
5964#[target_feature(enable = "neon,dotprod")]
5965unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5966 unsafe {
5969 use core::arch::aarch64::*;
5970 use core::arch::asm;
5971 let lomask = vdupq_n_u8(0x0F);
5972 let eight = vdupq_n_s8(8);
5973 let mut acc = 0f32;
5974 for gi in 0..gpr {
5975 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5976 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5977 let b = vld1q_u8(t.add(2));
5978 let lo = vandq_u8(b, lomask);
5979 let hi = vshrq_n_u8::<4>(b);
5980 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5981 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5982 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5983 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5984 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5985 asm!(
5986 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5987 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5988 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5989 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5990 options(pure, nomem, nostack),
5991 );
5992 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5993 }
5994 acc
5995 }
5996}
5997
5998#[cfg(target_arch = "x86_64")]
5999#[target_feature(enable = "avx2")]
6000unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6001 unsafe {
6003 use core::arch::x86_64::*;
6004 let lomask = _mm_set1_epi8(0x0F);
6005 let eight = _mm256_set1_epi8(8);
6006 let ones = _mm256_set1_epi16(1);
6007 let mut acc = 0f32;
6008 for gi in 0..gpr {
6009 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6010 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6011 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
6012 let lo = _mm_and_si128(b, lomask);
6013 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6014 let w = _mm256_sub_epi8(
6015 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6016 eight,
6017 );
6018 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6019 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6020 let d = _mm256_madd_epi16(p16, ones);
6021 let hi128 = _mm256_extracti128_si256::<1>(d);
6022 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6023 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6024 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6025 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6026 }
6027 acc
6028 }
6029}
6030
6031#[cfg(target_arch = "x86_64")]
6035#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6036unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6037 unsafe {
6039 use core::arch::x86_64::*;
6040 let lomask = _mm_set1_epi8(0x0F);
6041 let eight = _mm256_set1_epi8(8);
6042 let mut acc = 0f32;
6043 for gi in 0..gpr {
6044 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6045 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6046 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
6047 let lo = _mm_and_si128(b, lomask);
6048 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6049 let w = _mm256_sub_epi8(
6050 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6051 eight,
6052 );
6053 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6054 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6055 acc += d as f32 * s;
6056 }
6057 acc
6058 }
6059}
6060
6061#[cfg(target_arch = "x86_64")]
6066#[target_feature(enable = "avx2,fma")]
6071unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6072 unsafe {
6074 use core::arch::x86_64::*;
6075 let lomask = _mm_set1_epi8(0x0F);
6076 let eight = _mm256_set1_epi8(8);
6077 let ones = _mm256_set1_epi16(1);
6078 let mut f0 = _mm256_setzero_ps();
6091 let mut f1 = _mm256_setzero_ps();
6092 let mut f2 = _mm256_setzero_ps();
6093 let mut f3 = _mm256_setzero_ps();
6094 for gi in 0..gpr {
6095 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6096 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6097 let sv = _mm256_set1_ps(s);
6098 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6099 let lo = _mm_and_si128(bb, lomask);
6100 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6101 let w = _mm256_sub_epi8(
6102 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6103 eight,
6104 );
6105 let aw = _mm256_abs_epi8(w);
6106 let off = gi * GROUP_SIZE;
6107 let dot = |xq: &[i8]| {
6108 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6109 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6110 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
6111 };
6112 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6113 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6114 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6115 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6116 }
6117 [
6118 hsum256_ps(f0),
6119 hsum256_ps(f1),
6120 hsum256_ps(f2),
6121 hsum256_ps(f3),
6122 ]
6123 }
6124}
6125
6126#[cfg(target_arch = "x86_64")]
6129#[target_feature(enable = "avx2")]
6130#[inline]
6131unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
6132 unsafe {
6134 use core::arch::x86_64::*;
6135 let hi = _mm256_extractf128_ps::<1>(v);
6136 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
6137 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
6138 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
6139 _mm_cvtss_f32(s)
6140 }
6141}
6142
6143#[cfg(target_arch = "x86_64")]
6145#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
6146unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6147 unsafe {
6149 use core::arch::x86_64::*;
6150 let lomask = _mm_set1_epi8(0x0F);
6151 let eight = _mm256_set1_epi8(8);
6152 let mut f0 = _mm256_setzero_ps();
6155 let mut f1 = _mm256_setzero_ps();
6156 let mut f2 = _mm256_setzero_ps();
6157 let mut f3 = _mm256_setzero_ps();
6158 for gi in 0..gpr {
6159 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6160 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6161 let sv = _mm256_set1_ps(s);
6162 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6163 let lo = _mm_and_si128(bb, lomask);
6164 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6165 let w = _mm256_sub_epi8(
6166 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6167 eight,
6168 );
6169 let aw = _mm256_abs_epi8(w);
6170 let off = gi * GROUP_SIZE;
6171 let dot = |xq: &[i8]| {
6172 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6173 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
6174 _mm256_setzero_si256(),
6175 aw,
6176 _mm256_sign_epi8(x, w),
6177 ))
6178 };
6179 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6180 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6181 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6182 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6183 }
6184 let acc = [
6185 hsum256_ps(f0),
6186 hsum256_ps(f1),
6187 hsum256_ps(f2),
6188 hsum256_ps(f3),
6189 ];
6190 acc
6191 }
6192}
6193
6194#[cfg(target_arch = "aarch64")]
6199#[target_feature(enable = "neon,dotprod")]
6200unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6201 unsafe {
6203 use core::arch::aarch64::*;
6204 use core::arch::asm;
6205 let lomask = vdupq_n_u8(0x0F);
6206 let eight = vdupq_n_s8(8);
6207 let mut acc = [0f32; 4];
6208 for gi in 0..gpr {
6209 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6210 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6211 let b = vld1q_u8(t.add(2));
6212 let lo = vandq_u8(b, lomask);
6213 let hi = vshrq_n_u8::<4>(b);
6214 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6215 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6216 for (k, xq) in xs.iter().enumerate() {
6217 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6218 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6219 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6220 asm!(
6221 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6222 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6223 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6224 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6225 options(pure, nomem, nostack),
6226 );
6227 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6228 }
6229 }
6230 acc
6231 }
6232}
6233
6234#[inline]
6236fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6237 let gi = j / GROUP_SIZE;
6238 let k = j % GROUP_SIZE;
6239 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6240 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6241 let byte = tile[2 + k / 2];
6242 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6243 ((nib as i32 - 8) as f32, s)
6244}
6245
6246#[inline]
6249fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6250 let mut acc = 0f32;
6251 for gi in 0..gpr {
6252 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6253 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6254 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6255 let mut ga = 0f32;
6256 for (k, &b) in tile[2..].iter().enumerate() {
6257 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6258 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6259 }
6260 acc += ga * s;
6261 }
6262 acc
6263}
6264
6265fn row_span_offsets(
6273 dtype: TensorDtype,
6274 rows: usize,
6275 cols: usize,
6276 r: usize,
6277) -> Option<[(usize, usize); 3]> {
6278 if r >= rows || cols % GROUP_SIZE != 0 {
6279 return None;
6280 }
6281 let gpr = cols / GROUP_SIZE;
6282 match dtype {
6283 TensorDtype::Q4Tiled => Some([(r * gpr * Q4_TILE, gpr * Q4_TILE), (0, 0), (0, 0)]),
6284 TensorDtype::Q4TiledP => {
6285 let (params, codes, stride) = q4tp_sections(rows, cols);
6286 Some([
6287 (r * gpr * Q4TP_NIB, gpr * Q4TP_NIB),
6288 (params + r * 4, 4),
6289 (codes + r * stride, stride),
6290 ])
6291 }
6292 _ => None,
6293 }
6294}
6295
6296struct Q4tpView<'a> {
6297 nib: &'a [u8],
6298 params: &'a [u8],
6299 codes: &'a [u8],
6300 stride: usize,
6301 zero_rung: bool,
6303}
6304
6305impl<'a> Q4tpView<'a> {
6306 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6307 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
6308 Self {
6309 nib: &bytes[..params_off],
6310 params: &bytes[params_off..codes_off],
6311 codes: &bytes[codes_off..],
6312 stride,
6313 zero_rung: false,
6314 }
6315 }
6316
6317 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6319 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
6320 Self {
6321 nib: &bytes[..params_off],
6322 params: &bytes[params_off..codes_off],
6323 codes: &bytes[codes_off..],
6324 stride,
6325 zero_rung: true,
6326 }
6327 }
6328
6329 #[cfg(target_arch = "x86_64")]
6352 #[inline]
6353 fn ladder_via_avx2(&self) -> bool {
6354 avx2_enabled()
6355 }
6356 #[cfg(not(target_arch = "x86_64"))]
6357 #[inline]
6358 fn ladder_via_avx2(&self) -> bool {
6359 false
6360 }
6361
6362 #[inline]
6363 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
6364 let tab = if self.zero_rung {
6365 q2tp_ladder(self.params, r)
6366 } else if self.ladder_via_avx2() {
6367 #[cfg(target_arch = "x86_64")]
6374 let tab = {
6375 let p = &self.params[r * 4..r * 4 + 4];
6376 let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
6377 let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
6378 let mut tabc = [0f32; 128];
6379 unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
6380 let mut t = [0f32; 32];
6381 for (g, o) in t.iter_mut().enumerate() {
6382 *o = tabc[g * 4];
6383 }
6384 t
6385 };
6386 #[cfg(not(target_arch = "x86_64"))]
6390 let tab = q4tp_ladder(self.params, r);
6391 tab
6392 } else {
6393 q4tp_ladder(self.params, r)
6394 };
6395 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
6396 let out = &mut out[..gpr];
6397 let mut chunks = out.chunks_exact_mut(8);
6398 let mut ci = 0usize;
6399 for c in &mut chunks {
6400 let w = u64::from(codes[ci])
6401 | u64::from(codes[ci + 1]) << 8
6402 | u64::from(codes[ci + 2]) << 16
6403 | u64::from(codes[ci + 3]) << 24
6404 | u64::from(codes[ci + 4]) << 32;
6405 for (k, o) in c.iter_mut().enumerate() {
6406 *o = tab[((w >> (5 * k)) & 31) as usize];
6407 }
6408 ci += 5;
6409 }
6410 let tail = &codes[ci..];
6413 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6414 *o = tab[q4tp_code(tail, k)];
6415 }
6416 }
6417}
6418
6419#[inline]
6420fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6421 #[cfg(target_arch = "aarch64")]
6422 unsafe {
6423 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
6424 }
6425 #[cfg(target_arch = "x86_64")]
6426 unsafe {
6427 if vnni_tiles_enabled() {
6428 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
6429 }
6430 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
6431 }
6432 #[allow(unreachable_code)]
6433 {
6434 let mut acc = 0f32;
6435 for gi in 0..gpr {
6436 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6437 let s = scales[gi];
6438 let mut d = 0i32;
6439 for (k, &b) in tile.iter().enumerate() {
6440 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6441 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6442 }
6443 acc += d as f32 * s;
6444 }
6445 acc
6446 }
6447}
6448
6449#[cfg(target_arch = "aarch64")]
6452#[target_feature(enable = "neon,dotprod")]
6453unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6454 unsafe {
6457 use core::arch::aarch64::*;
6458 use core::arch::asm;
6459 let lomask = vdupq_n_u8(0x0F);
6460 let eight = vdupq_n_s8(8);
6461 let mut acc = 0f32;
6462 for gi in 0..gpr {
6463 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6464 let s = *scales.get_unchecked(gi);
6465 let b = vld1q_u8(t);
6466 let lo = vandq_u8(b, lomask);
6467 let hi = vshrq_n_u8::<4>(b);
6468 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6469 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6470 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6471 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6472 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6473 asm!(
6474 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6475 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6476 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6477 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6478 options(pure, nomem, nostack),
6479 );
6480 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6481 }
6482 acc
6483 }
6484}
6485
6486#[cfg(target_arch = "x86_64")]
6487#[target_feature(enable = "avx2")]
6488unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6489 unsafe {
6491 use core::arch::x86_64::*;
6492 let lomask = _mm_set1_epi8(0x0F);
6493 let eight = _mm256_set1_epi8(8);
6494 let ones = _mm256_set1_epi16(1);
6495 let mut acc = 0f32;
6496 for gi in 0..gpr {
6497 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6498 let s = *scales.get_unchecked(gi);
6499 let b = _mm_loadu_si128(t as *const __m128i);
6500 let lo = _mm_and_si128(b, lomask);
6501 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6502 let w = _mm256_sub_epi8(
6503 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6504 eight,
6505 );
6506 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6507 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6508 let d = _mm256_madd_epi16(p16, ones);
6509 let hi128 = _mm256_extracti128_si256::<1>(d);
6510 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6511 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6512 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6513 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6514 }
6515 acc
6516 }
6517}
6518
6519#[cfg(target_arch = "x86_64")]
6522#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6523unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6524 unsafe {
6526 use core::arch::x86_64::*;
6527 let lomask = _mm_set1_epi8(0x0F);
6528 let eight = _mm256_set1_epi8(8);
6529 let mut acc = 0f32;
6530 for gi in 0..gpr {
6531 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6532 let s = *scales.get_unchecked(gi);
6533 let b = _mm_loadu_si128(t as *const __m128i);
6534 let lo = _mm_and_si128(b, lomask);
6535 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6536 let w = _mm256_sub_epi8(
6537 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6538 eight,
6539 );
6540 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6541 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6542 }
6543 acc
6544 }
6545}
6546
6547#[cfg(target_arch = "x86_64")]
6554#[target_feature(enable = "avx2")]
6555#[inline]
6556unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6557 unsafe {
6558 use core::arch::x86_64::*;
6559 let tr = _mm_cvttps_epi32(x);
6561 let below = _mm_and_si128(
6564 _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6565 _mm_set1_epi32(1),
6566 );
6567 let n = _mm_sub_epi32(tr, below); let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); #[allow(clippy::excessive_precision, clippy::approx_constant)]
6578 let a = [
6579 6.931_471_805_599_453e-1,
6580 2.402_265_069_591_007e-1,
6581 5.550_410_866_482_158e-2,
6582 9.618_129_107_628_477e-3,
6583 1.333_355_814_642_844e-3,
6584 1.540_353_039_338_161e-4,
6585 1.525_273_380_405_984e-5,
6586 1.321_548_679_014_431e-6,
6587 1.017_808_600_923_969_6e-7,
6588 ];
6589 let mut p = _mm_set1_ps(a[8]);
6590 for &c in a[..8].iter().rev() {
6591 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6592 }
6593 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6594 let ncl = _mm_min_epi32(
6600 _mm_max_epi32(n, _mm_set1_epi32(-126)),
6601 _mm_set1_epi32(127),
6602 );
6603 let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6604 let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6605 bits = _mm_or_si128(
6606 _mm_andnot_si128(over, bits),
6607 _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6608 );
6609 let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6610 _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6611 }
6612}
6613
6614#[cfg(target_arch = "x86_64")]
6616#[inline]
6617unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6618 use core::arch::x86_64::*;
6619 unsafe {
6620 [
6621 _mm_cvtss_f32(v),
6622 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6623 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6624 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6625 ]
6626 }
6627}
6628
6629#[cfg(target_arch = "x86_64")]
6637#[target_feature(enable = "avx2")]
6638#[inline]
6639unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6640 unsafe {
6641 use core::arch::x86_64::*;
6642 let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6643 let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6644 tabc[0..4].copy_from_slice(&lanes4(t));
6645 for c in 1..32 {
6646 t = _mm_mul_ps(t, ratio);
6647 tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6648 }
6649 }
6650}
6651
6652#[cfg(target_arch = "x86_64")]
6667#[target_feature(enable = "avx2")]
6668unsafe fn q4tp_scales_lanes(
6669 vs: &[&Q4tpView<'_>],
6670 rows: &[usize],
6671 gpr: usize,
6672 out: &mut [f32],
6673) {
6674 use core::arch::x86_64::*;
6675 debug_assert_eq!(vs.len(), rows.len());
6676 debug_assert!(out.len() >= rows.len() * gpr);
6677 let mut done = 0usize;
6678 for (ci, chunk) in rows.chunks(4).enumerate() {
6679 let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6680 if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6681 for (l, &r) in chunk.iter().enumerate() {
6682 vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6683 }
6684 done += chunk.len();
6685 continue;
6686 }
6687 let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6688 let v0 = vchunk[0];
6689 let v1 = vchunk[1];
6690 let v2 = vchunk[2];
6691 let v3 = vchunk[3];
6692 let mut lo = [0f32; 4];
6694 let mut st = [0f32; 4];
6695 for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6696 let o = *r * 4;
6697 let vv = vchunk[l];
6698 lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6699 vv.params[o],
6700 vv.params[o + 1],
6701 ]));
6702 st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6703 vv.params[o + 2],
6704 vv.params[o + 3],
6705 ]));
6706 }
6707 let lo_v = _mm_loadu_ps(lo.as_ptr());
6708 let st_v = _mm_loadu_ps(st.as_ptr());
6709 let ratio = exp2_128(st_v);
6710 let mut t = exp2_128(lo_v);
6711 let mut tabc = [0f32; 32 * 4];
6713 tabc[0..4].copy_from_slice(&[
6714 _mm_cvtss_f32(t),
6715 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6716 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6717 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6718 ]);
6719 for c in 1..32 {
6720 t = _mm_mul_ps(t, ratio);
6721 tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6722 _mm_cvtss_f32(t),
6723 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6724 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6725 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6726 ]);
6727 }
6728 let vs4 = [v0, v1, v2, v3];
6730 for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6731 let v = vs4[l];
6732 let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6733 let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6734 let mut chunks = dst.chunks_exact_mut(8);
6735 let mut ci = 0usize;
6736 for c in &mut chunks {
6737 let w = u64::from(codes[ci])
6738 | u64::from(codes[ci + 1]) << 8
6739 | u64::from(codes[ci + 2]) << 16
6740 | u64::from(codes[ci + 3]) << 24
6741 | u64::from(codes[ci + 4]) << 32;
6742 for (k, o) in c.iter_mut().enumerate() {
6743 *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6744 }
6745 ci += 5;
6746 }
6747 let tail = &codes[ci..];
6748 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6749 *o = tabc[q4tp_code(tail, k) * 4 + l];
6750 }
6751 }
6752 done += 4;
6753 }
6754}
6755
6756fn q4tp_lanes() -> usize {
6764 #[cfg(not(target_arch = "x86_64"))]
6765 {
6766 1
6767 }
6768 #[cfg(target_arch = "x86_64")]
6769 {
6770 static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6771 *ON.get_or_init(|| {
6772 let n: usize = std::env::var("CMF_Q4TP_LANES")
6773 .ok()
6774 .and_then(|v| v.parse().ok())
6775 .unwrap_or(4);
6776 match n {
6777 1 => 1,
6778 2 => 2,
6779 8 => 8,
6780 _ => 4,
6781 }
6782 })
6783 }
6784}
6785
6786#[cfg(target_arch = "x86_64")]
6797#[target_feature(enable = "avx2")]
6798#[inline]
6799unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6800 unsafe {
6801 use core::arch::x86_64::*;
6802 let lomask = _mm_set1_epi8(0x0F);
6803 let b = _mm_loadu_si128(tile as *const __m128i);
6804 let lo = _mm_and_si128(b, lomask);
6805 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6806 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6807 }
6808}
6809
6810#[cfg(target_arch = "x86_64")]
6817#[target_feature(enable = "avx2")]
6818#[inline]
6819unsafe fn q4tp_madd_wx(
6820 w: core::arch::x86_64::__m256i,
6821 xp: *const i8,
6822) -> core::arch::x86_64::__m256i {
6823 unsafe {
6824 use core::arch::x86_64::*;
6825 let ones = _mm256_set1_epi16(1);
6826 let x = _mm256_loadu_si256(xp as *const __m256i);
6827 _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6828 }
6829}
6830
6831#[cfg(target_arch = "x86_64")]
6836#[target_feature(enable = "avx2")]
6837#[inline]
6838unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6839 unsafe {
6840 use core::arch::x86_64::*;
6841 let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6842 let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6843 let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6844 _mm_cvtss_f32(c)
6845 }
6846}
6847
6848#[cfg(target_arch = "x86_64")]
6852#[inline]
6853fn q4tp_xg8_corr(xg8: &[f32], sc: &[f32]) -> f32 {
6854 let n = sc.len().min(xg8.len());
6855 8.0 * crate::attention::dot_f32(&xg8[..n], &sc[..n])
6856}
6857
6858#[cfg(target_arch = "x86_64")]
6866#[target_feature(enable = "avx2,fma")]
6867#[inline]
6868unsafe fn q4tp_cell_dot_v(
6869 nib: *const u8,
6870 r: usize,
6871 gpr: usize,
6872 xp: *const i8,
6873 sc: *const f32,
6874 xg8: *const f32,
6875) -> f32 {
6876 unsafe {
6877 use core::arch::x86_64::*;
6878 let mut acc = _mm256_setzero_ps();
6879 for gi in 0..gpr {
6880 let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6881 let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6882 acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*sc.add(gi)), acc);
6883 }
6884 q4tp_hsum256(acc)
6885 - q4tp_xg8_corr(
6886 std::slice::from_raw_parts(xg8, gpr),
6887 std::slice::from_raw_parts(sc, gpr),
6888 )
6889 }
6890}
6891
6892#[cfg(target_arch = "x86_64")]
6905#[target_feature(enable = "avx2,fma")]
6906unsafe fn q4tp_matmat_row_cols_v(
6907 v: &Q4tpView<'_>,
6908 r: usize,
6909 gpr: usize,
6910 acts: &[SplitAct],
6911 sc: &[f32],
6912 abase: usize,
6913 rows: usize,
6914 out: SendMut,
6915) {
6916 use core::arch::x86_64::*;
6917 debug_assert_eq!(sc.len(), gpr);
6918 let n = acts.len();
6919 let mut bi = 0usize;
6920 while bi + 4 <= n {
6921 let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6922 let mut accs = [_mm256_setzero_ps(); 4];
6923 for gi in 0..gpr {
6924 let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6925 let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6926 for (k, act) in a.iter().enumerate() {
6927 let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6928 let d = q4tp_madd_wx(w, xp);
6929 accs[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[k]);
6930 }
6931 }
6932 for k in 0..4 {
6933 let act = a[k];
6934 let mut acc = (q4tp_hsum256(accs[k]) - q4tp_xg8_corr(&act.xg8, sc)) * act.sx;
6935 for &(j, xv) in &act.outliers {
6936 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6937 acc += w * s * xv;
6938 }
6939 unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6941 }
6942 bi += 4;
6943 }
6944 while bi < n {
6946 let act = &acts[bi];
6947 let mut acc = q4tp_cell_dot_v(
6948 v.nib.as_ptr(),
6949 r,
6950 gpr,
6951 act.xq.as_ptr(),
6952 sc.as_ptr(),
6953 act.xg8.as_ptr(),
6954 ) * act.sx;
6955 for &(j, xv) in &act.outliers {
6956 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6957 acc += w * s * xv;
6958 }
6959 unsafe { *out.at((abase + bi) * rows + r) = acc };
6961 bi += 1;
6962 }
6963}
6964
6965fn q4tp_panel_cols(cols: usize) -> usize {
6971 static ENV: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
6972 if let Some(v) = *ENV.get_or_init(|| {
6973 std::env::var("CMF_Q4TP_PANEL")
6974 .ok()
6975 .and_then(|v| v.parse().ok())
6976 .filter(|v: &usize| *v > 0)
6977 }) {
6978 return v;
6979 }
6980 #[cfg(target_arch = "x86_64")]
6981 {
6982 ((l2_bytes() / 2) / cols.max(1)).clamp(8, 256) & !3
6983 }
6984 #[cfg(not(target_arch = "x86_64"))]
6985 {
6986 let _ = cols;
6987 256
6988 }
6989}
6990
6991#[cfg(target_arch = "x86_64")]
6993fn l2_bytes() -> usize {
6994 static L2: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6995 *L2.get_or_init(|| {
6996 let s = std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cache/index2/size").unwrap_or_default();
6997 let s = s.trim();
6998 let (num, mul) = match s.strip_suffix('K') {
6999 Some(n) => (n, 1024),
7000 None => match s.strip_suffix('M') {
7001 Some(n) => (n, 1024 * 1024),
7002 None => (s, 1),
7003 },
7004 };
7005 num.parse::<usize>().ok().map(|n| n * mul).filter(|&b| b >= 64 * 1024).unwrap_or(1 << 20)
7006 })
7007}
7008
7009#[cfg(target_arch = "x86_64")]
7012fn q4tp_fma_x86() -> bool {
7013 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7014 *ON.get_or_init(|| {
7015 std::is_x86_feature_detected!("avx2")
7016 && std::is_x86_feature_detected!("fma")
7017 && std::env::var("CMF_Q4TP_FMA").map(|v| v != "0").unwrap_or(true)
7018 })
7019}
7020
7021#[cfg(target_arch = "x86_64")]
7037#[target_feature(enable = "avx2,fma")]
7038#[allow(clippy::too_many_arguments)]
7039unsafe fn q4tp_matmat_2rows_fma(
7040 v: &Q4tpView<'_>,
7041 r: usize,
7042 gpr: usize,
7043 acts: &[SplitAct],
7044 sc0: &[f32],
7045 sc1: &[f32],
7046 abase: usize,
7047 rows: usize,
7048 out: SendMut,
7049) {
7050 use core::arch::x86_64::*;
7051 debug_assert!(sc0.len() == gpr && sc1.len() == gpr);
7052 let n = acts.len();
7053 let nib = v.nib.as_ptr();
7054 let ones = _mm256_set1_epi16(1);
7055 let mut bi = 0usize;
7056 unsafe {
7057 while bi + 4 <= n {
7058 let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
7059 let xp: [*const i8; 4] = [
7060 a[0].xq.as_ptr(),
7061 a[1].xq.as_ptr(),
7062 a[2].xq.as_ptr(),
7063 a[3].xq.as_ptr(),
7064 ];
7065 let mut c0 = [_mm256_setzero_ps(); 4];
7066 let mut c1 = [_mm256_setzero_ps(); 4];
7067 for gi in 0..gpr {
7068 let w0 = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
7069 let w1 = q4tp_unpack_w(nib.add(((r + 1) * gpr + gi) * Q4TP_NIB));
7070 let s0 = _mm256_set1_ps(*sc0.get_unchecked(gi));
7071 let s1 = _mm256_set1_ps(*sc1.get_unchecked(gi));
7072 for k in 0..4 {
7073 let x = _mm256_loadu_si256(xp[k].add(gi * GROUP_SIZE) as *const __m256i);
7074 let d0 = _mm256_madd_epi16(_mm256_maddubs_epi16(w0, x), ones);
7075 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(w1, x), ones);
7076 c0[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d0), s0, c0[k]);
7077 c1[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d1), s1, c1[k]);
7078 }
7079 }
7080 for (row, cr, scr) in [(r, &c0, sc0), (r + 1, &c1, sc1)] {
7081 for k in 0..4 {
7082 let act = a[k];
7083 let mut acc = (q4tp_hsum256(cr[k]) - q4tp_xg8_corr(&act.xg8, scr)) * act.sx;
7084 for &(j, xv) in &act.outliers {
7085 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7086 acc += w * s * xv;
7087 }
7088 *out.at((abase + bi + k) * rows + row) = acc;
7090 }
7091 }
7092 bi += 4;
7093 }
7094 if bi < n {
7095 q4tp_matmat_row_cols_v(v, r, gpr, &acts[bi..], sc0, abase + bi, rows, out);
7096 q4tp_matmat_row_cols_v(v, r + 1, gpr, &acts[bi..], sc1, abase + bi, rows, out);
7097 }
7098 }
7099}
7100
7101#[cfg(target_arch = "x86_64")]
7110#[target_feature(enable = "avx2,fma")]
7111unsafe fn q4tp_matvec_lanes<const L: usize>(
7112 v: &Q4tpView<'_>,
7113 start: usize,
7114 end: usize,
7115 gpr: usize,
7116 xq: &[i8],
7117 xg8: &[f32],
7118 sx: f32,
7119 outliers: &[(usize, f32)],
7120 sc: &mut [f32],
7121 out: SendMut,
7122) {
7123 use core::arch::x86_64::*;
7124 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7125 crate::cpuprof::count_lane_rows((end - start) as u64);
7126 let mut base = start;
7127 while base + L <= end {
7128 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7129 let mut rows = [0usize; L];
7130 for l in 0..L {
7131 rows[l] = base + l;
7132 }
7133 let v4: [&Q4tpView; L] = [v; L];
7134 q4tp_scales_lanes(&v4, &rows, gpr, sc);
7135 drop(_t0);
7136 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7137 let mut accs = [_mm256_setzero_ps(); L];
7138 for gi in 0..gpr {
7139 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7140 for l in 0..L {
7141 let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7142 let d = q4tp_madd_wx(w, xp);
7143 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7144 accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7145 }
7146 }
7147 drop(_t1);
7148 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7149 for l in 0..L {
7150 let r = rows[l];
7151 let scl = &sc[l * gpr..(l + 1) * gpr];
7152 let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7153 for &(j, xv) in outliers {
7154 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7155 acc += w * s * xv;
7156 }
7157 unsafe { *out.at(r) = acc };
7159 }
7160 drop(_t2);
7161 base += L;
7162 }
7163 for r in base..end {
7167 let scl = &mut sc[..gpr];
7168 v.scales_into(r, gpr, scl);
7169 let scl: &[f32] = scl;
7170 let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
7171 * sx;
7172 for &(j, xv) in outliers {
7173 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7174 acc += w * s * xv;
7175 }
7176 unsafe { *out.at(r) = acc };
7178 }
7179}
7180
7181#[cfg(target_arch = "x86_64")]
7189#[target_feature(enable = "avx2,fma")]
7190unsafe fn q4tp_gateup_lanes<const L: usize>(
7191 gv: &Q4tpView<'_>,
7192 uv: &Q4tpView<'_>,
7193 start: usize,
7194 end: usize,
7195 gpr: usize,
7196 xq: &[i8],
7197 xg8: &[f32],
7198 sx: f32,
7199 limit: f32,
7200 outliers: &[(usize, f32)],
7201 gsc: &mut [f32],
7202 usc: &mut [f32],
7203 out: SendMut,
7204) {
7205 use core::arch::x86_64::*;
7206 debug_assert!(
7207 L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
7208 );
7209 crate::cpuprof::count_lane_rows((end - start) as u64);
7210 let mut base = start;
7211 while base + L <= end {
7212 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7213 let mut rows = [0usize; L];
7214 for l in 0..L {
7215 rows[l] = base + l;
7216 }
7217 let gv4: [&Q4tpView; L] = [gv; L];
7218 let uv4: [&Q4tpView; L] = [uv; L];
7219 q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
7220 q4tp_scales_lanes(&uv4, &rows, gpr, usc);
7221 drop(_t0);
7222 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7223 let mut ag = [_mm256_setzero_ps(); L];
7224 let mut au = [_mm256_setzero_ps(); L];
7225 for gi in 0..gpr {
7226 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7227 for l in 0..L {
7228 let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7229 let d = q4tp_madd_wx(w, xp);
7230 let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
7231 ag[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, ag[l]);
7232 }
7233 for l in 0..L {
7234 let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7235 let d = q4tp_madd_wx(w, xp);
7236 let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
7237 au[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, au[l]);
7238 }
7239 }
7240 drop(_t1);
7241 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7242 for l in 0..L {
7243 let r = rows[l];
7244 let gs = &gsc[l * gpr..(l + 1) * gpr];
7245 let us = &usc[l * gpr..(l + 1) * gpr];
7246 let mut g = (q4tp_hsum256(ag[l]) - q4tp_xg8_corr(xg8, gs)) * sx;
7247 let mut u = (q4tp_hsum256(au[l]) - q4tp_xg8_corr(xg8, us)) * sx;
7248 for &(j, xv) in outliers {
7249 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
7250 g += w * s * xv;
7251 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
7252 u += w * s * xv;
7253 }
7254 unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
7256 }
7257 drop(_t2);
7258 base += L;
7259 }
7260 for r in base..end {
7262 let g = &mut gsc[..gpr];
7263 let u = &mut usc[..gpr];
7264 gv.scales_into(r, gpr, g);
7265 uv.scales_into(r, gpr, u);
7266 let (g, u): (&[f32], &[f32]) = (g, u);
7267 let mut gv_ =
7268 q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
7269 let mut uv_ =
7270 q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
7271 for &(j, xv) in outliers {
7272 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
7273 gv_ += w * s * xv;
7274 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
7275 uv_ += w * s * xv;
7276 }
7277 unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
7279 }
7280}
7281
7282#[cfg(target_arch = "x86_64")]
7290#[target_feature(enable = "avx2,fma")]
7291unsafe fn q4tp_many_lanes<const L: usize>(
7292 views: &[Q4tpView<'_>],
7293 locate: &(dyn Fn(usize) -> (usize, usize)),
7294 start: usize,
7295 end: usize,
7296 gpr: usize,
7297 xq: &[i8],
7298 xg8: &[f32],
7299 sx: f32,
7300 outliers: &[(usize, f32)],
7301 sc: &mut [f32],
7302 outs: &[SendMut],
7303) {
7304 use core::arch::x86_64::*;
7305 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7306 crate::cpuprof::count_lane_rows((end - start) as u64);
7307 let mut base = start;
7308 while base + L <= end {
7309 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7310 let mut tk = [0usize; L];
7311 let mut rows = [0usize; L];
7312 for l in 0..L {
7313 let (t, r) = locate(base + l);
7314 tk[l] = t;
7315 rows[l] = r;
7316 }
7317 let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
7318 q4tp_scales_lanes(&v4, &rows, gpr, sc);
7319 drop(_t0);
7320 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7321 let mut accs = [_mm256_setzero_ps(); L];
7322 for gi in 0..gpr {
7323 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7324 for l in 0..L {
7325 let w = q4tp_unpack_w(
7326 views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
7327 );
7328 let d = q4tp_madd_wx(w, xp);
7329 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7330 accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7331 }
7332 }
7333 drop(_t1);
7334 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7335 for l in 0..L {
7336 let (t, r) = (tk[l], rows[l]);
7337 let scl = &sc[l * gpr..(l + 1) * gpr];
7338 let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7339 for &(j, xv) in outliers {
7340 let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
7341 acc += w * s * xv;
7342 }
7343 unsafe { *outs[t].at(r) = acc };
7345 }
7346 drop(_t2);
7347 base += L;
7348 }
7349 for flat in base..end {
7351 let (t, r) = locate(flat);
7352 let v = &views[t];
7353 let scl = &mut sc[..gpr];
7354 v.scales_into(r, gpr, scl);
7355 let scl: &[f32] = scl;
7356 let mut acc =
7357 q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
7358 for &(j, xv) in outliers {
7359 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7360 acc += w * s * xv;
7361 }
7362 unsafe { *outs[t].at(r) = acc };
7364 }
7365}
7366
7367#[inline]
7370fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7371 #[cfg(target_arch = "x86_64")]
7372 if avx2_enabled() {
7373 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
7375 }
7376 q4tp_row_float_scalar(nib, r, gpr, x, scales)
7377}
7378
7379#[inline]
7380fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7381 let mut acc = 0f32;
7382 for gi in 0..gpr {
7383 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7384 let s = scales[gi];
7385 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7386 let mut ga = 0f32;
7387 for (k, &b) in tile.iter().enumerate() {
7388 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7389 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7390 }
7391 acc += ga * s;
7392 }
7393 acc
7394}
7395
7396#[cfg(target_arch = "x86_64")]
7400#[target_feature(enable = "avx2")]
7401unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7402 unsafe {
7405 use core::arch::x86_64::*;
7406 let mask = _mm_set1_epi8(15);
7407 let eight = _mm_set1_epi8(8);
7408 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
7409 let mut acc = 0.0f32;
7410 for gi in 0..gpr {
7411 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
7412 let lo = _mm_and_si128(packed, mask);
7413 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
7414 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
7415 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
7416 let xp = x.as_ptr().add(gi * GROUP_SIZE);
7417 let a = _mm256_mul_ps(
7418 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
7419 _mm256_loadu_ps(xp),
7420 );
7421 let b = _mm256_mul_ps(
7422 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
7423 _mm256_loadu_ps(xp.add(8)),
7424 );
7425 let c = _mm256_mul_ps(
7426 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
7427 _mm256_loadu_ps(xp.add(16)),
7428 );
7429 let d = _mm256_mul_ps(
7430 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
7431 _mm256_loadu_ps(xp.add(24)),
7432 );
7433 let mut pairs = [0.0f32; 16];
7434 _mm256_storeu_ps(
7435 pairs.as_mut_ptr(),
7436 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
7437 );
7438 _mm256_storeu_ps(
7439 pairs.as_mut_ptr().add(8),
7440 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
7441 );
7442 let mut ga = 0.0f32;
7443 for v in pairs {
7444 ga += v;
7445 }
7446 acc += ga * scales[gi];
7447 }
7448 acc
7449 }
7450}
7451
7452#[inline]
7455fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7456 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7457 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
7458 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
7459 ((n as i32 - 8) as f32, scales[gi])
7460}
7461
7462fn q4tp_matvec(
7464 bytes: &[u8],
7465 x: &[f32],
7466 rows: usize,
7467 cols: usize,
7468 out: &mut [f32],
7469 pool: Option<&Pool>,
7470) {
7471 debug_assert_eq!(out.len(), rows);
7472 let gpr = cols / GROUP_SIZE;
7473 let v = Q4tpView::new(bytes, rows, cols);
7474 let out_addr = SendMut(out.as_mut_ptr());
7475 if a8w8_enabled() {
7476 let act = split_act(x);
7477 #[cfg(target_arch = "x86_64")]
7478 let lanes = q4tp_lanes();
7479 let run = |start: usize, end: usize| {
7480 #[cfg(target_arch = "x86_64")]
7481 if lanes > 1 && act.xg8.len() == gpr {
7482 with_krow(gpr * lanes, |sc| {
7485 unsafe {
7487 match lanes {
7488 2 => q4tp_matvec_lanes::<2>(
7489 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7490 &act.outliers, sc, out_addr,
7491 ),
7492 8 => q4tp_matvec_lanes::<8>(
7493 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7494 &act.outliers, sc, out_addr,
7495 ),
7496 _ => q4tp_matvec_lanes::<4>(
7497 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7498 &act.outliers, sc, out_addr,
7499 ),
7500 }
7501 }
7502 });
7503 return;
7504 }
7505 crate::cpuprof::count_plain_rows((end - start) as u64);
7506 with_krow(gpr, |sc| {
7508 for r in start..end {
7509 v.scales_into(r, gpr, sc);
7510 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
7511 for &(j, xv) in &act.outliers {
7512 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
7513 acc += w * s * xv;
7514 }
7515 unsafe { *out_addr.at(r) = acc }
7517 }
7518 })
7519 };
7520 dispatch_rows(pool, rows, &run);
7521 return;
7522 }
7523 let run = |start: usize, end: usize| {
7524 with_krow(gpr, |sc| {
7525 for r in start..end {
7526 v.scales_into(r, gpr, sc);
7527 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
7529 }
7530 })
7531 };
7532 dispatch_rows(pool, rows, &run);
7533}
7534
7535#[allow(clippy::too_many_arguments)]
7538fn q4tp_matvec2(
7539 bytes: &[u8],
7540 x1: &[f32],
7541 x2: &[f32],
7542 rows: usize,
7543 cols: usize,
7544 o1: &mut [f32],
7545 o2: &mut [f32],
7546 pool: Option<&Pool>,
7547) {
7548 let gpr = cols / GROUP_SIZE;
7549 let v = Q4tpView::new(bytes, rows, cols);
7550 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7551 let run = |start: usize, end: usize| {
7552 let mut sc = vec![0f32; gpr];
7553 for r in start..end {
7554 v.scales_into(r, gpr, &mut sc);
7555 unsafe {
7557 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
7558 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
7559 }
7560 }
7561 };
7562 dispatch_rows(pool, rows, &run);
7563}
7564
7565#[inline]
7568fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7569 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7570 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
7571 let c = (byte >> (2 * (k % 4))) & 3;
7572 (c as f32 - 1.5, scales[gi])
7573}
7574
7575#[cfg(target_arch = "x86_64")]
7576const Q2TP_DECODE_U32: [u32; 256] = {
7577 let mut tab = [0u32; 256];
7578 let mut b = 0usize;
7579 while b < 256 {
7580 tab[b] = ((b as u32) & 3)
7581 | ((((b as u32) >> 2) & 3) << 8)
7582 | ((((b as u32) >> 4) & 3) << 16)
7583 | ((((b as u32) >> 6) & 3) << 24);
7584 b += 1;
7585 }
7586 tab
7587};
7588
7589#[cfg(target_arch = "x86_64")]
7593#[target_feature(enable = "avx2")]
7594unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
7595 use core::arch::x86_64::*;
7596 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
7597 let codes = _mm256_setr_epi32(
7598 Q2TP_DECODE_U32[ch[0] as usize] as i32,
7599 Q2TP_DECODE_U32[ch[1] as usize] as i32,
7600 Q2TP_DECODE_U32[ch[2] as usize] as i32,
7601 Q2TP_DECODE_U32[ch[3] as usize] as i32,
7602 Q2TP_DECODE_U32[ch[4] as usize] as i32,
7603 Q2TP_DECODE_U32[ch[5] as usize] as i32,
7604 Q2TP_DECODE_U32[ch[6] as usize] as i32,
7605 Q2TP_DECODE_U32[ch[7] as usize] as i32,
7606 );
7607 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
7608 let pair = _mm256_maddubs_epi16(codes, xv);
7609 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
7610 let sum128 = _mm_add_epi32(
7611 _mm256_castsi256_si128(quad),
7612 _mm256_extracti128_si256(quad, 1),
7613 );
7614 let sum64 = _mm_hadd_epi32(sum128, sum128);
7615 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
7616}
7617
7618#[cfg(target_arch = "x86_64")]
7619#[inline]
7620fn q2tp_avx2_enabled() -> bool {
7621 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7622 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
7623}
7624
7625#[inline]
7632fn dot_q2tp_row_i8(
7633 chunks: &[u8],
7634 r: usize,
7635 gpr: usize,
7636 xq: &[i8],
7637 gsum: &[i32],
7638 scales: &[f32],
7639) -> f32 {
7640 let mut acc = 0f32;
7641 let base = r * gpr * Q2TP_CHUNK;
7642 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7643 let mut codes = [0i8; GROUP_SIZE];
7644 #[cfg(target_arch = "x86_64")]
7645 let avx2 = q2tp_avx2_enabled();
7650 for gi in 0..gpr {
7651 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7652 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7653 #[cfg(target_arch = "aarch64")]
7654 let dot = unsafe {
7660 use core::arch::aarch64::*;
7661 let b = vld1_u8(ch.as_ptr());
7662 let three = vdup_n_u8(3);
7663 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7664 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7665 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7666 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7667 let x4 = vld4_s8(xg.as_ptr());
7668 let mut acc4 = vdupq_n_s32(0);
7669 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7670 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7671 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7672 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7673 vaddvq_s32(acc4)
7674 };
7675 #[cfg(target_arch = "x86_64")]
7676 let dot: i32 = if avx2 {
7677 unsafe { q2tp_code_dot_avx2(ch, xg) }
7680 } else {
7681 ch.iter()
7682 .enumerate()
7683 .map(|(k, &b)| {
7684 ((b & 3) as i32) * xg[k * 4] as i32
7685 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7686 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7687 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7688 })
7689 .sum()
7690 };
7691 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7692 let dot: i32 = {
7693 for (k, &b) in ch.iter().enumerate() {
7694 codes[k * 4] = (b & 3) as i8;
7695 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7696 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7697 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7698 }
7699 codes
7700 .iter()
7701 .zip(xg)
7702 .map(|(&c, &x)| c as i32 * x as i32)
7703 .sum()
7704 };
7705 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7706 }
7707 acc
7708}
7709
7710fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7714 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7715}
7716
7717#[inline]
7721fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7722 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7723}
7724
7725#[inline]
7726fn q2tp_row_exact_center(
7727 chunks: &[u8],
7728 r: usize,
7729 gpr: usize,
7730 x: &[f32],
7731 scales: &[f32],
7732 center: f32,
7733) -> f32 {
7734 let mut acc = 0f32;
7735 for gi in 0..gpr {
7736 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7737 let s = scales[gi];
7738 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7739 let mut g = 0f32;
7740 for (k, &b) in ch.iter().enumerate() {
7741 g += ((b & 3) as f32 - center) * xb[k * 4]
7742 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7743 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7744 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7745 }
7746 acc += s * g;
7747 }
7748 acc
7749}
7750
7751fn q2tp_matvec(
7752 bytes: &[u8],
7753 x: &[f32],
7754 rows: usize,
7755 cols: usize,
7756 out: &mut [f32],
7757 pool: Option<&Pool>,
7758) {
7759 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7760}
7761
7762fn q2tp_affine_matvec(
7763 bytes: &[u8],
7764 x: &[f32],
7765 rows: usize,
7766 cols: usize,
7767 out: &mut [f32],
7768 pool: Option<&Pool>,
7769) {
7770 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7771}
7772
7773fn q2tp_matvec_mode(
7774 bytes: &[u8],
7775 x: &[f32],
7776 rows: usize,
7777 cols: usize,
7778 out: &mut [f32],
7779 pool: Option<&Pool>,
7780 affine: bool,
7781) {
7782 debug_assert_eq!(out.len(), rows);
7783 let gpr = cols / GROUP_SIZE;
7784 let v = Q4tpView::new_q2(bytes, rows, cols);
7785 let out_addr = SendMut(out.as_mut_ptr());
7786 if !affine && a8w8_enabled() {
7791 let act = split_act(x);
7792 let gsum = q1_group_sums(&act.xq, gpr);
7793 let (act, gsum) = (&act, &gsum);
7794 let run = move |start: usize, end: usize| {
7795 with_krow(gpr, |sc| {
7796 for r in start..end {
7797 v.scales_into(r, gpr, sc);
7798 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7799 for &(j, xv) in &act.outliers {
7800 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7801 acc += w * s * xv;
7802 }
7803 unsafe { *out_addr.at(r) = acc };
7805 }
7806 })
7807 };
7808 dispatch_rows(pool, rows, &run);
7809 return;
7810 }
7811 let run = |start: usize, end: usize| {
7812 with_krow(gpr, |sc| {
7813 for r in start..end {
7814 v.scales_into(r, gpr, sc);
7815 unsafe {
7817 *out_addr.at(r) = if affine {
7818 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7819 } else {
7820 q2tp_row_exact(v.nib, r, gpr, x, sc)
7821 }
7822 };
7823 }
7824 })
7825 };
7826 dispatch_rows(pool, rows, &run);
7827}
7828
7829#[allow(clippy::too_many_arguments)]
7831fn q2tp_matvec2(
7832 bytes: &[u8],
7833 x1: &[f32],
7834 x2: &[f32],
7835 rows: usize,
7836 cols: usize,
7837 o1: &mut [f32],
7838 o2: &mut [f32],
7839 pool: Option<&Pool>,
7840) {
7841 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7842}
7843
7844#[allow(clippy::too_many_arguments)]
7845fn q2tp_affine_matvec2(
7846 bytes: &[u8],
7847 x1: &[f32],
7848 x2: &[f32],
7849 rows: usize,
7850 cols: usize,
7851 o1: &mut [f32],
7852 o2: &mut [f32],
7853 pool: Option<&Pool>,
7854) {
7855 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7856}
7857
7858#[allow(clippy::too_many_arguments)]
7859fn q2tp_matvec2_mode(
7860 bytes: &[u8],
7861 x1: &[f32],
7862 x2: &[f32],
7863 rows: usize,
7864 cols: usize,
7865 o1: &mut [f32],
7866 o2: &mut [f32],
7867 pool: Option<&Pool>,
7868 affine: bool,
7869) {
7870 let gpr = cols / GROUP_SIZE;
7871 let v = Q4tpView::new_q2(bytes, rows, cols);
7872 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7873 let run = |start: usize, end: usize| {
7874 let mut sc = vec![0f32; gpr];
7875 for r in start..end {
7876 v.scales_into(r, gpr, &mut sc);
7877 unsafe {
7879 *p1.at(r) = if affine {
7880 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7881 } else {
7882 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7883 };
7884 *p2.at(r) = if affine {
7885 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7886 } else {
7887 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7888 };
7889 }
7890 }
7891 };
7892 dispatch_rows(pool, rows, &run);
7893}
7894
7895pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7902 let gpr = cols / GROUP_SIZE;
7907 let v = Q4tpView::new_q2(bytes, rows, cols);
7908 with_krow(gpr, |sc| {
7909 for r in 0..rows {
7910 v.scales_into(r, gpr, sc);
7911 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7912 }
7913 });
7914}
7915
7916pub fn q2tp_affine_matvec_for_test(
7919 bytes: &[u8],
7920 x: &[f32],
7921 rows: usize,
7922 cols: usize,
7923 out: &mut [f32],
7924) {
7925 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7926}
7927
7928pub fn q2tp_matmat_for_test(
7929 bytes: &[u8],
7930 xs_all: &[f32],
7931 b: usize,
7932 rows: usize,
7933 cols: usize,
7934 out: &mut [f32],
7935) {
7936 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7937}
7938
7939fn q2tp_matmat(
7940 bytes: &[u8],
7941 xs_all: &[f32],
7942 b: usize,
7943 rows: usize,
7944 cols: usize,
7945 out: &mut [f32],
7946 pool: Option<&Pool>,
7947) {
7948 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7949}
7950
7951fn q2tp_affine_matmat(
7952 bytes: &[u8],
7953 xs_all: &[f32],
7954 b: usize,
7955 rows: usize,
7956 cols: usize,
7957 out: &mut [f32],
7958 pool: Option<&Pool>,
7959) {
7960 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7961}
7962
7963fn q2tp_matmat_mode(
7964 bytes: &[u8],
7965 xs_all: &[f32],
7966 b: usize,
7967 rows: usize,
7968 cols: usize,
7969 out: &mut [f32],
7970 pool: Option<&Pool>,
7971 affine: bool,
7972) {
7973 debug_assert_eq!(out.len(), b * rows);
7974 let gpr = cols / GROUP_SIZE;
7975 let v = Q4tpView::new_q2(bytes, rows, cols);
7976 let out_addr = SendMut(out.as_mut_ptr());
7977 let run = |start: usize, end: usize| {
7978 let mut sc = vec![0f32; gpr];
7979 for r in start..end {
7980 v.scales_into(r, gpr, &mut sc);
7981 for bi in 0..b {
7982 let x = &xs_all[bi * cols..(bi + 1) * cols];
7983 unsafe {
7985 *out_addr.at(bi * rows + r) = if affine {
7986 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7987 } else {
7988 q2tp_row_exact(v.nib, r, gpr, x, &sc)
7989 }
7990 };
7991 }
7992 }
7993 };
7994 dispatch_rows(pool, rows, &run);
7995}
7996
7997#[cfg(target_arch = "aarch64")]
8005#[target_feature(enable = "neon,dotprod")]
8006unsafe fn dot_q4tp_row_1x4_sdot_v1(
8007 nib: &[u8],
8008 r: usize,
8009 gpr: usize,
8010 xs: [&[i8]; 4],
8011 scales: &[f32],
8012) -> [f32; 4] {
8013 unsafe {
8014 use core::arch::aarch64::*;
8015 use core::arch::asm;
8016 let lomask = vdupq_n_u8(0x0F);
8017 let eight = vdupq_n_s8(8);
8018 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
8019 for gi in 0..gpr {
8020 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8021 let s = *scales.get_unchecked(gi);
8022 let bb = vld1q_u8(t);
8023 let lo = vandq_u8(bb, lomask);
8024 let hi = vshrq_n_u8::<4>(bb);
8025 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8026 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8027 let mut d = [0f32; 4];
8028 for (k, dk) in d.iter_mut().enumerate() {
8029 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
8030 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
8031 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8032 asm!(
8033 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8034 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8035 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8036 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8037 options(pure, nomem, nostack),
8038 );
8039 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8040 }
8041 f0 += d[0];
8042 f1 += d[1];
8043 f2 += d[2];
8044 f3 += d[3];
8045 }
8046 [f0, f1, f2, f3]
8047 }
8048}
8049
8050#[allow(dead_code)]
8057static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
8058
8059#[cfg(test)]
8062static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
8063
8064#[cfg(target_arch = "x86_64")]
8070fn q4tp_blocked_x86() -> bool {
8071 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8072 1 => false,
8073 2 => avx512vnni_enabled(),
8078 _ => blocked_enabled() && avx512vnni_enabled(),
8082 }
8083}
8084
8085#[cfg(target_arch = "aarch64")]
8087#[allow(dead_code)]
8088fn q4tp_v1() -> bool {
8089 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8090 1 => true,
8091 2 => false,
8092 _ => {
8093 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8094 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
8095 }
8096 }
8097}
8098
8099#[cfg(target_arch = "x86_64")]
8111#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8112unsafe fn dot_q4tp_2x8_avx512(
8113 nib: &[u8],
8114 r0: usize,
8115 gpr: usize,
8116 xs: [&[i8]; 8],
8117 sc0: &[f32],
8118 sc1: &[f32],
8119) -> [[f32; 8]; 2] {
8120 unsafe {
8123 use core::arch::x86_64::*;
8124 let lomask = _mm256_set1_epi8(0x0F);
8125 let eight = _mm256_set1_epi8(8);
8126 let zero = _mm512_setzero_si512();
8127 let mut v0 = [_mm512_setzero_ps(); 8];
8128 let mut v1 = [_mm512_setzero_ps(); 8];
8129 let pairs = gpr / 2;
8130 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
8131 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8132 let bb = _mm256_loadu_si256(t as *const __m256i);
8133 let lo = _mm256_and_si256(bb, lomask);
8134 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8135 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8136 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8137 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8138 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8139 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
8140 };
8141 for gp in 0..pairs {
8142 let gi = gp * 2;
8143 let (wa0, neg0) = unpack(r0, gi);
8144 let (wa1, neg1) = unpack(r0 + 1, gi);
8145 let off = gi * GROUP_SIZE;
8146 let sv = |sc: &[f32]| {
8147 _mm512_insertf32x8::<1>(
8148 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
8149 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
8150 )
8151 };
8152 let s0 = sv(sc0);
8153 let s1 = sv(sc1);
8154 for k in 0..8 {
8155 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
8156 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8157 zero,
8158 wa0,
8159 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
8160 ));
8161 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8162 zero,
8163 wa1,
8164 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
8165 ));
8166 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
8167 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
8168 }
8169 }
8170 let mut acc = [[0f32; 8]; 2];
8171 for k in 0..8 {
8172 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
8173 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
8174 }
8175 if gpr % 2 == 1 {
8176 let off = (gpr - 1) * GROUP_SIZE;
8177 for j in off..off + GROUP_SIZE {
8178 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
8179 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
8180 for k in 0..8 {
8181 let x = *xs[k].get_unchecked(j) as f32;
8182 acc[0][k] += w0 * sa * x;
8183 acc[1][k] += w1 * sb * x;
8184 }
8185 }
8186 }
8187 acc
8188 }
8189}
8190
8191#[cfg(target_arch = "x86_64")]
8196#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8197unsafe fn dot_q4tp_row_1x8_avx512(
8198 nib: &[u8],
8199 r: usize,
8200 gpr: usize,
8201 xs: [&[i8]; 8],
8202 scales: &[f32],
8203) -> [f32; 8] {
8204 unsafe {
8206 use core::arch::x86_64::*;
8207 let lomask = _mm256_set1_epi8(0x0F);
8208 let eight = _mm256_set1_epi8(8);
8209 let zero = _mm512_setzero_si512();
8210 let (mut v0, mut v1, mut v2, mut v3) = (
8211 _mm512_setzero_ps(),
8212 _mm512_setzero_ps(),
8213 _mm512_setzero_ps(),
8214 _mm512_setzero_ps(),
8215 );
8216 let (mut v4, mut v5, mut v6, mut v7) = (
8217 _mm512_setzero_ps(),
8218 _mm512_setzero_ps(),
8219 _mm512_setzero_ps(),
8220 _mm512_setzero_ps(),
8221 );
8222 let pairs = gpr / 2;
8223 for gp in 0..pairs {
8224 let gi = gp * 2;
8225 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8226 let bb = _mm256_loadu_si256(t as *const __m256i);
8227 let lo = _mm256_and_si256(bb, lomask);
8228 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8229 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8234 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8235 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8236 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8237 let wabs = _mm512_abs_epi8(w);
8238 let neg = _mm512_movepi8_mask(w);
8239 let off = gi * GROUP_SIZE;
8240 let sv = _mm512_insertf32x8::<1>(
8241 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8242 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8243 );
8244 let dot = |x: &[i8]| -> __m512 {
8245 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8246 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8247 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8248 };
8249 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8250 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8251 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8252 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8253 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
8254 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
8255 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
8256 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
8257 }
8258 let mut acc = [
8259 _mm512_reduce_add_ps(v0),
8260 _mm512_reduce_add_ps(v1),
8261 _mm512_reduce_add_ps(v2),
8262 _mm512_reduce_add_ps(v3),
8263 _mm512_reduce_add_ps(v4),
8264 _mm512_reduce_add_ps(v5),
8265 _mm512_reduce_add_ps(v6),
8266 _mm512_reduce_add_ps(v7),
8267 ];
8268 if gpr % 2 == 1 {
8271 let off = (gpr - 1) * GROUP_SIZE;
8272 for j in off..off + GROUP_SIZE {
8273 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8274 let ws = w * s;
8275 for k in 0..8 {
8276 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8277 }
8278 }
8279 }
8280 acc
8281 }
8282}
8283
8284#[cfg(target_arch = "x86_64")]
8297#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8298unsafe fn dot_q4tp_row_1x4_avx512(
8299 nib: &[u8],
8300 r: usize,
8301 gpr: usize,
8302 xs: [&[i8]; 4],
8303 scales: &[f32],
8304) -> [f32; 4] {
8305 unsafe {
8307 use core::arch::x86_64::*;
8308 let lomask = _mm256_set1_epi8(0x0F);
8309 let eight = _mm256_set1_epi8(8);
8310 let zero = _mm512_setzero_si512();
8311 let (mut v0, mut v1, mut v2, mut v3) = (
8312 _mm512_setzero_ps(),
8313 _mm512_setzero_ps(),
8314 _mm512_setzero_ps(),
8315 _mm512_setzero_ps(),
8316 );
8317 let pairs = gpr / 2;
8318 for gp in 0..pairs {
8319 let gi = gp * 2;
8320 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8321 let bb = _mm256_loadu_si256(t as *const __m256i);
8322 let lo = _mm256_and_si256(bb, lomask);
8323 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8324 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8329 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8330 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8331 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8332 let wabs = _mm512_abs_epi8(w);
8333 let neg = _mm512_movepi8_mask(w);
8334 let off = gi * GROUP_SIZE;
8335 let sv = _mm512_insertf32x8::<1>(
8336 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8337 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8338 );
8339 let dot = |x: &[i8]| -> __m512 {
8340 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8341 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8342 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8343 };
8344 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8345 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8346 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8347 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8348 }
8349 let mut acc = [
8350 _mm512_reduce_add_ps(v0),
8351 _mm512_reduce_add_ps(v1),
8352 _mm512_reduce_add_ps(v2),
8353 _mm512_reduce_add_ps(v3),
8354 ];
8355 if gpr % 2 == 1 {
8358 let off = (gpr - 1) * GROUP_SIZE;
8359 for j in off..off + GROUP_SIZE {
8360 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8361 let ws = w * s;
8362 for k in 0..4 {
8363 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8364 }
8365 }
8366 }
8367 acc
8368 }
8369}
8370
8371#[cfg(target_arch = "aarch64")]
8375#[target_feature(enable = "neon,dotprod")]
8376unsafe fn dot_q4tp_row_1x4_sdot(
8377 nib: &[u8],
8378 r: usize,
8379 gpr: usize,
8380 xs: [&[i8]; 4],
8381 scales: &[f32],
8382) -> [f32; 4] {
8383 unsafe {
8385 use core::arch::aarch64::*;
8386 use core::arch::asm;
8387 let lomask = vdupq_n_u8(0x0F);
8388 let eight = vdupq_n_s8(8);
8389 let (mut v0, mut v1, mut v2, mut v3) = (
8405 vdupq_n_f32(0.0),
8406 vdupq_n_f32(0.0),
8407 vdupq_n_f32(0.0),
8408 vdupq_n_f32(0.0),
8409 );
8410 for gi in 0..gpr {
8411 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8412 let s = *scales.get_unchecked(gi);
8413 let bb = vld1q_u8(t);
8414 let lo = vandq_u8(bb, lomask);
8415 let hi = vshrq_n_u8::<4>(bb);
8416 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8417 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8418 let off = gi * GROUP_SIZE;
8419 let dot4 = |x: &[i8]| -> int32x4_t {
8420 let x0 = vld1q_s8(x.as_ptr().add(off));
8421 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
8422 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8423 asm!(
8424 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8425 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8426 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8427 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8428 options(pure, nomem, nostack),
8429 );
8430 vaddq_s32(a0, a1)
8431 };
8432 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
8433 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
8434 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
8435 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
8436 }
8437 [
8438 vaddvq_f32(v0),
8439 vaddvq_f32(v1),
8440 vaddvq_f32(v2),
8441 vaddvq_f32(v3),
8442 ]
8443 }
8444}
8445
8446fn q4tp_matmat(
8454 bytes: &[u8],
8455 xs_all: &[f32],
8456 b: usize,
8457 rows: usize,
8458 cols: usize,
8459 out: &mut [f32],
8460 pool: Option<&Pool>,
8461) {
8462 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
8463}
8464
8465#[allow(clippy::too_many_arguments)]
8469fn q4tp_matmat_with(
8470 bytes: &[u8],
8471 xs_all: &[f32],
8472 b: usize,
8473 rows: usize,
8474 cols: usize,
8475 out: &mut [f32],
8476 pool: Option<&Pool>,
8477 exact: bool,
8478) {
8479 debug_assert_eq!(out.len(), b * rows);
8480 let gpr = cols / GROUP_SIZE;
8481 let v = Q4tpView::new(bytes, rows, cols);
8482
8483 #[cfg(target_os = "macos")]
8487 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8488 dequant_matmat_accel(
8489 &|r, dst| {
8490 let mut sc = [0f32; 32];
8491 let mut scv;
8492 let s: &[f32] = if gpr <= 32 {
8493 v.scales_into(r, gpr, &mut sc);
8494 &sc[..gpr]
8495 } else {
8496 scv = vec![0f32; gpr];
8497 v.scales_into(r, gpr, &mut scv);
8498 &scv
8499 };
8500 for gi in 0..gpr {
8501 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
8502 for (k, &bb) in tile.iter().enumerate() {
8503 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
8504 dst[gi * GROUP_SIZE + k * 2 + 1] =
8505 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
8506 }
8507 }
8508 },
8509 xs_all,
8510 b,
8511 rows,
8512 cols,
8513 out,
8514 pool,
8515 );
8516 return;
8517 }
8518
8519 let out_addr = SendMut(out.as_mut_ptr());
8520 if a8w8_enabled() {
8521 let acts = split_acts_pooled(xs_all, b, cols, pool);
8522 let acts = &acts;
8523 #[cfg(target_arch = "aarch64")]
8530 let blocked_ok = sdot_enabled() && blocked_enabled();
8531 #[cfg(target_arch = "x86_64")]
8538 let blocked_ok = q4tp_blocked_x86() && !exact;
8539 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8540 let blocked_ok = {
8541 let _ = exact;
8542 false
8543 };
8544 let panel_cols = q4tp_panel_cols(cols);
8554 let run = |start: usize, end: usize| {
8555 for abase in (0..acts.len()).step_by(panel_cols) {
8556 let alen = (acts.len() - abase).min(panel_cols);
8557 let mut sc = vec![0f32; gpr];
8558 #[cfg(target_arch = "x86_64")]
8559 let mut r_lo = start;
8560 #[cfg(target_arch = "x86_64")]
8561 if blocked_ok && alen >= 8 {
8562 let mut sc1 = vec![0f32; gpr];
8563 while r_lo + 2 <= end {
8564 v.scales_into(r_lo, gpr, &mut sc);
8565 v.scales_into(r_lo + 1, gpr, &mut sc1);
8566 let mut bi = 0usize;
8567 while bi + 8 <= alen {
8568 let xs = [
8569 acts[abase + bi].xq.as_slice(),
8570 acts[abase + bi + 1].xq.as_slice(),
8571 acts[abase + bi + 2].xq.as_slice(),
8572 acts[abase + bi + 3].xq.as_slice(),
8573 acts[abase + bi + 4].xq.as_slice(),
8574 acts[abase + bi + 5].xq.as_slice(),
8575 acts[abase + bi + 6].xq.as_slice(),
8576 acts[abase + bi + 7].xq.as_slice(),
8577 ];
8578 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
8579 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
8580 for k in 0..8 {
8581 let act = &acts[abase + bi + k];
8582 let mut acc = dr[k] * act.sx;
8583 for &(j, xv) in &act.outliers {
8584 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8585 acc += w * s * xv;
8586 }
8587 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
8589 }
8590 }
8591 bi += 8;
8592 }
8593 for row in [r_lo, r_lo + 1] {
8596 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
8597 for b2 in bi..alen {
8598 let act = &acts[abase + b2];
8599 let xs4 = [
8600 act.xq.as_slice(),
8601 act.xq.as_slice(),
8602 act.xq.as_slice(),
8603 act.xq.as_slice(),
8604 ];
8605 let d =
8606 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
8607 let mut acc = d[0] * act.sx;
8608 for &(j, xv) in &act.outliers {
8609 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8610 acc += w * s * xv;
8611 }
8612 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
8614 }
8615 }
8616 r_lo += 2;
8617 }
8618 }
8619 #[cfg(target_arch = "x86_64")]
8624 if !blocked_ok
8625 && !exact
8626 && alen >= 4
8627 && q4tp_fma_x86()
8628 && acts[abase..abase + alen].iter().all(|a| a.xg8.len() == gpr)
8629 {
8630 let mut sc1 = vec![0f32; gpr];
8631 while r_lo + 2 <= end {
8632 v.scales_into(r_lo, gpr, &mut sc);
8633 v.scales_into(r_lo + 1, gpr, &mut sc1);
8634 unsafe {
8637 q4tp_matmat_2rows_fma(
8638 &v,
8639 r_lo,
8640 gpr,
8641 &acts[abase..abase + alen],
8642 &sc,
8643 &sc1,
8644 abase,
8645 rows,
8646 out_addr,
8647 )
8648 };
8649 r_lo += 2;
8650 }
8651 }
8652 #[cfg(target_arch = "x86_64")]
8653 let row_start = r_lo;
8654 #[cfg(not(target_arch = "x86_64"))]
8655 let row_start = start;
8656 for r in row_start..end {
8657 v.scales_into(r, gpr, &mut sc);
8658 let mut bi = 0usize;
8659 #[cfg(target_arch = "x86_64")]
8660 if blocked_ok {
8661 while bi + 8 <= alen {
8662 let xs = [
8663 acts[abase + bi].xq.as_slice(),
8664 acts[abase + bi + 1].xq.as_slice(),
8665 acts[abase + bi + 2].xq.as_slice(),
8666 acts[abase + bi + 3].xq.as_slice(),
8667 acts[abase + bi + 4].xq.as_slice(),
8668 acts[abase + bi + 5].xq.as_slice(),
8669 acts[abase + bi + 6].xq.as_slice(),
8670 acts[abase + bi + 7].xq.as_slice(),
8671 ];
8672 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8673 for k in 0..8 {
8674 let act = &acts[abase + bi + k];
8675 let mut acc = d[k] * act.sx;
8676 for &(j, xv) in &act.outliers {
8677 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8678 acc += w * s * xv;
8679 }
8680 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8682 }
8683 bi += 8;
8684 }
8685 while bi + 4 <= alen {
8686 let xs = [
8687 acts[abase + bi].xq.as_slice(),
8688 acts[abase + bi + 1].xq.as_slice(),
8689 acts[abase + bi + 2].xq.as_slice(),
8690 acts[abase + bi + 3].xq.as_slice(),
8691 ];
8692 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8693 for k in 0..4 {
8694 let act = &acts[abase + bi + k];
8695 let mut acc = d[k] * act.sx;
8696 for &(j, xv) in &act.outliers {
8697 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8698 acc += w * s * xv;
8699 }
8700 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8702 }
8703 bi += 4;
8704 }
8705 }
8706 #[cfg(target_arch = "aarch64")]
8707 if blocked_ok {
8708 while bi + 4 <= alen {
8709 let xs = [
8710 acts[abase + bi].xq.as_slice(),
8711 acts[abase + bi + 1].xq.as_slice(),
8712 acts[abase + bi + 2].xq.as_slice(),
8713 acts[abase + bi + 3].xq.as_slice(),
8714 ];
8715 let d = unsafe {
8716 if exact || q4tp_v1() {
8717 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8718 } else {
8719 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8720 }
8721 };
8722 for k in 0..4 {
8723 let act = &acts[abase + bi + k];
8724 let mut acc = d[k] * act.sx;
8725 for &(j, xv) in &act.outliers {
8726 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8727 acc += w * s * xv;
8728 }
8729 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8731 }
8732 bi += 4;
8733 }
8734 }
8735 let _ = blocked_ok;
8736 #[cfg(target_arch = "x86_64")]
8742 if acts[abase..abase + alen]
8743 .iter()
8744 .all(|a| a.xg8.len() == gpr)
8745 {
8746 unsafe {
8749 q4tp_matmat_row_cols_v(
8750 &v,
8751 r,
8752 gpr,
8753 &acts[abase + bi..abase + alen],
8754 &sc,
8755 abase + bi,
8756 rows,
8757 out_addr,
8758 )
8759 };
8760 bi = alen;
8761 }
8762 while bi < alen {
8763 let act = &acts[abase + bi];
8764 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8765 for &(j, xv) in &act.outliers {
8766 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8767 acc += w * s * xv;
8768 }
8769 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8771 bi += 1;
8772 }
8773 }
8774 }
8775 };
8776 dispatch_rows(pool, rows, &run);
8777 return;
8778 }
8779
8780 let run = |start: usize, end: usize| {
8781 let mut sc = vec![0f32; gpr];
8782 for r in start..end {
8783 v.scales_into(r, gpr, &mut sc);
8784 for bi in 0..b {
8785 let x = &xs_all[bi * cols..(bi + 1) * cols];
8786 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8788 }
8789 }
8790 };
8791 dispatch_rows(pool, rows, &run);
8792}
8793
8794fn q4t_matvec(
8796 bytes: &[u8],
8797 x: &[f32],
8798 rows: usize,
8799 cols: usize,
8800 out: &mut [f32],
8801 pool: Option<&Pool>,
8802) {
8803 debug_assert_eq!(out.len(), rows);
8804 let gpr = cols / GROUP_SIZE;
8805 let out_addr = SendMut(out.as_mut_ptr());
8806 if a8w8_enabled() {
8807 let act = split_act(x);
8808 let run = move |start: usize, end: usize| {
8809 for r in start..end {
8810 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8811 for &(j, xv) in &act.outliers {
8812 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8813 acc += w * s * xv;
8814 }
8815 unsafe { *out_addr.at(r) = acc };
8817 }
8818 };
8819 dispatch_rows(pool, rows, &run);
8820 return;
8821 }
8822 let run = move |start: usize, end: usize| {
8823 for r in start..end {
8824 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8826 }
8827 };
8828 dispatch_rows(pool, rows, &run);
8829}
8830
8831#[allow(clippy::too_many_arguments)]
8833fn q4t_matvec2(
8834 bytes: &[u8],
8835 x1: &[f32],
8836 x2: &[f32],
8837 rows: usize,
8838 cols: usize,
8839 o1: &mut [f32],
8840 o2: &mut [f32],
8841 pool: Option<&Pool>,
8842) {
8843 let gpr = cols / GROUP_SIZE;
8844 let p1 = SendMut(o1.as_mut_ptr());
8845 let p2 = SendMut(o2.as_mut_ptr());
8846 if a8w8_enabled() {
8847 let a1 = split_act(x1);
8848 let a2 = split_act(x2);
8849 let run = move |start: usize, end: usize| {
8850 for r in start..end {
8851 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8852 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8853 for &(j, xv) in &a1.outliers {
8854 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8855 v1 += w * s * xv;
8856 }
8857 for &(j, xv) in &a2.outliers {
8858 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8859 v2 += w * s * xv;
8860 }
8861 unsafe {
8863 *p1.at(r) = v1;
8864 *p2.at(r) = v2;
8865 }
8866 }
8867 };
8868 dispatch_rows(pool, rows, &run);
8869 return;
8870 }
8871 let run = move |start: usize, end: usize| {
8872 for r in start..end {
8873 unsafe {
8875 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8876 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8877 }
8878 }
8879 };
8880 dispatch_rows(pool, rows, &run);
8881}
8882
8883#[allow(clippy::too_many_arguments)]
8885#[cfg(target_os = "macos")]
8891fn dequant_matmat_accel(
8892 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8893 xs_all: &[f32],
8894 b: usize,
8895 rows: usize,
8896 cols: usize,
8897 out: &mut [f32],
8898 pool: Option<&Pool>,
8899) {
8900 const TR: usize = 2048;
8901 thread_local! {
8902 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8903 }
8904 WTILE.with(|wt| {
8905 let mut wtile = wt.borrow_mut();
8906 wtile.resize(TR * cols, 0.0);
8907 let mut r0 = 0usize;
8908 while r0 < rows {
8909 let tr = TR.min(rows - r0);
8910 let wt_addr = SendMut(wtile.as_mut_ptr());
8911 let run = |start: usize, end: usize| {
8912 for r in start..end {
8913 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8915 dequant_row(r0 + r, dst);
8916 }
8917 };
8918 dispatch_rows(pool, tr, &run);
8919 unsafe {
8920 accel_blas::cblas_sgemm(
8921 101, 111, 112, b as i32,
8925 tr as i32,
8926 cols as i32,
8927 1.0,
8928 xs_all.as_ptr(),
8929 cols as i32,
8930 wtile.as_ptr(),
8931 cols as i32,
8932 0.0,
8933 out.as_mut_ptr().add(r0),
8934 rows as i32,
8935 );
8936 }
8937 r0 += tr;
8938 }
8939 });
8940}
8941
8942fn q4t_matmat(
8943 bytes: &[u8],
8944 xs_all: &[f32],
8945 b: usize,
8946 rows: usize,
8947 cols: usize,
8948 out: &mut [f32],
8949 pool: Option<&Pool>,
8950) {
8951 debug_assert_eq!(out.len(), b * rows);
8952 let gpr = cols / GROUP_SIZE;
8953 #[cfg(target_os = "macos")]
8957 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8958 dequant_matmat_accel(
8959 &|r, dst| {
8960 for gi in 0..gpr {
8961 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8962 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8963 for (k, &bb) in tile[2..].iter().enumerate() {
8964 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8965 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8966 }
8967 }
8968 },
8969 xs_all,
8970 b,
8971 rows,
8972 cols,
8973 out,
8974 pool,
8975 );
8976 return;
8977 }
8978 let out_addr = SendMut(out.as_mut_ptr());
8979 if a8w8_enabled() {
8980 let acts: Vec<SplitAct> = (0..b)
8981 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8982 .collect();
8983 let acts = &acts;
8984 #[cfg(target_arch = "x86_64")]
8985 let blocked_ok = avx2_enabled() && blocked_enabled();
8986 #[cfg(target_arch = "aarch64")]
8987 let blocked_ok = sdot_enabled() && blocked_enabled();
8988 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8989 let blocked_ok = false;
8990 let run = move |start: usize, end: usize| {
8991 for r in start..end {
8992 let mut bi = 0usize;
8993 #[cfg(target_arch = "aarch64")]
8994 if blocked_ok {
8995 while bi + 4 <= acts.len() {
8996 let xs = [
8997 acts[bi].xq.as_slice(),
8998 acts[bi + 1].xq.as_slice(),
8999 acts[bi + 2].xq.as_slice(),
9000 acts[bi + 3].xq.as_slice(),
9001 ];
9002 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
9003 for k in 0..4 {
9004 let act = &acts[bi + k];
9005 let mut acc = d[k] * act.sx;
9006 for &(j, xv) in &act.outliers {
9007 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
9008 acc += w * sc * xv;
9009 }
9010 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9012 }
9013 bi += 4;
9014 }
9015 }
9016 #[cfg(target_arch = "x86_64")]
9017 if blocked_ok {
9018 while bi + 4 <= acts.len() {
9019 let xs = [
9020 acts[bi].xq.as_slice(),
9021 acts[bi + 1].xq.as_slice(),
9022 acts[bi + 2].xq.as_slice(),
9023 acts[bi + 3].xq.as_slice(),
9024 ];
9025 let d = unsafe {
9026 if vnni_tiles_enabled() {
9027 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
9028 } else {
9029 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
9030 }
9031 };
9032 for k in 0..4 {
9033 let act = &acts[bi + k];
9034 let mut acc = d[k] * act.sx;
9035 for &(j, xv) in &act.outliers {
9036 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
9037 acc += w * sc * xv;
9038 }
9039 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9041 }
9042 bi += 4;
9043 }
9044 }
9045 let _ = blocked_ok;
9046 while bi < acts.len() {
9047 let act = &acts[bi];
9048 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9049 for &(j, xv) in &act.outliers {
9050 let (w, s) = q4t_outlier(bytes, r, gpr, j);
9051 acc += w * s * xv;
9052 }
9053 unsafe { *out_addr.at(bi * rows + r) = acc };
9055 bi += 1;
9056 }
9057 }
9058 };
9059 dispatch_rows(pool, rows, &run);
9060 return;
9061 }
9062 let run = move |start: usize, end: usize| {
9063 for r in start..end {
9064 for bi in 0..b {
9065 let x = &xs_all[bi * cols..(bi + 1) * cols];
9066 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
9068 }
9069 }
9070 };
9071 dispatch_rows(pool, rows, &run);
9072}
9073
9074fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
9083 (0..gpr)
9084 .map(|gi| {
9085 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
9086 .iter()
9087 .map(|&v| v as i32)
9088 .sum()
9089 })
9090 .collect()
9091}
9092
9093#[inline]
9097#[allow(unreachable_code)]
9098#[cfg(target_arch = "x86_64")]
9103#[target_feature(enable = "avx2")]
9104unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9105 unsafe {
9107 use core::arch::x86_64::*;
9108 let expand = _mm256_setr_epi8(
9110 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9111 3, 3, 3,
9112 );
9113 let bitsel = _mm256_setr_epi8(
9114 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9115 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9116 );
9117 let ones8 = _mm256_set1_epi8(1);
9118 let ones16 = _mm256_set1_epi16(1);
9119 let mut acc = 0f32;
9120 for gi in 0..gpr {
9121 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9122 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9123 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9124 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9125 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9126 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9127 let sel = _mm256_and_si256(x, mask);
9128 let p16 = _mm256_maddubs_epi16(ones8, sel);
9130 let d32 = _mm256_madd_epi16(p16, ones16);
9131 let hi128 = _mm256_extracti128_si256::<1>(d32);
9132 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9133 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9134 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9135 let msum = _mm_cvtsi128_si32(s32);
9136 let d = 2 * msum - gsum[gi];
9139 acc += d as f32 * s;
9140 }
9141 acc
9142 }
9143}
9144
9145#[cfg(target_arch = "x86_64")]
9148#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9149unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9150 unsafe {
9152 use core::arch::x86_64::*;
9153 let expand = _mm256_setr_epi8(
9154 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9155 3, 3, 3,
9156 );
9157 let bitsel = _mm256_setr_epi8(
9158 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9159 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9160 );
9161 let ones8 = _mm256_set1_epi8(1);
9162 let mut acc = 0f32;
9163 for gi in 0..gpr {
9164 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9165 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9166 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9167 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9168 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9169 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9170 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9171 let d = 2 * msum - gsum[gi];
9172 acc += d as f32 * s;
9173 }
9174 acc
9175 }
9176}
9177
9178#[cfg(target_arch = "x86_64")]
9180#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9181unsafe fn dot_q1_row_1x4_vnni(
9182 bytes: &[u8],
9183 r: usize,
9184 gpr: usize,
9185 xs: [&[i8]; 4],
9186 gsums: [&[i32]; 4],
9187) -> [f32; 4] {
9188 unsafe {
9190 use core::arch::x86_64::*;
9191 let expand = _mm256_setr_epi8(
9192 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9193 3, 3, 3,
9194 );
9195 let bitsel = _mm256_setr_epi8(
9196 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9197 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9198 );
9199 let ones8 = _mm256_set1_epi8(1);
9200 let mut acc = [0f32; 4];
9201 for gi in 0..gpr {
9202 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9203 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9204 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9205 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9206 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9207 for (k, xq) in xs.iter().enumerate() {
9208 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9209 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9210 let d = 2 * msum - gsums[k][gi];
9211 acc[k] += d as f32 * s;
9212 }
9213 }
9214 acc
9215 }
9216}
9217
9218#[cfg(target_arch = "x86_64")]
9221#[target_feature(enable = "avx2")]
9222unsafe fn dot_q1_row_1x4_avx2(
9223 bytes: &[u8],
9224 r: usize,
9225 gpr: usize,
9226 xs: [&[i8]; 4],
9227 gsums: [&[i32]; 4],
9228) -> [f32; 4] {
9229 unsafe {
9231 use core::arch::x86_64::*;
9232 let expand = _mm256_setr_epi8(
9233 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9234 3, 3, 3,
9235 );
9236 let bitsel = _mm256_setr_epi8(
9237 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9238 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9239 );
9240 let ones8 = _mm256_set1_epi8(1);
9241 let ones16 = _mm256_set1_epi16(1);
9242 let mut acc = [0f32; 4];
9243 for gi in 0..gpr {
9244 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9245 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9246 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9247 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9248 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9249 for (k, xq) in xs.iter().enumerate() {
9250 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9251 let sel = _mm256_and_si256(x, mask);
9252 let p16 = _mm256_maddubs_epi16(ones8, sel);
9253 let d32 = _mm256_madd_epi16(p16, ones16);
9254 let hi128 = _mm256_extracti128_si256::<1>(d32);
9255 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9256 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9257 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9258 let msum = _mm_cvtsi128_si32(s32);
9259 let d = 2 * msum - gsums[k][gi];
9260 acc[k] += d as f32 * s;
9261 }
9262 }
9263 acc
9264 }
9265}
9266
9267#[allow(unreachable_code)]
9268fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9269 #[cfg(target_arch = "aarch64")]
9270 unsafe {
9271 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
9272 }
9273 #[cfg(target_arch = "x86_64")]
9274 if avx2_enabled() {
9275 unsafe {
9276 if vnni_tiles_enabled() {
9277 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
9278 }
9279 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
9280 }
9281 }
9282 let _ = gsum;
9283 let mut acc = 0f32;
9284 for gi in 0..gpr {
9285 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9286 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9287 let mut d = 0i32;
9288 for (j, &b) in tile[2..].iter().enumerate() {
9289 for k in 0..8 {
9290 let w = ((b >> k) & 1) as i32 * 2 - 1;
9291 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
9292 }
9293 }
9294 acc += d as f32 * s;
9295 }
9296 acc
9297}
9298
9299#[cfg(target_arch = "aarch64")]
9308#[target_feature(enable = "neon,dotprod")]
9309unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9310 unsafe {
9313 use core::arch::aarch64::*;
9314 use core::arch::asm;
9315 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9316 let m = vld1q_u8(MASKS.as_ptr());
9317 macro_rules! tile_dot {
9319 ($t:expr, $x:expr) => {{
9320 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
9321 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
9322 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9323 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9324 let x0 = vld1q_s8($x);
9325 let x1 = vld1q_s8($x.add(16));
9326 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9327 asm!(
9328 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9329 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9330 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9331 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9332 options(pure, nomem, nostack),
9333 );
9334 vaddq_s32(a0, a1)
9335 }};
9336 }
9337 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9346 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9347 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9348 const IW11: [u8; 16] = [
9349 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9350 ];
9351 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9352 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9353 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9354 let isc = vld1_u8(ISC.as_ptr());
9355 macro_rules! tile_dot_tbl {
9357 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
9358 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
9359 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
9360 let x0 = vld1q_s8($x);
9361 let x1 = vld1q_s8($x.add(16));
9362 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9363 asm!(
9364 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9365 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9366 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9367 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9368 options(pure, nomem, nostack),
9369 );
9370 vaddq_s32(a0, a1)
9371 }};
9372 }
9373 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9374 let row_base = r * gpr * Q1_TILE;
9375 let abs_end = bytes.len();
9376 let xp = xq.as_ptr();
9377 let gp = gsum.as_ptr();
9378 let mut accv = vdupq_n_f32(0.0);
9379 let mut gi = 0;
9380 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9383 let t0 = base.add(gi * Q1_TILE);
9384 let ld_a = vld1q_u8(t0);
9385 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9386 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
9387 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
9388 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
9389 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
9390 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9392 let g = vld1q_s32(gp.add(gi));
9393 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9394 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9395 let scf: float32x4_t;
9396 asm!(
9397 "fcvtl {o:v}.4s, {i:v}.4h",
9398 o = out(vreg) scf, i = in(vreg) sc16,
9399 options(pure, nomem, nostack),
9400 );
9401 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9402 gi += 4;
9403 }
9404 let mut acc = vaddvq_f32(accv);
9405 while gi < gpr {
9406 let t = base.add(gi * Q1_TILE);
9407 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9408 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
9409 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9410 gi += 1;
9411 }
9412 acc
9413 }
9414}
9415
9416#[cfg(target_arch = "aarch64")]
9421#[target_feature(enable = "neon,dotprod")]
9422unsafe fn dot_q1_row_1x4_sdot(
9423 bytes: &[u8],
9424 r: usize,
9425 gpr: usize,
9426 xs: [&[i8]; 4],
9427 gs: [&[i32]; 4],
9428) -> [f32; 4] {
9429 unsafe {
9431 use core::arch::aarch64::*;
9432 use core::arch::asm;
9433 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9434 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9435 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9436 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9437 const IW11: [u8; 16] = [
9438 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9439 ];
9440 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9441 let m = vld1q_u8(MASKS.as_ptr());
9442 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9443 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9444 let isc = vld1_u8(ISC.as_ptr());
9445 macro_rules! sdot2 {
9446 ($w0:expr, $w1:expr, $x:expr) => {{
9447 let x0 = vld1q_s8($x);
9448 let x1 = vld1q_s8($x.add(16));
9449 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9450 asm!(
9451 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9452 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9453 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9454 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9455 options(pure, nomem, nostack),
9456 );
9457 vaddq_s32(a0, a1)
9458 }};
9459 }
9460 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9461 let row_base = r * gpr * Q1_TILE;
9462 let abs_end = bytes.len();
9463 let mut accv = [vdupq_n_f32(0.0); 4];
9464 let mut gi = 0;
9465 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9466 let t0 = base.add(gi * Q1_TILE);
9467 let ld_a = vld1q_u8(t0);
9468 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9469 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
9471 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
9472 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
9473 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
9474 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
9475 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
9476 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
9477 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
9478 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9479 let scf: float32x4_t;
9480 asm!(
9481 "fcvtl {o:v}.4s, {i:v}.4h",
9482 o = out(vreg) scf, i = in(vreg) sc16,
9483 options(pure, nomem, nostack),
9484 );
9485 for k in 0..4 {
9486 let xp = xs[k].as_ptr();
9487 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
9488 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
9489 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
9490 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
9491 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9492 let g = vld1q_s32(gs[k].as_ptr().add(gi));
9493 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9494 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9495 }
9496 gi += 4;
9497 }
9498 let mut acc = [
9499 vaddvq_f32(accv[0]),
9500 vaddvq_f32(accv[1]),
9501 vaddvq_f32(accv[2]),
9502 vaddvq_f32(accv[3]),
9503 ];
9504 while gi < gpr {
9505 let t = base.add(gi * Q1_TILE);
9506 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9507 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9508 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
9509 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9510 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9511 for k in 0..4 {
9512 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9513 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9514 }
9515 gi += 1;
9516 }
9517 acc
9518 }
9519}
9520
9521#[inline]
9523fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
9524 let gi = j / GROUP_SIZE;
9525 let k = j % GROUP_SIZE;
9526 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9527 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9528 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
9529 ((bit as i32 * 2 - 1) as f32, s)
9530}
9531
9532#[inline]
9534fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
9535 let mut acc = 0f32;
9536 for gi in 0..gpr {
9537 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9538 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9539 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9540 let mut ga = 0f32;
9541 for (j, &b) in tile[2..].iter().enumerate() {
9542 for k in 0..8 {
9543 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9544 }
9545 }
9546 acc += ga * s;
9547 }
9548 acc
9549}
9550
9551#[allow(clippy::too_many_arguments)]
9554fn bin_range_a8w8(
9555 w: BinW,
9556 gpr: usize,
9557 act: &SplitAct,
9558 gsum: &[i32],
9559 out: SendMut,
9560 start: usize,
9561 end: usize,
9562) {
9563 for r in start..end {
9564 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9565 for &(j, xv) in &act.outliers {
9566 let (wv, s) = w.outlier(r, gpr, j);
9567 acc += wv * s * xv;
9568 }
9569 unsafe { *out.at(r) = acc };
9571 }
9572}
9573
9574fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
9576 for r in start..end {
9577 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
9579 }
9580}
9581
9582#[derive(Clone, Copy)]
9590pub(crate) enum BinW<'a> {
9591 Q1(&'a [u8]),
9592 Q1g {
9593 bytes: &'a [u8],
9594 sc_off: usize,
9596 rows: usize,
9599 },
9600}
9601
9602impl<'a> BinW<'a> {
9603 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
9604 Self::Q1g {
9605 bytes,
9606 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
9607 rows,
9608 }
9609 }
9610
9611 #[inline]
9612 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9613 match self {
9614 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
9615 Self::Q1g {
9616 bytes,
9617 sc_off,
9618 rows,
9619 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
9620 }
9621 }
9622
9623 #[inline]
9624 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
9625 match self {
9626 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
9627 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
9628 }
9629 }
9630
9631 #[inline]
9632 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9633 match self {
9634 Self::Q1(b) => q1_outlier(b, r, gpr, j),
9635 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
9636 }
9637 }
9638
9639 #[cfg(target_arch = "aarch64")]
9644 #[inline]
9645 unsafe fn dot_1x4_sdot(
9646 self,
9647 r: usize,
9648 gpr: usize,
9649 xs: [&[i8]; 4],
9650 gs: [&[i32]; 4],
9651 ) -> [f32; 4] {
9652 unsafe {
9654 match self {
9655 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
9656 Self::Q1g {
9657 bytes,
9658 sc_off,
9659 rows,
9660 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
9661 }
9662 }
9663 }
9664
9665 #[cfg(target_arch = "x86_64")]
9670 #[inline]
9671 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9672 unsafe {
9674 match (self, vnni_tiles_enabled()) {
9675 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9676 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9677 (Self::Q1g { bytes, sc_off, .. }, true) => {
9678 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9679 }
9680 (Self::Q1g { bytes, sc_off, .. }, false) => {
9681 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9682 }
9683 }
9684 }
9685 }
9686}
9687
9688#[inline]
9698fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9699 u32::from_le_bytes([
9700 bytes[t * 4],
9701 bytes[t * 4 + 1],
9702 bytes[t * 4 + 2],
9703 bytes[t * 4 + 3],
9704 ])
9705}
9706
9707#[inline]
9709fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9710 let o = sc_off + (t / 4) * 2;
9711 u16::from_le_bytes([bytes[o], bytes[o + 1]])
9712}
9713
9714#[allow(unreachable_code)]
9716#[allow(clippy::too_many_arguments)]
9717fn dot_q1g_row_i8(
9718 bytes: &[u8],
9719 sc_off: usize,
9720 rows: usize,
9721 r: usize,
9722 gpr: usize,
9723 xq: &[i8],
9724 gsum: &[i32],
9725) -> f32 {
9726 #[cfg(target_arch = "aarch64")]
9727 unsafe {
9728 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9729 }
9730 #[cfg(target_arch = "x86_64")]
9731 if avx2_enabled() {
9732 unsafe {
9733 if vnni_tiles_enabled() {
9734 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9735 }
9736 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9737 }
9738 }
9739 let _ = (gsum, rows);
9740 let mut acc = 0f32;
9741 for gi in 0..gpr {
9742 let t = r * gpr + gi;
9743 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9744 let bits = q1g_bits(bytes, t);
9745 let mut d = 0i32;
9746 for k in 0..GROUP_SIZE {
9747 let w = ((bits >> k) & 1) as i32 * 2 - 1;
9748 d += w * xq[gi * GROUP_SIZE + k] as i32;
9749 }
9750 acc += d as f32 * s;
9751 }
9752 acc
9753}
9754
9755#[cfg(target_arch = "x86_64")]
9760#[target_feature(enable = "avx2")]
9761#[inline]
9762unsafe fn q1g_hsum4(
9763 d0: core::arch::x86_64::__m256i,
9764 d1: core::arch::x86_64::__m256i,
9765 d2: core::arch::x86_64::__m256i,
9766 d3: core::arch::x86_64::__m256i,
9767) -> core::arch::x86_64::__m128i {
9768 unsafe {
9770 use core::arch::x86_64::*;
9771 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9772 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9773 }
9774}
9775
9776#[cfg(target_arch = "x86_64")]
9780#[target_feature(enable = "avx2")]
9781#[inline]
9782unsafe fn q1g_group_acc(
9783 acc: f32,
9784 msum4: core::arch::x86_64::__m128i,
9785 gsum4: &[i32],
9786 s: f32,
9787) -> f32 {
9788 unsafe {
9790 use core::arch::x86_64::*;
9791 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9792 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9793 let mut f = [0f32; 4];
9794 _mm_storeu_ps(f.as_mut_ptr(), d);
9795 let mut acc = acc;
9796 for v in f {
9797 acc += v * s;
9798 }
9799 acc
9800 }
9801}
9802
9803#[cfg(target_arch = "x86_64")]
9808#[target_feature(enable = "avx2")]
9809unsafe fn dot_q1g_row_avx2(
9810 bytes: &[u8],
9811 sc_off: usize,
9812 r: usize,
9813 gpr: usize,
9814 xq: &[i8],
9815 gsum: &[i32],
9816) -> f32 {
9817 unsafe {
9820 use core::arch::x86_64::*;
9821 let expand = _mm256_setr_epi8(
9822 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9823 3, 3, 3,
9824 );
9825 let bitsel = _mm256_setr_epi8(
9826 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9827 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9828 );
9829 let ones8 = _mm256_set1_epi8(1);
9830 let ones16 = _mm256_set1_epi16(1);
9831 let bp = bytes.as_ptr().add(r * gpr * 4);
9832 let mut acc = 0f32;
9833 let mut gi = 0;
9834 while gi < gpr {
9835 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9836 let mut d = [_mm256_setzero_si256(); 4];
9837 for (k, dk) in d.iter_mut().enumerate() {
9838 let t = gi + k;
9839 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9840 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9841 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9842 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9843 let sel = _mm256_and_si256(x, mask);
9844 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9845 }
9846 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9847 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9848 gi += 4;
9849 }
9850 acc
9851 }
9852}
9853
9854#[cfg(target_arch = "x86_64")]
9857#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9858unsafe fn dot_q1g_row_vnni(
9859 bytes: &[u8],
9860 sc_off: usize,
9861 r: usize,
9862 gpr: usize,
9863 xq: &[i8],
9864 gsum: &[i32],
9865) -> f32 {
9866 unsafe {
9868 use core::arch::x86_64::*;
9869 let expand = _mm256_setr_epi8(
9870 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9871 3, 3, 3,
9872 );
9873 let bitsel = _mm256_setr_epi8(
9874 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9875 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9876 );
9877 let ones8 = _mm256_set1_epi8(1);
9878 let bp = bytes.as_ptr().add(r * gpr * 4);
9879 let mut acc = 0f32;
9880 let mut gi = 0;
9881 while gi < gpr {
9882 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9883 let mut d = [_mm256_setzero_si256(); 4];
9884 for (k, dk) in d.iter_mut().enumerate() {
9885 let t = gi + k;
9886 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9887 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9888 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9889 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9890 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9891 }
9892 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9893 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9894 gi += 4;
9895 }
9896 acc
9897 }
9898}
9899
9900#[cfg(target_arch = "x86_64")]
9904#[target_feature(enable = "avx2")]
9905unsafe fn dot_q1g_row_1x4_avx2(
9906 bytes: &[u8],
9907 sc_off: usize,
9908 r: usize,
9909 gpr: usize,
9910 xs: [&[i8]; 4],
9911 gsums: [&[i32]; 4],
9912) -> [f32; 4] {
9913 unsafe {
9915 use core::arch::x86_64::*;
9916 let expand = _mm256_setr_epi8(
9917 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9918 3, 3, 3,
9919 );
9920 let bitsel = _mm256_setr_epi8(
9921 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9922 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9923 );
9924 let ones8 = _mm256_set1_epi8(1);
9925 let ones16 = _mm256_set1_epi16(1);
9926 let bp = bytes.as_ptr().add(r * gpr * 4);
9927 let mut acc = [0f32; 4];
9928 let mut gi = 0;
9929 while gi < gpr {
9930 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9931 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9933 for t in 0..4 {
9934 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9935 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9936 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9937 for (k, xq) in xs.iter().enumerate() {
9938 let x = _mm256_loadu_si256(
9939 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9940 );
9941 let sel = _mm256_and_si256(x, mask);
9942 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9943 }
9944 }
9945 for k in 0..4 {
9946 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9947 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9948 }
9949 gi += 4;
9950 }
9951 acc
9952 }
9953}
9954
9955#[cfg(target_arch = "x86_64")]
9957#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9958unsafe fn dot_q1g_row_1x4_vnni(
9959 bytes: &[u8],
9960 sc_off: usize,
9961 r: usize,
9962 gpr: usize,
9963 xs: [&[i8]; 4],
9964 gsums: [&[i32]; 4],
9965) -> [f32; 4] {
9966 unsafe {
9968 use core::arch::x86_64::*;
9969 let expand = _mm256_setr_epi8(
9970 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
9971 3, 3, 3,
9972 );
9973 let bitsel = _mm256_setr_epi8(
9974 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9975 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9976 );
9977 let ones8 = _mm256_set1_epi8(1);
9978 let bp = bytes.as_ptr().add(r * gpr * 4);
9979 let mut acc = [0f32; 4];
9980 let mut gi = 0;
9981 while gi < gpr {
9982 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9983 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9984 for t in 0..4 {
9985 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9986 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9987 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9988 for (k, xq) in xs.iter().enumerate() {
9989 let x = _mm256_loadu_si256(
9990 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9991 );
9992 d[k][t] = _mm256_dpbusd_epi32(
9993 _mm256_setzero_si256(),
9994 ones8,
9995 _mm256_and_si256(x, mask),
9996 );
9997 }
9998 }
9999 for k in 0..4 {
10000 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
10001 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
10002 }
10003 gi += 4;
10004 }
10005 acc
10006 }
10007}
10008
10009#[cfg(target_arch = "aarch64")]
10020#[target_feature(enable = "neon,dotprod")]
10021#[allow(clippy::too_many_arguments)]
10022unsafe fn dot_q1g_row_sdot(
10023 bytes: &[u8],
10024 sc_off: usize,
10025 rows: usize,
10026 r: usize,
10027 gpr: usize,
10028 xq: &[i8],
10029 gsum: &[i32],
10030) -> f32 {
10031 unsafe {
10034 use core::arch::aarch64::*;
10035 use core::arch::asm;
10036 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10037 const IW: [[u8; 16]; 8] = [
10039 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10040 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10041 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10042 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10043 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10044 [
10045 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10046 ],
10047 [
10048 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10049 ],
10050 [
10051 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10052 ],
10053 ];
10054 let m = vld1q_u8(MASKS.as_ptr());
10055 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10056 macro_rules! tile_dot {
10057 ($w0:expr, $w1:expr, $x:expr) => {{
10058 let x0 = vld1q_s8($x);
10059 let x1 = vld1q_s8($x.add(16));
10060 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10061 asm!(
10062 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10063 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10064 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10065 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10066 options(pure, nomem, nostack),
10067 );
10068 vaddq_s32(a0, a1)
10069 }};
10070 }
10071 macro_rules! tbl_mask {
10072 ($ld:expr, $i:expr) => {
10073 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
10074 };
10075 }
10076 let bp = bytes.as_ptr().add(r * gpr * 4);
10077 let xp = xq.as_ptr();
10078 let gp = gsum.as_ptr();
10079 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10080 let mut accv = vdupq_n_f32(0.0);
10081 let mut gi = 0;
10082 while gi + 4 <= vec_end {
10083 let ld = vld1q_u8(bp.add(gi * 4));
10084 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
10085 let d1 = tile_dot!(
10086 tbl_mask!(ld, 2),
10087 tbl_mask!(ld, 3),
10088 xp.add((gi + 1) * GROUP_SIZE)
10089 );
10090 let d2 = tile_dot!(
10091 tbl_mask!(ld, 4),
10092 tbl_mask!(ld, 5),
10093 xp.add((gi + 2) * GROUP_SIZE)
10094 );
10095 let d3 = tile_dot!(
10096 tbl_mask!(ld, 6),
10097 tbl_mask!(ld, 7),
10098 xp.add((gi + 3) * GROUP_SIZE)
10099 );
10100 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10101 let g = vld1q_s32(gp.add(gi));
10102 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10103 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10104 let scf: float32x4_t;
10105 asm!(
10106 "fcvtl {o:v}.4s, {i:v}.4h",
10107 o = out(vreg) scf, i = in(vreg) h,
10108 options(pure, nomem, nostack),
10109 );
10110 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
10111 gi += 4;
10112 }
10113 let mut acc = vaddvq_f32(accv);
10114 while gi < gpr {
10115 let t = bp.add(gi * 4);
10116 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10117 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10118 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10119 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10120 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10121 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
10122 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
10123 gi += 1;
10124 }
10125 acc
10126 }
10127}
10128
10129#[cfg(target_arch = "aarch64")]
10133#[target_feature(enable = "neon,dotprod")]
10134#[allow(clippy::too_many_arguments)]
10135unsafe fn dot_q1g_row_1x4_sdot(
10136 bytes: &[u8],
10137 sc_off: usize,
10138 rows: usize,
10139 r: usize,
10140 gpr: usize,
10141 xs: [&[i8]; 4],
10142 gs: [&[i32]; 4],
10143) -> [f32; 4] {
10144 unsafe {
10146 use core::arch::aarch64::*;
10147 use core::arch::asm;
10148 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10149 const IW: [[u8; 16]; 8] = [
10150 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10151 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10152 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10153 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10154 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10155 [
10156 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10157 ],
10158 [
10159 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10160 ],
10161 [
10162 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10163 ],
10164 ];
10165 let m = vld1q_u8(MASKS.as_ptr());
10166 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10167 macro_rules! sdot2 {
10168 ($w0:expr, $w1:expr, $x:expr) => {{
10169 let x0 = vld1q_s8($x);
10170 let x1 = vld1q_s8($x.add(16));
10171 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10172 asm!(
10173 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10174 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10175 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10176 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10177 options(pure, nomem, nostack),
10178 );
10179 vaddq_s32(a0, a1)
10180 }};
10181 }
10182 let bp = bytes.as_ptr().add(r * gpr * 4);
10183 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10184 let mut accv = [vdupq_n_f32(0.0); 4];
10185 let mut gi = 0;
10186 while gi + 4 <= vec_end {
10187 let ld = vld1q_u8(bp.add(gi * 4));
10188 let w: [int8x16_t; 8] =
10190 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
10191 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10192 let scf: float32x4_t;
10193 asm!(
10194 "fcvtl {o:v}.4s, {i:v}.4h",
10195 o = out(vreg) scf, i = in(vreg) h,
10196 options(pure, nomem, nostack),
10197 );
10198 for k in 0..4 {
10199 let xp = xs[k].as_ptr();
10200 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
10201 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
10202 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
10203 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
10204 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10205 let g = vld1q_s32(gs[k].as_ptr().add(gi));
10206 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10207 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
10208 }
10209 gi += 4;
10210 }
10211 let mut acc = [
10212 vaddvq_f32(accv[0]),
10213 vaddvq_f32(accv[1]),
10214 vaddvq_f32(accv[2]),
10215 vaddvq_f32(accv[3]),
10216 ];
10217 while gi < gpr {
10218 let t = bp.add(gi * 4);
10219 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10220 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10221 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10222 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10223 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10224 for k in 0..4 {
10225 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
10226 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
10227 }
10228 gi += 1;
10229 }
10230 acc
10231 }
10232}
10233
10234#[inline]
10236fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
10237 let t = r * gpr + j / GROUP_SIZE;
10238 let k = j % GROUP_SIZE;
10239 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10240 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
10241 ((bit as i32 * 2 - 1) as f32, s)
10242}
10243
10244#[inline]
10246fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
10247 let mut acc = 0f32;
10248 for gi in 0..gpr {
10249 let t = r * gpr + gi;
10250 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10251 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10252 let mut ga = 0f32;
10253 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
10254 for k in 0..8 {
10255 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
10256 }
10257 }
10258 acc += ga * s;
10259 }
10260 acc
10261}
10262
10263fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
10268 let entries = base_len + (rows + 1) * 4;
10269 (base_len, entries, entries <= bytes.len())
10270}
10271
10272#[inline]
10274fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
10275 let o = rp_off + r * 4;
10276 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
10277}
10278
10279const SIGN5: [[f32; 5]; 256] = {
10283 let mut lut = [[0.0f32; 5]; 256];
10284 let pow3 = [1u16, 3, 9, 27, 81];
10285 let mut byte = 0usize;
10286 while byte < 256 {
10287 let mut i = 0usize;
10288 while i < 5 {
10289 let code = (byte as u16 / pow3[i]) % 3;
10290 lut[byte][i] = if code == 1 {
10291 1.0
10292 } else if code == 2 {
10293 -1.0
10294 } else {
10295 0.0
10296 };
10297 i += 1;
10298 }
10299 byte += 1;
10300 }
10301 lut
10302};
10303
10304const SIGN5_I8: [[i8; 5]; 256] = {
10306 let mut lut = [[0i8; 5]; 256];
10307 let pow3 = [1u16, 3, 9, 27, 81];
10308 let mut byte = 0usize;
10309 while byte < 256 {
10310 let mut i = 0usize;
10311 while i < 5 {
10312 let code = (byte as u16 / pow3[i]) % 3;
10313 lut[byte][i] = if code == 1 {
10314 1
10315 } else if code == 2 {
10316 -1
10317 } else {
10318 0
10319 };
10320 i += 1;
10321 }
10322 byte += 1;
10323 }
10324 lut
10325};
10326
10327const SIGN5_U64: [u64; 256] = {
10333 let mut lut = [0u64; 256];
10334 let pow3 = [1u16, 3, 9, 27, 81];
10335 let mut byte = 0usize;
10336 while byte < 256 {
10337 let mut v = 0u64;
10338 let mut i = 0usize;
10339 while i < 5 {
10340 let code = (byte as u16 / pow3[i]) % 3;
10341 let s: u8 = if code == 1 {
10342 1
10343 } else if code == 2 {
10344 0xFF
10345 } else {
10346 0
10347 };
10348 v |= (s as u64) << (i * 8);
10349 i += 1;
10350 }
10351 lut[byte] = v;
10352 byte += 1;
10353 }
10354 lut
10355};
10356
10357#[inline]
10362fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
10363 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10364 let off = (r * gpr + j / GROUP_SIZE) * TILE;
10365 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10366 let within = j % GROUP_SIZE;
10367 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
10368}
10369
10370#[cfg(target_arch = "aarch64")]
10373#[target_feature(enable = "neon,dotprod")]
10374#[inline]
10375unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
10376 unsafe {
10378 use core::arch::aarch64::*;
10379 use core::arch::asm;
10380 let w0 = vld1q_s8(w);
10381 let w1 = vld1q_s8(w.add(16));
10382 let x0 = vld1q_s8(x);
10383 let x1 = vld1q_s8(x.add(16));
10384 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10385 asm!(
10386 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10387 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10388 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10389 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10390 options(pure, nomem, nostack),
10391 );
10392 vaddvq_s32(vaddq_s32(a0, a1))
10393 }
10394}
10395
10396#[cfg(target_arch = "x86_64")]
10399#[target_feature(enable = "avx2")]
10400#[inline]
10401unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
10402 unsafe {
10404 use core::arch::x86_64::*;
10405 let wv = _mm256_loadu_si256(w as *const __m256i);
10406 let xv = _mm256_loadu_si256(x as *const __m256i);
10407 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10408 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
10409 let hi128 = _mm256_extracti128_si256::<1>(d);
10410 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10411 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10412 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10413 _mm_cvtsi128_si32(s32)
10414 }
10415}
10416
10417#[inline]
10422fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
10423 debug_assert!(dst.len() >= 40);
10424 unsafe {
10427 let p = dst.as_mut_ptr();
10428 for bi in 0..7 {
10429 core::ptr::write_unaligned(
10430 p.add(bi * 5) as *mut u64,
10431 SIGN5_U64[*codes.add(bi) as usize],
10432 );
10433 }
10434 }
10435}
10436
10437#[inline]
10442fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
10443 #[cfg(target_arch = "aarch64")]
10444 unsafe {
10445 return sdot32_i8(w, x);
10446 }
10447 #[cfg(target_arch = "x86_64")]
10448 unsafe {
10449 return i8dot32_avx2(w, x);
10450 }
10451 #[allow(unreachable_code)]
10452 unsafe {
10453 let mut s = 0i32;
10454 for k in 0..GROUP_SIZE {
10455 s += *w.add(k) as i32 * *x.add(k) as i32;
10456 }
10457 s
10458 }
10459}
10460
10461#[inline]
10462unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
10463 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
10464 (
10465 SIGN5_U64[*codes as usize],
10466 SIGN5_U64[*codes.add(1) as usize],
10467 SIGN5_U64[*codes.add(2) as usize],
10468 SIGN5_U64[*codes.add(3) as usize],
10469 SIGN5_U64[*codes.add(4) as usize],
10470 SIGN5_U64[*codes.add(5) as usize],
10471 SIGN5_U64[*codes.add(6) as usize],
10472 )
10473 };
10474
10475 let u0 = s0 | (s1 << 40);
10476 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
10477 let u2 = (s3 >> 8) | (s4 << 32);
10478 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
10479
10480 (u0, u1, u2, u3)
10481}
10482
10483#[cfg(target_arch = "aarch64")]
10487#[target_feature(enable = "neon,dotprod")]
10488unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10489 use core::arch::aarch64::*;
10490 use core::arch::asm;
10491 unsafe {
10492 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10493 let mut acc = 0f32;
10494 let bytes_ptr = bytes.as_ptr();
10495 let xq_ptr = xq.as_ptr();
10496 let row_off = r * gpr * TILE;
10497
10498 let gpr2 = gpr & !1;
10499 let mut gi = 0;
10500 while gi < gpr2 {
10501 let off0 = row_off + gi * TILE;
10502 let off1 = off0 + TILE;
10503 let s0 = f16_to_f32(u16::from_le_bytes([
10504 *bytes_ptr.add(off0),
10505 *bytes_ptr.add(off0 + 1),
10506 ]));
10507 let s1 = f16_to_f32(u16::from_le_bytes([
10508 *bytes_ptr.add(off1),
10509 *bytes_ptr.add(off1 + 1),
10510 ]));
10511
10512 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10513 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10514
10515 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10516 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10517 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10518 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10519
10520 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10521 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10522 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
10523 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
10524
10525 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
10526 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10527 asm!(
10528 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
10529 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
10530 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
10531 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
10532 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
10533 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
10534 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
10535 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
10536 options(pure, nomem, nostack),
10537 );
10538 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
10539 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
10540 acc += d0 as f32 * s0 + d1 as f32 * s1;
10541 gi += 2;
10542 }
10543
10544 if gi < gpr {
10545 let off = row_off + gi * TILE;
10546 let s = f16_to_f32(u16::from_le_bytes([
10547 *bytes_ptr.add(off),
10548 *bytes_ptr.add(off + 1),
10549 ]));
10550 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10551 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10552 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10553 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10554 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10555 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10556 asm!(
10557 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10558 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10559 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10560 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10561 options(pure, nomem, nostack),
10562 );
10563 let d = vaddvq_s32(vaddq_s32(a0, a1));
10564 acc += d as f32 * s;
10565 }
10566 acc
10567 }
10568}
10569
10570#[cfg(target_arch = "x86_64")]
10572#[target_feature(enable = "avx2")]
10573unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10574 use core::arch::x86_64::*;
10575 unsafe {
10576 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10577 let mut acc = 0f32;
10578 let bytes_ptr = bytes.as_ptr();
10579 let xq_ptr = xq.as_ptr();
10580 let row_off = r * gpr * TILE;
10581
10582 let ones = _mm256_set1_epi16(1);
10583 for gi in 0..gpr {
10584 let off = row_off + gi * TILE;
10585 let s = f16_to_f32(u16::from_le_bytes([
10586 *bytes_ptr.add(off),
10587 *bytes_ptr.add(off + 1),
10588 ]));
10589 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10590 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10591 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10592 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10593 let d256 = _mm256_madd_epi16(p16, ones);
10594 let d128 = _mm_add_epi32(
10595 _mm256_castsi256_si128(d256),
10596 _mm256_extracti128_si256(d256, 1),
10597 );
10598 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
10599 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
10600 acc += d32 as f32 * s;
10601 }
10602 acc
10603 }
10604}
10605
10606#[cfg(target_arch = "x86_64")]
10608#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10609unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10610 use core::arch::x86_64::*;
10611 unsafe {
10613 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10614 let mut acc = 0f32;
10615 let bytes_ptr = bytes.as_ptr();
10616 let xq_ptr = xq.as_ptr();
10617 let row_off = r * gpr * TILE;
10618 for gi in 0..gpr {
10619 let off = row_off + gi * TILE;
10620 let s = f16_to_f32(u16::from_le_bytes([
10621 *bytes_ptr.add(off),
10622 *bytes_ptr.add(off + 1),
10623 ]));
10624 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10625 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10626 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10627 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10628 acc += d as f32 * s;
10629 }
10630 acc
10631 }
10632}
10633
10634#[inline]
10638fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10639 #[cfg(target_arch = "aarch64")]
10640 unsafe {
10641 return q1t_dot_row_sdot(bytes, r, gpr, xq);
10642 }
10643 #[cfg(target_arch = "x86_64")]
10644 unsafe {
10645 if vnni_tiles_enabled() {
10646 return q1t_dot_row_vnni(bytes, r, gpr, xq);
10647 }
10648 return q1t_dot_row_avx2(bytes, r, gpr, xq);
10649 }
10650 #[allow(unreachable_code)]
10651 {
10652 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10653 let mut acc = 0f32;
10654 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
10656 let off = (r * gpr + gi) * TILE;
10657 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10658 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
10659 let mut d = 0i32;
10660 for k in 0..GROUP_SIZE {
10661 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
10662 }
10663 acc += d as f32 * s;
10664 }
10665 acc
10666 }
10667}
10668
10669fn q1t_row_outlier_correction(
10676 bytes: &[u8],
10677 r: usize,
10678 rp_off: usize,
10679 entries_off: usize,
10680 has_ov: bool,
10681 x: &[f32],
10682) -> f32 {
10683 if !has_ov {
10684 return 0.0;
10685 }
10686 let (c0, c1) = (
10687 q1t_rowptr(bytes, rp_off, r),
10688 q1t_rowptr(bytes, rp_off, r + 1),
10689 );
10690 let mut corr = 0f32;
10691 for p in c0..c1 {
10692 let e = entries_off + p * 4;
10693 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10694 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10695 corr += val * x[col];
10696 }
10697 corr
10698}
10699
10700fn q1t_dequant_row(
10704 bytes: &[u8],
10705 r: usize,
10706 gpr: usize,
10707 rp_off: usize,
10708 entries_off: usize,
10709 has_ov: bool,
10710 buf: &mut [f32],
10711) {
10712 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10713 for g in 0..gpr {
10714 let off = (r * gpr + g) * TILE;
10715 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10716 let codes = &bytes[off + 2..off + TILE];
10717 let bc = g * GROUP_SIZE;
10718 for bi in 0..6 {
10720 let lut = &SIGN5[codes[bi] as usize];
10721 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10722 for i in 0..5 {
10723 d[i] = lut[i] * s;
10724 }
10725 }
10726 let lut = &SIGN5[codes[6] as usize];
10727 buf[bc + 30] = lut[0] * s;
10728 buf[bc + 31] = lut[1] * s;
10729 }
10730 if !has_ov {
10731 return;
10732 }
10733 let (c0, c1) = (
10734 q1t_rowptr(bytes, rp_off, r),
10735 q1t_rowptr(bytes, rp_off, r + 1),
10736 );
10737 for p in c0..c1 {
10738 let e = entries_off + p * 4;
10739 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10740 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10741 }
10742}
10743
10744fn q1t_add_overlay(
10748 bytes: &[u8],
10749 x: &[f32],
10750 rows: usize,
10751 cols: usize,
10752 out: &mut [f32],
10753 pool: Option<&Pool>,
10754) {
10755 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10756 let gpr = cols / GROUP_SIZE;
10757 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10758 if !has_ov {
10759 return;
10760 }
10761 let out_addr = SendMut(out.as_mut_ptr());
10762 let run = move |start: usize, end: usize| {
10763 for r in start..end {
10764 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10765 unsafe { *out_addr.at(r) += corr };
10767 }
10768 };
10769 dispatch_rows(pool, rows, &run);
10770}
10771
10772#[allow(clippy::too_many_arguments)]
10775fn q1t_range_a8w8(
10776 bytes: &[u8],
10777 gpr: usize,
10778 rp_off: usize,
10779 ent_off: usize,
10780 has_ov: bool,
10781 act: &SplitAct,
10782 x: &[f32],
10783 out: SendMut,
10784 start: usize,
10785 end: usize,
10786) {
10787 for r in start..end {
10788 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10789 for &(j, xv) in &act.outliers {
10790 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10791 }
10792 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10793 unsafe { *out.at(r) = acc };
10795 }
10796}
10797
10798#[allow(clippy::too_many_arguments)]
10801fn q1t_range_f32_batch(
10802 bytes: &[u8],
10803 gpr: usize,
10804 rp_off: usize,
10805 ent_off: usize,
10806 has_ov: bool,
10807 x: &[f32],
10808 out: SendMut,
10809 start: usize,
10810 end: usize,
10811) {
10812 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10813 let mut sg = [0f32; GROUP_SIZE];
10814 for r in start..end {
10815 let mut acc = 0f32;
10816 for g in 0..gpr {
10817 let off = (r * gpr + g) * TILE;
10818 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10819 let codes = &bytes[off + 2..off + TILE];
10820 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10821 for bi in 0..6 {
10822 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10823 }
10824 let lut = &SIGN5[codes[6] as usize];
10825 sg[30] = lut[0];
10826 sg[31] = lut[1];
10827 let mut gsum = 0f32;
10828 for k in 0..GROUP_SIZE {
10829 gsum += sg[k] * xg[k];
10830 }
10831 acc += s * gsum;
10832 }
10833 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10834 unsafe { *out.at(r) = acc };
10836 }
10837}
10838
10839fn q1t_matvec(
10843 bytes: &[u8],
10844 x: &[f32],
10845 rows: usize,
10846 cols: usize,
10847 out: &mut [f32],
10848 pool: Option<&Pool>,
10849) {
10850 debug_assert_eq!(out.len(), rows);
10851 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10852 let gpr = cols / GROUP_SIZE;
10853 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10854 let out_addr = SendMut(out.as_mut_ptr());
10855 if a8w8_enabled() {
10859 let act = split_act(x);
10860 let act = &act;
10861 let run = move |start: usize, end: usize| {
10862 for r in start..end {
10863 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10864 for &(j, xv) in &act.outliers {
10865 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10866 }
10867 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10868 unsafe { *out_addr.at(r) = acc };
10870 }
10871 };
10872 dispatch_rows(pool, rows, &run);
10873 return;
10874 }
10875 let run = move |start: usize, end: usize| {
10876 let mut sg = [0f32; GROUP_SIZE];
10880 for r in start..end {
10881 let mut acc = 0f32;
10882 for g in 0..gpr {
10883 let off = (r * gpr + g) * TILE;
10884 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10885 let codes = &bytes[off + 2..off + TILE];
10886 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10887 for bi in 0..6 {
10888 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10889 }
10890 let lut = &SIGN5[codes[6] as usize];
10891 sg[30] = lut[0];
10892 sg[31] = lut[1];
10893 let mut gsum = 0f32;
10894 for k in 0..GROUP_SIZE {
10895 gsum += sg[k] * xg[k];
10896 }
10897 acc += s * gsum;
10898 }
10899 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10900 unsafe { *out_addr.at(r) = acc };
10901 }
10902 };
10903 dispatch_rows(pool, rows, &run);
10904}
10905
10906#[cfg(target_arch = "aarch64")]
10912#[target_feature(enable = "neon,dotprod")]
10913unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10914 use core::arch::aarch64::*;
10915 use core::arch::asm;
10916 unsafe {
10918 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10919 let bytes_ptr = bytes.as_ptr();
10920 let row_off = r * gpr * TILE;
10921 let xp = [xa.as_ptr(), xb.as_ptr()];
10922 let mut acc = [0f32; 2];
10923 macro_rules! sdot2 {
10924 ($w0:expr, $w1:expr, $x:expr) => {{
10925 let x0 = vld1q_s8($x);
10926 let x1 = vld1q_s8($x.add(16));
10927 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10928 asm!(
10929 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10930 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10931 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10932 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10933 options(pure, nomem, nostack),
10934 );
10935 vaddvq_s32(vaddq_s32(a0, a1))
10936 }};
10937 }
10938 let gpr2 = gpr & !1;
10939 let mut gi = 0;
10940 while gi < gpr2 {
10941 let off0 = row_off + gi * TILE;
10942 let off1 = off0 + TILE;
10943 let s0 = f16_to_f32(u16::from_le_bytes([
10944 *bytes_ptr.add(off0),
10945 *bytes_ptr.add(off0 + 1),
10946 ]));
10947 let s1 = f16_to_f32(u16::from_le_bytes([
10948 *bytes_ptr.add(off1),
10949 *bytes_ptr.add(off1 + 1),
10950 ]));
10951 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10952 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10953 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10954 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10955 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10956 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10957 for k in 0..2 {
10958 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10959 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10960 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10961 }
10962 gi += 2;
10963 }
10964 if gi < gpr {
10965 let off = row_off + gi * TILE;
10966 let s = f16_to_f32(u16::from_le_bytes([
10967 *bytes_ptr.add(off),
10968 *bytes_ptr.add(off + 1),
10969 ]));
10970 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10971 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10972 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10973 for k in 0..2 {
10974 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10975 acc[k] += d as f32 * s;
10976 }
10977 }
10978 acc
10979 }
10980}
10981
10982fn q1t_matvec2(
10988 bytes: &[u8],
10989 x1: &[f32],
10990 x2: &[f32],
10991 rows: usize,
10992 cols: usize,
10993 o1: &mut [f32],
10994 o2: &mut [f32],
10995 pool: Option<&Pool>,
10996) {
10997 debug_assert_eq!(o1.len(), rows);
10998 debug_assert_eq!(o2.len(), rows);
10999 const TILE: usize = cortiq_core::quant::Q1T_TILE;
11000 let gpr = cols / GROUP_SIZE;
11001 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11002 let out1 = SendMut(o1.as_mut_ptr());
11003 let out2 = SendMut(o2.as_mut_ptr());
11004 if a8w8_enabled() {
11005 let a1 = split_act(x1);
11006 let a2 = split_act(x2);
11007 let (a1, a2) = (&a1, &a2);
11008 let run = move |start: usize, end: usize| {
11009 for r in start..end {
11010 #[cfg(target_arch = "aarch64")]
11011 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
11014 #[cfg(not(target_arch = "aarch64"))]
11015 let ds = [
11016 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
11017 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
11018 ];
11019 let mut acc1 = ds[0] * a1.sx;
11020 for &(j, xv) in &a1.outliers {
11021 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
11022 }
11023 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11024 let mut acc2 = ds[1] * a2.sx;
11025 for &(j, xv) in &a2.outliers {
11026 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
11027 }
11028 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11029 unsafe {
11031 *out1.at(r) = acc1;
11032 *out2.at(r) = acc2;
11033 }
11034 }
11035 };
11036 dispatch_rows(pool, rows, &run);
11037 return;
11038 }
11039 let run = move |start: usize, end: usize| {
11040 let mut sg = [0f32; GROUP_SIZE];
11043 for r in start..end {
11044 let mut acc1 = 0f32;
11045 let mut acc2 = 0f32;
11046 for g in 0..gpr {
11047 let off = (r * gpr + g) * TILE;
11048 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11049 let codes = &bytes[off + 2..off + TILE];
11050 for bi in 0..6 {
11051 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
11052 }
11053 let lut = &SIGN5[codes[6] as usize];
11054 sg[30] = lut[0];
11055 sg[31] = lut[1];
11056 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
11057 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
11058 let mut gsum1 = 0f32;
11059 for k in 0..GROUP_SIZE {
11060 gsum1 += sg[k] * xg1[k];
11061 }
11062 acc1 += s * gsum1;
11063 let mut gsum2 = 0f32;
11064 for k in 0..GROUP_SIZE {
11065 gsum2 += sg[k] * xg2[k];
11066 }
11067 acc2 += s * gsum2;
11068 }
11069 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11070 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11071 unsafe {
11073 *out1.at(r) = acc1;
11074 *out2.at(r) = acc2;
11075 }
11076 }
11077 };
11078 dispatch_rows(pool, rows, &run);
11079}
11080
11081fn q1t_matmat(
11084 bytes: &[u8],
11085 xs: &[f32],
11086 b: usize,
11087 rows: usize,
11088 cols: usize,
11089 out: &mut [f32],
11090 pool: Option<&Pool>,
11091) {
11092 debug_assert_eq!(out.len(), b * rows);
11093 const TILE: usize = cortiq_core::quant::Q1T_TILE;
11094 let gpr = cols / GROUP_SIZE;
11095 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11096 let out_addr = SendMut(out.as_mut_ptr());
11097 if a8w8_enabled() {
11101 let acts: Vec<SplitAct> = (0..b)
11102 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
11103 .collect();
11104 let acts = &acts;
11105 let run = move |start: usize, end: usize| {
11106 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
11110 for g in 0..gpr {
11111 let off = (r * gpr + g) * TILE;
11112 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11113 q1t_unpack_group_i8(
11114 bytes.as_ptr().wrapping_add(off + 2),
11115 &mut sg[g * GROUP_SIZE..],
11116 );
11117 }
11118 for bi in 0..b {
11119 let act = &acts[bi];
11120 let mut isum = 0f32;
11121 for g in 0..gpr {
11122 let d = q1t_i8dot32(
11123 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
11124 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
11125 );
11126 isum += d as f32 * sc[g];
11127 }
11128 let mut acc = isum * act.sx;
11129 for &(j, xv) in &act.outliers {
11130 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
11131 }
11132 accs[bi] = acc;
11133 }
11134 if has_ov {
11138 let (c0, c1) = (
11139 q1t_rowptr(bytes, rp_off, r),
11140 q1t_rowptr(bytes, rp_off, r + 1),
11141 );
11142 for p in c0..c1 {
11143 let e = ent_off + p * 4;
11144 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
11145 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
11146 for bi in 0..b {
11147 accs[bi] += val * xs[bi * cols + col];
11148 }
11149 }
11150 }
11151 for bi in 0..b {
11152 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
11153 }
11154 }
11155 };
11156 dispatch_rows(pool, rows, &run);
11157 return;
11158 }
11159 let run = move |start: usize, end: usize| {
11160 let mut buf = vec![0f32; cols];
11161 for r in start..end {
11162 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
11163 for bi in 0..b {
11164 let xr = &xs[bi * cols..(bi + 1) * cols];
11165 let mut acc = 0f32;
11166 for j in 0..cols {
11167 acc += buf[j] * xr[j];
11168 }
11169 unsafe { *out_addr.at(bi * rows + r) = acc };
11170 }
11171 }
11172 };
11173 dispatch_rows(pool, rows, &run);
11174}
11175
11176#[cfg(test)]
11179fn q1_matvec(
11180 bytes: &[u8],
11181 x: &[f32],
11182 rows: usize,
11183 cols: usize,
11184 out: &mut [f32],
11185 pool: Option<&Pool>,
11186) {
11187 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
11188}
11189
11190fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
11193 debug_assert_eq!(out.len(), rows);
11194 let gpr = cols / GROUP_SIZE;
11195 let out_addr = SendMut(out.as_mut_ptr());
11196 if a8w8_enabled() {
11197 let act = split_act(x);
11198 let gsum = q1_group_sums(&act.xq, gpr);
11199 let (act, gsum) = (&act, &gsum);
11200 let run =
11201 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
11202 dispatch_rows(pool, rows, &run);
11203 return;
11204 }
11205 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
11206 dispatch_rows(pool, rows, &run);
11207}
11208
11209#[cfg(test)]
11211#[allow(clippy::too_many_arguments)]
11212fn q1_matvec2(
11213 bytes: &[u8],
11214 x1: &[f32],
11215 x2: &[f32],
11216 rows: usize,
11217 cols: usize,
11218 o1: &mut [f32],
11219 o2: &mut [f32],
11220 pool: Option<&Pool>,
11221) {
11222 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
11223}
11224
11225#[allow(clippy::too_many_arguments)]
11227fn bin_matvec2(
11228 w: BinW,
11229 x1: &[f32],
11230 x2: &[f32],
11231 rows: usize,
11232 cols: usize,
11233 o1: &mut [f32],
11234 o2: &mut [f32],
11235 pool: Option<&Pool>,
11236) {
11237 let gpr = cols / GROUP_SIZE;
11238 let p1 = SendMut(o1.as_mut_ptr());
11239 let p2 = SendMut(o2.as_mut_ptr());
11240 if a8w8_enabled() {
11241 let a1 = split_act(x1);
11242 let a2 = split_act(x2);
11243 let g1 = q1_group_sums(&a1.xq, gpr);
11244 let g2 = q1_group_sums(&a2.xq, gpr);
11245 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
11246 let run = move |start: usize, end: usize| {
11247 for r in start..end {
11248 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
11249 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
11250 for &(j, xv) in &a1.outliers {
11251 let (wv, s) = w.outlier(r, gpr, j);
11252 v1 += wv * s * xv;
11253 }
11254 for &(j, xv) in &a2.outliers {
11255 let (wv, s) = w.outlier(r, gpr, j);
11256 v2 += wv * s * xv;
11257 }
11258 unsafe {
11260 *p1.at(r) = v1;
11261 *p2.at(r) = v2;
11262 }
11263 }
11264 };
11265 dispatch_rows(pool, rows, &run);
11266 return;
11267 }
11268 let run = move |start: usize, end: usize| {
11269 for r in start..end {
11270 unsafe {
11272 *p1.at(r) = w.row_exact(r, gpr, x1);
11273 *p2.at(r) = w.row_exact(r, gpr, x2);
11274 }
11275 }
11276 };
11277 dispatch_rows(pool, rows, &run);
11278}
11279
11280#[cfg(test)]
11282#[allow(clippy::too_many_arguments)]
11283fn q1_matmat(
11284 bytes: &[u8],
11285 xs_all: &[f32],
11286 b: usize,
11287 rows: usize,
11288 cols: usize,
11289 out: &mut [f32],
11290 pool: Option<&Pool>,
11291) {
11292 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
11293}
11294
11295#[allow(clippy::too_many_arguments)]
11297fn bin_matmat(
11298 w: BinW,
11299 xs_all: &[f32],
11300 b: usize,
11301 rows: usize,
11302 cols: usize,
11303 out: &mut [f32],
11304 pool: Option<&Pool>,
11305) {
11306 debug_assert_eq!(out.len(), b * rows);
11307 let gpr = cols / GROUP_SIZE;
11308 let out_addr = SendMut(out.as_mut_ptr());
11309 if a8w8_enabled() {
11310 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
11311 .map(|bi| {
11312 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
11313 let gsum = q1_group_sums(&act.xq, gpr);
11314 (act, gsum)
11315 })
11316 .collect();
11317 let acts = &acts;
11318 #[cfg(target_arch = "x86_64")]
11319 let blocked_ok = avx2_enabled() && blocked_enabled();
11320 #[cfg(target_arch = "aarch64")]
11321 let blocked_ok = sdot_enabled() && blocked_enabled();
11322 let run = move |start: usize, end: usize| {
11323 for r in start..end {
11324 let mut bi = 0usize;
11325 #[cfg(target_arch = "aarch64")]
11328 if blocked_ok {
11329 while bi + 4 <= acts.len() {
11330 let xs = [
11331 acts[bi].0.xq.as_slice(),
11332 acts[bi + 1].0.xq.as_slice(),
11333 acts[bi + 2].0.xq.as_slice(),
11334 acts[bi + 3].0.xq.as_slice(),
11335 ];
11336 let gs = [
11337 acts[bi].1.as_slice(),
11338 acts[bi + 1].1.as_slice(),
11339 acts[bi + 2].1.as_slice(),
11340 acts[bi + 3].1.as_slice(),
11341 ];
11342 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
11344 for k in 0..4 {
11345 let (act, _) = &acts[bi + k];
11346 let mut acc = d[k] * act.sx;
11347 for &(j, xv) in &act.outliers {
11348 let (wv, sc) = w.outlier(r, gpr, j);
11349 acc += wv * sc * xv;
11350 }
11351 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11353 }
11354 bi += 4;
11355 }
11356 }
11357 #[cfg(target_arch = "x86_64")]
11358 if blocked_ok {
11359 while bi + 4 <= acts.len() {
11360 let xs = [
11361 acts[bi].0.xq.as_slice(),
11362 acts[bi + 1].0.xq.as_slice(),
11363 acts[bi + 2].0.xq.as_slice(),
11364 acts[bi + 3].0.xq.as_slice(),
11365 ];
11366 let gs = [
11367 acts[bi].1.as_slice(),
11368 acts[bi + 1].1.as_slice(),
11369 acts[bi + 2].1.as_slice(),
11370 acts[bi + 3].1.as_slice(),
11371 ];
11372 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
11374 for k in 0..4 {
11375 let (act, _) = &acts[bi + k];
11376 let mut acc = d[k] * act.sx;
11377 for &(j, xv) in &act.outliers {
11378 let (wv, sc) = w.outlier(r, gpr, j);
11379 acc += wv * sc * xv;
11380 }
11381 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11383 }
11384 bi += 4;
11385 }
11386 }
11387 while bi < acts.len() {
11388 let (act, gsum) = &acts[bi];
11389 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
11390 for &(j, xv) in &act.outliers {
11391 let (wv, s) = w.outlier(r, gpr, j);
11392 acc += wv * s * xv;
11393 }
11394 unsafe { *out_addr.at(bi * rows + r) = acc };
11396 bi += 1;
11397 }
11398 }
11399 };
11400 dispatch_rows(pool, rows, &run);
11401 return;
11402 }
11403 let run = move |start: usize, end: usize| {
11404 for r in start..end {
11405 for bi in 0..b {
11406 let x = &xs_all[bi * cols..(bi + 1) * cols];
11407 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
11409 }
11410 }
11411 };
11412 dispatch_rows(pool, rows, &run);
11413}
11414
11415fn q4matvec(
11421 bytes: &[u8],
11422 x: &[f32],
11423 rows: usize,
11424 cols: usize,
11425 out: &mut [f32],
11426 pool: Option<&Pool>,
11427) {
11428 debug_assert_eq!(out.len(), rows);
11429 let (packed, scales) = q4_split(bytes, rows, cols);
11430 let gpr = cols / GROUP_SIZE;
11431 let out_addr = SendMut(out.as_mut_ptr());
11432
11433 if a8w8_enabled() {
11434 let act = split_act(x);
11435 let run = move |start: usize, end: usize| {
11436 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
11437 };
11438 dispatch_rows(pool, rows, &run);
11439 return;
11440 }
11441
11442 let run =
11443 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
11444 dispatch_rows(pool, rows, &run);
11445}
11446
11447#[inline]
11450#[allow(unreachable_code)]
11451#[cfg(target_arch = "x86_64")]
11456#[target_feature(enable = "avx2")]
11457unsafe fn dot_q4b_row_1x4_avx2(
11458 buf: &[u8],
11459 scales: &[u8],
11460 g0: usize,
11461 gpr: usize,
11462 xs: [&[i8]; 4],
11463) -> [f32; 4] {
11464 unsafe {
11466 use core::arch::x86_64::*;
11467 let ones = _mm256_set1_epi16(1);
11468 let mut acc = [0f32; 4];
11469 for gi in 0..gpr {
11470 let s = f16_to_f32(u16::from_le_bytes([
11471 scales[(g0 + gi) * 2],
11472 scales[(g0 + gi) * 2 + 1],
11473 ]));
11474 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11475 let aw = _mm256_abs_epi8(w);
11476 for (k, xq) in xs.iter().enumerate() {
11477 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11478 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11479 let d = _mm256_madd_epi16(p16, ones);
11480 let hi128 = _mm256_extracti128_si256::<1>(d);
11481 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11482 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11483 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11484 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
11485 }
11486 }
11487 acc
11488 }
11489}
11490
11491#[cfg(target_arch = "x86_64")]
11493#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11494unsafe fn dot_q4b_row_1x4_vnni(
11495 buf: &[u8],
11496 scales: &[u8],
11497 g0: usize,
11498 gpr: usize,
11499 xs: [&[i8]; 4],
11500) -> [f32; 4] {
11501 unsafe {
11503 use core::arch::x86_64::*;
11504 let mut acc = [0f32; 4];
11505 for gi in 0..gpr {
11506 let s = f16_to_f32(u16::from_le_bytes([
11507 scales[(g0 + gi) * 2],
11508 scales[(g0 + gi) * 2 + 1],
11509 ]));
11510 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11511 let aw = _mm256_abs_epi8(w);
11512 for (k, xq) in xs.iter().enumerate() {
11513 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11514 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11515 acc[k] += d as f32 * s;
11516 }
11517 }
11518 acc
11519 }
11520}
11521
11522#[cfg(target_arch = "x86_64")]
11528#[target_feature(enable = "avx2")]
11529unsafe fn dot_q4b_row_1x4_sx_avx2(
11530 buf: &[u8],
11531 scales: &[u8],
11532 g0: usize,
11533 gpr: usize,
11534 xs: [&[i8]; 4],
11535 sxs: [f32; 4],
11536) -> [f32; 4] {
11537 unsafe {
11539 use core::arch::x86_64::*;
11540 let ones = _mm256_set1_epi16(1);
11541 let mut acc = [0f32; 4];
11542 for gi in 0..gpr {
11543 let s = f16_to_f32(u16::from_le_bytes([
11544 scales[(g0 + gi) * 2],
11545 scales[(g0 + gi) * 2 + 1],
11546 ]));
11547 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11548 let aw = _mm256_abs_epi8(w);
11549 for (k, xq) in xs.iter().enumerate() {
11550 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11551 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11552 let d = _mm256_madd_epi16(p16, ones);
11553 let hi128 = _mm256_extracti128_si256::<1>(d);
11554 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11555 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11556 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11557 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
11558 }
11559 }
11560 acc
11561 }
11562}
11563
11564#[cfg(target_arch = "x86_64")]
11567#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11568unsafe fn dot_q4b_row_1x4_sx_vnni(
11569 buf: &[u8],
11570 scales: &[u8],
11571 g0: usize,
11572 gpr: usize,
11573 xs: [&[i8]; 4],
11574 sxs: [f32; 4],
11575) -> [f32; 4] {
11576 unsafe {
11578 use core::arch::x86_64::*;
11579 let mut acc = [0f32; 4];
11580 for gi in 0..gpr {
11581 let s = f16_to_f32(u16::from_le_bytes([
11582 scales[(g0 + gi) * 2],
11583 scales[(g0 + gi) * 2 + 1],
11584 ]));
11585 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11586 let aw = _mm256_abs_epi8(w);
11587 for (k, xq) in xs.iter().enumerate() {
11588 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11589 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11590 acc[k] += (d as f32 * sxs[k]) * s;
11591 }
11592 }
11593 acc
11594 }
11595}
11596
11597#[allow(unreachable_code)]
11598fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11599 #[cfg(target_arch = "aarch64")]
11600 unsafe {
11601 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
11602 }
11603 #[cfg(target_arch = "x86_64")]
11604 unsafe {
11605 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
11606 }
11607 let mut acc = 0f32;
11608 for gi in 0..gpr {
11609 let g = g0 + gi;
11610 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11611 let mut d = 0i32;
11612 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11613 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
11614 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
11615 }
11616 acc += d as f32 * s;
11617 }
11618 acc
11619}
11620
11621#[inline]
11623#[allow(unreachable_code)]
11624fn dot_q4_row_i8_2(
11625 packed: &[u8],
11626 scales: &[u8],
11627 g0: usize,
11628 gpr: usize,
11629 xq1: &[i8],
11630 xq2: &[i8],
11631) -> (f32, f32) {
11632 #[cfg(target_arch = "aarch64")]
11633 unsafe {
11634 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
11635 }
11636 #[cfg(target_arch = "x86_64")]
11637 unsafe {
11638 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
11639 }
11640 (
11641 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
11642 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
11643 )
11644}
11645
11646#[allow(clippy::too_many_arguments)]
11649fn q4_range_a8w8(
11650 packed: &[u8],
11651 scales: &[u8],
11652 gpr: usize,
11653 cols: usize,
11654 act: &SplitAct,
11655 out: SendMut,
11656 start: usize,
11657 end: usize,
11658) {
11659 for r in start..end {
11660 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
11661 for &(j, xv) in &act.outliers {
11663 let flat = r * cols + j;
11664 let byte = packed[flat / 2];
11665 let nib = if flat & 1 == 0 {
11666 byte & 0x0F
11667 } else {
11668 byte >> 4
11669 };
11670 let s = f16_to_f32(u16::from_le_bytes([
11671 scales[(flat / GROUP_SIZE) * 2],
11672 scales[(flat / GROUP_SIZE) * 2 + 1],
11673 ]));
11674 acc += ((nib as i32 - 8) as f32) * s * xv;
11675 }
11676 unsafe { *out.at(r) = acc };
11678 }
11679}
11680
11681#[allow(clippy::too_many_arguments)]
11684fn q4_range2_a8w8(
11685 packed: &[u8],
11686 scales: &[u8],
11687 gpr: usize,
11688 cols: usize,
11689 a1: &SplitAct,
11690 a2: &SplitAct,
11691 p1: SendMut,
11692 p2: SendMut,
11693 start: usize,
11694 end: usize,
11695) {
11696 for r in start..end {
11697 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11698 let mut acc1 = s1 * a1.sx;
11699 let mut acc2 = s2 * a2.sx;
11700 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11702 for &(j, xv) in outliers {
11703 let flat = r * cols + j;
11704 let byte = packed[flat / 2];
11705 let nib = if flat & 1 == 0 {
11706 byte & 0x0F
11707 } else {
11708 byte >> 4
11709 };
11710 let s = f16_to_f32(u16::from_le_bytes([
11711 scales[(flat / GROUP_SIZE) * 2],
11712 scales[(flat / GROUP_SIZE) * 2 + 1],
11713 ]));
11714 *acc += ((nib as i32 - 8) as f32) * s * xv;
11715 }
11716 };
11717 fix(&a1.outliers, &mut acc1);
11718 fix(&a2.outliers, &mut acc2);
11719 unsafe {
11721 *p1.at(r) = acc1;
11722 *p2.at(r) = acc2;
11723 }
11724 }
11725}
11726
11727fn q4_range_f32(
11729 packed: &[u8],
11730 scales: &[u8],
11731 gpr: usize,
11732 x: &[f32],
11733 out: SendMut,
11734 start: usize,
11735 end: usize,
11736) {
11737 for r in start..end {
11738 let mut acc = 0f32;
11739 for gi in 0..gpr {
11740 let g = r * gpr + gi;
11741 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11742 let pk = &packed[g * 16..(g + 1) * 16];
11743 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11744 let mut ga = 0f32;
11745 for (k, &b) in pk.iter().enumerate() {
11746 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11747 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11748 }
11749 acc += ga * s;
11750 }
11751 unsafe { *out.at(r) = acc };
11753 }
11754}
11755
11756#[allow(clippy::too_many_arguments)]
11760fn q4matvec2(
11761 bytes: &[u8],
11762 x1: &[f32],
11763 x2: &[f32],
11764 rows: usize,
11765 cols: usize,
11766 o1: &mut [f32],
11767 o2: &mut [f32],
11768 pool: Option<&Pool>,
11769) {
11770 debug_assert_eq!(o1.len(), rows);
11771 debug_assert_eq!(o2.len(), rows);
11772 let (packed, scales) = q4_split(bytes, rows, cols);
11773 let gpr = cols / GROUP_SIZE;
11774
11775 if a8w8_enabled() {
11776 let a1 = split_act(x1);
11777 let a2 = split_act(x2);
11778 let p1 = SendMut(o1.as_mut_ptr());
11779 let p2 = SendMut(o2.as_mut_ptr());
11780 let run = move |start: usize, end: usize| {
11781 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11782 };
11783 dispatch_rows(pool, rows, &run);
11784 return;
11785 }
11786
11787 let p1 = SendMut(o1.as_mut_ptr());
11788 let p2 = SendMut(o2.as_mut_ptr());
11789 let run = move |start: usize, end: usize| {
11790 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11791 };
11792 dispatch_rows(pool, rows, &run);
11793}
11794
11795#[allow(clippy::too_many_arguments)]
11797fn q4_range2_f32(
11798 packed: &[u8],
11799 scales: &[u8],
11800 gpr: usize,
11801 x1: &[f32],
11802 x2: &[f32],
11803 p1: SendMut,
11804 p2: SendMut,
11805 start: usize,
11806 end: usize,
11807) {
11808 for r in start..end {
11809 let (mut acc1, mut acc2) = (0f32, 0f32);
11810 for gi in 0..gpr {
11811 let g = r * gpr + gi;
11812 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11813 let pk = &packed[g * 16..(g + 1) * 16];
11814 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11815 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11816 let (mut g1, mut g2) = (0f32, 0f32);
11817 for (k, &b) in pk.iter().enumerate() {
11818 let wl = (b & 0x0F) as f32 - 8.0;
11819 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11820 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11821 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11822 }
11823 acc1 += g1 * s;
11824 acc2 += g2 * s;
11825 }
11826 unsafe {
11828 *p1.at(r) = acc1;
11829 *p2.at(r) = acc2;
11830 }
11831 }
11832}
11833
11834thread_local! {
11835 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11838 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11839}
11840
11841#[allow(clippy::too_many_arguments)]
11847fn q4matmat(
11848 bytes: &[u8],
11849 xs_all: &[f32],
11850 b: usize,
11851 rows: usize,
11852 cols: usize,
11853 out: &mut [f32],
11854 pool: Option<&Pool>,
11855) {
11856 debug_assert_eq!(xs_all.len(), b * cols);
11857 debug_assert_eq!(out.len(), b * rows);
11858 let (packed, scales) = q4_split(bytes, rows, cols);
11859 let gpr = cols / GROUP_SIZE;
11860 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11861
11862 if a8w8_enabled() {
11863 let acts: Vec<SplitAct> = (0..b)
11864 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11865 .collect();
11866 let acts = &acts;
11867 let out_addr = SendMut(out.as_mut_ptr());
11868 let run = move |start: usize, end: usize| {
11869 ROW_I8.with(|rb| {
11870 let mut buf = rb.borrow_mut();
11871 buf.resize(cols, 0);
11872 for r in start..end {
11873 for gi in 0..gpr {
11877 let g = r * gpr + gi;
11878 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11879 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11880 buf[gi * GROUP_SIZE + k * 2 + 1] =
11881 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11882 }
11883 }
11884 let mut bi = 0usize;
11885 #[cfg(target_arch = "x86_64")]
11886 if avx2_enabled() && blocked_enabled() {
11887 while bi + 4 <= acts.len() {
11888 let xs = [
11889 acts[bi].xq.as_slice(),
11890 acts[bi + 1].xq.as_slice(),
11891 acts[bi + 2].xq.as_slice(),
11892 acts[bi + 3].xq.as_slice(),
11893 ];
11894 let d = unsafe {
11895 if vnni_tiles_enabled() {
11896 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11897 } else {
11898 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11899 }
11900 };
11901 for k in 0..4 {
11902 let act = &acts[bi + k];
11903 let mut acc = d[k] * act.sx;
11904 for &(j, xv) in &act.outliers {
11905 acc += (buf[j] as i8) as f32
11906 * gscale((r * cols + j) / GROUP_SIZE)
11907 * xv;
11908 }
11909 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11911 }
11912 bi += 4;
11913 }
11914 }
11915 while bi < acts.len() {
11916 let act = &acts[bi];
11917 let mut acc = 0f32;
11918 for gi in 0..gpr {
11919 let d = dot_i8_i8(
11920 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11921 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11922 );
11923 acc += d as f32 * gscale(r * gpr + gi);
11924 }
11925 acc *= act.sx;
11926 for &(j, xv) in &act.outliers {
11928 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11929 }
11930 unsafe { *out_addr.at(bi * rows + r) = acc };
11932 bi += 1;
11933 }
11934 }
11935 })
11936 };
11937 dispatch_rows(pool, rows, &run);
11938 return;
11939 }
11940
11941 let out_addr = SendMut(out.as_mut_ptr());
11942 let run = move |start: usize, end: usize| {
11943 ROW_F32.with(|rb| {
11944 let mut buf = rb.borrow_mut();
11945 buf.resize(cols, 0.0);
11946 for r in start..end {
11947 for gi in 0..gpr {
11950 let g = r * gpr + gi;
11951 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11952 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11953 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11954 }
11955 }
11956 for bi in 0..b {
11957 let x = &xs_all[bi * cols..(bi + 1) * cols];
11958 let mut acc = 0f32;
11959 for gi in 0..gpr {
11960 let mut ga = 0f32;
11961 for k in 0..GROUP_SIZE / 2 {
11966 let e = gi * GROUP_SIZE + k * 2;
11967 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11968 }
11969 acc += ga * gscale(r * gpr + gi);
11970 }
11971 unsafe { *out_addr.at(bi * rows + r) = acc };
11973 }
11974 }
11975 })
11976 };
11977 dispatch_rows(pool, rows, &run);
11978}
11979
11980#[allow(clippy::too_many_arguments)]
11985fn vbitmatmat(
11986 bytes: &[u8],
11987 offsets: &[usize],
11988 xs_all: &[f32],
11989 b: usize,
11990 rows: usize,
11991 cols: usize,
11992 out: &mut [f32],
11993 pool: Option<&Pool>,
11994) {
11995 debug_assert_eq!(xs_all.len(), b * cols);
11996 debug_assert_eq!(out.len(), b * rows);
11997 debug_assert_eq!(offsets.len(), rows + 1);
11998 let ng = cols / GROUP_SIZE;
11999 let bits = &bytes[..rows];
12000 let sc_off = rows;
12001 let gscale = |r: usize, g: usize| {
12002 let so = (r * ng + g) * 2;
12003 f16_to_f32(u16::from_le_bytes([
12004 bytes[sc_off + so],
12005 bytes[sc_off + so + 1],
12006 ]))
12007 };
12008
12009 let decode_f32 = |r: usize, dst: &mut [f32]| {
12011 let bw = bits[r] as usize;
12012 let l = ((1i32 << (bw - 1)) - 1) as f32;
12013 let data = &bytes[offsets[r]..offsets[r + 1]];
12014 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
12015 for d in dst.iter_mut() {
12016 while nbits < bw {
12017 acc = (acc << 8) | data[idx] as u64;
12018 idx += 1;
12019 nbits += 8;
12020 }
12021 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
12022 nbits -= bw;
12023 *d = u - l;
12024 }
12025 };
12026
12027 if a8w8_enabled() {
12028 let acts: Vec<SplitAct> = (0..b)
12029 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
12030 .collect();
12031 let acts = &acts;
12032 let out_addr = SendMut(out.as_mut_ptr());
12033 let run = move |start: usize, end: usize| {
12034 for r in start..end {
12035 let bw = bits[r] as usize;
12036 if bw == 8 {
12037 ROW_F32.with(|rb| {
12040 let mut buf = rb.borrow_mut();
12041 buf.resize(cols, 0.0);
12042 decode_f32(r, &mut buf);
12043 for bi in 0..b {
12044 let x = &xs_all[bi * cols..(bi + 1) * cols];
12045 let mut dot = 0f32;
12046 for g in 0..ng {
12047 let mut gd = 0f32;
12048 for k in 0..GROUP_SIZE {
12049 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12050 }
12051 dot += gd * gscale(r, g);
12052 }
12053 unsafe { *out_addr.at(bi * rows + r) = dot };
12055 }
12056 });
12057 continue;
12058 }
12059 let l = (1i32 << (bw - 1)) - 1;
12060 let data = &bytes[offsets[r]..offsets[r + 1]];
12061 ROW_I8.with(|rb| {
12062 let mut buf = rb.borrow_mut();
12063 buf.resize(cols, 0);
12064 #[inline(always)]
12065 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
12066 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
12067 let u = unpack8::<B>(&data[blk * B..]);
12068 for k in 0..8 {
12069 chunk[k] = (u[k] - l) as i8 as u8;
12070 }
12071 }
12072 }
12073 match bw {
12074 3 => fill::<3>(data, l, &mut buf),
12075 4 => vbit_fill4(data, &mut buf),
12076 5 => fill::<5>(data, l, &mut buf),
12077 6 => fill::<6>(data, l, &mut buf),
12078 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
12079 }
12080 let mut bi = 0usize;
12081 #[cfg(target_arch = "x86_64")]
12085 if avx2_enabled() && blocked_enabled() {
12086 while bi + 4 <= acts.len() {
12087 let xs = [
12088 acts[bi].xq.as_slice(),
12089 acts[bi + 1].xq.as_slice(),
12090 acts[bi + 2].xq.as_slice(),
12091 acts[bi + 3].xq.as_slice(),
12092 ];
12093 let sxs = [
12094 acts[bi].sx,
12095 acts[bi + 1].sx,
12096 acts[bi + 2].sx,
12097 acts[bi + 3].sx,
12098 ];
12099 let d = unsafe {
12100 if vnni_tiles_enabled() {
12101 dot_q4b_row_1x4_sx_vnni(
12102 &buf,
12103 &bytes[sc_off..],
12104 r * ng,
12105 ng,
12106 xs,
12107 sxs,
12108 )
12109 } else {
12110 dot_q4b_row_1x4_sx_avx2(
12111 &buf,
12112 &bytes[sc_off..],
12113 r * ng,
12114 ng,
12115 xs,
12116 sxs,
12117 )
12118 }
12119 };
12120 for k in 0..4 {
12121 let act = &acts[bi + k];
12122 let mut dot = d[k];
12123 for &(j, xv) in &act.outliers {
12124 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12125 }
12126 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
12128 }
12129 bi += 4;
12130 }
12131 }
12132 while bi < acts.len() {
12133 let act = &acts[bi];
12134 let mut dot = 0f32;
12135 for g in 0..ng {
12136 let d = dot_i8_i8(
12137 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12138 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12139 ) as f32
12140 * act.sx;
12141 dot += d * gscale(r, g);
12142 }
12143 for &(j, xv) in &act.outliers {
12144 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12145 }
12146 unsafe { *out_addr.at(bi * rows + r) = dot };
12148 bi += 1;
12149 }
12150 });
12151 }
12152 };
12153 dispatch_rows(pool, rows, &run);
12154 return;
12155 }
12156
12157 let out_addr = SendMut(out.as_mut_ptr());
12158 let run = move |start: usize, end: usize| {
12159 ROW_F32.with(|rb| {
12160 let mut buf = rb.borrow_mut();
12161 buf.resize(cols, 0.0);
12162 for r in start..end {
12163 decode_f32(r, &mut buf);
12164 for bi in 0..b {
12165 let x = &xs_all[bi * cols..(bi + 1) * cols];
12166 let mut dot = 0f32;
12167 for g in 0..ng {
12168 let mut gd = 0f32;
12169 for k in 0..GROUP_SIZE {
12170 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12171 }
12172 dot += gd * gscale(r, g);
12173 }
12174 unsafe { *out_addr.at(bi * rows + r) = dot };
12176 }
12177 }
12178 })
12179 };
12180 dispatch_rows(pool, rows, &run);
12181}
12182
12183pub(crate) fn gpu_batch_job<'a>(
12187 t: &'a QTensor,
12188 x: &[f32],
12189) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
12190 match t {
12191 QTensor::Mapped {
12192 model,
12193 idx,
12194 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
12195 rows,
12196 cols,
12197 row_scale,
12198 col_field,
12199 ..
12200 } => Some((
12201 model.clone(),
12202 crate::gpu::BatchJob {
12203 idx: *idx,
12204 rows: *rows,
12205 cols: *cols,
12206 row_scale,
12207 xs: prescale(x, col_field, *dt).into_owned(),
12208 layout: crate::gpu::BatchLayout::Q8,
12209 },
12210 )),
12211 QTensor::Mapped {
12213 model,
12214 idx,
12215 dtype: TensorDtype::Q1,
12216 rows,
12217 cols,
12218 ..
12219 } => Some((
12220 model.clone(),
12221 crate::gpu::BatchJob {
12222 idx: *idx,
12223 rows: *rows,
12224 cols: *cols,
12225 row_scale: &[],
12226 xs: x.to_vec(),
12227 layout: crate::gpu::BatchLayout::Q1,
12228 },
12229 )),
12230 QTensor::Mapped {
12232 model,
12233 idx,
12234 dtype: TensorDtype::Q1G,
12235 rows,
12236 cols,
12237 ..
12238 } => Some((
12239 model.clone(),
12240 crate::gpu::BatchJob {
12241 idx: *idx,
12242 rows: *rows,
12243 cols: *cols,
12244 row_scale: &[],
12245 xs: x.to_vec(),
12246 layout: crate::gpu::BatchLayout::Q1g,
12247 },
12248 )),
12249 QTensor::Mapped {
12254 model,
12255 idx,
12256 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
12257 rows,
12258 cols,
12259 ..
12260 } => Some((
12261 model.clone(),
12262 crate::gpu::BatchJob {
12263 idx: *idx,
12264 rows: *rows,
12265 cols: *cols,
12266 row_scale: &[],
12267 xs: x.to_vec(),
12268 layout: if *dt == TensorDtype::Q4Tiled {
12269 crate::gpu::BatchLayout::Q4t
12270 } else {
12271 crate::gpu::BatchLayout::Q4tp
12272 },
12273 },
12274 )),
12275 _ => None,
12276 }
12277}
12278
12279thread_local! {
12280 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12281 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12282}
12283
12284pub(crate) fn prescale<'a>(
12285 x: &'a [f32],
12286 col_field: &[f32],
12287 dtype: TensorDtype,
12288) -> std::borrow::Cow<'a, [f32]> {
12289 if dtype == TensorDtype::Q8_2f {
12290 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
12291 } else {
12292 std::borrow::Cow::Borrowed(x)
12293 }
12294}
12295
12296pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
12299 x: &[f32],
12300 col_field: &[f32],
12301 dtype: TensorDtype,
12302 buf_id: u8,
12303 f: F,
12304) -> R {
12305 if dtype == TensorDtype::Q8_2f {
12306 if buf_id == 1 {
12307 PRESCALE_BUF1.with(|b| {
12308 let mut buf = b.borrow_mut();
12309 buf.clear();
12310 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12311 f(&buf)
12312 })
12313 } else {
12314 PRESCALE_BUF2.with(|b| {
12315 let mut buf = b.borrow_mut();
12316 buf.clear();
12317 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12318 f(&buf)
12319 })
12320 }
12321 } else {
12322 f(x)
12323 }
12324}
12325
12326#[cfg(target_arch = "x86_64")]
12331pub(crate) fn avx2_enabled() -> bool {
12332 use std::sync::OnceLock;
12333 static ON: OnceLock<bool> = OnceLock::new();
12334 *ON.get_or_init(|| {
12335 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
12336 && std::arch::is_x86_feature_detected!("avx2")
12337 && std::arch::is_x86_feature_detected!("fma")
12338 })
12339}
12340
12341#[cfg(target_arch = "x86_64")]
12346fn avx2_a8w8_enabled() -> bool {
12347 if FLOAT_ACTIVATIONS.get() {
12348 return false;
12349 }
12350 use std::sync::OnceLock;
12351 static ON: OnceLock<bool> = OnceLock::new();
12352 *ON.get_or_init(|| {
12353 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
12354 })
12355}
12356
12357thread_local! {
12358 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12359}
12360
12361pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
12364 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
12365 impl Drop for Restore {
12366 fn drop(&mut self) {
12367 FULL_GPU_Q8.set(self.0);
12368 }
12369 }
12370 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
12371}
12372
12373thread_local! {
12377 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12378}
12379
12380fn exact_cpu_default_enabled() -> bool {
12389 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12390 *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
12391}
12392
12393impl QTensor {
12394 #[inline]
12395 fn exact_cpu_default(&self) -> bool {
12396 !FLOAT_ACTIVATIONS.get()
12397 && exact_cpu_default_enabled()
12398 && matches!(
12399 self,
12400 Self::Mapped { model, dtype, .. } if matches!(model.header.arch.arch_name.as_str(), "flint" | "cortiq_embryo")
12401 || (model.header.arch.arch_name == "dory" && *dtype == TensorDtype::Q8_2f)
12404 )
12405 }
12406
12407 #[inline]
12408 fn any_exact_cpu_default(ts: &[&QTensor]) -> bool {
12409 ts.iter().any(|t| t.exact_cpu_default())
12410 }
12411
12412 #[inline]
12413 fn any_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
12414 pairs
12415 .iter()
12416 .any(|(gate, up)| gate.exact_cpu_default() || up.exact_cpu_default())
12417 }
12418}
12419
12420pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
12425 float_activations_scope(f)
12426}
12427
12428#[inline]
12432pub(crate) fn float_activations_active() -> bool {
12433 FLOAT_ACTIVATIONS.get()
12434}
12435
12436pub(crate) struct FloatActivationsScope {
12440 previous: bool,
12441 _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
12442}
12443
12444impl Drop for FloatActivationsScope {
12445 fn drop(&mut self) {
12446 FLOAT_ACTIVATIONS.set(self.previous);
12447 }
12448}
12449
12450#[inline]
12451pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
12452 FloatActivationsScope {
12453 previous: FLOAT_ACTIVATIONS.replace(true),
12454 _not_send: std::marker::PhantomData,
12455 }
12456}
12457
12458#[doc(hidden)]
12464pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
12465 let _scope = enter_float_activations_scope();
12466 f()
12467}
12468
12469static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
12481
12482pub(crate) fn row_exact() -> bool {
12483 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
12484}
12485
12486fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
12487 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
12488 impl Drop for Restore<'_> {
12489 fn drop(&mut self) {
12490 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
12491 }
12492 }
12493 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
12494 let _restore = Restore(active);
12495 f()
12496}
12497
12498pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
12500 counted_row_exact_scope(&ROW_EXACT, f)
12501}
12502
12503#[inline]
12507pub(crate) fn a8w8_enabled() -> bool {
12508 #[cfg(target_arch = "aarch64")]
12509 {
12510 sdot_enabled()
12511 }
12512 #[cfg(target_arch = "x86_64")]
12513 {
12514 avx2_a8w8_enabled()
12515 }
12516 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
12517 {
12518 false
12519 }
12520}
12521
12522#[inline]
12525#[allow(unreachable_code)]
12526fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
12527 #[cfg(target_arch = "aarch64")]
12528 unsafe {
12529 return dot_i8_sdot(w, xq);
12530 }
12531 #[cfg(target_arch = "x86_64")]
12532 unsafe {
12533 if avx512vnni_enabled() {
12534 return dot_i8_i8_vnni(w, xq);
12535 }
12536 return dot_i8_i8_avx2(w, xq);
12537 }
12538 w.iter()
12539 .zip(xq)
12540 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
12541 .sum()
12542}
12543
12544#[cfg(target_arch = "x86_64")]
12548fn avx512vnni_enabled() -> bool {
12549 use std::sync::OnceLock;
12550 static ON: OnceLock<bool> = OnceLock::new();
12551 *ON.get_or_init(|| {
12552 std::env::var("CMF_AVX512")
12553 .map(|v| v != "0")
12554 .unwrap_or(true)
12555 && std::arch::is_x86_feature_detected!("avx512f")
12556 && std::arch::is_x86_feature_detected!("avx512bw")
12557 && std::arch::is_x86_feature_detected!("avx512vl")
12558 && std::arch::is_x86_feature_detected!("avx512vnni")
12559 })
12560}
12561
12562#[cfg(target_arch = "x86_64")]
12570fn vnni_tiles_enabled() -> bool {
12571 use std::sync::OnceLock;
12572 static ON: OnceLock<bool> = OnceLock::new();
12573 *ON.get_or_init(|| {
12574 std::env::var("CMF_VNNI_TILES")
12575 .map(|v| v != "0")
12576 .unwrap_or(true)
12577 && avx512vnni_enabled()
12578 })
12579}
12580
12581#[cfg(target_arch = "x86_64")]
12586#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12587#[inline]
12588unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
12589 unsafe {
12591 use core::arch::x86_64::*;
12592 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
12593 let hi128 = _mm256_extracti128_si256::<1>(d);
12594 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12595 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12596 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12597 _mm_cvtsi128_si32(s32)
12598 }
12599}
12600
12601#[cfg(target_arch = "x86_64")]
12606#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12607unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12608 unsafe {
12610 use core::arch::x86_64::*;
12611 let n = w.len();
12612 let mut j = 0usize;
12613 let mut total: i32;
12614 {
12619 #[inline(always)]
12620 unsafe fn step(
12621 w: *const u8,
12622 x: *const i8,
12623 acc: core::arch::x86_64::__m512i,
12624 ) -> core::arch::x86_64::__m512i {
12625 unsafe {
12626 use core::arch::x86_64::*;
12627 let wv = _mm512_loadu_si512(w as *const _);
12628 let xv = _mm512_loadu_si512(x as *const _);
12629 let aw = _mm512_abs_epi8(wv);
12630 let neg = _mm512_movepi8_mask(wv);
12631 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
12632 _mm512_dpbusd_epi32(acc, aw, sx)
12633 }
12634 }
12635 let (mut a0, mut a1, mut a2, mut a3) = (
12636 _mm512_setzero_si512(),
12637 _mm512_setzero_si512(),
12638 _mm512_setzero_si512(),
12639 _mm512_setzero_si512(),
12640 );
12641 while j + 256 <= n {
12642 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12643 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
12644 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
12645 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
12646 j += 256;
12647 }
12648 while j + 64 <= n {
12649 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12650 j += 64;
12651 }
12652 let s01 = _mm512_add_epi32(a0, a1);
12653 let s23 = _mm512_add_epi32(a2, a3);
12654 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
12655 }
12656 if j + 32 <= n {
12658 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12659 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12660 let d = _mm256_dpbusd_epi32(
12661 _mm256_setzero_si256(),
12662 _mm256_abs_epi8(wv),
12663 _mm256_sign_epi8(xv, wv),
12664 );
12665 let hi128 = _mm256_extracti128_si256::<1>(d);
12666 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12667 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12668 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12669 total += _mm_cvtsi128_si32(s32);
12670 j += 32;
12671 }
12672 while j < n {
12673 total += (w[j] as i8) as i32 * xq[j] as i32;
12674 j += 1;
12675 }
12676 total
12677 }
12678}
12679
12680#[cfg(target_arch = "x86_64")]
12682#[target_feature(enable = "avx2,fma")]
12683unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12684 unsafe {
12686 use core::arch::x86_64::*;
12687 let n = x.len();
12688 let wp = w.as_ptr();
12689 let xp = x.as_ptr();
12690 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12691 let mut j = 0usize;
12692 while j + 16 <= n {
12693 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12694 let lo = _mm256_cvtepi8_epi32(wb);
12695 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12696 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12697 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12698 j += 16;
12699 }
12700 let acc = _mm256_add_ps(a0, a1);
12701 let hi128 = _mm256_extractf128_ps::<1>(acc);
12702 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12703 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12704 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12705 let mut sum = _mm_cvtss_f32(s32);
12706 while j < n {
12707 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12708 j += 1;
12709 }
12710 sum
12711 }
12712}
12713
12714#[cfg(target_arch = "x86_64")]
12719#[target_feature(enable = "avx2")]
12720unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12721 unsafe {
12723 use core::arch::x86_64::*;
12724 let n = w.len();
12725 let ones = _mm256_set1_epi16(1);
12726 let mut acc = _mm256_setzero_si256();
12727 let mut j = 0usize;
12728 while j + 32 <= n {
12729 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12730 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12731 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12732 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12733 j += 32;
12734 }
12735 let hi128 = _mm256_extracti128_si256::<1>(acc);
12736 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12737 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12738 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12739 let mut s = _mm_cvtsi128_si32(s32);
12740 while j < n {
12741 s += (w[j] as i8) as i32 * xq[j] as i32;
12742 j += 1;
12743 }
12744 s
12745 }
12746}
12747
12748#[cfg(target_arch = "aarch64")]
12752#[target_feature(enable = "neon,i8mm")]
12753unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12754 unsafe {
12756 use core::arch::aarch64::*;
12757 use core::arch::asm;
12758 let n = w0.len();
12759 let w0p = w0.as_ptr() as *const i8;
12760 let w1p = w1.as_ptr() as *const i8;
12761 let mut acc01 = vdupq_n_s32(0);
12764 let mut acc23 = vdupq_n_s32(0);
12765 let mut i = 0usize;
12766 while i + 8 <= n {
12767 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12768 let xb01 = vcombine_s8(
12769 vld1_s8(xs[0].as_ptr().add(i)),
12770 vld1_s8(xs[1].as_ptr().add(i)),
12771 );
12772 let xb23 = vcombine_s8(
12773 vld1_s8(xs[2].as_ptr().add(i)),
12774 vld1_s8(xs[3].as_ptr().add(i)),
12775 );
12776 asm!(
12777 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12778 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12779 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12780 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12781 options(pure, nomem, nostack),
12782 );
12783 i += 8;
12784 }
12785 let mut out = [[0i32; 4]; 2];
12786 let a01: [i32; 4] = core::mem::transmute(acc01);
12787 let a23: [i32; 4] = core::mem::transmute(acc23);
12788 out[0][0] = a01[0];
12789 out[0][1] = a01[1];
12790 out[1][0] = a01[2];
12791 out[1][1] = a01[3];
12792 out[0][2] = a23[0];
12793 out[0][3] = a23[1];
12794 out[1][2] = a23[2];
12795 out[1][3] = a23[3];
12796 if i < n {
12797 for (k, x) in xs.iter().enumerate() {
12798 for j in i..n {
12799 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12800 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12801 }
12802 }
12803 }
12804 out
12805 }
12806}
12807
12808#[cfg(target_arch = "aarch64")]
12812#[target_feature(enable = "neon,dotprod")]
12813unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12814 unsafe {
12816 use core::arch::aarch64::*;
12817 use core::arch::asm;
12818 let n = w0.len();
12819 let w0p = w0.as_ptr() as *const i8;
12820 let w1p = w1.as_ptr() as *const i8;
12821 let mut acc = [[vdupq_n_s32(0); 4]; 2];
12822 let mut i = 0usize;
12823 while i + 16 <= n {
12824 let wv0 = vld1q_s8(w0p.add(i));
12825 let wv1 = vld1q_s8(w1p.add(i));
12826 for (k, x) in xs.iter().enumerate() {
12827 let xv = vld1q_s8(x.as_ptr().add(i));
12828 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12829 asm!(
12830 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12831 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12832 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12833 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12834 options(pure, nomem, nostack),
12835 );
12836 acc[0][k] = a0;
12837 acc[1][k] = a1;
12838 }
12839 i += 16;
12840 }
12841 let mut out = [[0i32; 4]; 2];
12842 for r in 0..2 {
12843 for k in 0..4 {
12844 out[r][k] = vaddvq_s32(acc[r][k]);
12845 }
12846 }
12847 if i < n {
12848 for (k, x) in xs.iter().enumerate() {
12849 for j in i..n {
12850 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12851 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12852 }
12853 }
12854 }
12855 out
12856 }
12857}
12858
12859#[cfg(target_arch = "x86_64")]
12865#[target_feature(enable = "avx2")]
12866unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12867 unsafe {
12869 use core::arch::x86_64::*;
12870 let n = w0.len();
12871 let ones = _mm256_set1_epi16(1);
12872 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12873 let mut j = 0usize;
12874 while j + 32 <= n {
12875 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12876 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12877 let aw0 = _mm256_abs_epi8(wv0);
12878 let aw1 = _mm256_abs_epi8(wv1);
12879 for (k, x) in xs.iter().enumerate() {
12880 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12881 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12882 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12883 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12884 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12885 }
12886 j += 32;
12887 }
12888 let mut out = [[0i32; 4]; 2];
12889 for r in 0..2 {
12890 for k in 0..4 {
12891 let a = acc[r][k];
12892 let hi128 = _mm256_extracti128_si256::<1>(a);
12893 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12894 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12895 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12896 out[r][k] = _mm_cvtsi128_si32(s32);
12897 }
12898 }
12899 if j < n {
12900 for (k, x) in xs.iter().enumerate() {
12901 for i in j..n {
12902 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12903 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12904 }
12905 }
12906 }
12907 out
12908 }
12909}
12910
12911#[cfg(target_arch = "x86_64")]
12916#[inline]
12917fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12918 let dot = if avx512vnni_enabled() && row.len() >= 64 {
12919 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12920 } else {
12921 unsafe { dot_i8_i8_avx2(row, &act.xq) }
12922 };
12923 let mut acc = dot as f32 * act.sx;
12924 for &(j, xv) in &act.outliers {
12925 acc += (row[j] as i8) as f32 * xv;
12926 }
12927 acc
12928}
12929
12930#[cfg(target_arch = "x86_64")]
12934#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12935unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12936 unsafe {
12938 use core::arch::x86_64::*;
12939 let n = w.len();
12940 let flip = _mm512_set1_epi8(-128); #[inline(always)]
12942 unsafe fn step(
12943 w: *const u8,
12944 x: *const i8,
12945 flip: core::arch::x86_64::__m512i,
12946 acc: core::arch::x86_64::__m512i,
12947 ) -> core::arch::x86_64::__m512i {
12948 unsafe {
12949 use core::arch::x86_64::*;
12950 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12951 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12952 }
12953 }
12954 let (mut a0, mut a1, mut a2, mut a3) = (
12955 _mm512_setzero_si512(),
12956 _mm512_setzero_si512(),
12957 _mm512_setzero_si512(),
12958 _mm512_setzero_si512(),
12959 );
12960 let mut j = 0usize;
12961 while j + 256 <= n {
12962 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12963 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12964 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12965 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12966 j += 256;
12967 }
12968 while j + 64 <= n {
12969 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12970 j += 64;
12971 }
12972 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12973 _mm512_add_epi32(a0, a1),
12974 _mm512_add_epi32(a2, a3),
12975 ));
12976 while j < n {
12978 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12979 j += 1;
12980 }
12981 total
12982 }
12983}
12984
12985#[cfg(target_arch = "x86_64")]
12991#[target_feature(enable = "avx2")]
12992unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12993 unsafe {
12996 use core::arch::x86_64::*;
12997 let lomask = _mm_set1_epi8(0x0F);
12998 let eight = _mm256_set1_epi8(8);
12999 let ones = _mm256_set1_epi16(1);
13000 let mut acc = 0f32;
13001 for gi in 0..gpr {
13002 let g = g0 + gi;
13003 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13004 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
13005 let lo = _mm_and_si128(b, lomask);
13006 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
13007 let w = _mm256_sub_epi8(
13008 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
13009 eight,
13010 );
13011 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13012 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
13013 let d = _mm256_madd_epi16(p16, ones);
13014 let hi128 = _mm256_extracti128_si256::<1>(d);
13015 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
13016 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
13017 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
13018 acc += _mm_cvtsi128_si32(s32) as f32 * s;
13019 }
13020 acc
13021 }
13022}
13023
13024#[cfg(target_arch = "x86_64")]
13027#[target_feature(enable = "avx2")]
13028unsafe fn dot_q4_row_avx2_2(
13029 packed: &[u8],
13030 scales: &[u8],
13031 g0: usize,
13032 gpr: usize,
13033 xq1: &[i8],
13034 xq2: &[i8],
13035) -> (f32, f32) {
13036 unsafe {
13038 use core::arch::x86_64::*;
13039 let lomask = _mm_set1_epi8(0x0F);
13040 let eight = _mm256_set1_epi8(8);
13041 let ones = _mm256_set1_epi16(1);
13042 let (mut acc1, mut acc2) = (0f32, 0f32);
13043 #[inline(always)]
13044 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
13045 unsafe {
13046 use core::arch::x86_64::*;
13047 let hi128 = _mm256_extracti128_si256::<1>(d);
13048 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
13049 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
13050 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
13051 _mm_cvtsi128_si32(s32)
13052 }
13053 }
13054 for gi in 0..gpr {
13055 let g = g0 + gi;
13056 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13057 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
13058 let lo = _mm_and_si128(b, lomask);
13059 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
13060 let w = _mm256_sub_epi8(
13061 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
13062 eight,
13063 );
13064 let aw = _mm256_abs_epi8(w);
13065 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13066 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13067 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
13068 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
13069 acc1 += hsum(d1) as f32 * s;
13070 acc2 += hsum(d2) as f32 * s;
13071 }
13072 (acc1, acc2)
13073 }
13074}
13075
13076#[cfg(target_arch = "x86_64")]
13078fn q8_range_avx2(
13079 q: &[u8],
13080 row_scale: &[f32],
13081 act: &SplitAct,
13082 cols: usize,
13083 out_addr: SendMut,
13084 start: usize,
13085 end: usize,
13086) {
13087 for o in start..end {
13088 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13089 unsafe { *out_addr.at(o) = v };
13091 }
13092}
13093
13094#[cfg(target_arch = "x86_64")]
13096#[allow(clippy::too_many_arguments)]
13097fn q8_range2_avx2(
13098 q: &[u8],
13099 row_scale: &[f32],
13100 a1: &SplitAct,
13101 a2: &SplitAct,
13102 cols: usize,
13103 p1: SendMut,
13104 p2: SendMut,
13105 start: usize,
13106 end: usize,
13107) {
13108 for o in start..end {
13109 let row = &q[o * cols..(o + 1) * cols];
13110 unsafe {
13112 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
13113 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
13114 }
13115 }
13116}
13117
13118#[cfg(target_arch = "aarch64")]
13129fn i8mm_enabled() -> bool {
13130 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13131 *ON.get_or_init(|| {
13132 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
13133 && std::arch::is_aarch64_feature_detected!("i8mm")
13134 })
13135}
13136
13137#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
13141fn sdot_enabled() -> bool {
13142 if FLOAT_ACTIVATIONS.get() {
13143 return false;
13144 }
13145 use std::sync::OnceLock;
13146 static ON: OnceLock<bool> = OnceLock::new();
13147 *ON.get_or_init(|| {
13148 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
13149 if !want {
13150 return false;
13151 }
13152
13153 #[cfg(target_arch = "aarch64")]
13154 {
13155 if std::arch::is_aarch64_feature_detected!("dotprod") {
13156 return true;
13157 }
13158 #[cfg(target_os = "android")]
13159 {
13160 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
13161 if cpuinfo.lines().any(|l| {
13162 (l.starts_with("Features") || l.starts_with("features"))
13163 && l.contains("asimddp")
13164 }) {
13165 return true;
13166 }
13167 }
13168 }
13169 false
13170 }
13171 #[cfg(not(target_arch = "aarch64"))]
13172 {
13173 false
13174 }
13175 })
13176}
13177
13178struct SplitAct {
13183 xq: Vec<i8>,
13184 sx: f32,
13185 outliers: Vec<(usize, f32)>,
13186 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13189 xsum: i32,
13190 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13200 xg8: Vec<f32>,
13201}
13202
13203thread_local! {
13204 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
13207 const { std::cell::RefCell::new(Vec::new()) };
13208 static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
13210 const { std::cell::RefCell::new(Vec::new()) };
13211}
13212
13213fn xg8_take() -> Vec<f32> {
13215 let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13216 v.clear();
13217 v
13218}
13219
13220impl Drop for SplitAct {
13221 fn drop(&mut self) {
13222 let buf = std::mem::take(&mut self.xq);
13223 if buf.capacity() > 0 {
13224 XQ_FREE.with(|f| {
13225 let mut f = f.borrow_mut();
13226 if f.len() < 16 {
13227 f.push(buf);
13228 }
13229 });
13230 }
13231 let buf = std::mem::take(&mut self.xg8);
13232 if buf.capacity() > 0 {
13233 XG8_FREE.with(|f| {
13234 let mut f = f.borrow_mut();
13235 if f.len() < 16 {
13236 f.push(buf);
13237 }
13238 });
13239 }
13240 }
13241}
13242
13243fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
13247 out.clear();
13248 out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
13249 let mut s = 0i32;
13250 for &v in c {
13251 s += v as i32;
13252 }
13253 s as f32
13254 }));
13255}
13256
13257thread_local! {
13258 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
13265 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
13270 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
13271}
13272
13273#[inline]
13276fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
13277 KROW.with(|s| {
13278 let b = unsafe { &mut *s.get() };
13281 if b.len() < n {
13282 b.resize(n, 0.0);
13283 }
13284 f(&mut b[..n])
13285 })
13286}
13287
13288#[inline(always)]
13299fn q8_round(t: f32) -> i8 {
13300 let t = t.clamp(-127.0, 127.0);
13301 let i = t as i32;
13302 let f = t - i as f32;
13303 let r = if f >= 0.5 {
13304 i + 1
13305 } else if f <= -0.5 {
13306 i - 1
13307 } else {
13308 i
13309 };
13310 r as i8
13311}
13312
13313#[inline]
13314fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
13315 KROWS.with(|s| {
13316 let b = unsafe { &mut *s.get() };
13319 for row in b.iter_mut() {
13320 if row.len() < n {
13321 row.resize(n, 0.0);
13322 }
13323 }
13324 let (a, b) = b.split_at_mut(1);
13325 f(&mut a[0][..n], &mut b[0][..n])
13326 })
13327}
13328
13329#[inline]
13330fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
13331 if limit > 0.0 {
13332 up = up.clamp(-limit, limit);
13333 gate = gate.min(limit);
13334 }
13335 gate / (1.0 + (-gate).exp()) * up
13336}
13337
13338fn split_acts_pooled(xs_all: &[f32], b: usize, cols: usize, pool: Option<&Pool>) -> Vec<SplitAct> {
13343 match pool {
13344 Some(p) if b >= 16 => {
13345 let slots: Vec<std::sync::Mutex<Option<SplitAct>>> =
13346 (0..b).map(|_| std::sync::Mutex::new(None)).collect();
13347 p.run_rows(b, &|start, end| {
13348 for bi in start..end {
13349 let a = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
13350 *slots[bi].lock().unwrap_or_else(|e| e.into_inner()) = Some(a);
13351 }
13352 });
13353 slots
13354 .into_iter()
13355 .map(|m| {
13356 m.into_inner()
13357 .unwrap_or_else(|e| e.into_inner())
13358 .expect("every column split")
13359 })
13360 .collect()
13361 }
13362 _ => (0..b)
13363 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
13364 .collect(),
13365 }
13366}
13367
13368fn split_act(x: &[f32]) -> SplitAct {
13369 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13370 let n = x.len();
13371 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
13372 let thr = 8.0 * rms;
13373 let mut outliers: Vec<(usize, f32)> = Vec::new();
13377 let mut amax = 0f32;
13378 for (j, &v) in x.iter().enumerate() {
13379 let a = v.abs();
13380 if a > thr {
13381 outliers.push((j, v));
13382 } else if a > amax {
13383 amax = a;
13384 }
13385 }
13386 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13387 let inv = 1.0 / sx;
13388 crate::cpuprof::count_outliers(outliers.len() as u64);
13389 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13390 xq.clear();
13391 xq.reserve(n);
13392 if outliers.is_empty() {
13393 xq.extend(
13394 x.iter()
13395 .map(|&v| q8_round(v * inv)),
13396 );
13397 } else {
13398 xq.extend(x.iter().map(|&v| {
13400 if v.abs() > thr {
13401 0
13402 } else {
13403 q8_round(v * inv)
13404 }
13405 }));
13406 }
13407 let xsum = xq.iter().map(|&v| v as i32).sum();
13408 let mut xg8 = xg8_take();
13409 xg8_fill(&xq, &mut xg8);
13410 SplitAct {
13411 xq,
13412 sx,
13413 outliers,
13414 xsum,
13415 xg8,
13416 }
13417}
13418
13419fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
13420 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13421 let n = x.len();
13422 let rms = (x
13423 .iter()
13424 .zip(col)
13425 .map(|(&a, &c)| {
13426 let v = a * c;
13427 (v * v) as f64
13428 })
13429 .sum::<f64>()
13430 / n.max(1) as f64)
13431 .sqrt() as f32;
13432 let thr = 8.0 * rms;
13433
13434 let mut outliers = Vec::new();
13435 let mut amax = 0f32;
13436 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
13437 let v = a * c;
13438 let s = v.abs();
13439 if s > thr {
13440 outliers.push((j, v));
13441 } else if s > amax {
13442 amax = s;
13443 }
13444 }
13445
13446 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13447 let inv = 1.0 / sx;
13448 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13449 xq.clear();
13450 xq.reserve(n);
13451 if outliers.is_empty() {
13452 xq.extend(
13453 x.iter()
13454 .zip(col)
13455 .map(|(&a, &c)| q8_round((a * c) * inv)),
13456 );
13457 } else {
13458 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
13459 let v = a * c;
13460 if v.abs() > thr {
13461 0
13462 } else {
13463 q8_round(v * inv)
13464 }
13465 }));
13466 }
13467 let xsum = xq.iter().map(|&v| v as i32).sum();
13468 let mut xg8 = xg8_take();
13469 xg8_fill(&xq, &mut xg8);
13470 SplitAct {
13471 xq,
13472 sx,
13473 outliers,
13474 xsum,
13475 xg8,
13476 }
13477}
13478
13479#[cfg(target_arch = "aarch64")]
13482#[target_feature(enable = "neon,dotprod")]
13483unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
13484 unsafe {
13486 use core::arch::aarch64::*;
13487 use core::arch::asm;
13488 let wp = w.as_ptr() as *const i8;
13489 let n = w.len();
13490 let (mut a0, mut a1, mut a2, mut a3) = (
13491 vdupq_n_s32(0),
13492 vdupq_n_s32(0),
13493 vdupq_n_s32(0),
13494 vdupq_n_s32(0),
13495 );
13496 let mut i = 0;
13497 while i + 64 <= n {
13498 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13499 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
13500 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
13501 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
13502 asm!(
13503 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
13504 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
13505 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
13506 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
13507 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13508 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
13509 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
13510 options(pure, nomem, nostack),
13511 );
13512 i += 64;
13513 }
13514 while i + 16 <= n {
13515 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13516 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
13517 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
13518 i += 16;
13519 }
13520 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
13521 while i < n {
13522 s += (*wp.add(i)) as i32 * xq[i] as i32;
13523 i += 1;
13524 }
13525 s
13526 }
13527}
13528
13529#[cfg(target_arch = "aarch64")]
13533#[target_feature(enable = "neon,dotprod")]
13534unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
13535 unsafe {
13537 use core::arch::aarch64::*;
13538 use core::arch::asm;
13539 let n = xq.len();
13540 let px = xq.as_ptr();
13541 let (p0, p1, p2, p3) = (
13542 w0.as_ptr() as *const i8,
13543 w1.as_ptr() as *const i8,
13544 w2.as_ptr() as *const i8,
13545 w3.as_ptr() as *const i8,
13546 );
13547 let (mut a0, mut a1, mut a2, mut a3) = (
13548 vdupq_n_s32(0),
13549 vdupq_n_s32(0),
13550 vdupq_n_s32(0),
13551 vdupq_n_s32(0),
13552 );
13553 let mut i = 0;
13554 while i + 16 <= n {
13555 let x = vld1q_s8(px.add(i));
13556 let v0 = vld1q_s8(p0.add(i));
13557 let v1 = vld1q_s8(p1.add(i));
13558 let v2 = vld1q_s8(p2.add(i));
13559 let v3 = vld1q_s8(p3.add(i));
13560 asm!(
13561 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13562 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13563 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13564 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13565 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13566 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13567 options(pure, nomem, nostack),
13568 );
13569 i += 16;
13570 }
13571 let mut r = [
13572 vaddvq_s32(a0),
13573 vaddvq_s32(a1),
13574 vaddvq_s32(a2),
13575 vaddvq_s32(a3),
13576 ];
13577 while i < n {
13578 let xi = *px.add(i) as i32;
13579 r[0] += (*p0.add(i)) as i32 * xi;
13580 r[1] += (*p1.add(i)) as i32 * xi;
13581 r[2] += (*p2.add(i)) as i32 * xi;
13582 r[3] += (*p3.add(i)) as i32 * xi;
13583 i += 1;
13584 }
13585 r
13586 }
13587}
13588
13589#[cfg(target_arch = "aarch64")]
13596#[target_feature(enable = "neon,dotprod")]
13597unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
13598 unsafe {
13601 use core::arch::aarch64::*;
13602 use core::arch::asm;
13603 let n = xq.len();
13604 let px = xq.as_ptr();
13605 let pg = g.as_ptr() as *const i8;
13606 let (mut a0, mut a1, mut a2, mut a3) = (
13607 vdupq_n_s32(0),
13608 vdupq_n_s32(0),
13609 vdupq_n_s32(0),
13610 vdupq_n_s32(0),
13611 );
13612 let mut i = 0;
13613 while i + 16 <= n {
13614 let x = vld1q_s8(px.add(i));
13615 let base = pg.add(4 * i);
13616 let v0 = vld1q_s8(base);
13617 let v1 = vld1q_s8(base.add(16));
13618 let v2 = vld1q_s8(base.add(32));
13619 let v3 = vld1q_s8(base.add(48));
13620 asm!(
13621 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13622 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13623 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13624 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13625 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13626 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13627 options(pure, nomem, nostack),
13628 );
13629 i += 16;
13630 }
13631 [
13632 vaddvq_s32(a0),
13633 vaddvq_s32(a1),
13634 vaddvq_s32(a2),
13635 vaddvq_s32(a3),
13636 ]
13637 }
13638}
13639
13640#[cfg(target_arch = "aarch64")]
13646fn q8_range_sdot(
13647 q: &[u8],
13648 rep: &[u8],
13649 row_scale: &[f32],
13650 act: &SplitAct,
13651 cols: usize,
13652 out_addr: SendMut,
13653 start: usize,
13654 end: usize,
13655) {
13656 let mut o = start;
13657 if !rep.is_empty() {
13660 while o < end && o % 4 != 0 {
13661 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13662 unsafe { *out_addr.at(o) = v };
13663 o += 1;
13664 }
13665 }
13666 while o + 4 <= end {
13667 let r = if rep.is_empty() {
13668 unsafe {
13669 dot_i8_sdot_4rows(
13670 &q[o * cols..(o + 1) * cols],
13671 &q[(o + 1) * cols..(o + 2) * cols],
13672 &q[(o + 2) * cols..(o + 3) * cols],
13673 &q[(o + 3) * cols..(o + 4) * cols],
13674 &act.xq,
13675 )
13676 }
13677 } else {
13678 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
13679 };
13680 for k in 0..4 {
13681 let mut acc = r[k] as f32 * act.sx;
13682 for &(j, xv) in &act.outliers {
13683 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
13684 }
13685 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
13687 }
13688 o += 4;
13689 }
13690 while o < end {
13691 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13692 unsafe { *out_addr.at(o) = v };
13693 o += 1;
13694 }
13695}
13696
13697#[cfg(target_arch = "aarch64")]
13700#[allow(clippy::too_many_arguments)]
13701fn q8_range2_sdot(
13702 q: &[u8],
13703 row_scale: &[f32],
13704 a1: &SplitAct,
13705 a2: &SplitAct,
13706 cols: usize,
13707 p1: SendMut,
13708 p2: SendMut,
13709 start: usize,
13710 end: usize,
13711) {
13712 for o in start..end {
13713 let row = &q[o * cols..(o + 1) * cols];
13714 unsafe {
13716 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13717 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13718 }
13719 }
13720}
13721
13722#[allow(clippy::too_many_arguments)]
13724fn q8_range2_f32(
13725 q: &[u8],
13726 row_scale: &[f32],
13727 x1: &[f32],
13728 x2: &[f32],
13729 cols: usize,
13730 p1: SendMut,
13731 p2: SendMut,
13732 start: usize,
13733 end: usize,
13734) {
13735 for o in start..end {
13736 let row = &q[o * cols..(o + 1) * cols];
13737 unsafe {
13739 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13740 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13741 }
13742 }
13743}
13744
13745fn q8_range_f32(
13747 q: &[u8],
13748 row_scale: &[f32],
13749 xs: &[f32],
13750 cols: usize,
13751 out_addr: SendMut,
13752 start: usize,
13753 end: usize,
13754) {
13755 for o in start..end {
13756 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13757 unsafe { *out_addr.at(o) = v };
13759 }
13760}
13761
13762#[inline]
13766fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13767 #[cfg(target_arch = "aarch64")]
13768 return row_dot_sdot(row, act);
13769 #[cfg(target_arch = "x86_64")]
13770 return row_dot_avx2(row, act);
13771 #[allow(unreachable_code)]
13772 q8_row_dot_scalar(row, act)
13773}
13774
13775#[allow(dead_code)]
13776fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13777 let mut acc = 0i32;
13778 for (k, &b) in row.iter().enumerate() {
13779 acc += (b as i8) as i32 * act.xq[k] as i32;
13780 }
13781 let mut acc = acc as f32 * act.sx;
13782 for &(j, xv) in &act.outliers {
13783 acc += (row[j] as i8) as f32 * xv;
13784 }
13785 acc
13786}
13787
13788#[cfg(target_arch = "aarch64")]
13791#[inline]
13792fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13793 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13794 for &(j, xv) in &act.outliers {
13795 acc += (row[j] as i8) as f32 * xv;
13796 }
13797 acc
13798}
13799
13800#[cfg(target_arch = "aarch64")]
13808#[target_feature(enable = "neon,dotprod")]
13809unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13810 unsafe {
13813 use core::arch::aarch64::*;
13814 use core::arch::asm;
13815 let lomask = vdupq_n_u8(0x0F);
13816 let eight = vdupq_n_s8(8);
13817 let mut acc = 0f32;
13818 for gi in 0..gpr {
13819 let g = g0 + gi;
13820 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13821 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13822 let lo = vandq_u8(b, lomask);
13823 let hi = vshrq_n_u8::<4>(b);
13824 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13825 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13826 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13827 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13828 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13829 asm!(
13830 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13831 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13832 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13833 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13834 options(pure, nomem, nostack),
13835 );
13836 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13837 }
13838 acc
13839 }
13840}
13841
13842#[cfg(target_arch = "aarch64")]
13847#[target_feature(enable = "neon,dotprod")]
13848unsafe fn dot_q4_row_sdot2(
13849 packed: &[u8],
13850 scales: &[u8],
13851 g0: usize,
13852 gpr: usize,
13853 xq1: &[i8],
13854 xq2: &[i8],
13855) -> (f32, f32) {
13856 unsafe {
13859 use core::arch::aarch64::*;
13860 use core::arch::asm;
13861 let lomask = vdupq_n_u8(0x0F);
13862 let eight = vdupq_n_s8(8);
13863 let (mut acc1, mut acc2) = (0f32, 0f32);
13864 for gi in 0..gpr {
13865 let g = g0 + gi;
13866 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13867 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13868 let lo = vandq_u8(b, lomask);
13869 let hi = vshrq_n_u8::<4>(b);
13870 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13871 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13872 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13873 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13874 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13875 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13876 let (mut a0, mut a1, mut b0, mut b1) = (
13877 vdupq_n_s32(0),
13878 vdupq_n_s32(0),
13879 vdupq_n_s32(0),
13880 vdupq_n_s32(0),
13881 );
13882 asm!(
13883 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13884 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13885 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13886 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13887 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13888 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13889 e0 = in(vreg) e0, e1 = in(vreg) e1,
13890 x10 = in(vreg) x10, x11 = in(vreg) x11,
13891 x20 = in(vreg) x20, x21 = in(vreg) x21,
13892 options(pure, nomem, nostack),
13893 );
13894 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13895 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13896 }
13897 (acc1, acc2)
13898 }
13899}
13900
13901#[inline]
13906pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13907 #[cfg(target_arch = "aarch64")]
13908 unsafe {
13909 return axpy_i8_f32_neon(acc, row, w);
13910 }
13911 #[cfg(target_arch = "x86_64")]
13912 if avx2_enabled() {
13913 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13914 }
13915 #[allow(unreachable_code)]
13916 {
13917 for (a, &b) in acc.iter_mut().zip(row) {
13918 *a += w * b as f32;
13919 }
13920 }
13921}
13922
13923#[cfg(target_arch = "x86_64")]
13925#[target_feature(enable = "avx2,fma")]
13926unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13927 unsafe {
13929 use core::arch::x86_64::*;
13930 let n = acc.len().min(row.len());
13931 let ap = acc.as_mut_ptr();
13932 let rp = row.as_ptr();
13933 let wv = _mm256_set1_ps(w);
13934 let mut j = 0usize;
13935 while j + 16 <= n {
13936 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13937 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13938 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13939 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13940 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13941 _mm256_storeu_ps(ap.add(j), v0);
13942 _mm256_storeu_ps(ap.add(j + 8), v1);
13943 j += 16;
13944 }
13945 while j < n {
13946 *ap.add(j) += w * (*rp.add(j)) as f32;
13947 j += 1;
13948 }
13949 }
13950}
13951
13952#[cfg(target_arch = "aarch64")]
13953#[target_feature(enable = "neon")]
13954unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13955 unsafe {
13957 use core::arch::aarch64::*;
13958 let n = acc.len().min(row.len());
13959 let ap = acc.as_mut_ptr();
13960 let rp = row.as_ptr();
13961 let wv = vdupq_n_f32(w);
13962 let mut j = 0usize;
13963 while j + 16 <= n {
13964 let rb = vld1q_s8(rp.add(j));
13965 let lo = vmovl_s8(vget_low_s8(rb));
13966 let hi = vmovl_s8(vget_high_s8(rb));
13967 for (off, half) in [(0, lo), (8, hi)] {
13968 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13969 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13970 let o = j + off;
13971 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13972 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13973 }
13974 j += 16;
13975 }
13976 while j < n {
13977 *ap.add(j) += w * (*rp.add(j)) as f32;
13978 j += 1;
13979 }
13980 }
13981}
13982
13983#[inline]
13986pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13987 #[cfg(target_arch = "aarch64")]
13988 unsafe {
13989 return dot_i8_f32_neon(w, x);
13990 }
13991 #[cfg(target_arch = "x86_64")]
13992 if avx2_enabled() {
13993 return unsafe { dot_i8_f32_avx2(w, x) };
13994 }
13995 #[allow(unreachable_code)]
13996 {
13997 let mut sum = 0.0f32;
13998 for (j, &b) in w.iter().enumerate() {
13999 sum += (b as i8) as f32 * x[j];
14000 }
14001 sum
14002 }
14003}
14004
14005#[inline]
14009fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
14010 #[cfg(target_arch = "aarch64")]
14011 unsafe {
14012 return dot_i8_col_f32_neon(w, x, col);
14013 }
14014 #[allow(unreachable_code)]
14015 {
14016 let mut sum = 0.0f32;
14017 for (j, &b) in w.iter().enumerate() {
14018 sum += (b as i8) as f32 * x[j] * col[j];
14019 }
14020 sum
14021 }
14022}
14023
14024#[cfg(target_arch = "aarch64")]
14025#[target_feature(enable = "neon")]
14026unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
14027 unsafe {
14029 use core::arch::aarch64::*;
14030 let n = x.len();
14031 let wp = w.as_ptr() as *const i8;
14032 let xp = x.as_ptr();
14033 let cp = col.as_ptr();
14034 let (mut a0, mut a1, mut a2, mut a3) = (
14035 vdupq_n_f32(0.0),
14036 vdupq_n_f32(0.0),
14037 vdupq_n_f32(0.0),
14038 vdupq_n_f32(0.0),
14039 );
14040 let mut j = 0usize;
14041 while j + 16 <= n {
14042 let wb = vld1q_s8(wp.add(j));
14043 let lo = vmovl_s8(vget_low_s8(wb));
14044 let hi = vmovl_s8(vget_high_s8(wb));
14045 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14046 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14047 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14048 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14049 a0 = vfmaq_f32(
14050 a0,
14051 w0,
14052 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
14053 );
14054 a1 = vfmaq_f32(
14055 a1,
14056 w1,
14057 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
14058 );
14059 a2 = vfmaq_f32(
14060 a2,
14061 w2,
14062 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
14063 );
14064 a3 = vfmaq_f32(
14065 a3,
14066 w3,
14067 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
14068 );
14069 j += 16;
14070 }
14071 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14072 while j < n {
14073 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
14074 j += 1;
14075 }
14076 sum
14077 }
14078}
14079
14080#[cfg(target_arch = "aarch64")]
14081#[target_feature(enable = "neon")]
14082unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
14083 unsafe {
14085 use core::arch::aarch64::*;
14086 let n = x.len();
14087 let wp = w.as_ptr() as *const i8;
14088 let xp = x.as_ptr();
14089 let (mut a0, mut a1, mut a2, mut a3) = (
14090 vdupq_n_f32(0.0),
14091 vdupq_n_f32(0.0),
14092 vdupq_n_f32(0.0),
14093 vdupq_n_f32(0.0),
14094 );
14095 let mut j = 0usize;
14096 while j + 16 <= n {
14097 let wb = vld1q_s8(wp.add(j));
14098 let lo = vmovl_s8(vget_low_s8(wb));
14099 let hi = vmovl_s8(vget_high_s8(wb));
14100 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14101 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14102 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14103 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14104 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
14105 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
14106 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
14107 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
14108 j += 16;
14109 }
14110 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14111 while j < n {
14112 sum += (*wp.add(j)) as f32 * *xp.add(j);
14113 j += 1;
14114 }
14115 sum
14116 }
14117}
14118
14119#[allow(clippy::too_many_arguments)]
14120fn qmatvec(
14121 q: &[u8],
14122 rep: &[u8],
14123 row_scale: &[f32],
14124 x: &[f32],
14125 col_field: &[f32],
14126 dtype: TensorDtype,
14127 rows: usize,
14128 cols: usize,
14129 out: &mut [f32],
14130 pool: Option<&Pool>,
14131) {
14132 debug_assert_eq!(out.len(), rows);
14133 #[cfg(not(target_arch = "aarch64"))]
14134 let _ = rep;
14135
14136 #[cfg(target_arch = "aarch64")]
14137 if sdot_enabled() {
14138 let act = if dtype == TensorDtype::Q8_2f {
14139 split_act_q8_2f(x, col_field)
14140 } else {
14141 split_act(x)
14142 };
14143 let out_addr = SendMut(out.as_mut_ptr());
14144 let run_range = |start: usize, end: usize| {
14145 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
14146 };
14147 match pool {
14148 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14149 _ => run_range(0, rows),
14150 }
14151 return;
14152 }
14153 #[cfg(target_arch = "x86_64")]
14156 if avx2_a8w8_enabled() {
14157 let act = if dtype == TensorDtype::Q8_2f {
14158 split_act_q8_2f(x, col_field)
14159 } else {
14160 split_act(x)
14161 };
14162 let out_addr = SendMut(out.as_mut_ptr());
14163 let run_range = |start: usize, end: usize| {
14164 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
14165 };
14166 match pool {
14167 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14168 _ => run_range(0, rows),
14169 }
14170 return;
14171 }
14172
14173 prescale_with(x, col_field, dtype, 1, |xs| {
14174 let out_addr = SendMut(out.as_mut_ptr());
14175 let run_range = move |start: usize, end: usize| {
14176 for o in start..end {
14177 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
14178 unsafe { *out_addr.at(o) = v };
14180 }
14181 };
14182 match pool {
14183 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14184 _ => run_range(0, rows),
14185 }
14186 });
14187}
14188
14189#[allow(clippy::too_many_arguments)]
14190fn qmatvec2(
14191 q: &[u8],
14192 row_scale: &[f32],
14193 x1: &[f32],
14194 x2: &[f32],
14195 col_field: &[f32],
14196 dtype: TensorDtype,
14197 rows: usize,
14198 cols: usize,
14199 o1: &mut [f32],
14200 o2: &mut [f32],
14201 pool: Option<&Pool>,
14202) {
14203 #[cfg(target_arch = "aarch64")]
14204 if sdot_enabled() {
14205 let a1s = if dtype == TensorDtype::Q8_2f {
14206 split_act_q8_2f(x1, col_field)
14207 } else {
14208 split_act(x1)
14209 };
14210 let a2s = if dtype == TensorDtype::Q8_2f {
14211 split_act_q8_2f(x2, col_field)
14212 } else {
14213 split_act(x2)
14214 };
14215 let p1 = SendMut(o1.as_mut_ptr());
14216 let p2 = SendMut(o2.as_mut_ptr());
14217 let run_range = |start: usize, end: usize| {
14218 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14219 };
14220 match pool {
14221 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14222 _ => run_range(0, rows),
14223 }
14224 return;
14225 }
14226 #[cfg(target_arch = "x86_64")]
14227 if avx2_a8w8_enabled() {
14228 let a1s = if dtype == TensorDtype::Q8_2f {
14229 split_act_q8_2f(x1, col_field)
14230 } else {
14231 split_act(x1)
14232 };
14233 let a2s = if dtype == TensorDtype::Q8_2f {
14234 split_act_q8_2f(x2, col_field)
14235 } else {
14236 split_act(x2)
14237 };
14238 let p1 = SendMut(o1.as_mut_ptr());
14239 let p2 = SendMut(o2.as_mut_ptr());
14240 let run_range = |start: usize, end: usize| {
14241 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14242 };
14243 match pool {
14244 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14245 _ => run_range(0, rows),
14246 }
14247 return;
14248 }
14249
14250 prescale_with(x1, col_field, dtype, 1, |x1s| {
14251 prescale_with(x2, col_field, dtype, 2, |x2s| {
14252 let p1 = SendMut(o1.as_mut_ptr());
14253 let p2 = SendMut(o2.as_mut_ptr());
14254 let run_range = move |start: usize, end: usize| {
14255 for o in start..end {
14256 let row = &q[o * cols..(o + 1) * cols];
14257 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
14258 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
14259 unsafe {
14261 *p1.at(o) = s1;
14262 *p2.at(o) = s2;
14263 }
14264 }
14265 };
14266 match pool {
14267 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14268 _ => run_range(0, rows),
14269 }
14270 });
14271 });
14272}
14273
14274#[derive(Clone, Copy)]
14275struct SendMut(*mut f32);
14276unsafe impl Send for SendMut {}
14277unsafe impl Sync for SendMut {}
14278
14279impl SendMut {
14280 #[inline]
14281 fn at(self, i: usize) -> *mut f32 {
14282 unsafe { self.0.add(i) }
14283 }
14284}
14285
14286#[cfg(test)]
14287mod tests {
14288
14289 #[test]
14293 fn q8_round_is_round_clamp() {
14294 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
14295 let mut probe = vec![
14296 0.0f32,
14297 -0.0,
14298 f32::NAN,
14299 f32::INFINITY,
14300 f32::NEG_INFINITY,
14301 f32::MAX,
14302 f32::MIN,
14303 1e30,
14304 -1e30,
14305 f32::MIN_POSITIVE,
14306 -f32::MIN_POSITIVE,
14307 ];
14308 for k in -300i32..=300 {
14309 let h = k as f32 * 0.5;
14310 let up = f32::from_bits(h.to_bits() + 1);
14311 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
14312 for t in [h, up, down] {
14313 probe.push(t);
14314 probe.push(-t);
14315 }
14316 }
14317 let mut t = -140.0f32;
14318 while t < 140.0 {
14319 probe.push(t);
14320 t += 0.000_731;
14321 }
14322 for t in probe {
14323 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
14324 }
14325 }
14326
14327 use super::*;
14328
14329 #[test]
14333 fn f32_matmat_equals_matvec_bitwise() {
14334 let (rows, cols, b) = (5usize, 37usize, 7usize);
14335 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14336 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14337 let t = QTensor::from_f32(w, rows, cols);
14338 let mut out = vec![0f32; b * rows];
14339 t.matmat(&xs, b, &mut out, None);
14340 for bi in 0..b {
14341 let mut one = vec![0f32; rows];
14342 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
14343 for o in 0..rows {
14344 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
14345 }
14346 }
14347 }
14348
14349 #[test]
14352 fn f32_matmat_row_split_on_pool_is_bitwise() {
14353 let (rows, cols, b) = (300usize, 37usize, 7usize);
14354 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14355 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14356 let t = QTensor::from_f32(w, rows, cols);
14357 let mut alone = vec![0f32; b * rows];
14358 t.matmat(&xs, b, &mut alone, None);
14359 let pool = Pool::new(3);
14360 let mut split = vec![0f32; b * rows];
14361 t.matmat(&xs, b, &mut split, Some(&pool));
14362 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14363 assert_eq!(bits(&alone), bits(&split));
14364 }
14365
14366 #[test]
14367 fn q2tp_i8_dot_matches_exact_on_grid() {
14368 let (rows, cols) = (5, 64);
14372 let gpr = cols / GROUP_SIZE;
14373 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
14377 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
14378 .collect();
14379 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
14380 let x: Vec<f32> = (0..cols)
14381 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
14382 .collect();
14383 let act = split_act(&x);
14384 assert!(
14385 act.outliers.is_empty(),
14386 "on-grid input must have no outliers"
14387 );
14388 let gsum = q1_group_sums(&act.xq, gpr);
14389 for r in 0..rows {
14390 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
14391 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
14392 assert!(
14393 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
14394 "row {r}: exact {exact} vs i8 {fast}"
14395 );
14396 }
14397 }
14398
14399 #[cfg(target_arch = "x86_64")]
14400 #[test]
14401 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
14402 if !std::arch::is_x86_feature_detected!("avx2") {
14407 return;
14408 }
14409 let mut seed = 0x9e3779b9u32;
14410 let mut next = || {
14411 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
14412 seed
14413 };
14414 for _ in 0..20_000 {
14415 let mut ch = [0u8; Q2TP_CHUNK];
14416 let mut x = [0i8; GROUP_SIZE];
14417 for b in &mut ch {
14418 *b = next() as u8;
14419 }
14420 for v in &mut x {
14421 *v = (next() >> 24) as i8;
14422 }
14423 let mut reference = 0i32;
14424 for (k, &b) in ch.iter().enumerate() {
14425 reference += (b & 3) as i32 * x[k * 4] as i32;
14426 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
14427 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
14428 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
14429 }
14430 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
14433 assert_eq!(got, reference, "packed q2 lane mismatch");
14434 }
14435 }
14436
14437 #[test]
14438 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
14439 let (rows, cols) = (1usize, GROUP_SIZE);
14440 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
14441 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
14445 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
14446 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
14448 let mut raw = vec![0.0f32; rows];
14449 let mut affine = vec![0.0f32; rows];
14450 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
14451 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
14452 assert_eq!(raw, vec![-24.0]);
14453 assert_eq!(affine, vec![-8.0]);
14454 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
14455 }
14456
14457 #[test]
14458 fn q8_row_dot_fast_matches_scalar() {
14459 let cols = 96;
14462 let row: Vec<u8> = (0..cols)
14463 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
14464 .collect();
14465 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
14466 let act = split_act(&x);
14467 let fast = q8_row_dot(&row, &act);
14468 let scalar = q8_row_dot_scalar(&row, &act);
14469 assert!(
14470 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
14471 "fast {fast} vs scalar {scalar}"
14472 );
14473 }
14474
14475 #[test]
14476 fn f32_matvec_matches_matvec_rows_bitexact() {
14477 let (rows, cols) = (300, 40);
14478 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
14479 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
14480 let qt = QTensor::from_f32(w.clone(), rows, cols);
14481
14482 let mut a = vec![0.0f32; rows];
14483 matvec_rows(None, &w, &x, &mut a);
14484 let mut b = vec![0.0f32; rows];
14485 qt.matvec(&x, &mut b, None);
14486 assert_eq!(a, b);
14487 }
14488
14489 #[test]
14490 fn sdot_kernel_exact_on_grid() {
14491 eprintln!("sdot_enabled = {}", sdot_enabled());
14496 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
14498 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14499 .collect();
14500 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
14501 let x: Vec<f32> = (0..cols)
14502 .map(|i| match i % 3 {
14503 0 => 1.0,
14504 1 => -1.0,
14505 _ => 0.0,
14506 })
14507 .collect();
14508 let mut a = vec![0.0f32; rows];
14509 qmatvec(
14510 &w,
14511 &[],
14512 &scales,
14513 &x,
14514 &[],
14515 TensorDtype::Q8Row,
14516 rows,
14517 cols,
14518 &mut a,
14519 None,
14520 );
14521 for o in 0..rows {
14522 let mut acc = 0.0f32;
14523 for j in 0..cols {
14524 acc += (w[o * cols + j] as i8) as f32 * x[j];
14525 }
14526 let expect = acc * scales[o];
14527 assert!(
14528 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14529 "row {o}: {} vs {expect}",
14530 a[o]
14531 );
14532 }
14533 }
14534
14535 #[test]
14536 fn q1_tbl_fast_path_matches_reference() {
14537 let (rows, cols) = (5, 256);
14542 let gpr = cols / GROUP_SIZE;
14543 let mut bytes = Vec::new();
14544 for t in 0..rows * gpr {
14545 let s = 0.007 + (t % 11) as f32 * 0.004;
14546 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14547 for j in 0..4 {
14548 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
14549 }
14550 }
14551 let x: Vec<f32> = (0..cols)
14552 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
14553 .collect();
14554 let mut w = vec![0.0f32; rows * cols];
14555 cortiq_core::quant::dequant_q1(&bytes, &mut w);
14556 let mut got = vec![0.0f32; rows];
14557 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14558 for o in 0..rows {
14559 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14560 assert!(
14561 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14562 "row {o}: {} vs {expect}",
14563 got[o]
14564 );
14565 }
14566 let b = 5usize;
14569 let mut xs_all = Vec::new();
14570 for bi in 0..b {
14571 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
14572 }
14573 let mut mm = vec![0.0f32; b * rows];
14574 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
14575 for bi in 0..b {
14576 let mut single = vec![0.0f32; rows];
14577 q1_matvec(
14578 &bytes,
14579 &xs_all[bi * cols..(bi + 1) * cols],
14580 rows,
14581 cols,
14582 &mut single,
14583 None,
14584 );
14585 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
14586 }
14587 }
14588
14589 #[test]
14590 fn q1_kernels_match_exact_reference() {
14591 let (rows, cols) = (7, 96);
14593 let gpr = cols / GROUP_SIZE;
14594 let mut bytes = Vec::new();
14595 for t in 0..rows * gpr {
14596 let s = 0.01 + (t % 13) as f32 * 0.003;
14597 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14598 for j in 0..4 {
14599 bytes.push(((t * 31 + j * 97) % 251) as u8);
14600 }
14601 }
14602 let x: Vec<f32> = (0..cols)
14604 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
14605 .collect();
14606 let mut w = vec![0.0f32; rows * cols];
14608 cortiq_core::quant::dequant_q1(&bytes, &mut w);
14609 let mut expect = vec![0.0f32; rows];
14610 for o in 0..rows {
14611 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14612 }
14613 let mut got = vec![0.0f32; rows];
14614 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14615 for o in 0..rows {
14616 assert!(
14617 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
14618 "row {o}: {} vs {}",
14619 got[o],
14620 expect[o]
14621 );
14622 }
14623 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
14625 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
14626 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14627 assert_eq!(a1, got);
14628 let mut xs = x.clone();
14629 xs.extend_from_slice(&x2);
14630 let mut mm = vec![0.0f32; 2 * rows];
14631 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
14632 assert_eq!(&mm[..rows], got.as_slice());
14633 assert_eq!(&mm[rows..], a2.as_slice());
14634 }
14635
14636 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
14639 let tiles = rows * cols / GROUP_SIZE;
14640 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
14641 for t in 0..tiles {
14642 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
14643 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14644 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
14645 q1.extend_from_slice(&bits.to_le_bytes());
14646 }
14647 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
14648 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
14649 (q1, rp.data)
14650 }
14651
14652 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
14655 let mut x: Vec<f32> = (0..cols)
14656 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
14657 .collect();
14658 x[3 % cols] = 41.0;
14659 x[(cols / 2 + 5) % cols] = -37.5;
14660 x
14661 }
14662
14663 #[test]
14670 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
14671 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
14672 let gpr = cols / GROUP_SIZE;
14673 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
14674 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14675 let x = q1g_test_x(cols, rows);
14676 let act = split_act(&x);
14677 assert!(!act.outliers.is_empty(), "test wants the outlier path");
14678 let gsum = q1_group_sums(&act.xq, gpr);
14679 for r in 0..rows {
14680 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
14681 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
14682 let (a, b) = (
14683 wa.dot_i8(r, gpr, &act.xq, &gsum),
14684 wb.dot_i8(r, gpr, &act.xq, &gsum),
14685 );
14686 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
14687 for j in [0, 31, 32, 127, cols - 1] {
14688 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
14689 assert_eq!(
14690 (a.0.to_bits(), a.1.to_bits()),
14691 (b.0.to_bits(), b.1.to_bits())
14692 );
14693 }
14694 #[cfg(target_arch = "x86_64")]
14695 if avx2_enabled() {
14696 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
14697 unsafe {
14699 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
14700 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
14701 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
14702 }
14703 if std::is_x86_feature_detected!("avx512vnni")
14704 && std::is_x86_feature_detected!("avx512bw")
14705 && std::is_x86_feature_detected!("avx512vl")
14706 {
14707 unsafe {
14709 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
14710 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
14711 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
14712 }
14713 }
14714 }
14715 }
14716 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
14718 bin_matvec(wa, &x, rows, cols, &mut ga, None);
14719 bin_matvec(wb, &x, rows, cols, &mut gb, None);
14720 assert_eq!(ga, gb, "{rows}x{cols} matvec");
14721 let x2 = q1g_test_x(cols, rows + 1);
14722 let (mut a1, mut a2, mut b1, mut b2) = (
14723 vec![0f32; rows],
14724 vec![0f32; rows],
14725 vec![0f32; rows],
14726 vec![0f32; rows],
14727 );
14728 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14729 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14730 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14731 let b = 5usize;
14732 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14733 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14734 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14735 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14736 assert_eq!(ma, mb, "{rows}x{cols} matmat");
14737 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14740 .map(|k| {
14741 let a = split_act(&xs[k * cols..(k + 1) * cols]);
14742 let g = q1_group_sums(&a.xq, gpr);
14743 (a, g)
14744 })
14745 .collect();
14746 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14747 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14748 for r in 0..rows {
14749 #[cfg(target_arch = "x86_64")]
14750 if avx2_enabled() {
14751 let (a, b) = unsafe {
14753 (
14754 wa.dot_1x4_x86(r, gpr, xq4, gs4),
14755 wb.dot_1x4_x86(r, gpr, xq4, gs4),
14756 )
14757 };
14758 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14759 }
14760 #[cfg(target_arch = "aarch64")]
14761 if sdot_enabled() {
14762 let (a, b) = unsafe {
14764 (
14765 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14766 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14767 )
14768 };
14769 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14770 }
14771 let _ = (r, &xq4, &gs4);
14772 }
14773 }
14774 }
14775
14776 #[test]
14782 #[ignore]
14783 fn q1g_cpu_matvec_speed() {
14784 use crate::pool::Pool;
14785 let threads = std::env::var("CMF_THREADS")
14786 .ok()
14787 .and_then(|v| v.parse().ok())
14788 .unwrap_or(8usize);
14789 let pool = Pool::new(threads);
14790 for &(name, rows, cols, reps) in &[
14791 ("gate/up_proj", 6144usize, 2048usize, 41usize),
14792 ("down_proj", 2048, 6144, 41),
14793 ("lm_head", 151_669, 2048, 9),
14794 ] {
14795 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14796 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14797 let x = q1g_test_x(cols, 1);
14798 let mut out = vec![0f32; rows];
14799 let mut time = |w: BinW| {
14800 let t0 = std::time::Instant::now();
14801 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14802 t0.elapsed().as_secs_f64()
14803 };
14804 let _ = (time(wa), time(wb));
14805 let (mut ta, mut tb) = (Vec::new(), Vec::new());
14806 for r in 0..reps {
14807 if r % 2 == 0 {
14808 ta.push(time(wa));
14809 tb.push(time(wb));
14810 } else {
14811 tb.push(time(wb));
14812 ta.push(time(wa));
14813 }
14814 }
14815 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14816 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14817 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14818 eprintln!(
14819 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14820 a8w8_enabled(),
14821 ma * 1e6,
14822 mb * 1e6,
14823 mb / ma
14824 );
14825 }
14826 }
14827
14828 #[test]
14833 fn q1g_mapped_tensor_paths_equal_q1() {
14834 use crate::pool::Pool;
14835 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14836 let (rows, cols) = (300usize, 256usize);
14837 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14838 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14839 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14840 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14841 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14842 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14843 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14844 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14845 "max_position_embeddings": 8},
14846 "quant_type": "VBIT"
14847 }))
14848 .unwrap();
14849 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14850 name: name.into(),
14851 dtype,
14852 shape: vec![rows, cols],
14853 data: data.clone(),
14854 };
14855 let specs = [
14856 spec("g1", TensorDtype::Q1, &gq1),
14857 spec("u1", TensorDtype::Q1, &uq1),
14858 spec("gg", TensorDtype::Q1G, &gq1g),
14859 spec("ug", TensorDtype::Q1G, &uq1g),
14860 ];
14861 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14862 std::fs::create_dir_all(&dir).unwrap();
14863 let path = dir.join("m.cmf");
14864 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14865 let model = Arc::new(CmfModel::open(&path).unwrap());
14866 assert!(model.verify().is_empty(), "{:?}", model.verify());
14867 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14868 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14869 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14870 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14871 let pool = Pool::new(3);
14872 let x = q1g_test_x(cols, 4);
14873 let run = |a: &QTensor| {
14874 let mut o = vec![0f32; rows];
14875 a.matvec(&x, &mut o, Some(&pool));
14876 o
14877 };
14878 assert_eq!(run(&g1), run(&gg), "matvec");
14879 let many = |a: &QTensor, b: &QTensor| {
14880 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14881 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14882 (oa, ob)
14883 };
14884 let want = many(&g1, &u1);
14885 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14886 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14887 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14888 let fused = |g: &QTensor, u: &QTensor| {
14889 let mut o = vec![0f32; rows];
14890 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14891 (ok, o)
14892 };
14893 let (ok_a, fa) = fused(&g1, &u1);
14894 let (ok_b, fb) = fused(&gg, &ug);
14895 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14896 assert_eq!(fa, fb, "fused gate/up");
14897 let b = 6usize;
14898 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14899 let mm = |a: &QTensor| {
14900 let mut o = vec![0f32; b * rows];
14901 a.matmat(&xs, b, &mut o, Some(&pool));
14902 o
14903 };
14904 assert_eq!(mm(&g1), mm(&gg), "matmat");
14905 let x2 = q1g_test_x(cols, 77);
14906 let pair = |a: &QTensor| {
14907 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14908 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14909 (o1, o2)
14910 };
14911 assert_eq!(pair(&g1), pair(&gg), "matvec2");
14912 for r in [0, 1, rows / 2, rows - 1] {
14913 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14914 g1.row_f32(r, &mut ra);
14915 gg.row_f32(r, &mut rb);
14916 assert_eq!(ra, rb, "embedding row {r}");
14917 }
14918 drop((g1, u1, gg, ug, model));
14919 let _ = std::fs::remove_dir_all(&dir);
14920 }
14921
14922 #[test]
14929 fn blr_mapped_tensor_paths_equal_dense() {
14930 use crate::pool::Pool;
14931 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14932 let (rows, cols, r) = (96usize, 128usize, 16usize);
14933 let mut rng = crate::blr::tests::Rng(41);
14934 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14935 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14936 let mut q8: Vec<u8> = (0..rows * cols)
14938 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14939 .collect();
14940 for o in 0..rows {
14941 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14942 }
14943 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14944 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14945 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14946 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14947 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14948 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14949 "max_position_embeddings": 8},
14950 "quant_type": "Q8_2F"
14951 }))
14952 .unwrap();
14953 let specs = [
14954 TensorSpec {
14955 name: "g".into(),
14956 dtype: TensorDtype::BinaryLowRank,
14957 shape: vec![rows, cols],
14958 data: gate.clone(),
14959 },
14960 TensorSpec {
14961 name: "u".into(),
14962 dtype: TensorDtype::BinaryLowRank,
14963 shape: vec![rows, cols],
14964 data: up.clone(),
14965 },
14966 TensorSpec {
14967 name: "q".into(),
14968 dtype: TensorDtype::Q8Row,
14969 shape: vec![rows, cols],
14970 data: q8,
14971 },
14972 ];
14973 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14974 std::fs::create_dir_all(&dir).unwrap();
14975 let path = dir.join("m.cmf");
14976 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14977 let model = Arc::new(CmfModel::open(&path).unwrap());
14978 assert!(model.verify().is_empty(), "{:?}", model.verify());
14979 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14980 let (g, u, q) = (t("g"), t("u"), t("q"));
14981 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14982 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14983 assert!(g.as_f32().is_none() && g.model_idx().is_none());
14984 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14985 assert!(!g.sparse_col_ok());
14986 assert_eq!((g.rows(), g.cols()), (rows, cols));
14987 let deq = |bytes: &[u8]| {
14988 let mut w = vec![0f32; rows * cols];
14989 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14990 w
14991 };
14992 let (wg, wu) = (deq(&gate), deq(&up));
14993 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14994 (0..rows)
14995 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14996 .collect()
14997 };
14998 let close = |a: &[f32], b: &[f32], what: &str| {
14999 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
15000 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
15001 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
15002 };
15003 let pool = Pool::new(3);
15004 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
15005 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
15006 let mut o = vec![0f32; rows];
15007 g.matvec(&x, &mut o, Some(&pool));
15008 close(&o, &mv(&wg, &x), "matvec");
15009 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15010 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
15011 close(&o1, &mv(&wg, &x), "matvec2 a");
15012 close(&o2, &mv(&wg, &x2), "matvec2 b");
15013 for b in [2usize, 5] {
15014 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
15015 let mut om = vec![0f32; b * rows];
15016 g.matmat(&xs, b, &mut om, Some(&pool));
15017 for t in 0..b {
15018 close(
15019 &om[t * rows..(t + 1) * rows],
15020 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
15021 "matmat",
15022 );
15023 }
15024 }
15025 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
15028 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
15029 close(&a, &mv(&wg, &x), "matvec_many blr gate");
15030 close(&c, &mv(&wu, &x), "matvec_many blr up");
15031 let mut qo = vec![0f32; rows];
15032 q.matvec(&x, &mut qo, Some(&pool));
15033 assert_eq!(bq, qo, "matvec_many q8_row sibling");
15034 let mut f = vec![0f32; rows];
15035 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
15036 let mut row = vec![0f32; cols];
15037 g.row_f32(rows / 3, &mut row);
15038 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
15039 let mut scratch = vec![0f32; cols];
15040 let d = g.row_dot(7, &x, &mut scratch);
15041 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
15042 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
15043 let mut col = vec![0f32; rows];
15044 g.add_col_scaled(cols / 2, 1.5, &mut col);
15045 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
15046 close(&col, &want, "add_col_scaled");
15047 drop((g, u, q, model));
15048 let _ = std::fs::remove_dir_all(&dir);
15049 }
15050
15051 #[test]
15052 fn repack_is_bit_identical() {
15053 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
15059 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
15060 .collect();
15061 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
15062 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
15063 let rep = q8_repack_layout(&w, rows, cols);
15064 for g in 0..rows / 4 {
15066 for c in 0..cols / 16 {
15067 for lane in 0..4 {
15068 assert_eq!(
15069 &rep[g * 4 * cols + c * 64 + lane * 16
15070 ..g * 4 * cols + c * 64 + lane * 16 + 16],
15071 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
15072 );
15073 }
15074 }
15075 }
15076 let mut a = vec![0.0f32; rows];
15077 qmatvec(
15078 &w,
15079 &[],
15080 &scales,
15081 &x,
15082 &[],
15083 TensorDtype::Q8Row,
15084 rows,
15085 cols,
15086 &mut a,
15087 None,
15088 );
15089 let mut b = vec![0.0f32; rows];
15090 qmatvec(
15091 &w,
15092 &rep,
15093 &scales,
15094 &x,
15095 &[],
15096 TensorDtype::Q8Row,
15097 rows,
15098 cols,
15099 &mut b,
15100 None,
15101 );
15102 assert_eq!(a, b, "full-range repack output diverged");
15103
15104 #[cfg(target_arch = "aarch64")]
15105 if sdot_enabled() {
15106 let act = split_act(&x);
15108 let mut c1 = vec![0.0f32; rows];
15109 let mut c2 = vec![0.0f32; rows];
15110 q8_range_sdot(
15111 &w,
15112 &[],
15113 &scales,
15114 &act,
15115 cols,
15116 SendMut(c1.as_mut_ptr()),
15117 3,
15118 rows - 2,
15119 );
15120 q8_range_sdot(
15121 &w,
15122 &rep,
15123 &scales,
15124 &act,
15125 cols,
15126 SendMut(c2.as_mut_ptr()),
15127 3,
15128 rows - 2,
15129 );
15130 assert_eq!(c1, c2, "unaligned-range repack output diverged");
15131 }
15132 }
15133
15134 #[test]
15135 fn sdot_a8w8_noise_is_bounded() {
15136 let (rows, cols) = (16, 512);
15140 let w: Vec<u8> = (0..rows * cols)
15141 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
15142 .collect();
15143 let scales = vec![0.01f32; rows];
15144 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
15145 let mut a = vec![0.0f32; rows];
15146 qmatvec(
15147 &w,
15148 &[],
15149 &scales,
15150 &x,
15151 &[],
15152 TensorDtype::Q8Row,
15153 rows,
15154 cols,
15155 &mut a,
15156 None,
15157 );
15158 let (mut num, mut den) = (0f64, 0f64);
15159 for o in 0..rows {
15160 let mut acc = 0.0f32;
15161 for j in 0..cols {
15162 acc += (w[o * cols + j] as i8) as f32 * x[j];
15163 }
15164 let expect = acc * scales[o];
15165 num += ((a[o] - expect) as f64).powi(2);
15166 den += (expect as f64).powi(2);
15167 }
15168 let rel = (num / den.max(1e-12)).sqrt();
15169 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
15170 }
15171
15172 #[test]
15173 fn i8_dot_neon_matches_scalar() {
15174 let n = 100;
15175 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
15176 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
15177 let mut scalar = 0.0f32;
15178 for j in 0..n {
15179 scalar += (w[j] as i8) as f32 * x[j];
15180 }
15181 let fast = dot_i8_f32(&w, &x);
15182 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
15183 }
15184
15185 #[test]
15187 fn vbitmatvec_matches_full_dequant() {
15188 let (rows, cols) = (6, 64);
15189 let ng = cols / GROUP_SIZE;
15190 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15192 let mut bytes = bits.clone();
15193 for g in 0..rows * ng {
15194 let s = 0.02 + 0.001 * g as f32;
15195 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15196 }
15197 for r in 0..rows {
15198 let b = bits[r] as usize;
15199 let (mut acc, mut nb) = (0u64, 0usize);
15200 let mut rowbytes = Vec::new();
15201 for i in 0..cols {
15202 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15203 acc = (acc << b) | v;
15204 nb += b;
15205 while nb >= 8 {
15206 nb -= 8;
15207 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15208 }
15209 }
15210 if nb > 0 {
15211 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15212 }
15213 bytes.extend_from_slice(&rowbytes);
15214 }
15215 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15216
15217 let mut reference = vec![0f32; rows * cols];
15218 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
15219 let mut expect = vec![0f32; rows];
15220 for r in 0..rows {
15221 expect[r] = reference[r * cols..(r + 1) * cols]
15222 .iter()
15223 .zip(&x)
15224 .map(|(w, xv)| w * xv)
15225 .sum();
15226 }
15227 let mut got = vec![0f32; rows];
15228 let offsets = vbit_row_offsets(&bytes, rows, cols);
15229 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
15230 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15234 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15235 for r in 0..rows {
15236 assert!(
15237 (got[r] - expect[r]).abs() < tol * scale,
15238 "row {r}: {} vs {}",
15239 got[r],
15240 expect[r]
15241 );
15242 }
15243 }
15244
15245 #[test]
15250 #[cfg(target_arch = "x86_64")]
15251 fn vbit_matmat_blocked_matches_per_row() {
15252 let (rows, cols, b) = (64usize, 128usize, 9usize);
15253 let ng = cols / GROUP_SIZE;
15254 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
15255 let mut bytes = bits.clone();
15256 for g in 0..rows * ng {
15257 let sc = 0.02 + 0.0005 * g as f32;
15258 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15259 }
15260 for r in 0..rows {
15261 let bw = bits[r] as usize;
15262 let (mut acc, mut nb) = (0u64, 0usize);
15263 let mut rowbytes = Vec::new();
15264 for i in 0..cols {
15265 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15266 acc = (acc << bw) | v;
15267 nb += bw;
15268 while nb >= 8 {
15269 nb -= 8;
15270 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15271 }
15272 }
15273 if nb > 0 {
15274 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15275 }
15276 bytes.extend_from_slice(&rowbytes);
15277 }
15278 let x: Vec<f32> = (0..b * cols)
15279 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15280 .collect();
15281 let offsets = vbit_row_offsets(&bytes, rows, cols);
15282 let mut y_a = vec![0f32; b * rows];
15283 let mut y_b = vec![0f32; b * rows];
15284 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15285 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
15286 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15287 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
15288 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15289 let max_d = y_a
15290 .iter()
15291 .zip(&y_b)
15292 .map(|(p, q)| (p - q).abs())
15293 .fold(0.0f32, f32::max);
15294 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
15295 }
15296
15297 #[test]
15305 fn q4t_matmat_blocked_matches_per_row() {
15306 let (rows, cols, b) = (16usize, 64usize, 9usize);
15307 let gpr = cols / GROUP_SIZE;
15308 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15309 for r in 0..rows {
15310 for g in 0..gpr {
15311 let t = (r * gpr + g) * Q4_TILE;
15312 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
15313 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15314 for k in 0..16 {
15315 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15316 }
15317 }
15318 }
15319 let x: Vec<f32> = (0..b * cols)
15320 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15321 .collect();
15322 let mut y_blk = vec![0f32; b * rows];
15323 let mut y_row = vec![0f32; b * rows];
15324 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15325 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
15326 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15327 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
15328 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15329 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
15330 }
15331
15332 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
15339 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
15340 let gpr = cols / GROUP_SIZE;
15341 let stride = q4tp_code_stride(gpr);
15342 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
15343 let mut b = vec![0u8; codes_off + rows * stride];
15344 for r in 0..rows {
15345 for g in 0..gpr {
15346 let t = (r * gpr + g) * Q4TP_NIB;
15347 for k in 0..16 {
15348 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15349 }
15350 }
15351 let lo = -6.0 - 0.03 * (r % 17) as f32;
15352 let step = 0.01 + 0.004 * (r % 11) as f32;
15353 let p = params_off + r * 4;
15354 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
15355 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
15356 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
15357 for g in 0..gpr {
15358 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
15359 }
15360 }
15361 b
15362 }
15363
15364 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
15368 let gpr = cols / GROUP_SIZE;
15369 let v = Q4tpView::new(bytes, rows, cols);
15370 let mut out = vec![0u8; rows * gpr * Q4_TILE];
15371 let mut sc = vec![0f32; gpr];
15372 for r in 0..rows {
15373 v.scales_into(r, gpr, &mut sc);
15374 for g in 0..gpr {
15375 let t = (r * gpr + g) * Q4_TILE;
15376 let s = sc[g];
15377 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15378 let src = (r * gpr + g) * Q4TP_NIB;
15379 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
15380 }
15381 }
15382 out
15383 }
15384
15385 #[test]
15391 fn q4tp_exact_path_matches_dequant_reference() {
15392 let (rows, cols) = (256usize, 512usize);
15393 let gpr = cols / GROUP_SIZE;
15394 let bytes = synth_q4tp(rows, cols);
15395 let mut w = vec![0f32; rows * cols];
15396 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15397
15398 let x: Vec<f32> = (0..cols)
15399 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15400 .collect();
15401 let v = Q4tpView::new(&bytes, rows, cols);
15402 let mut sc = vec![0f32; gpr];
15403 for r in 0..rows {
15404 v.scales_into(r, gpr, &mut sc);
15405 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
15406 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
15407 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15411 assert!(
15412 (got - want).abs() <= 1e-5 * mag,
15413 "row {r}: kernel {got} vs dequant {want}"
15414 );
15415 }
15416 }
15417
15418 #[test]
15424 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
15425 let (rows, cols) = (256usize, 512usize);
15426 let bytes = synth_q4tp(rows, cols);
15427 let twin = q4tp_as_q4t(&bytes, rows, cols);
15428 let x: Vec<f32> = (0..cols)
15429 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15430 .collect();
15431
15432 let mut got = vec![0f32; rows];
15433 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
15434 let mut want = vec![0f32; rows];
15435 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
15436
15437 let mut w = vec![0f32; rows * cols];
15440 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15441 for r in 0..rows {
15442 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15443 assert!(
15444 (got[r] - want[r]).abs() <= 1e-3 * mag,
15445 "row {r}: q4tp {} vs q4t {}",
15446 got[r],
15447 want[r]
15448 );
15449 }
15450 }
15451
15452 #[test]
15457 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
15458 let (rows, cols, b) = (256usize, 512usize, 5usize);
15459 let bytes = synth_q4tp(rows, cols);
15460 let twin = q4tp_as_q4t(&bytes, rows, cols);
15461 let xs: Vec<f32> = (0..b * cols)
15462 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15463 .collect();
15464
15465 let mut got = vec![0f32; b * rows];
15466 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
15467 let mut want = vec![0f32; b * rows];
15468 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
15469
15470 let mut w = vec![0f32; rows * cols];
15471 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15472 for t in 0..b {
15473 for r in 0..rows {
15474 let mag: f32 = (0..cols)
15475 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
15476 .sum();
15477 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
15478 assert!(
15479 (g - wa).abs() <= 1e-3 * mag,
15480 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
15481 );
15482 }
15483 }
15484 }
15485
15486 #[test]
15487 fn q4tp_matvec2_matches_the_single_stream_kernel() {
15488 let (rows, cols) = (128usize, 256usize);
15489 let gpr = cols / GROUP_SIZE;
15490 let bytes = synth_q4tp(rows, cols);
15491 let xs: Vec<f32> = (0..2 * cols)
15492 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15493 .collect();
15494
15495 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15496 q4tp_matvec2(
15497 &bytes,
15498 &xs[..cols],
15499 &xs[cols..],
15500 rows,
15501 cols,
15502 &mut o1,
15503 &mut o2,
15504 None,
15505 );
15506
15507 let v = Q4tpView::new(&bytes, rows, cols);
15510 let mut sc = vec![0f32; gpr];
15511 for r in 0..rows {
15512 v.scales_into(r, gpr, &mut sc);
15513 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
15514 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
15515 }
15516 }
15517
15518 #[test]
15525 fn q4tp_matvec_keeps_pace_with_q4t() {
15526 let (rows, cols) = (4096usize, 3072usize);
15527 let bytes = synth_q4tp(rows, cols);
15528 let twin = q4tp_as_q4t(&bytes, rows, cols);
15529 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
15530 let mut o = vec![0f32; rows];
15531 let n = 12;
15532 let mut best = (f64::MAX, f64::MAX);
15533 for _ in 0..3 {
15536 let t0 = std::time::Instant::now();
15537 for _ in 0..n {
15538 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
15539 }
15540 best.0 = best.0.min(t0.elapsed().as_secs_f64());
15541 let t0 = std::time::Instant::now();
15542 for _ in 0..n {
15543 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
15544 }
15545 best.1 = best.1.min(t0.elapsed().as_secs_f64());
15546 }
15547 let ratio = best.1 / best.0;
15548 println!(
15549 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
15550 best.0 * 1e3 / n as f64,
15551 best.1 * 1e3 / n as f64
15552 );
15553 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
15554 }
15555
15556 #[cfg(target_os = "macos")]
15557 #[test]
15558 fn q4t_matmat_accel_matches_dequant_reference() {
15559 if !accel_gemm_enabled() {
15560 return; }
15562 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
15564 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15565 for r in 0..rows {
15566 for g in 0..gpr {
15567 let t = (r * gpr + g) * Q4_TILE;
15568 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
15569 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15570 for k in 0..16 {
15571 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15572 }
15573 }
15574 }
15575 let x: Vec<f32> = (0..b * cols)
15576 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15577 .collect();
15578 let mut got = vec![0f32; b * rows];
15579 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
15580 let mut w = vec![0f32; rows * cols];
15582 for r in 0..rows {
15583 for g in 0..gpr {
15584 let t = (r * gpr + g) * Q4_TILE;
15585 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
15586 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
15587 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
15588 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
15589 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
15590 }
15591 }
15592 }
15593 for bi in 0..b {
15594 for r in 0..rows {
15595 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
15596 let d = (got[bi * rows + r] - want).abs();
15597 assert!(
15598 d <= want.abs().max(1.0) * 1e-4,
15599 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
15600 got[bi * rows + r]
15601 );
15602 }
15603 }
15604 }
15605
15606 #[test]
15607 fn q4matvec_matches_full_dequant() {
15608 let (rows, cols) = (8, 64);
15609 let groups = rows * cols / GROUP_SIZE;
15610 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15612 for i in 0..groups * 16 {
15613 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15614 }
15615 for g in 0..groups {
15616 let s = 0.01 + 0.003 * g as f32;
15617 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15618 }
15619 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15620
15621 let mut reference = vec![0.0f32; rows * cols];
15622 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15623 let mut expect = vec![0.0f32; rows];
15624 for r in 0..rows {
15625 expect[r] = reference[r * cols..(r + 1) * cols]
15626 .iter()
15627 .zip(&x)
15628 .map(|(w, xv)| w * xv)
15629 .sum();
15630 }
15631
15632 let mut got = vec![0.0f32; rows];
15633 q4matvec(&bytes, &x, rows, cols, &mut got, None);
15634 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15638 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15639 for r in 0..rows {
15640 assert!(
15641 (got[r] - expect[r]).abs() < tol * scale,
15642 "row {r}: {} vs {}",
15643 got[r],
15644 expect[r]
15645 );
15646 }
15647 }
15648
15649 #[test]
15652 fn vbitmatvec2_equals_two_singles() {
15653 let (rows, cols) = (6, 64);
15654 let ng = cols / GROUP_SIZE;
15655 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15656 let mut bytes = bits.clone();
15657 for g in 0..rows * ng {
15658 let s = 0.02 + 0.001 * g as f32;
15659 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15660 }
15661 for r in 0..rows {
15662 let b = bits[r] as usize;
15663 let (mut acc, mut nb) = (0u64, 0usize);
15664 let mut rowbytes = Vec::new();
15665 for i in 0..cols {
15666 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15667 acc = (acc << b) | v;
15668 nb += b;
15669 while nb >= 8 {
15670 nb -= 8;
15671 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15672 }
15673 }
15674 if nb > 0 {
15675 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15676 }
15677 bytes.extend_from_slice(&rowbytes);
15678 }
15679 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15680 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
15681 let offsets = vbit_row_offsets(&bytes, rows, cols);
15682
15683 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15684 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
15685 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
15686 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15687 vbitmatvec2(
15688 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
15689 );
15690 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
15691 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
15692 }
15693
15694 #[test]
15696 fn q4matvec2_equals_two_singles() {
15697 let (rows, cols) = (8, 128);
15698 let groups = rows * cols / GROUP_SIZE;
15699 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15700 for i in 0..groups * 16 {
15701 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15702 }
15703 for g in 0..groups {
15704 let s = 0.01 + 0.003 * g as f32;
15705 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15706 }
15707 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15710 x1[9] = 250.0;
15711 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15712
15713 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15714 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
15715 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
15716 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15717 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
15718 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
15719 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
15720 }
15721
15722 #[test]
15725 fn matvec_many_equals_separate_matvecs() {
15726 use crate::pool::Pool;
15727 let (r1, r2, cols) = (300, 200, 64);
15728 let mk = |salt: usize, rows: usize| {
15729 QTensor::from_f32(
15730 (0..rows * cols)
15731 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15732 .collect(),
15733 rows,
15734 cols,
15735 )
15736 };
15737 let (a, b) = (mk(1, r1), mk(5, r2));
15738 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15739 let pool = Pool::new(3);
15740
15741 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15742 a.matvec(&x, &mut ea, Some(&pool));
15743 b.matvec(&x, &mut eb, Some(&pool));
15744 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15745 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15746 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15747 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15748 }
15749
15750 #[test]
15755 fn q4tp_matvec_many_equals_separate_matvecs() {
15756 use crate::pool::Pool;
15757 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15758
15759 let (r1, r2, cols) = (300usize, 200usize, 64usize);
15760 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15761 "arch_name": "tiny-q4tp",
15762 "hidden_size": cols,
15763 "intermediate_size": cols * 2,
15764 "num_layers": 1,
15765 "num_attention_heads": 2,
15766 "num_kv_heads": 1,
15767 "head_dim": 32,
15768 "vocab_size": r1,
15769 "layer_types": ["FullAttention"],
15770 "rms_norm_eps": 1e-6,
15771 "max_position_embeddings": 8,
15772 "linear_conv_kernel_dim": 0,
15773 "linear_num_key_heads": 0,
15774 "linear_num_value_heads": 0
15775 }))
15776 .unwrap();
15777 let header = CmfHeader {
15778 format: "cmf".into(),
15779 version: CMF_VERSION,
15780 arch,
15781 quant_type: QuantType::Q4Block,
15782 provenance: None,
15783 tokenizer_config: None,
15784 section_hashes: None,
15785 skills: Vec::new(),
15786 shard: None,
15787 calibration: None,
15788 routing: None,
15789 genome: None,
15790 lineage: Vec::new(),
15791 router: None,
15792 segments: Vec::new(),
15793 };
15794 let specs = [
15795 TensorSpec {
15796 name: "q".into(),
15797 dtype: TensorDtype::Q4TiledP,
15798 shape: vec![r1, cols],
15799 data: synth_q4tp(r1, cols),
15800 },
15801 TensorSpec {
15802 name: "kv".into(),
15803 dtype: TensorDtype::Q4TiledP,
15804 shape: vec![r2, cols],
15805 data: synth_q4tp(r2, cols),
15806 },
15807 ];
15808 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15809 std::fs::create_dir_all(&dir).unwrap();
15810 let path = dir.join("m.cmf");
15811 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15812 let model = Arc::new(CmfModel::open(&path).unwrap());
15813 let (a, b) = (
15814 QTensor::from_model(&model, "q").unwrap(),
15815 QTensor::from_model(&model, "kv").unwrap(),
15816 );
15817 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15818 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15819 let x: Vec<f32> = (0..cols)
15820 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15821 .collect();
15822 let pool = Pool::new(3);
15823 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15824 a.matvec(&x, &mut ea, Some(&pool));
15825 b.matvec(&x, &mut eb, Some(&pool));
15826 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15827 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15828 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15829 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15830 let _ = std::fs::remove_dir_all(&dir);
15831 }
15832
15833 #[cfg(target_arch = "x86_64")]
15839 #[test]
15840 fn fast_expert_rows_match_the_row_kernels() {
15841 use crate::pool::Pool;
15842 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15843
15844 for (rows, cols) in [(32usize, 2560usize), (16, 640), (48, 192), (16, 64)] {
15846 let bytes = synth_q4tp(rows, cols);
15847 let v = Q4tpView::new(&bytes, rows, cols);
15848 let gpr = cols / GROUP_SIZE;
15849 for r0 in (0..rows).step_by(16) {
15850 let Some(got) = moe_fast::block_scales(&v, r0, gpr) else {
15851 eprintln!("fast_expert_rows: no AVX-512 VNNI here, skipped");
15852 return;
15853 };
15854 for i in 0..16 {
15855 let mut want = vec![0f32; gpr];
15856 v.scales_into(r0 + i, gpr, &mut want);
15857 assert_eq!(
15858 want.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15859 got[i * gpr..(i + 1) * gpr].iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15860 "scales of row {} ({rows}x{cols})",
15861 r0 + i
15862 );
15863 }
15864 }
15865 }
15866
15867 let (h, inter, ne) = (256usize, 640usize, 5usize);
15868 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15869 "arch_name": "tiny-q4tp-moe",
15870 "hidden_size": h,
15871 "intermediate_size": inter,
15872 "num_layers": 1,
15873 "num_attention_heads": h / 32,
15874 "num_kv_heads": 1,
15875 "head_dim": 32,
15876 "vocab_size": 8,
15877 "layer_types": ["FullAttention"],
15878 "rms_norm_eps": 1e-6,
15879 "max_position_embeddings": 8,
15880 "linear_conv_kernel_dim": 0,
15881 "linear_num_key_heads": 0,
15882 "linear_num_value_heads": 0
15883 }))
15884 .unwrap();
15885 let header = CmfHeader {
15886 format: "cmf".into(),
15887 version: CMF_VERSION,
15888 arch,
15889 quant_type: QuantType::Q4Block,
15890 provenance: None,
15891 tokenizer_config: None,
15892 section_hashes: None,
15893 skills: Vec::new(),
15894 shard: None,
15895 calibration: None,
15896 routing: None,
15897 genome: None,
15898 lineage: Vec::new(),
15899 router: None,
15900 segments: Vec::new(),
15901 };
15902 let mut specs = Vec::new();
15903 for e in 0..ne {
15904 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15905 .into_iter()
15906 .enumerate()
15907 {
15908 let mut data = synth_q4tp(r, c);
15909 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15910 .iter_mut()
15911 .enumerate()
15912 {
15913 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15914 }
15915 specs.push(TensorSpec {
15916 name: format!("{n}{e}"),
15917 dtype: TensorDtype::Q4TiledP,
15918 shape: vec![r, c],
15919 data,
15920 });
15921 }
15922 }
15923 let dir = std::env::temp_dir().join(format!("cmf-moe-fast-{}", std::process::id()));
15924 std::fs::create_dir_all(&dir).unwrap();
15925 let path = dir.join("m.cmf");
15926 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15927 let model = Arc::new(CmfModel::open(&path).unwrap());
15928 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15929 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15930 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15931 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15932 let b = 6usize;
15933 let mut xs: Vec<f32> = (0..b * h)
15934 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15935 .collect();
15936 xs[5] = 9.0; let groups: Vec<Vec<usize>> = vec![vec![0, 1, 2, 3, 4], vec![5], vec![], vec![2, 0, 5], vec![1, 3]];
15939 let lens: Vec<usize> = groups.iter().map(|v| v.len()).collect();
15940 let n_pairs: usize = lens.iter().sum();
15941 let pool = Pool::new(3);
15942 let close = |want: &[f32], got: &[f32], rel: f32, what: &str| {
15943 assert_eq!(want.len(), got.len(), "{what}: length");
15944 let top = want.iter().fold(0f32, |m, v| m.max(v.abs()));
15945 assert!(top > 0.0, "{what}: an all-zero reference proves nothing");
15946 for (i, (w, v)) in want.iter().zip(got).enumerate() {
15947 assert!((w - v).abs() <= rel * top, "{what}: [{i}] {w} vs {v} (scale {top})");
15948 }
15949 };
15950
15951 let pairs: Vec<(&QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e])).collect();
15953 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15954 assert!(QTensor::moe_gate_up_rows(&pairs, &groups, &xs, &mut gs, Some(&pool)));
15955 let downs: Vec<&QTensor> = d.iter().collect();
15956 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15957 assert!(QTensor::moe_down_rows(&downs, &lens, &gs, &mut ds, Some(&pool)));
15958
15959 let triples: Vec<(&QTensor, &QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e], &d[e])).collect();
15960 let mut out = vec![0f32; n_pairs * h];
15961 let mut gs_fast = Vec::new();
15962 if !moe_fast::run_stages(&triples, &groups, &xs, &mut out, Some(&pool), None, Some(&mut gs_fast)) {
15963 eprintln!("fast_expert_rows: the int8 activation kernels are off here, skipped");
15964 let _ = std::fs::remove_dir_all(&dir);
15965 return;
15966 }
15967 close(&gs.concat(), &gs_fast, 1e-4, "gate and up");
15968 close(&ds.concat(), &out, 1e-2, "whole experts");
15969
15970 let mut gs_out = gs.clone();
15973 let top = gs.iter().flatten().fold(0f32, |m, v| m.max(v.abs()));
15974 gs_out[0][3] = 40.0 * top;
15975 let mut ds_out: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15976 assert!(QTensor::moe_down_rows(&downs, &lens, &gs_out, &mut ds_out, Some(&pool)));
15977 let mut out2 = vec![0f32; n_pairs * h];
15978 assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out2, Some(&pool), Some(&gs_out.concat()), None));
15979 close(&ds_out.concat(), &out2, 1e-4, "down");
15980
15981 let mut out3 = vec![0f32; n_pairs * h];
15983 assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out3, None));
15984 assert_eq!(
15985 out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15986 out3.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15987 "pooled and serial runs must agree bit for bit"
15988 );
15989 pool.mark_two_nodes(true);
15991 let mut out4 = vec![0f32; n_pairs * h];
15992 assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out4, Some(&pool)));
15993 let mut out5 = vec![0f32; n_pairs * h];
15994 assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out5, Some(&pool), Some(&gs_out.concat()), None));
15995 pool.mark_two_nodes(false);
15996 assert_eq!(
15997 out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15998 out4.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15999 "the rows split over two sides must agree bit for bit"
16000 );
16001 assert_eq!(
16002 out2.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
16003 out5.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
16004 "the down rows split over two sides must agree bit for bit"
16005 );
16006 let mut short = vec![0f32; n_pairs * h - 1];
16007 assert!(!QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut short, Some(&pool)));
16008 let _ = std::fs::remove_dir_all(&dir);
16009 }
16010
16011 #[test]
16017 fn multi_token_moe_rows_equal_single_token_decode() {
16018 use crate::pool::Pool;
16019 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
16020
16021 let (h, inter, ne) = (64usize, 128usize, 3usize);
16022 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
16023 "arch_name": "tiny-q4tp-moe",
16024 "hidden_size": h,
16025 "intermediate_size": inter,
16026 "num_layers": 1,
16027 "num_attention_heads": 2,
16028 "num_kv_heads": 1,
16029 "head_dim": 32,
16030 "vocab_size": 8,
16031 "layer_types": ["FullAttention"],
16032 "rms_norm_eps": 1e-6,
16033 "max_position_embeddings": 8,
16034 "linear_conv_kernel_dim": 0,
16035 "linear_num_key_heads": 0,
16036 "linear_num_value_heads": 0
16037 }))
16038 .unwrap();
16039 let header = CmfHeader {
16040 format: "cmf".into(),
16041 version: CMF_VERSION,
16042 arch,
16043 quant_type: QuantType::Q4Block,
16044 provenance: None,
16045 tokenizer_config: None,
16046 section_hashes: None,
16047 skills: Vec::new(),
16048 shard: None,
16049 calibration: None,
16050 routing: None,
16051 genome: None,
16052 lineage: Vec::new(),
16053 router: None,
16054 segments: Vec::new(),
16055 };
16056 let mut specs = Vec::new();
16057 for e in 0..ne {
16058 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
16059 .into_iter()
16060 .enumerate()
16061 {
16062 let mut data = synth_q4tp(r, c);
16065 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
16066 .iter_mut()
16067 .enumerate()
16068 {
16069 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
16070 }
16071 specs.push(TensorSpec {
16072 name: format!("{n}{e}"),
16073 dtype: TensorDtype::Q4TiledP,
16074 shape: vec![r, c],
16075 data,
16076 });
16077 }
16078 }
16079 let dir = std::env::temp_dir().join(format!(
16080 "cmf-moe-rows-{}-{}",
16081 std::process::id(),
16082 FLOAT_ACTIVATIONS.get()
16083 ));
16084 std::fs::create_dir_all(&dir).unwrap();
16085 let path = dir.join("m.cmf");
16086 CmfModel::write(&path, &header, &specs, None, None).unwrap();
16087 let model = Arc::new(CmfModel::open(&path).unwrap());
16088 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
16089 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
16090 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
16091 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
16092 let b = 4usize;
16093 let mut xs: Vec<f32> = (0..b * h)
16094 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
16095 .collect();
16096 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
16099 (vec![2, 0], vec![0.6, 0.4]),
16100 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
16101 (vec![1], vec![1.0]),
16102 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
16103 ];
16104 let pool = Pool::new(3);
16105 let mut want = vec![0f32; b * h];
16107 for (tk, (idx, w)) in routes.iter().enumerate() {
16108 let x = &xs[tk * h..(tk + 1) * h];
16109 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
16110 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
16111 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
16112 if FLOAT_ACTIVATIONS.get() {
16113 for (slot, &e) in idx.iter().enumerate() {
16114 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
16115 g[e].matvec(x, &mut gate, Some(&pool));
16116 u[e].matvec(x, &mut up, Some(&pool));
16117 for (v, u) in gate.iter_mut().zip(up) {
16118 *v = (*v / (1.0 + (-*v).exp())) * u;
16119 }
16120 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
16121 }
16122 }
16123 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
16124 assert!(QTensor::moe_down_many(
16125 &downs,
16126 &gs,
16127 w,
16128 &mut want[tk * h..(tk + 1) * h],
16129 Some(&pool)
16130 ));
16131 }
16132 if FLOAT_ACTIVATIONS.get() {
16133 for (tk, (idx, w)) in routes.iter().enumerate() {
16134 let mut scalar = vec![0.0; h];
16135 for (&e, &weight) in idx.iter().zip(w) {
16136 let (mut gate, mut up, mut down) =
16137 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
16138 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
16139 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
16140 for (v, u) in gate.iter_mut().zip(up) {
16141 *v = (*v / (1.0 + (-*v).exp())) * u;
16142 }
16143 d[e].matvec(&gate, &mut down, Some(&pool));
16144 for (v, d) in scalar.iter_mut().zip(down) {
16145 *v += weight * d;
16146 }
16147 }
16148 assert_eq!(
16149 &want[tk * h..(tk + 1) * h],
16150 scalar,
16151 "float many equals scalar experts"
16152 );
16153 }
16154 }
16155 let mut experts: Vec<usize> = Vec::new();
16157 let mut groups: Vec<Vec<usize>> = Vec::new();
16158 for (tk, (idx, _)) in routes.iter().enumerate() {
16159 for &e in idx {
16160 match experts.iter().position(|&x| x == e) {
16161 Some(k) => groups[k].push(tk),
16162 None => {
16163 experts.push(e);
16164 groups.push(vec![tk]);
16165 }
16166 }
16167 }
16168 }
16169 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
16170 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
16171 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
16172 assert!(QTensor::moe_gate_up_rows(
16173 &pairs,
16174 &groups,
16175 &xs,
16176 &mut gs,
16177 Some(&pool)
16178 ));
16179 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
16180 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
16181 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
16182 assert!(QTensor::moe_down_rows(
16183 &downs,
16184 &lens,
16185 &gs,
16186 &mut ds,
16187 Some(&pool)
16188 ));
16189 let slot = |tk: usize, e: usize| {
16190 let k = experts.iter().position(|&x| x == e).unwrap();
16191 groups[..k].iter().map(|g| g.len()).sum::<usize>()
16192 + groups[k].iter().position(|&x| x == tk).unwrap()
16193 };
16194 let mut got = vec![0f32; b * h];
16195 for (tk, (idx, w)) in routes.iter().enumerate() {
16196 for i in 0..h {
16197 let mut acc = 0f32;
16198 for (&e, &we) in idx.iter().zip(w) {
16199 acc += we * ds[slot(tk, e)][i];
16200 }
16201 got[tk * h + i] = acc;
16202 }
16203 }
16204 assert!(want.iter().any(|v| *v != 0.0));
16205 assert_eq!(
16206 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16207 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16208 "multi-token MoE must equal decode bit for bit"
16209 );
16210
16211 let b5 = 5usize;
16214 let x5: Vec<f32> = (0..b5 * h)
16215 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
16216 .collect();
16217 let mut mm = vec![0f32; b5 * inter];
16218 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
16219 for tk in 0..b5 {
16220 let mut mv = vec![0f32; inter];
16221 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
16222 assert_eq!(
16223 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16224 mm[tk * inter..(tk + 1) * inter]
16225 .iter()
16226 .map(|v| v.to_bits())
16227 .collect::<Vec<_>>(),
16228 "row-exact matmat token {tk}"
16229 );
16230 }
16231 let _ = std::fs::remove_dir_all(&dir);
16234 }
16235
16236 #[test]
16244 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
16245 use crate::pool::Pool;
16246 use std::sync::atomic::Ordering::Relaxed;
16247 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
16248 Q4TP_ALT.store(2, Relaxed);
16251 let pool = Pool::new(3);
16252 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
16256 let bytes = synth_q4tp(rows, cols);
16257 let mut xs: Vec<f32> = (0..b * cols)
16258 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
16259 .collect();
16260 xs[3] = 7.5; let run = |exact: bool| {
16262 let mut out = vec![0f32; b * rows];
16263 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
16264 out
16265 };
16266 let (fast, exact) = (run(false), run(true));
16267 #[cfg(not(target_arch = "aarch64"))]
16268 let _ = fast;
16269 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
16270 let mut matvecs = vec![0f32; b * rows];
16271 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
16272 q4tp_matvec(
16273 &bytes,
16274 &xs[bi * cols..(bi + 1) * cols],
16275 rows,
16276 cols,
16277 o,
16278 Some(&pool),
16279 );
16280 }
16281 assert!(matvecs.iter().any(|v| *v != 0.0));
16282 assert_eq!(
16283 bits(&exact),
16284 bits(&matvecs),
16285 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
16286 );
16287 #[cfg(target_arch = "aarch64")]
16288 {
16289 let gpr = cols / GROUP_SIZE;
16291 let v = Q4tpView::new(&bytes, rows, cols);
16292 let mut old = vec![0f32; b * rows];
16293 let mut sc = vec![0f32; gpr];
16294 let a8w8 = a8w8_enabled();
16295 let blocked = sdot_enabled() && blocked_enabled();
16296 let acts: Vec<SplitAct> = (0..b)
16297 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
16298 .collect();
16299 for r in 0..rows {
16300 v.scales_into(r, gpr, &mut sc);
16301 if !a8w8 {
16302 for bi in 0..b {
16303 let x = &xs[bi * cols..(bi + 1) * cols];
16304 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
16305 }
16306 continue;
16307 }
16308 let finish = |d: f32, act: &SplitAct| {
16309 let mut acc = d * act.sx;
16310 for &(j, xv) in &act.outliers {
16311 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
16312 acc += w * s * xv;
16313 }
16314 acc
16315 };
16316 let mut bi = 0usize;
16317 while blocked && bi + 4 <= b {
16318 let xs4 = [
16319 acts[bi].xq.as_slice(),
16320 acts[bi + 1].xq.as_slice(),
16321 acts[bi + 2].xq.as_slice(),
16322 acts[bi + 3].xq.as_slice(),
16323 ];
16324 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
16325 for k in 0..4 {
16326 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
16327 }
16328 bi += 4;
16329 }
16330 for (bi, act) in acts.iter().enumerate().skip(bi) {
16331 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
16332 old[bi * rows + r] = finish(d, act);
16333 }
16334 }
16335 assert_eq!(
16336 bits(&fast),
16337 bits(&old),
16338 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
16339 );
16340 if blocked {
16343 assert_ne!(
16344 bits(&fast),
16345 bits(&matvecs),
16346 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
16347 );
16348 }
16349 }
16350 }
16351 Q4TP_ALT.store(0, Relaxed);
16352 }
16353
16354 #[test]
16355 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
16356 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
16357 let active = AtomicUsize::new(0);
16360 counted_row_exact_scope(&active, || {
16361 assert_eq!(active.load(Ordering::Acquire), 1);
16362 counted_row_exact_scope(&active, || {
16363 assert_eq!(active.load(Ordering::Acquire), 2);
16364 });
16365 assert_eq!(active.load(Ordering::Acquire), 1);
16366 });
16367 assert_eq!(active.load(Ordering::Acquire), 0);
16368
16369 let both_entered = Barrier::new(2);
16370 let release_last = Barrier::new(2);
16371 std::thread::scope(|s| {
16372 let first = s.spawn(|| counted_row_exact_scope(&active, || {
16373 both_entered.wait();
16374 }));
16375 let last = s.spawn(|| counted_row_exact_scope(&active, || {
16376 both_entered.wait();
16377 release_last.wait();
16378 }));
16379 first.join().unwrap();
16380 let after_first = active.load(Ordering::Acquire);
16381 release_last.wait();
16382 last.join().unwrap();
16383 assert_eq!(after_first, 1, "second request must remain exact");
16384 });
16385 assert_eq!(active.load(Ordering::Acquire), 0);
16386 let panic = std::panic::catch_unwind(|| {
16387 counted_row_exact_scope(&active, || panic!("scope unwind"));
16388 });
16389 assert!(panic.is_err());
16390 assert_eq!(active.load(Ordering::Acquire), 0);
16391 }
16392
16393 #[test]
16394 #[cfg(target_arch = "x86_64")]
16395 fn q4tp_float_avx2_is_bitwise_scalar() {
16396 if !avx2_enabled() {
16397 return;
16398 }
16399 for cols in [32, 64, 96, 2048, 4096] {
16400 let rows = 9;
16401 let bytes = synth_q4tp(rows, cols);
16402 let v = Q4tpView::new(&bytes, rows, cols);
16403 let gpr = cols / GROUP_SIZE;
16404 let mut sc = vec![0.0; gpr];
16405 for seed in 1..=5 {
16406 let xs: Vec<f32> = (0..cols)
16407 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
16408 .collect();
16409 for r in 0..rows {
16410 v.scales_into(r, gpr, &mut sc);
16411 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
16412 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
16413 assert_eq!(
16414 scalar.to_bits(),
16415 vector.to_bits(),
16416 "cols={cols} row={r} seed={seed}"
16417 );
16418 }
16419 }
16420 }
16421 }
16422
16423 #[test]
16424 fn multi_token_moe_rows_float_equal_single_token_decode() {
16425 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
16426 }
16427
16428 #[test]
16429 fn full_gpu_q8_scope_is_nested_and_thread_local() {
16430 assert!(!FULL_GPU_Q8.get());
16431 let before = gpu_split_frac();
16432 {
16433 let _guard = enter_full_gpu_q8_scope();
16434 assert_eq!(gpu_split_frac(), 1.0);
16435 {
16436 let _nested = enter_full_gpu_q8_scope();
16437 }
16438 assert_eq!(gpu_split_frac(), 1.0);
16439 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
16440 }
16441 assert!(!FULL_GPU_Q8.get());
16442 assert_eq!(gpu_split_frac(), before);
16443 }
16444
16445 #[test]
16446 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
16447 assert!(!FLOAT_ACTIVATIONS.get());
16448 let before = a8w8_enabled();
16449 float_activations_scope(|| {
16450 assert!(!a8w8_enabled());
16451 float_activations_scope(|| assert!(!a8w8_enabled()));
16452 assert!(FLOAT_ACTIVATIONS.get());
16453 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
16454 .join()
16455 .unwrap();
16456 });
16457 assert!(!FLOAT_ACTIVATIONS.get());
16458 assert_eq!(a8w8_enabled(), before);
16459 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
16460 assert!(!FLOAT_ACTIVATIONS.get());
16461 }
16462
16463 #[test]
16466 fn batched_matmat_equals_per_position_matvec() {
16467 let (rows, cols, b) = (8, 64, 5);
16468 let groups = rows * cols / GROUP_SIZE;
16470 let mut q4 = Vec::new();
16471 for i in 0..groups * 16 {
16472 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16473 }
16474 for g in 0..groups {
16475 q4.extend_from_slice(
16476 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16477 );
16478 }
16479 let ng = cols / GROUP_SIZE;
16481 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
16482 let mut vb = bits.clone();
16483 for g in 0..rows * ng {
16484 vb.extend_from_slice(
16485 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
16486 );
16487 }
16488 for r in 0..rows {
16489 let bw = bits[r] as usize;
16490 let (mut acc, mut nb) = (0u64, 0usize);
16491 let mut rowbytes = Vec::new();
16492 for i in 0..cols {
16493 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
16494 acc = (acc << bw) | v;
16495 nb += bw;
16496 while nb >= 8 {
16497 nb -= 8;
16498 rowbytes.push(((acc >> nb) & 0xFF) as u8);
16499 }
16500 }
16501 if nb > 0 {
16502 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
16503 }
16504 vb.extend_from_slice(&rowbytes);
16505 }
16506 let offsets = vbit_row_offsets(&vb, rows, cols);
16507
16508 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16509
16510 let mut got = vec![0f32; b * rows];
16512 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
16513 for bi in 0..b {
16514 let mut expect = vec![0f32; rows];
16515 q4matvec(
16516 &q4,
16517 &xs[bi * cols..(bi + 1) * cols],
16518 rows,
16519 cols,
16520 &mut expect,
16521 None,
16522 );
16523 assert_eq!(
16524 &got[bi * rows..(bi + 1) * rows],
16525 &expect[..],
16526 "q4 batch pos {bi}"
16527 );
16528 }
16529
16530 let mut got = vec![0f32; b * rows];
16532 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
16533 for bi in 0..b {
16534 let mut expect = vec![0f32; rows];
16535 vbitmatvec(
16536 &vb,
16537 &offsets,
16538 &xs[bi * cols..(bi + 1) * cols],
16539 rows,
16540 cols,
16541 &mut expect,
16542 None,
16543 );
16544 assert_eq!(
16545 &got[bi * rows..(bi + 1) * rows],
16546 &expect[..],
16547 "vbit batch pos {bi}"
16548 );
16549 }
16550 }
16551
16552 #[test]
16556 fn q4_tiled_matches_q4_block_bitexact() {
16557 let (rows, cols, b) = (8usize, 128usize, 3usize);
16558 let groups = rows * cols / GROUP_SIZE;
16559 let mut split = Vec::with_capacity(groups * 18);
16560 for i in 0..groups * 16 {
16561 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16562 }
16563 for g in 0..groups {
16564 split.extend_from_slice(
16565 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16566 );
16567 }
16568 let (packed, scales) = split.split_at(groups * 16);
16570 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
16571 for g in 0..groups {
16572 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
16573 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
16574 }
16575
16576 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
16577 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
16579
16580 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
16581 q4matvec(&split, &x1, rows, cols, &mut a, None);
16582 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
16583 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
16584
16585 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
16586 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
16587 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
16588 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
16589 assert_eq!(a1, t1);
16590 assert_eq!(a2, t2);
16591
16592 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16593 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
16594 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
16595 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
16596 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
16597 }
16598
16599 #[test]
16606 fn q4matvec_sdot_outlier_exact() {
16607 let (rows, cols) = (4, 128);
16608 let groups = rows * cols / GROUP_SIZE;
16609 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
16610 for i in 0..groups * 16 {
16611 bytes.push(((i * 11 + 5) % 256) as u8);
16612 }
16613 for g in 0..groups {
16614 let s = 0.02 + 0.002 * g as f32;
16615 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
16616 }
16617 let mut x: Vec<f32> = (0..cols)
16618 .map(|i| match i % 3 {
16619 0 => 1.0,
16620 1 => -1.0,
16621 _ => 0.0,
16622 })
16623 .collect();
16624 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
16627 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
16628 let mut expect = vec![0.0f32; rows];
16629 for r in 0..rows {
16630 expect[r] = reference[r * cols..(r + 1) * cols]
16631 .iter()
16632 .zip(&x)
16633 .map(|(w, xv)| w * xv)
16634 .sum();
16635 }
16636 let mut got = vec![0.0f32; rows];
16637 q4matvec(&bytes, &x, rows, cols, &mut got, None);
16638 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
16639 for r in 0..rows {
16640 assert!(
16641 (got[r] - expect[r]).abs() < 2e-3 * scale,
16642 "row {r}: {} vs {} (outlier term must be exact)",
16643 got[r],
16644 expect[r]
16645 );
16646 }
16647 }
16648
16649 #[test]
16653 fn q1t_matvec_matches_reference() {
16654 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
16655 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
16657 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
16658 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
16660 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
16661 let mut bytes = Vec::new();
16662 for r in 0..rows {
16663 for g in 0..gpr {
16664 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
16665 let mut c = [0u8; 7];
16666 for k in 0..GROUP_SIZE {
16667 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
16669 0
16670 } else {
16671 ((k + r * 3 + g) % 3) as u8 };
16673 cortiq_core::quant::q1t_pack(&mut c, k, code);
16674 }
16675 bytes.extend_from_slice(&c);
16676 }
16677 }
16678 let mut row_ptr = vec![0u32; rows + 1];
16681 for &(idx, _) in &outliers {
16682 row_ptr[idx as usize / cols + 1] += 1;
16683 }
16684 for r in 0..rows {
16685 row_ptr[r + 1] += row_ptr[r];
16686 }
16687 for &p in &row_ptr {
16688 bytes.extend_from_slice(&p.to_le_bytes());
16689 }
16690 for &(idx, v) in &outliers {
16691 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
16692 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
16693 }
16694
16695 let mut refw = vec![0f32; rows * cols];
16696 dequant_q1t(&bytes, rows, cols, &mut refw);
16697 let x: Vec<f32> = (0..cols)
16700 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16701 .collect();
16702 let mut expect = vec![0f32; rows];
16703 for r in 0..rows {
16704 let mut a = 0.0f32;
16705 for j in 0..cols {
16706 a += refw[r * cols + j] * x[j];
16707 }
16708 expect[r] = a;
16709 }
16710 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
16711 let mut got = vec![0f32; rows];
16712 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
16713 for r in 0..rows {
16714 assert!(
16715 (got[r] - expect[r]).abs() < tol(expect[r]),
16716 "row {r}: {} vs {}",
16717 got[r],
16718 expect[r]
16719 );
16720 }
16721 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
16723 let mut gm = vec![0f32; 2 * rows];
16724 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
16725 for r in 0..rows {
16726 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
16727 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
16728 }
16729 let xb: Vec<f32> = (0..cols)
16733 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
16734 .collect();
16735 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16736 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
16737 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
16738 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16739 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
16740 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
16741 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
16742 }
16743
16744 #[test]
16747 fn q1t_matvec2_odd_gpr_matches_singles() {
16748 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16749 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
16751 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16752 for r in 0..rows {
16753 for g in 0..gpr {
16754 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
16755 let mut c = [0u8; 7];
16756 for k in 0..GROUP_SIZE {
16757 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
16758 }
16759 bytes.extend_from_slice(&c);
16760 }
16761 }
16762 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16763 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16764 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16765 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16766 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16767 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16768 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16769 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
16770 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
16771 }
16772
16773 #[test]
16777 #[ignore]
16778 fn q1t_matvec2_speed() {
16779 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16780 use std::time::Instant;
16781 let (rows, cols) = (8192usize, 4096usize);
16782 let gpr = cols / GROUP_SIZE;
16783 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16784 for r in 0..rows {
16785 for g in 0..gpr {
16786 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16787 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16788 let mut c = [0u8; 7];
16789 for k in 0..GROUP_SIZE {
16790 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16791 }
16792 bytes.extend_from_slice(&c);
16793 }
16794 }
16795 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16796 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16797 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16798 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16799 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16801 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16802 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
16803 for _ in 0..8 {
16804 let t0 = Instant::now();
16805 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16806 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
16807 let t1 = Instant::now();
16808 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16809 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16810 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
16811 }
16812 assert_eq!(p1, s1);
16813 assert_eq!(p2, s2);
16814 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
16815 }
16816
16817 #[test]
16821 #[ignore]
16822 fn q1t_matvec_speed() {
16823 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
16824 use std::time::Instant;
16825 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
16827 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
16828 for r in 0..rows {
16829 for g in 0..gpr {
16830 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16831 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16832 let mut c = [0u8; 7];
16833 for k in 0..GROUP_SIZE {
16834 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16835 }
16836 bytes.extend_from_slice(&c);
16837 }
16838 }
16839 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
16841 let mut idx = 0usize;
16842 while idx < n {
16843 row_ptr[idx / cols + 1] += 1;
16844 idx += stride;
16845 }
16846 for r in 0..rows {
16847 row_ptr[r + 1] += row_ptr[r];
16848 }
16849 for &p in &row_ptr {
16850 bytes.extend_from_slice(&p.to_le_bytes());
16851 }
16852 let mut idx = 0usize;
16853 while idx < n {
16854 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
16855 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
16856 idx += stride;
16857 }
16858 let x: Vec<f32> = (0..cols)
16861 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16862 .collect();
16863 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
16864
16865 let slow = |out: &mut [f32]| {
16867 let mut buf = vec![0f32; cols];
16868 for r in 0..rows {
16869 for g in 0..gpr {
16870 let off = (r * gpr + g) * Q1T_TILE;
16871 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
16872 let codes = &bytes[off + 2..off + Q1T_TILE];
16873 for k in 0..GROUP_SIZE {
16874 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
16875 1 => s,
16876 2 => -s,
16877 _ => 0.0,
16878 };
16879 }
16880 }
16881 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
16882 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
16883 }
16884 };
16885 let iters = 5;
16886 let mut a = vec![0f32; rows];
16887 slow(&mut a); let t = Instant::now();
16889 for _ in 0..iters {
16890 slow(&mut a);
16891 }
16892 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16893
16894 let mut b = vec![0f32; rows];
16895 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
16897 for _ in 0..iters {
16898 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
16899 }
16900 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16901
16902 for r in 0..rows {
16903 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16904 }
16905 println!(
16906 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
16907 slow_ms / fast_ms
16908 );
16909 }
16910}
16911
16912#[cfg(test)]
16913mod row_span_tests {
16914 use cortiq_core::TensorDtype;
16915
16916 #[test]
16920 fn q4tp_row_spans_cover_the_row() {
16921 let (rows, cols) = (37usize, 160usize);
16922 let gpr = cols / 32;
16923 let total = cortiq_core::quant::expected_nbytes(TensorDtype::Q4TiledP, &[rows, cols]).unwrap();
16924 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16925 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16926 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16927 let step = cortiq_core::quant::f32_to_f16(0.1);
16928 for r in 0..rows {
16929 let o = params_off + r * 4;
16930 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16931 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16932 }
16933 for r in [0usize, 1, 17, rows - 1] {
16934 let spans = super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, r).unwrap();
16935 assert_eq!(spans.iter().map(|s| s.1).sum::<usize>(), total / rows);
16936 let mut keep = vec![false; total];
16937 for &(o, n) in &spans {
16938 keep[o..o + n].fill(true);
16939 }
16940 let other: Vec<u8> =
16941 bytes.iter().zip(&keep).map(|(&b, &k)| if k { b } else { 0x5A }).collect();
16942 let (a, b) = (super::Q4tpView::new(&bytes, rows, cols), super::Q4tpView::new(&other, rows, cols));
16943 let (mut sa, mut sb) = (vec![0f32; gpr], vec![0f32; gpr]);
16944 a.scales_into(r, gpr, &mut sa);
16945 b.scales_into(r, gpr, &mut sb);
16946 assert_eq!(sa, sb, "row {r} scales");
16947 let nib = r * gpr * super::Q4TP_NIB..(r + 1) * gpr * super::Q4TP_NIB;
16948 assert_eq!(a.nib[nib.clone()], b.nib[nib], "row {r} nibbles");
16949 }
16950 assert!(super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, rows).is_none());
16951 assert!(super::row_span_offsets(TensorDtype::Q8_2f, rows, cols, 0).is_none());
16952 let t = super::row_span_offsets(TensorDtype::Q4Tiled, rows, cols, 3).unwrap();
16953 assert_eq!(t[0], (3 * gpr * super::Q4_TILE, gpr * super::Q4_TILE));
16954 }
16955}
16956
16957#[cfg(test)]
16958mod gemm_bench {
16959 #[test]
16970 #[ignore]
16971 fn q4tp_matmat_throughput() {
16972 let _alt = super::Q4TP_ALT_TEST_LOCK
16973 .lock()
16974 .unwrap_or_else(|e| e.into_inner());
16975 let b: usize = std::env::var("CMF_BENCH_B")
16979 .ok()
16980 .and_then(|v| v.parse().ok())
16981 .unwrap_or(296);
16982 let (rows, cols) = (9216usize, 2304usize);
16983 let (_, _, _) = (rows, cols, b);
16984 let total =
16985 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16986 .unwrap();
16987 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16992 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16993 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16994 let step = cortiq_core::quant::f32_to_f16(0.1);
16995 for r in 0..rows {
16996 let o = params_off + r * 4;
16997 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16998 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16999 }
17000 let _ = codes_off;
17001 let xs: Vec<f32> = (0..b * cols)
17002 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17003 .collect();
17004 let mut out = vec![0f32; b * rows];
17005 let pool = crate::pool::Pool::from_env();
17006 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17012 let reps: usize = std::env::var("CMF_BENCH_REPS")
17013 .ok()
17014 .and_then(|v| v.parse().ok())
17015 .unwrap_or(10);
17016 let mut best = [f64::MAX; 2];
17017 let mut sums = [0f32; 2];
17018 for _ in 0..reps {
17019 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
17020 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
17021 let t = std::time::Instant::now();
17022 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17023 best[k] = best[k].min(t.elapsed().as_secs_f64());
17024 sums[k] = out.iter().take(64).sum::<f32>();
17025 }
17026 }
17027 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17028 for (k, name) in ["previous", "tuned "].iter().enumerate() {
17029 println!(
17030 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
17031 best[k] * 1e3,
17032 flops / best[k] / 1e9,
17033 sums[k]
17034 );
17035 }
17036 assert!(
17037 (sums[0] - sums[1]).abs() < 1e-2,
17038 "the tuned kernel changed the result: {} vs {}",
17039 sums[0],
17040 sums[1]
17041 );
17042 }
17043
17044 #[test]
17050 fn q4tp_matmat_blocked_matches_scalar() {
17051 use std::sync::atomic::Ordering::Relaxed;
17052 let _alt = super::Q4TP_ALT_TEST_LOCK
17053 .lock()
17054 .unwrap_or_else(|e| e.into_inner());
17055 for &(rows, cols, b) in &[
17063 (64usize, 128usize, 7usize),
17064 (33, 96, 4),
17065 (16, 256, 9),
17066 (192, 2304, 37),
17067 ] {
17068 let total = cortiq_core::quant::expected_nbytes(
17069 cortiq_core::TensorDtype::Q4TiledP,
17070 &[rows, cols],
17071 )
17072 .unwrap();
17073 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17074 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17075 let lo = cortiq_core::quant::f32_to_f16(-4.0);
17076 let step = cortiq_core::quant::f32_to_f16(0.1);
17077 for r in 0..rows {
17078 let o = params_off + r * 4;
17079 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17080 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17081 }
17082 let xs: Vec<f32> = (0..b * cols)
17083 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
17084 .collect();
17085 let mut got = vec![0f32; b * rows];
17086 let mut want = vec![0f32; b * rows];
17087 let gpr = cols / 32;
17088 let view = super::Q4tpView::new(&bytes, rows, cols);
17089 let pool = crate::pool::Pool::from_env();
17090 super::Q4TP_ALT.store(2, Relaxed);
17091 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
17092 super::Q4TP_ALT.store(1, Relaxed);
17093 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
17094 super::Q4TP_ALT.store(0, Relaxed);
17095 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17102 let (mut worst, mut at) = (0f32, 0usize);
17103 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
17104 if (g - w).abs() > worst {
17105 worst = (g - w).abs();
17106 at = i;
17107 }
17108 }
17109 assert!(
17110 worst <= 1e-4 * scale,
17111 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
17112 (scale {scale:.3e}) at cell {at}: {} vs {}",
17113 got[at],
17114 want[at]
17115 );
17116
17117 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
17124 for bi in 0..b {
17125 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
17126 for r in 0..rows {
17127 let mut sc = vec![0f32; gpr];
17128 view.scales_into(r, gpr, &mut sc);
17129 let mut exact = 0f64;
17130 for j in 0..cols {
17131 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17132 exact += w as f64 * sq as f64 * act.xq[j] as f64;
17133 }
17134 exact *= act.sx as f64;
17135 for &(j, xv) in &act.outliers {
17136 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17137 exact += w as f64 * sq as f64 * xv as f64;
17138 }
17139 let i = bi * rows + r;
17140 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
17141 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
17142 }
17143 }
17144 println!(
17145 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
17146 per-column {e_scalar:.3e}"
17147 );
17148 assert!(
17153 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
17154 "{rows}x{cols} b={b}: error against f64 too large — blocked \
17155 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
17156 );
17157 }
17158 }
17159
17160 #[test]
17166 #[ignore]
17167 fn q4tp_matmat_row_exact_cost() {
17168 let _alt = super::Q4TP_ALT_TEST_LOCK
17169 .lock()
17170 .unwrap_or_else(|e| e.into_inner());
17171 let pool = crate::pool::Pool::from_env();
17172 let reps: usize = std::env::var("CMF_BENCH_REPS")
17173 .ok()
17174 .and_then(|v| v.parse().ok())
17175 .unwrap_or(30);
17176 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
17177 let total = cortiq_core::quant::expected_nbytes(
17178 cortiq_core::TensorDtype::Q4TiledP,
17179 &[rows, cols],
17180 )
17181 .unwrap();
17182 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17183 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17184 let lo = cortiq_core::quant::f32_to_f16(-4.0);
17185 let step = cortiq_core::quant::f32_to_f16(0.1);
17186 for r in 0..rows {
17187 let o = params_off + r * 4;
17188 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17189 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17190 }
17191 for &b in &[2usize, 4, 5, 8] {
17192 let xs: Vec<f32> = (0..b * cols)
17193 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17194 .collect();
17195 let mut out = vec![0f32; b * rows];
17196 let mut best = [f64::MAX; 3];
17197 for _ in 0..reps {
17198 for (k, best_k) in best.iter_mut().enumerate() {
17199 let t = std::time::Instant::now();
17200 match k {
17201 0 | 1 => super::q4tp_matmat_with(
17202 &bytes,
17203 &xs,
17204 b,
17205 rows,
17206 cols,
17207 &mut out,
17208 pool.as_deref(),
17209 k == 1,
17210 ),
17211 _ => {
17212 for (bi, o) in out.chunks_mut(rows).enumerate() {
17213 super::q4tp_matvec(
17214 &bytes,
17215 &xs[bi * cols..(bi + 1) * cols],
17216 rows,
17217 cols,
17218 o,
17219 pool.as_deref(),
17220 );
17221 }
17222 }
17223 }
17224 *best_k = best_k.min(t.elapsed().as_secs_f64());
17225 }
17226 }
17227 println!(
17228 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
17229 {b} matvecs {:.3} ms",
17230 best[0] * 1e3,
17231 best[1] * 1e3,
17232 best[2] * 1e3
17233 );
17234 }
17235 }
17236 }
17237
17238 #[test]
17242 #[ignore]
17243 fn q4t_matmat_throughput() {
17244 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
17245 let total =
17246 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
17247 .unwrap();
17248 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17251 let sc = cortiq_core::quant::f32_to_f16(0.02);
17252 for t in bytes.chunks_mut(super::Q4_TILE) {
17253 t[..2].copy_from_slice(&sc.to_le_bytes());
17254 }
17255 let xs: Vec<f32> = (0..b * cols)
17256 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17257 .collect();
17258 let mut out = vec![0f32; b * rows];
17259 let pool = crate::pool::Pool::from_env();
17260 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17261 let reps: usize = std::env::var("CMF_BENCH_REPS")
17262 .ok()
17263 .and_then(|v| v.parse().ok())
17264 .unwrap_or(10);
17265 let mut best = f64::MAX;
17266 for _ in 0..reps {
17267 let t = std::time::Instant::now();
17268 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17269 best = best.min(t.elapsed().as_secs_f64());
17270 }
17271 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17272 println!(
17273 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
17274 best * 1e3,
17275 flops / best / 1e9,
17276 out.iter().take(64).sum::<f32>()
17277 );
17278 }
17279}
17280
17281#[cfg(all(test, target_arch = "x86_64"))]
17282mod x86_gemm_tests {
17283 use super::x86_gemm_rm;
17284
17285 #[test]
17292 fn q4tp_fma_batch_matches_the_row_exact_cell() {
17293 if !super::q4tp_fma_x86() || !super::a8w8_enabled() {
17294 return;
17295 }
17296 let pool = crate::pool::Pool::from_env();
17297 for &(rows, cols, b) in &[(33usize, 96usize, 7usize), (64, 256, 13), (257, 2560, 37)] {
17298 let total = cortiq_core::quant::expected_nbytes(
17299 cortiq_core::TensorDtype::Q4TiledP,
17300 &[rows, cols],
17301 )
17302 .unwrap();
17303 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17304 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17305 let lo = cortiq_core::quant::f32_to_f16(-4.0);
17306 let step = cortiq_core::quant::f32_to_f16(0.1);
17307 for r in 0..rows {
17308 let o = params_off + r * 4;
17309 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17310 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17311 }
17312 let mut xs: Vec<f32> = (0..b * cols)
17313 .map(|i| ((i * 37 % 89) as f32 - 44.0) / 44.0)
17314 .collect();
17315 for bi in (0..b).step_by(3) {
17316 xs[bi * cols + (bi * 7) % cols] = 40.0; }
17318 let mut got = vec![0f32; b * rows];
17319 let mut want = vec![0f32; b * rows];
17320 super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref(), false);
17321 super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref(), true);
17322 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17323 let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
17324 assert!(
17325 worst <= 1e-5 * scale,
17326 "{rows}x{cols} b={b}: FMA batch and row-exact cell disagree by {worst:.3e} (scale {scale:.3e})"
17327 );
17328 }
17329 }
17330
17331 #[allow(clippy::too_many_arguments)]
17332 fn naive(
17333 m: usize,
17334 n: usize,
17335 k: usize,
17336 alpha: f32,
17337 a: &[f32],
17338 lda: usize,
17339 b: &[f32],
17340 ldb: usize,
17341 b_rows_are_n: bool,
17342 ) -> Vec<f64> {
17343 let mut c = vec![0f64; m * n];
17344 for i in 0..m {
17345 for j in 0..n {
17346 let mut acc = 0f64;
17347 for p in 0..k {
17348 let bv = if b_rows_are_n { b[j * ldb + p] } else { b[p * ldb + j] };
17349 acc += a[i * lda + p] as f64 * bv as f64;
17350 }
17351 c[i * n + j] = acc * alpha as f64;
17352 }
17353 }
17354 c
17355 }
17356
17357 #[test]
17362 fn x86_gemm_matches_naive_on_odd_shapes() {
17363 for &(m, n, k) in &[
17364 (1usize, 1usize, 1usize),
17365 (4, 16, 256),
17366 (7, 37, 19),
17367 (9, 300, 257),
17368 (64, 1000, 256),
17369 (5, 256, 1001),
17370 (13, 129, 512),
17371 (8, 16, 0),
17372 ] {
17373 for &bt in &[false, true] {
17374 let (lda, ldb, ldc) = (k + 3, if bt { k + 2 } else { n + 5 }, n + 1);
17375 let a: Vec<f32> = (0..m * lda)
17376 .map(|i| ((i * 7919) % 1000) as f32 / 500.0 - 1.0)
17377 .collect();
17378 let brows = if bt { n } else { k.max(1) };
17379 let b: Vec<f32> = (0..brows * ldb)
17380 .map(|i| ((i * 104_729) % 997) as f32 / 498.5 - 1.0)
17381 .collect();
17382 let mut c = vec![f32::NAN; m * ldc];
17383 x86_gemm_rm(m, n, k, 0.37, &a, lda, &b, ldb, bt, &mut c, ldc);
17384 let r = naive(m, n, k, 0.37, &a, lda, &b, ldb, bt);
17385 for i in 0..m {
17386 for j in 0..n {
17387 let (x, y) = (c[i * ldc + j] as f64, r[i * n + j]);
17388 assert!(
17389 (x - y).abs() <= 1e-3 * (1.0 + y.abs()),
17390 "m{m} n{n} k{k} bt{bt} [{i},{j}]: {x} vs {y}"
17391 );
17392 }
17393 assert!(c[i * ldc + n].is_nan(), "m{m} n{n} k{k}: padding column written");
17394 }
17395 }
17396 }
17397 }
17398}