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;
22
23mod mxfp4;
25
26pub enum QTensor {
27 F32 {
28 data: Vec<f32>,
29 rows: usize,
30 cols: usize,
31 },
32 Mapped {
33 model: Arc<CmfModel>,
34 idx: usize,
36 dtype: TensorDtype,
37 rows: usize,
38 cols: usize,
39 row_scale: Vec<f32>,
41 col_field: Vec<f32>,
43 vbit_offsets: Vec<usize>,
47 repack: Vec<u8>,
55 },
56 Blr(Box<crate::blr::BlrTensor>),
62}
63
64fn repack_enabled() -> bool {
71 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
72 *ON.get_or_init(|| {
73 std::env::var("CMF_REPACK")
74 .map(|v| v == "1")
75 .unwrap_or(cfg!(target_os = "android"))
76 })
77}
78
79fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
83 #[cfg(target_arch = "aarch64")]
84 let arch_ok = sdot_enabled();
85 #[cfg(not(target_arch = "aarch64"))]
86 let arch_ok = false;
87 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
88 return Vec::new();
89 }
90 q8_repack_layout(bytes, rows, cols)
91}
92
93fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
96 let groups = rows / 4;
97 let mut rep = vec![0u8; groups * 4 * cols];
98 for g in 0..groups {
99 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
100 for c in 0..cols / 16 {
101 for lane in 0..4 {
102 let src = (g * 4 + lane) * cols + c * 16;
103 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
104 .copy_from_slice(&bytes[src..src + 16]);
105 }
106 }
107 }
108 rep
109}
110
111fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
114 let ng = cols / GROUP_SIZE;
115 let bits = &bytes[..rows];
116 let mut offsets = Vec::with_capacity(rows + 1);
117 let mut off = rows + rows * ng * 2;
118 for r in 0..rows {
119 offsets.push(off);
120 off += (cols * bits[r] as usize).div_ceil(8);
121 }
122 offsets.push(off);
123 offsets
124}
125
126fn blocked_enabled() -> bool {
132 use std::sync::atomic::Ordering::Relaxed;
133 match BLOCKED_OVERRIDE.load(Relaxed) {
134 1 => false,
135 2 => true,
136 _ => {
137 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
138 *ON.get_or_init(|| {
139 std::env::var("CMF_X86_BLOCKED")
140 .map(|v| v != "0")
141 .unwrap_or(true)
142 })
143 }
144 }
145}
146
147static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
148
149pub fn set_blocked_override(on: Option<bool>) {
156 let v = match on {
157 None => 0,
158 Some(false) => 1,
159 Some(true) => 2,
160 };
161 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
162}
163
164fn gpu_lmhead_enabled() -> bool {
165 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
166 *ON.get_or_init(|| {
167 std::env::var("CMF_GPU_LMHEAD")
168 .map(|v| v != "0")
169 .unwrap_or(true)
170 })
171}
172
173fn gpu_split_frac() -> f32 {
174 if FULL_GPU_Q8.get() {
177 return 1.0;
178 }
179 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
180 *F.get_or_init(|| {
181 std::env::var("CMF_GPU_SPLIT")
182 .ok()
183 .and_then(|v| v.parse::<f32>().ok())
184 .unwrap_or(0.5)
185 .clamp(0.0, 1.0)
186 })
187}
188
189impl QTensor {
190 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
191 debug_assert_eq!(data.len(), rows * cols);
192 Self::F32 { data, rows, cols }
193 }
194
195 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
198 let idx = model
201 .tensor_index(name)
202 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
203 let entry = &model.tensors[idx];
204 if entry.shape.len() != 2 {
205 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
206 }
207 let (rows, cols) = (entry.shape[0], entry.shape[1]);
208 let bytes = model.entry_bytes(entry);
209
210 match entry.dtype {
211 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
212 let n = rows * cols;
213 let scales_off = n;
214 let row_scale: Vec<f32> = (0..rows)
215 .map(|o| {
216 f16_to_f32(u16::from_le_bytes([
217 bytes[scales_off + o * 2],
218 bytes[scales_off + o * 2 + 1],
219 ]))
220 })
221 .collect();
222 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
223 let col_off = n + rows * 2;
224 (0..cols)
225 .map(|i| {
226 f16_to_f32(u16::from_le_bytes([
227 bytes[col_off + i * 2],
228 bytes[col_off + i * 2 + 1],
229 ]))
230 })
231 .collect()
232 } else {
233 Vec::new()
234 };
235 Ok(Self::Mapped {
236 model: model.clone(),
237 idx,
238 dtype: entry.dtype,
239 rows,
240 cols,
241 row_scale,
242 col_field,
243 vbit_offsets: Vec::new(),
244 repack: q8_repack(bytes, rows, cols),
245 })
246 }
247 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
249 model: model.clone(),
250 idx,
251 dtype: entry.dtype,
252 rows,
253 cols,
254 row_scale: Vec::new(),
255 col_field: Vec::new(),
256 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
257 repack: Vec::new(),
258 }),
259 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
263 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
264 let offsets: Vec<usize> = (0..=rows)
265 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
266 .collect();
267 Ok(Self::Mapped {
268 model: model.clone(),
269 idx,
270 dtype: entry.dtype,
271 rows,
272 cols,
273 row_scale: Vec::new(),
274 col_field: Vec::new(),
275 vbit_offsets: offsets,
276 repack: Vec::new(),
277 })
278 }
279 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
285 model: model.clone(),
286 idx,
287 dtype: entry.dtype,
288 rows,
289 cols,
290 row_scale: Vec::new(),
291 col_field: Vec::new(),
292 vbit_offsets: Vec::new(),
293 repack: Vec::new(),
294 }),
295 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
298 model: model.clone(),
299 idx,
300 dtype: entry.dtype,
301 rows,
302 cols,
303 row_scale: Vec::new(),
304 col_field: Vec::new(),
305 vbit_offsets: Vec::new(),
306 repack: Vec::new(),
307 }),
308 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
310 model: model.clone(),
311 idx,
312 dtype: entry.dtype,
313 rows,
314 cols,
315 row_scale: Vec::new(),
316 col_field: Vec::new(),
317 vbit_offsets: Vec::new(),
318 repack: Vec::new(),
319 }),
320 TensorDtype::Q4Block 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::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
334 model: model.clone(),
335 idx,
336 dtype: entry.dtype,
337 rows,
338 cols,
339 row_scale: Vec::new(),
340 col_field: Vec::new(),
341 vbit_offsets: Vec::new(),
342 repack: Vec::new(),
343 }),
344 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 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::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
358 model: model.clone(),
359 idx,
360 dtype: entry.dtype,
361 rows,
362 cols,
363 row_scale: Vec::new(),
364 col_field: Vec::new(),
365 vbit_offsets: Vec::new(),
366 repack: Vec::new(),
367 }),
368 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
372 model: model.clone(),
373 idx,
374 dtype: entry.dtype,
375 rows,
376 cols,
377 row_scale: Vec::new(),
378 col_field: Vec::new(),
379 vbit_offsets: Vec::new(),
380 repack: Vec::new(),
381 }),
382 TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
385 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
386 model, idx,
387 )?)))
388 }
389 _ => {
391 let mut data = vec![0.0f32; rows * cols];
392 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
393 Ok(Self::from_f32(data, rows, cols))
394 }
395 }
396 }
397
398 pub fn is_blr(&self) -> bool {
401 matches!(self, Self::Blr(_))
402 }
403
404 pub fn is_mxfp4(&self) -> bool {
407 matches!(
408 self,
409 Self::Mapped {
410 dtype: TensorDtype::Mxfp4,
411 ..
412 }
413 )
414 }
415
416 fn per_tensor_only(&self) -> bool {
419 self.is_blr() || self.is_mxfp4()
420 }
421
422 pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
428 match self {
429 Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
430 _ => None,
431 }
432 }
433
434 pub(crate) fn graph_weight_or_blr(
438 &self,
439 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
440 self.graph_weight()
441 .or_else(|| self.graph_weight_descriptor())
442 .or_else(|| {
443 self.graph_blr()
444 .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
445 })
446 }
447
448 pub(crate) fn is_q1(&self) -> bool {
452 matches!(
453 self,
454 Self::Mapped {
455 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
456 ..
457 }
458 )
459 }
460
461 fn bin_w(&self) -> Option<BinW<'_>> {
464 match self {
465 Self::Mapped {
466 dtype: TensorDtype::Q1,
467 ..
468 } => Some(BinW::Q1(self.quant_bytes())),
469 Self::Mapped {
470 dtype: TensorDtype::Q1G,
471 rows,
472 cols,
473 ..
474 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
475 _ => None,
476 }
477 }
478
479 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
482 match self {
483 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
484 _ => None,
485 }
486 }
487
488 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
495 if self.has_prism_contract() {
496 return None;
497 }
498 match self {
499 #[cfg(target_os = "macos")]
500 Self::Mapped {
501 dtype: TensorDtype::Q1T,
502 ..
503 } if !crate::gpu::metal_q1t_enabled() => None,
504 Self::Mapped {
505 idx,
506 dtype:
507 TensorDtype::Q1
508 | TensorDtype::Q1G
509 | TensorDtype::Q1T
510 | TensorDtype::Q4Block
511 | TensorDtype::Q4Tiled
512 | TensorDtype::Q4TiledP
516 | TensorDtype::Q8Row
517 | TensorDtype::Q8_2f,
518 rows,
519 cols,
520 ..
521 } => Some((*idx, *rows, *cols)),
522 _ => None,
523 }
524 }
525
526 #[cfg(target_os = "macos")]
533 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
534 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
535 let name = &model.tensors[idx].name;
536 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
537 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
538 if forward && affine {
539 let e = model.tensors.get(idx)?;
540 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
541 }
542 }
543 self.q1_parts()
544 }
545
546 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
552 if self.has_prism_contract() {
553 return None;
554 }
555 match self {
556 Self::Mapped {
557 idx,
558 dtype: TensorDtype::Q4Tiled,
559 rows,
560 cols,
561 ..
562 } => Some((*idx, *rows, *cols)),
563 _ => None,
564 }
565 }
566
567 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
571 if self.has_prism_contract() {
572 return None;
573 }
574 match self {
575 Self::Mapped {
576 idx,
577 dtype: TensorDtype::Q4TiledP,
578 rows,
579 cols,
580 ..
581 } => Some((*idx, *rows, *cols)),
582 _ => None,
583 }
584 }
585
586 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
591 if self.has_prism_contract() {
592 return None;
593 }
594 match self {
595 Self::Mapped {
596 idx,
597 dtype: TensorDtype::Q8Row,
598 rows,
599 cols,
600 row_scale,
601 col_field,
602 ..
603 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
604 _ => None,
605 }
606 }
607
608 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
611 if self.has_prism_contract() {
612 return None;
613 }
614 match self {
615 Self::Mapped {
616 idx,
617 dtype: TensorDtype::Q8_2f,
618 rows,
619 cols,
620 row_scale,
621 col_field,
622 ..
623 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
624 _ => None,
625 }
626 }
627
628 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
632 match self {
633 Self::Mapped { dtype, .. } => Some(*dtype),
634 _ => None,
635 }
636 }
637
638 pub fn model_idx(&self) -> Option<usize> {
643 match self {
644 Self::Mapped { idx, .. } => Some(*idx),
645 _ => None,
646 }
647 }
648
649 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
654 match self {
655 Self::Mapped { model, .. } => Some(model.clone()),
656 _ => None,
657 }
658 }
659
660 pub(crate) fn has_prism_contract(&self) -> bool {
665 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
666 }
667
668 pub fn rows(&self) -> usize {
669 match self {
670 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
671 Self::Blr(t) => t.rows,
672 }
673 }
674
675 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
678 if self.has_prism_contract() {
679 return None;
680 }
681 match self {
682 Self::Mapped {
683 model,
684 idx,
685 dtype: TensorDtype::Q4Tiled,
686 ..
687 } => Some((model, *idx)),
688 _ => None,
689 }
690 }
691
692 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
695 if self.has_prism_contract() {
696 return None;
697 }
698 match self {
699 Self::Mapped {
700 model,
701 idx,
702 dtype: TensorDtype::Q4TiledP,
703 ..
704 } => Some((model, *idx)),
705 _ => None,
706 }
707 }
708
709 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
717 if self.has_prism_contract() {
718 return None;
719 }
720 match self {
721 Self::Mapped {
722 model,
723 idx,
724 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
725 ..
726 } => Some((model, *idx)),
727 _ => None,
728 }
729 }
730
731 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
734 if self.has_prism_contract() {
735 return None;
736 }
737 match self {
738 Self::Mapped {
739 model,
740 idx,
741 dtype: TensorDtype::Q8_2f,
742 ..
743 } => Some((model, *idx)),
744 _ => None,
745 }
746 }
747
748 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
751 if self.has_prism_contract() {
752 return None;
753 }
754 match self {
755 Self::Mapped {
756 model,
757 idx,
758 dtype: TensorDtype::Q2TiledP,
759 ..
760 } => Some((model, *idx)),
761 _ => None,
762 }
763 }
764
765 pub fn cols(&self) -> usize {
766 match self {
767 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
768 Self::Blr(t) => t.cols,
769 }
770 }
771
772 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
775 if self.has_prism_contract() {
776 return None;
777 }
778 match self {
779 Self::Mapped {
780 model,
781 idx,
782 dtype: TensorDtype::Q1,
783 ..
784 } => Some((model, *idx)),
785 _ => None,
786 }
787 }
788
789 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
800 if self.has_prism_contract() {
801 return None;
802 }
803 self.graph_weight_descriptor()
804 }
805
806 pub(crate) fn graph_weight_descriptor(
811 &self,
812 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
813 match self {
814 Self::Mapped {
815 model,
816 idx,
817 dtype: TensorDtype::Q8Row,
818 row_scale,
819 ..
820 } => Some((model, *idx, 0, row_scale.as_slice())),
821 Self::Mapped {
822 model,
823 idx,
824 dtype: TensorDtype::Q1,
825 ..
826 } => Some((model, *idx, 1, &[])),
827 Self::Mapped {
830 model,
831 idx,
832 dtype: TensorDtype::Q1G,
833 ..
834 } => Some((model, *idx, 10, &[])),
835 Self::Mapped {
840 model,
841 idx,
842 dtype: TensorDtype::Q4Tiled,
843 ..
844 } => Some((model, *idx, 5, &[])),
845 Self::Mapped {
849 model,
850 idx,
851 dtype: TensorDtype::Q4TiledP,
852 ..
853 } => Some((model, *idx, 6, &[])),
854 Self::Mapped {
855 model,
856 idx,
857 dtype: TensorDtype::Q4Block,
858 ..
859 } => Some((model, *idx, 2, &[])),
860 Self::Mapped {
865 model,
866 idx,
867 dtype: TensorDtype::Q8_2f,
868 ..
869 } => Some((model, *idx, 7, &[])),
870 Self::Mapped {
871 model,
872 idx,
873 dtype: TensorDtype::Q1T,
874 ..
875 } => Some((model, *idx, 3, &[])),
876 Self::Mapped {
880 model,
881 idx,
882 dtype: TensorDtype::Q2TiledP,
883 ..
884 } => Some((model, *idx, 9, &[])),
885 _ => None,
886 }
887 }
888
889 pub fn as_f32(&self) -> Option<&[f32]> {
892 match self {
893 Self::F32 { data, .. } => Some(data),
894 Self::Mapped { .. } | Self::Blr(_) => None,
895 }
896 }
897
898 fn quant_bytes(&self) -> &[u8] {
899 match self {
900 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
901 Self::Blr(t) => t.bytes(),
902 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
903 }
904 }
905
906 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
908 let cols = self.cols();
909 debug_assert_eq!(dst.len(), cols);
910 match self {
911 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
912 Self::Blr(t) => t.row_f32(r, dst),
913 Self::Mapped {
914 model,
915 idx,
916 dtype,
917 row_scale,
918 col_field,
919 vbit_offsets,
920 ..
921 } => {
922 if *dtype == TensorDtype::Q4Tiled {
923 let bytes = self.quant_bytes();
924 let gpr = cols / GROUP_SIZE;
925 for gi in 0..gpr {
926 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
927 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
928 for (k, &b) in tile[2..].iter().enumerate() {
929 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
930 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
931 }
932 }
933 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
934 crate::prism::inverse_embedding(model, dst);
935 }
936 return;
937 }
938 if *dtype == TensorDtype::Q4TiledP {
939 let bytes = self.quant_bytes();
940 let gpr = cols / GROUP_SIZE;
941 let v = Q4tpView::new(bytes, self.rows(), cols);
942 let mut sc = vec![0f32; gpr];
943 v.scales_into(r, gpr, &mut sc);
944 for gi in 0..gpr {
945 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
946 let s = sc[gi];
947 for (k, &b) in tile.iter().enumerate() {
948 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
949 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
950 }
951 }
952 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
953 crate::prism::inverse_embedding(model, dst);
954 }
955 return;
956 }
957 if *dtype == TensorDtype::Q2TiledP {
958 let bytes = self.quant_bytes();
959 let gpr = cols / GROUP_SIZE;
960 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
961 let mut sc = vec![0f32; gpr];
962 v.scales_into(r, gpr, &mut sc);
963 for gi in 0..gpr {
964 let ch =
965 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
966 let s = sc[gi];
967 for (k, &b) in ch.iter().enumerate() {
968 for j in 0..4 {
969 let center = if crate::prism::is_affine_target(
970 model,
971 &model.tensors[*idx].name,
972 ) {
973 1.0
974 } else {
975 1.5
976 };
977 dst[gi * GROUP_SIZE + k * 4 + j] =
978 (((b >> (2 * j)) & 3) as f32 - center) * s;
979 }
980 }
981 }
982 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
983 crate::prism::inverse_embedding(model, dst);
984 }
985 return;
986 }
987 if *dtype == TensorDtype::Q4Block {
988 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
989 let gpr = cols / GROUP_SIZE;
990 for gi in 0..gpr {
991 let g = r * gpr + gi;
992 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
993 for (k, &b) in packed[g * 16..(g + 1) * 16].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::Mxfp4 {
1004 mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1005 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1006 crate::prism::inverse_embedding(model, dst);
1007 }
1008 return;
1009 }
1010 if *dtype == TensorDtype::Q1 {
1011 let bytes = self.quant_bytes();
1012 let gpr = cols / GROUP_SIZE;
1013 for gi in 0..gpr {
1014 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1015 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1016 for (j, &b) in tile[2..].iter().enumerate() {
1017 for k in 0..8 {
1018 dst[gi * GROUP_SIZE + j * 8 + k] =
1019 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1020 }
1021 }
1022 }
1023 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1024 crate::prism::inverse_embedding(model, dst);
1025 }
1026 return;
1027 }
1028 if *dtype == TensorDtype::Q1G {
1029 let bytes = self.quant_bytes();
1032 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1033 let gpr = cols / GROUP_SIZE;
1034 for gi in 0..gpr {
1035 let t = r * gpr + gi;
1036 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1037 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1038 for k in 0..8 {
1039 dst[gi * GROUP_SIZE + j * 8 + k] =
1040 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1041 }
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::Q1T {
1050 let bytes = self.quant_bytes();
1051 let gpr = cols / GROUP_SIZE;
1052 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1053 for gi in 0..gpr {
1054 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1055 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1056 bytes[off],
1057 bytes[off + 1],
1058 ]));
1059 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1060 for k in 0..GROUP_SIZE {
1061 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1062 {
1063 1 => s,
1064 2 => -s,
1065 _ => 0.0,
1066 };
1067 }
1068 }
1069 let rows = self.rows();
1071 let entries = base_len + (rows + 1) * 4;
1072 if entries <= bytes.len() {
1073 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1074 let r0 = u32::from_le_bytes([
1075 ptrs[r * 4],
1076 ptrs[r * 4 + 1],
1077 ptrs[r * 4 + 2],
1078 ptrs[r * 4 + 3],
1079 ]) as usize;
1080 let r1 = u32::from_le_bytes([
1081 ptrs[(r + 1) * 4],
1082 ptrs[(r + 1) * 4 + 1],
1083 ptrs[(r + 1) * 4 + 2],
1084 ptrs[(r + 1) * 4 + 3],
1085 ]) as usize;
1086 let off = entries + r0 * 4;
1087 for i in 0..r1 - r0 {
1088 let item = &bytes[off + i * 4..off + i * 4 + 4];
1089 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1090 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1091 item[2], item[3],
1092 ]));
1093 if c < cols {
1094 dst[c] = v;
1095 }
1096 }
1097 }
1098 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1099 crate::prism::inverse_embedding(model, dst);
1100 }
1101 return;
1102 }
1103 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1104 let bytes = self.quant_bytes();
1105 let rows = self.rows();
1106 let ng = cols / GROUP_SIZE;
1107 let bits = &bytes[..rows];
1108 let sc_off = rows;
1109 let off = vbit_offsets[r];
1112 let b = bits[r] as usize;
1113 let l = ((1usize << (b - 1)) - 1) as f32;
1114 let data = &bytes[off..];
1115 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1116 for (i, d) in dst.iter_mut().enumerate() {
1117 while nbits < b {
1118 acc = (acc << 8) | data[byte_idx] as u64;
1119 byte_idx += 1;
1120 nbits += 8;
1121 }
1122 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1123 nbits -= b;
1124 let so = (r * ng + i / GROUP_SIZE) * 2;
1125 let sv = f16_to_f32(u16::from_le_bytes([
1126 bytes[sc_off + so],
1127 bytes[sc_off + so + 1],
1128 ]));
1129 *d = (u - l) * sv;
1130 }
1131 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1132 crate::prism::inverse_embedding(model, dst);
1133 }
1134 return;
1135 }
1136 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1137 let s = row_scale[r];
1138 match dtype {
1139 TensorDtype::Q8Row => {
1140 for (d, &b) in dst.iter_mut().zip(q) {
1141 *d = (b as i8) as f32 * s;
1142 }
1143 }
1144 TensorDtype::Q8_2f => {
1145 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1146 *d = (b as i8) as f32 * s * col_field[i];
1147 }
1148 }
1149 _ => unreachable!(),
1150 }
1151 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1152 crate::prism::inverse_embedding(model, dst);
1153 }
1154 }
1155 }
1156 }
1157
1158 pub fn sparse_col_ok(&self) -> bool {
1163 match self {
1164 Self::F32 { .. } => true,
1165 Self::Blr(_) => false,
1166 Self::Mapped { dtype, .. } => {
1167 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1168 }
1169 }
1170 }
1171
1172 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1176 let inter = self.cols();
1177 let hidden = self.rows();
1178 debug_assert_eq!(out.len(), hidden);
1179 match self {
1180 Self::F32 { data, .. } => {
1181 for (k, o) in out.iter_mut().enumerate() {
1182 *o += w * data[k * inter + c];
1183 }
1184 }
1185 Self::Blr(t) => t.add_col_scaled(c, w, out),
1186 Self::Mapped {
1187 dtype,
1188 row_scale,
1189 col_field,
1190 ..
1191 } => {
1192 let q = self.quant_bytes();
1193 let colf = if *dtype == TensorDtype::Q8_2f {
1194 col_field[c]
1195 } else {
1196 1.0
1197 };
1198 let wc = w * colf;
1199 for (k, o) in out.iter_mut().enumerate() {
1200 let b = q[k * inter + c] as i8 as f32;
1201 *o += wc * b * row_scale[k];
1202 }
1203 }
1204 }
1205 }
1206
1207 #[inline]
1216 pub fn prefetch_row(&self, r: usize) {
1217 let Self::Mapped { dtype, .. } = self else {
1218 return;
1219 };
1220 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1221 return;
1222 }
1223 let cols = self.cols();
1224 let q = self.quant_bytes();
1225 let (a, b) = (r * cols, (r + 1) * cols);
1226 if b > q.len() {
1227 return;
1228 }
1229 let mut j = a;
1230 while j < b {
1231 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1232 j += 512;
1233 }
1234 }
1235
1236 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1245 let cols = self.cols();
1246 debug_assert_eq!(out.len(), cols);
1247 match self {
1248 Self::F32 { data, .. } => {
1249 let row = &data[r * cols..(r + 1) * cols];
1250 for (o, v) in out.iter_mut().zip(row) {
1251 *o += w * v;
1252 }
1253 }
1254 Self::Blr(t) => {
1255 let tmp = &mut scratch[..cols];
1256 t.row_f32(r, tmp);
1257 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1258 *o += w * v;
1259 }
1260 }
1261 Self::Mapped {
1262 dtype,
1263 row_scale,
1264 col_field,
1265 ..
1266 } => match dtype {
1267 TensorDtype::Q8Row => {
1268 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1269 let ws = w * row_scale[r];
1270 let row: &[i8] =
1271 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1272 axpy_i8_f32(out, row, ws);
1273 }
1274 TensorDtype::Q8_2f => {
1275 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1276 let ws = w * row_scale[r];
1277 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1278 *o += ws * c * (*b as i8 as f32);
1279 }
1280 }
1281 _ => {
1282 self.row_f32(r, scratch);
1283 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1284 *o += w * v;
1285 }
1286 }
1287 },
1288 }
1289 }
1290
1291 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1295 let cols = self.cols();
1296 match self {
1297 Self::F32 { data, .. } => {
1298 let row = &data[r * cols..(r + 1) * cols];
1299 row.iter().zip(x).map(|(w, v)| w * v).sum()
1300 }
1301 Self::Blr(t) => {
1302 let tmp = &mut scratch[..cols];
1303 t.row_f32(r, tmp);
1304 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1305 }
1306 Self::Mapped {
1307 model,
1308 idx,
1309 dtype,
1310 row_scale,
1311 col_field,
1312 ..
1313 } => {
1314 let prism_forward =
1315 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1316 if prism_forward {
1317 let transformed = crate::prism::forward(model, &x[..cols]);
1318 let gpr = cols / GROUP_SIZE;
1319 match dtype {
1320 TensorDtype::Q2TiledP => {
1321 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1322 let mut sc = vec![0f32; gpr];
1323 v.scales_into(r, gpr, &mut sc);
1324 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1325 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1326 }
1327 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1328 }
1329 _ => {
1330 self.row_f32(r, scratch);
1331 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1332 }
1333 }
1334 }
1335 match dtype {
1336 TensorDtype::Q8Row => {
1337 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1338 dot_i8_f32(q, x) * row_scale[r]
1339 }
1340 TensorDtype::Q8_2f => {
1341 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1342 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1343 }
1344 _ => {
1345 self.row_f32(r, scratch);
1346 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1347 }
1348 }
1349 }
1350 }
1351 }
1352
1353 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1356 if self.flint_exact_cpu_default() {
1357 return float_activations_scope(|| self.matvec(x, out, pool));
1358 }
1359 match self {
1360 Self::F32 { data, .. } => {
1364 if !crate::f32_backend::matvec(data, x, out) {
1365 matvec_rows(pool, data, x, out);
1366 }
1367 }
1368 Self::Blr(t) => {
1369 if crate::gpu::enabled_here()
1370 && crate::blr::gpu_wanted()
1371 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1372 {
1373 return;
1374 }
1375 t.matvec(x, out, pool)
1376 }
1377 Self::Mapped {
1378 model,
1379 idx,
1380 dtype,
1381 rows,
1382 cols,
1383 row_scale,
1384 col_field,
1385 vbit_offsets,
1386 repack,
1387 } => {
1388 let _ = (model, idx);
1389 assert!(
1399 out.len() >= *rows && x.len() >= *cols,
1400 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1401 out.len(),
1402 x.len(),
1403 );
1404 let prism_forward =
1405 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1406 if *dtype == TensorDtype::Q2TiledP
1407 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1408 {
1409 use std::sync::atomic::{AtomicUsize, Ordering};
1410 static N: AtomicUsize = AtomicUsize::new(0);
1411 let n = N.fetch_add(1, Ordering::Relaxed);
1412 if n < 128 {
1413 eprintln!(
1414 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1415 model.tensors[*idx].name,
1416 prism_forward,
1417 rows,
1418 cols,
1419 crate::gpu::enabled_here(),
1420 crate::gpu::q2tp_gpu_opt_in(),
1421 crate::gpu::cur_layer(),
1422 );
1423 }
1424 }
1425 if prism_forward {
1430 let transformed = crate::prism::forward(model, &x[..*cols]);
1431 match dtype {
1432 TensorDtype::Q4Block => {
1433 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1434 }
1435 TensorDtype::Q4Tiled => {
1436 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1437 }
1438 TensorDtype::Q4TiledP => {
1439 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1440 }
1441 TensorDtype::Q2TiledP => {
1442 let affine =
1443 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1444 if *rows * *cols >= 8_388_608
1445 && crate::gpu::enabled_here()
1446 && crate::gpu::q2tp_gpu_opt_in()
1447 {
1448 let gpu_ok = if affine {
1449 crate::gpu::q2tp_affine_matvec(
1450 model,
1451 *idx,
1452 &transformed,
1453 *rows,
1454 *cols,
1455 out,
1456 )
1457 } else {
1458 crate::gpu::q2tp_matvec(
1459 model,
1460 *idx,
1461 &transformed,
1462 *rows,
1463 *cols,
1464 out,
1465 )
1466 };
1467 if gpu_ok {
1468 return;
1469 }
1470 }
1471 if affine {
1472 q2tp_affine_matvec(
1473 self.quant_bytes(),
1474 &transformed,
1475 *rows,
1476 *cols,
1477 out,
1478 pool,
1479 )
1480 } else {
1481 q2tp_matvec(
1482 self.quant_bytes(),
1483 &transformed,
1484 *rows,
1485 *cols,
1486 out,
1487 pool,
1488 )
1489 }
1490 }
1491 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1492 self.bin_w().expect("binary dtype"),
1493 &transformed,
1494 *rows,
1495 *cols,
1496 out,
1497 pool,
1498 ),
1499 TensorDtype::Q1T => {
1500 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1501 }
1502 TensorDtype::Mxfp4 => {
1503 mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1504 }
1505 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1506 self.quant_bytes(),
1507 vbit_offsets,
1508 &transformed,
1509 *rows,
1510 *cols,
1511 out,
1512 pool,
1513 ),
1514 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1515 self.quant_bytes(),
1516 repack,
1517 row_scale,
1518 &transformed,
1519 col_field,
1520 *dtype,
1521 *rows,
1522 *cols,
1523 out,
1524 pool,
1525 ),
1526 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1527 }
1528 return;
1529 }
1530 if *dtype == TensorDtype::Q4Block {
1531 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1535 let t0 = std::time::Instant::now();
1536 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1537 crate::gpu::ProbeArm::Gpu => {
1538 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1539 crate::gpu::probe_record(
1540 crate::gpu::OpClass::Matvec,
1541 true,
1542 t0.elapsed(),
1543 );
1544 return;
1545 }
1546 }
1547 crate::gpu::ProbeArm::CpuTimed => {
1548 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1549 crate::gpu::probe_record(
1550 crate::gpu::OpClass::Matvec,
1551 false,
1552 t0.elapsed(),
1553 );
1554 return;
1555 }
1556 crate::gpu::ProbeArm::Cpu => {}
1557 }
1558 }
1559 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1560 return;
1561 }
1562 if *dtype == TensorDtype::Q4Tiled {
1563 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1568 let t0 = std::time::Instant::now();
1569 let cls = crate::gpu::matvec_class(*rows, *cols);
1570 match crate::gpu::probe_arm(cls) {
1571 crate::gpu::ProbeArm::Gpu => {
1572 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1573 crate::gpu::probe_record(cls, true, t0.elapsed());
1574 return;
1575 }
1576 }
1577 crate::gpu::ProbeArm::CpuTimed => {
1578 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1579 crate::gpu::probe_record(cls, false, t0.elapsed());
1580 return;
1581 }
1582 crate::gpu::ProbeArm::Cpu => {}
1583 }
1584 }
1585 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1586 return;
1587 }
1588 if *dtype == TensorDtype::Q4TiledP {
1589 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1595 let t0 = std::time::Instant::now();
1596 let cls = crate::gpu::matvec_class(*rows, *cols);
1597 match crate::gpu::probe_arm(cls) {
1598 crate::gpu::ProbeArm::Gpu => {
1599 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1600 crate::gpu::probe_record(cls, true, t0.elapsed());
1601 return;
1602 }
1603 }
1604 crate::gpu::ProbeArm::CpuTimed => {
1605 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1606 crate::gpu::probe_record(cls, false, t0.elapsed());
1607 return;
1608 }
1609 crate::gpu::ProbeArm::Cpu => {}
1610 }
1611 }
1612 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1613 return;
1614 }
1615 if *dtype == TensorDtype::Q2TiledP {
1616 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1617 return;
1618 }
1619 if *dtype == TensorDtype::Mxfp4 {
1620 self.mxfp4_matvec(x, out, pool);
1621 return;
1622 }
1623 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1624 let w = self.bin_w().expect("binary dtype");
1629 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1630 let t0 = std::time::Instant::now();
1631 let arm = if crate::gpu::q1_force() {
1632 crate::gpu::ProbeArm::Gpu
1633 } else {
1634 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1635 };
1636 match arm {
1637 crate::gpu::ProbeArm::Gpu => {
1638 let ok = if *dtype == TensorDtype::Q1G {
1639 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1640 } else {
1641 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1642 };
1643 if ok {
1644 crate::gpu::probe_record(
1645 crate::gpu::OpClass::Matvec,
1646 true,
1647 t0.elapsed(),
1648 );
1649 return;
1650 }
1651 }
1652 crate::gpu::ProbeArm::CpuTimed => {
1653 bin_matvec(w, x, *rows, *cols, out, pool);
1654 crate::gpu::probe_record(
1655 crate::gpu::OpClass::Matvec,
1656 false,
1657 t0.elapsed(),
1658 );
1659 return;
1660 }
1661 crate::gpu::ProbeArm::Cpu => {}
1662 }
1663 }
1664 bin_matvec(w, x, *rows, *cols, out, pool);
1665 return;
1666 }
1667 if *dtype == TensorDtype::Q1T {
1668 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1672 let t0 = std::time::Instant::now();
1673 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1674 crate::gpu::ProbeArm::Gpu => {
1675 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1676 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1677 crate::gpu::probe_record(
1678 crate::gpu::OpClass::Matvec,
1679 true,
1680 t0.elapsed(),
1681 );
1682 return;
1683 }
1684 }
1685 crate::gpu::ProbeArm::CpuTimed => {
1686 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1687 crate::gpu::probe_record(
1688 crate::gpu::OpClass::Matvec,
1689 false,
1690 t0.elapsed(),
1691 );
1692 return;
1693 }
1694 crate::gpu::ProbeArm::Cpu => {}
1695 }
1696 }
1697 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1698 return;
1699 }
1700 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1701 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1702 return;
1703 }
1704 let xs = prescale(x, col_field, *dtype);
1705 if *rows >= crate::gpu::min_rows()
1710 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1711 && gpu_lmhead_enabled()
1712 && crate::gpu::enabled_here()
1713 {
1714 let t0 = std::time::Instant::now();
1717 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1718 crate::gpu::ProbeArm::Gpu => {}
1719 crate::gpu::ProbeArm::CpuTimed => {
1720 qmatvec(
1721 self.quant_bytes(),
1722 repack,
1723 row_scale,
1724 x,
1725 col_field,
1726 *dtype,
1727 *rows,
1728 *cols,
1729 out,
1730 pool,
1731 );
1732 crate::gpu::probe_record(
1733 crate::gpu::OpClass::Matvec,
1734 false,
1735 t0.elapsed(),
1736 );
1737 return;
1738 }
1739 crate::gpu::ProbeArm::Cpu => {
1740 qmatvec(
1741 self.quant_bytes(),
1742 repack,
1743 row_scale,
1744 x,
1745 col_field,
1746 *dtype,
1747 *rows,
1748 *cols,
1749 out,
1750 pool,
1751 );
1752 return;
1753 }
1754 }
1755 let frac = gpu_split_frac();
1756 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1757 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1758 let bytes = self.quant_bytes();
1759 let ok = std::thread::scope(|sc| {
1760 let g = sc.spawn(|| {
1761 crate::gpu::q8_matvec_range(
1762 model,
1763 *idx,
1764 cpu_rows,
1765 &row_scale[cpu_rows..],
1766 &xs,
1767 *rows - cpu_rows,
1768 *cols,
1769 out_gpu,
1770 )
1771 });
1772 if cpu_rows > 0 {
1773 let rep_cpu = if repack.is_empty() {
1776 &[][..]
1777 } else {
1778 &repack[..(cpu_rows / 4) * 4 * *cols]
1779 };
1780 qmatvec(
1781 &bytes[..cpu_rows * *cols],
1782 rep_cpu,
1783 &row_scale[..cpu_rows],
1784 x,
1785 col_field,
1786 *dtype,
1787 cpu_rows,
1788 *cols,
1789 out_cpu,
1790 pool,
1791 );
1792 }
1793 g.join().unwrap_or(false)
1794 });
1795 if ok {
1796 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1797 return;
1798 }
1799 qmatvec(
1802 &bytes[cpu_rows * *cols..(*rows) * *cols],
1803 &[],
1804 &row_scale[cpu_rows..],
1805 x,
1806 col_field,
1807 *dtype,
1808 *rows - cpu_rows,
1809 *cols,
1810 out_gpu,
1811 pool,
1812 );
1813 return;
1814 }
1815 qmatvec(
1816 self.quant_bytes(),
1817 repack,
1818 row_scale,
1819 x,
1820 col_field,
1821 *dtype,
1822 *rows,
1823 *cols,
1824 out,
1825 pool,
1826 );
1827 }
1828 }
1829 }
1830
1831 pub fn matvec2(
1833 &self,
1834 x1: &[f32],
1835 x2: &[f32],
1836 o1: &mut [f32],
1837 o2: &mut [f32],
1838 pool: Option<&Pool>,
1839 ) {
1840 if self.flint_exact_cpu_default() {
1841 return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1842 }
1843 match self {
1844 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1845 Self::Blr(_) => {
1846 self.matvec(x1, o1, pool);
1847 self.matvec(x2, o2, pool);
1848 }
1849 Self::Mapped {
1850 model,
1851 idx,
1852 dtype,
1853 rows,
1854 cols,
1855 row_scale,
1856 col_field,
1857 vbit_offsets,
1858 ..
1859 } => {
1860 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1861 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1862 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1863 match dtype {
1864 TensorDtype::Q4Block => {
1865 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1866 }
1867 TensorDtype::Q4Tiled => {
1868 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1869 }
1870 TensorDtype::Q4TiledP => {
1871 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1872 }
1873 TensorDtype::Q2TiledP => {
1874 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1875 q2tp_affine_matvec2(
1876 self.quant_bytes(),
1877 &tx1,
1878 &tx2,
1879 *rows,
1880 *cols,
1881 o1,
1882 o2,
1883 pool,
1884 )
1885 } else {
1886 q2tp_matvec2(
1887 self.quant_bytes(),
1888 &tx1,
1889 &tx2,
1890 *rows,
1891 *cols,
1892 o1,
1893 o2,
1894 pool,
1895 )
1896 }
1897 }
1898 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1899 self.bin_w().expect("binary dtype"),
1900 &tx1,
1901 &tx2,
1902 *rows,
1903 *cols,
1904 o1,
1905 o2,
1906 pool,
1907 ),
1908 TensorDtype::Q1T => {
1909 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1910 }
1911 TensorDtype::Mxfp4 => mxfp4::matvec2(
1912 self.quant_bytes(),
1913 &tx1,
1914 &tx2,
1915 *rows,
1916 *cols,
1917 o1,
1918 o2,
1919 pool,
1920 ),
1921 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1922 self.quant_bytes(),
1923 vbit_offsets,
1924 &tx1,
1925 &tx2,
1926 *rows,
1927 *cols,
1928 o1,
1929 o2,
1930 pool,
1931 ),
1932 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1933 self.quant_bytes(),
1934 row_scale,
1935 &tx1,
1936 &tx2,
1937 col_field,
1938 *dtype,
1939 *rows,
1940 *cols,
1941 o1,
1942 o2,
1943 pool,
1944 ),
1945 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1946 }
1947 return;
1948 }
1949 if *dtype == TensorDtype::Q4Block {
1950 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1951 return;
1952 }
1953 if *dtype == TensorDtype::Q4Tiled {
1954 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1955 return;
1956 }
1957 if *dtype == TensorDtype::Q4TiledP {
1958 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1959 return;
1960 }
1961 if *dtype == TensorDtype::Q2TiledP {
1962 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1963 return;
1964 }
1965 if *dtype == TensorDtype::Mxfp4 {
1966 mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1967 return;
1968 }
1969 if let Some(w) = self.bin_w() {
1970 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1971 return;
1972 }
1973 if *dtype == TensorDtype::Q1T {
1974 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1980 return;
1981 }
1982 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1983 vbitmatvec2(
1984 self.quant_bytes(),
1985 vbit_offsets,
1986 x1,
1987 x2,
1988 *rows,
1989 *cols,
1990 o1,
1991 o2,
1992 pool,
1993 );
1994 return;
1995 }
1996 qmatvec2(
1997 self.quant_bytes(),
1998 row_scale,
1999 x1,
2000 x2,
2001 col_field,
2002 *dtype,
2003 *rows,
2004 *cols,
2005 o1,
2006 o2,
2007 pool,
2008 );
2009 }
2010 }
2011 }
2012}
2013
2014impl QTensor {
2015 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2023 if self.has_prism_contract() {
2024 return None;
2025 }
2026 match self {
2027 Self::Mapped {
2028 model, idx, dtype, ..
2029 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2030 _ => None,
2031 }
2032 }
2033
2034 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2035 if self.flint_exact_cpu_default() {
2036 return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2037 }
2038 let cols = self.cols();
2039 let rows = self.rows();
2040 debug_assert_eq!(xs_all.len(), b * cols);
2041 debug_assert_eq!(out.len(), b * rows);
2042 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2043 if crate::gptq_capture::capturing() {
2047 if let Self::Mapped { model, idx, .. } = self {
2048 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2049 }
2050 if let Self::Blr(t) = self {
2051 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2052 }
2053 }
2054 match self {
2055 Self::Blr(t) => {
2056 if crate::gpu::enabled_here()
2057 && crate::blr::gpu_wanted()
2058 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2059 {
2060 return;
2061 }
2062 t.matmat(xs_all, b, out, pool)
2063 }
2064 Self::F32 { data, .. } => {
2065 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2066 return;
2067 }
2068 let out_addr = SendMut(out.as_mut_ptr());
2069 let run = |start: usize, end: usize| {
2076 let mut bi = start;
2077 while bi < end {
2078 let n = (end - bi).min(4);
2079 for o in 0..rows {
2080 let row = &data[o * cols..(o + 1) * cols];
2081 if n == 4 {
2082 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2083 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2084 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2085 for j in 0..cols {
2086 let w = row[j];
2087 a0 += w * x0[j];
2088 a1 += w * x1[j];
2089 a2 += w * x2[j];
2090 a3 += w * x3[j];
2091 }
2092 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2093 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2094 }
2095 } else {
2096 for k in 0..n {
2097 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2098 let mut acc = 0f32;
2099 for j in 0..cols {
2100 acc += row[j] * x[j];
2101 }
2102 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2103 }
2104 }
2105 }
2106 bi += n;
2107 }
2108 };
2109 dispatch_rows(pool, b, &run);
2110 }
2111 Self::Mapped {
2112 model,
2113 idx,
2114 dtype,
2115 row_scale,
2116 col_field,
2117 vbit_offsets,
2118 ..
2119 } => {
2120 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2121 let mut transformed = Vec::with_capacity(xs_all.len());
2122 for bi in 0..b {
2123 transformed.extend_from_slice(&crate::prism::forward(
2124 model,
2125 &xs_all[bi * cols..(bi + 1) * cols],
2126 ));
2127 }
2128 match dtype {
2129 TensorDtype::Q4Block => {
2130 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2131 }
2132 TensorDtype::Q4Tiled => {
2133 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2134 }
2135 TensorDtype::Q4TiledP => {
2136 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2137 }
2138 TensorDtype::Q2TiledP => {
2139 let affine =
2140 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2141 let gpu_batch_ok = if affine {
2147 b >= 2
2148 } else {
2149 b >= 32 && b * rows * cols >= 128_000_000
2150 };
2151 if gpu_batch_ok
2152 && cols % 32 == 0
2153 && crate::gpu::enabled_here()
2154 && crate::gpu::q2tp_gpu_opt_in()
2155 {
2156 let gpu_ok = if affine {
2157 crate::gpu::q2tp_affine_matmat(
2158 model,
2159 *idx,
2160 &transformed,
2161 b,
2162 rows,
2163 cols,
2164 out,
2165 )
2166 } else {
2167 crate::gpu::q2tp_matmat(
2168 model,
2169 *idx,
2170 &transformed,
2171 b,
2172 rows,
2173 cols,
2174 out,
2175 )
2176 };
2177 if gpu_ok {
2178 return;
2179 }
2180 }
2181 if affine
2185 && b == 1
2186 && cols % 32 == 0
2187 && crate::gpu::enabled_here()
2188 && crate::gpu::q2tp_gpu_opt_in()
2189 && crate::gpu::q2tp_affine_matvec(
2190 model,
2191 *idx,
2192 &transformed[..cols],
2193 rows,
2194 cols,
2195 &mut out[..rows],
2196 )
2197 {
2198 return;
2199 }
2200 if affine {
2201 q2tp_affine_matmat(
2202 self.quant_bytes(),
2203 &transformed,
2204 b,
2205 rows,
2206 cols,
2207 out,
2208 pool,
2209 )
2210 } else {
2211 q2tp_matmat(
2212 self.quant_bytes(),
2213 &transformed,
2214 b,
2215 rows,
2216 cols,
2217 out,
2218 pool,
2219 )
2220 }
2221 }
2222 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2223 self.bin_w().expect("binary dtype"),
2224 &transformed,
2225 b,
2226 rows,
2227 cols,
2228 out,
2229 pool,
2230 ),
2231 TensorDtype::Q1T => {
2232 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2233 }
2234 TensorDtype::Mxfp4 => {
2235 mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2236 }
2237 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2238 self.quant_bytes(),
2239 vbit_offsets,
2240 &transformed,
2241 b,
2242 rows,
2243 cols,
2244 out,
2245 pool,
2246 ),
2247 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2248 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2249 .map(|bi| {
2250 prescale(
2251 &transformed[bi * cols..(bi + 1) * cols],
2252 col_field,
2253 *dtype,
2254 )
2255 })
2256 .collect();
2257 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2258 }
2259 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2260 }
2261 return;
2262 }
2263 if *dtype == TensorDtype::Q4Block {
2264 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2265 return;
2266 }
2267 if *dtype == TensorDtype::Mxfp4 {
2268 self.mxfp4_matmat(xs_all, b, out, pool);
2269 return;
2270 }
2271 if *dtype == TensorDtype::Q4TiledP {
2272 if b >= 32
2286 && b * rows * cols >= 128_000_000
2287 && cols % 32 == 0
2288 && !row_exact()
2289 && !crate::gpu::mm_killed()
2290 && crate::gpu::enabled_here()
2291 {
2292 let class = if b >= 128 {
2293 crate::gpu::OpClass::MatmatWide
2294 } else {
2295 crate::gpu::OpClass::Matmat
2296 };
2297 if let Self::Mapped { model, idx, .. } = self {
2298 if crate::mm_ab::on() {
2310 let mut g = vec![0f32; b * rows];
2311 let t = std::time::Instant::now();
2312 let took = crate::gpu::q4tp_matmat(
2313 model, *idx, xs_all, b, rows, cols, &mut g,
2314 );
2315 let dg = t.elapsed();
2316 let t = std::time::Instant::now();
2317 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2318 let dc = t.elapsed();
2319 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2320 return;
2321 }
2322 let t0 = std::time::Instant::now();
2323 let resident = crate::gpu::weight_is_resident(model, *idx);
2327 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2328 crate::gpu::ProbeArm::Gpu => {
2329 if crate::gpu::q4tp_matmat(
2330 model, *idx, xs_all, b, rows, cols, out,
2331 ) {
2332 let el = t0.elapsed();
2333 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2343 let budget = std::time::Duration::from_secs_f64(
2344 flops / 1.5e12 * 8.0 + 0.020,
2345 );
2346 crate::gpu::mm_budget_check(
2347 "q4tp matmat",
2348 el,
2349 budget,
2350 crate::gpu::probe_was_cold() || !resident,
2351 );
2352 crate::gpu::probe_record(class, true, el);
2353 return;
2354 }
2355 }
2356 crate::gpu::ProbeArm::CpuTimed => {
2357 q4tp_matmat(
2358 self.quant_bytes(),
2359 xs_all,
2360 b,
2361 rows,
2362 cols,
2363 out,
2364 pool,
2365 );
2366 crate::gpu::probe_record(class, false, t0.elapsed());
2367 return;
2368 }
2369 crate::gpu::ProbeArm::Cpu => {}
2370 }
2371 }
2372 }
2373 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2374 return;
2375 }
2376 if *dtype == TensorDtype::Q2TiledP {
2377 if b >= 32
2383 && b * rows * cols >= 128_000_000
2384 && cols % 32 == 0
2385 && !crate::gpu::mm_killed()
2386 && crate::gpu::enabled_here()
2387 {
2388 let class = if b >= 128 {
2389 crate::gpu::OpClass::MatmatWide
2390 } else {
2391 crate::gpu::OpClass::Matmat
2392 };
2393 if let Self::Mapped { model, idx, .. } = self {
2394 let t0 = std::time::Instant::now();
2395 match crate::gpu::probe_arm(class) {
2396 crate::gpu::ProbeArm::Gpu => {
2397 if crate::gpu::q2tp_matmat(
2398 model, *idx, xs_all, b, rows, cols, out,
2399 ) {
2400 crate::gpu::probe_record(class, true, t0.elapsed());
2401 return;
2402 }
2403 }
2404 crate::gpu::ProbeArm::CpuTimed => {
2405 q2tp_matmat(
2406 self.quant_bytes(),
2407 xs_all,
2408 b,
2409 rows,
2410 cols,
2411 out,
2412 pool,
2413 );
2414 crate::gpu::probe_record(class, false, t0.elapsed());
2415 return;
2416 }
2417 crate::gpu::ProbeArm::Cpu => {}
2418 }
2419 }
2420 }
2421 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2426 return;
2427 }
2428 if *dtype == TensorDtype::Q4Tiled {
2429 if b >= 32
2441 && b * rows * cols >= 128_000_000
2442 && cols % 32 == 0
2443 && !crate::gpu::mm_killed()
2444 && crate::gpu::enabled_here()
2445 {
2446 let class = if b >= 128 {
2447 crate::gpu::OpClass::MatmatWide
2448 } else {
2449 crate::gpu::OpClass::Matmat
2450 };
2451 if let Self::Mapped { model, idx, .. } = self {
2452 let t0 = std::time::Instant::now();
2453 match crate::gpu::probe_arm(class) {
2454 crate::gpu::ProbeArm::Gpu => {
2455 if crate::gpu::q4t_matmat(
2456 model, *idx, xs_all, b, rows, cols, out,
2457 ) {
2458 let el = t0.elapsed();
2459 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2469 let budget = std::time::Duration::from_secs_f64(
2470 flops / 1.5e12 * 8.0 + 0.020,
2471 );
2472 crate::gpu::mm_budget_check(
2473 "q4t matmat",
2474 el,
2475 budget,
2476 crate::gpu::probe_was_cold(),
2477 );
2478 crate::gpu::probe_record(class, true, el);
2479 return;
2480 }
2481 }
2482 crate::gpu::ProbeArm::CpuTimed => {
2483 q4t_matmat(
2484 self.quant_bytes(),
2485 xs_all,
2486 b,
2487 rows,
2488 cols,
2489 out,
2490 pool,
2491 );
2492 crate::gpu::probe_record(class, false, t0.elapsed());
2493 return;
2494 }
2495 crate::gpu::ProbeArm::Cpu => {}
2496 }
2497 }
2498 }
2499 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2500 return;
2501 }
2502 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2503 let w = self.bin_w().expect("binary dtype");
2507 if b >= 32
2508 && b * rows * cols >= 128_000_000
2509 && cols % 64 == 0
2510 && crate::gpu::enabled_here()
2511 {
2512 if let Self::Mapped { model, idx, .. } = self {
2513 let t0 = std::time::Instant::now();
2514 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2515 crate::gpu::ProbeArm::Gpu => {
2516 let ok = if *dtype == TensorDtype::Q1G {
2517 crate::gpu::q1g_matmat(
2518 model, *idx, xs_all, b, rows, cols, out,
2519 )
2520 } else {
2521 crate::gpu::q1_matmat(
2522 model, *idx, xs_all, b, rows, cols, out,
2523 )
2524 };
2525 if ok {
2526 crate::gpu::probe_record(
2527 crate::gpu::OpClass::Matmat,
2528 true,
2529 t0.elapsed(),
2530 );
2531 return;
2532 }
2533 }
2534 crate::gpu::ProbeArm::CpuTimed => {
2535 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2536 crate::gpu::probe_record(
2537 crate::gpu::OpClass::Matmat,
2538 false,
2539 t0.elapsed(),
2540 );
2541 return;
2542 }
2543 crate::gpu::ProbeArm::Cpu => {}
2544 }
2545 }
2546 }
2547 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2548 return;
2549 }
2550 if *dtype == TensorDtype::Q1T {
2551 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2554 if let Self::Mapped { model, idx, .. } = self {
2555 let t0 = std::time::Instant::now();
2556 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2557 crate::gpu::ProbeArm::Gpu => {
2558 if crate::gpu::q1t_matmat(
2559 model, *idx, xs_all, b, rows, cols, out,
2560 ) {
2561 crate::gpu::probe_record(
2562 crate::gpu::OpClass::Matmat,
2563 true,
2564 t0.elapsed(),
2565 );
2566 return;
2567 }
2568 }
2569 crate::gpu::ProbeArm::CpuTimed => {
2570 q1t_matmat(
2571 self.quant_bytes(),
2572 xs_all,
2573 b,
2574 rows,
2575 cols,
2576 out,
2577 pool,
2578 );
2579 crate::gpu::probe_record(
2580 crate::gpu::OpClass::Matmat,
2581 false,
2582 t0.elapsed(),
2583 );
2584 return;
2585 }
2586 crate::gpu::ProbeArm::Cpu => {}
2587 }
2588 }
2589 }
2590 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2591 return;
2592 }
2593 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2594 vbitmatmat(
2595 self.quant_bytes(),
2596 vbit_offsets,
2597 xs_all,
2598 b,
2599 rows,
2600 cols,
2601 out,
2602 pool,
2603 );
2604 return;
2605 }
2606 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2607 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2608 .collect();
2609 if row_exact()
2615 && (1..=4).contains(&b)
2616 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2617 && crate::gpu::enabled_here()
2618 && crate::gpu::wgpu_active()
2619 {
2620 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2621 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2622 return;
2623 }
2624 }
2625 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2631 if let Self::Mapped { model, idx, .. } = self {
2632 let t0 = std::time::Instant::now();
2633 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2634 crate::gpu::ProbeArm::Gpu
2635 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2636 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2637 {
2638 let q = self.quant_bytes();
2642 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2643 return;
2644 }
2645 crate::gpu::ProbeArm::Gpu => {
2646 if *dtype == TensorDtype::Q8_2f
2654 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2655 && crate::gpu::q8_matmat_2f(
2656 model, *idx, row_scale, col_field, xs_all, b, rows,
2657 cols, out,
2658 )
2659 {
2660 crate::gpu::probe_record(
2661 crate::gpu::OpClass::Matmat,
2662 true,
2663 t0.elapsed(),
2664 );
2665 return;
2666 }
2667 let flat: Vec<f32> =
2668 pre.iter().flat_map(|v| v.iter().copied()).collect();
2669 if crate::gpu::q8_matmat(
2670 model, *idx, row_scale, &flat, b, rows, cols, out,
2671 ) {
2672 crate::gpu::probe_record(
2673 crate::gpu::OpClass::Matmat,
2674 true,
2675 t0.elapsed(),
2676 );
2677 return;
2678 }
2679 }
2680 crate::gpu::ProbeArm::CpuTimed => {
2681 let q = self.quant_bytes();
2682 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2683 crate::gpu::probe_record(
2684 crate::gpu::OpClass::Matmat,
2685 false,
2686 t0.elapsed(),
2687 );
2688 return;
2689 }
2690 crate::gpu::ProbeArm::Cpu => {}
2691 }
2692 }
2693 }
2694 let q = self.quant_bytes();
2695 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2696 }
2697 }
2698 }
2699}
2700
2701impl QTensor {
2702 fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2706 let (rows, cols) = (self.rows(), self.cols());
2707 if let Self::Mapped { model, idx, .. } = self {
2708 if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2709 let t0 = std::time::Instant::now();
2710 let cls = crate::gpu::matvec_class(rows, cols);
2711 match crate::gpu::probe_arm(cls) {
2712 crate::gpu::ProbeArm::Gpu => {
2713 if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2714 crate::gpu::probe_record(cls, true, t0.elapsed());
2715 return;
2716 }
2717 }
2718 crate::gpu::ProbeArm::CpuTimed => {
2719 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2720 crate::gpu::probe_record(cls, false, t0.elapsed());
2721 return;
2722 }
2723 crate::gpu::ProbeArm::Cpu => {}
2724 }
2725 }
2726 }
2727 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2728 }
2729
2730 fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2735 let (rows, cols) = (self.rows(), self.cols());
2736 if let Self::Mapped { model, idx, .. } = self {
2737 if b >= 2
2738 && b * rows * cols >= 33_554_432
2739 && !row_exact()
2740 && !crate::gpu::mm_killed()
2741 && crate::gpu::enabled_here()
2742 {
2743 let class = if b >= 128 {
2744 crate::gpu::OpClass::MatmatWide
2745 } else {
2746 crate::gpu::OpClass::Matmat
2747 };
2748 let t0 = std::time::Instant::now();
2749 match crate::gpu::probe_arm(class) {
2750 crate::gpu::ProbeArm::Gpu => {
2751 if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2752 crate::gpu::probe_record(class, true, t0.elapsed());
2753 return;
2754 }
2755 }
2756 crate::gpu::ProbeArm::CpuTimed => {
2757 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2758 crate::gpu::probe_record(class, false, t0.elapsed());
2759 return;
2760 }
2761 crate::gpu::ProbeArm::Cpu => {}
2762 }
2763 }
2764 }
2765 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2766 }
2767}
2768
2769impl QTensor {
2770 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2780 let (rows, cols) = (self.rows(), self.cols());
2781 let Self::Mapped {
2782 model,
2783 idx,
2784 dtype,
2785 row_scale,
2786 col_field,
2787 ..
2788 } = self
2789 else {
2790 return false;
2791 };
2792 if crate::prism::has_contract(model) {
2793 return false;
2794 }
2795 match *dtype {
2796 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2797 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2801 if *dtype == TensorDtype::Q8_2f
2804 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2805 && crate::gpu::q8_matmat_2f(
2806 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2807 )
2808 {
2809 return true;
2810 }
2811 let flat: Vec<f32> = (0..b)
2812 .flat_map(|bi| {
2813 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2814 })
2815 .collect();
2816 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2817 }
2818 _ => false,
2819 }
2820 }
2821
2822 pub fn matvec_many<const N: usize>(
2829 ts: [&QTensor; N],
2830 x: &[f32],
2831 mut outs: [&mut [f32]; N],
2832 pool: Option<&Pool>,
2833 ) {
2834 if Self::any_flint_exact_cpu_default(&ts) {
2835 return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2836 }
2837 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2838 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2839 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2844 t.matvec(x, o, pool);
2845 }
2846 return;
2847 }
2848 let uniform_q8 = ts.iter().all(|t| {
2849 matches!(
2850 t,
2851 Self::Mapped {
2852 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2853 ..
2854 }
2855 )
2856 });
2857 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2858 if uniform_f32 && crate::f32_backend::active() {
2859 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2860 t.matvec(x, o, pool);
2861 }
2862 return;
2863 }
2864 let uniform_q4 = ts.iter().all(|t| {
2865 matches!(
2866 t,
2867 Self::Mapped {
2868 dtype: TensorDtype::Q4Block,
2869 ..
2870 }
2871 )
2872 });
2873 let uniform_vbit = ts.iter().all(|t| {
2874 matches!(
2875 t,
2876 Self::Mapped {
2877 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2878 ..
2879 }
2880 )
2881 });
2882 let uniform_q1 =
2885 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2886 let uniform_q1t = ts.iter().all(|t| {
2887 matches!(
2888 t,
2889 Self::Mapped {
2890 dtype: TensorDtype::Q1T,
2891 ..
2892 }
2893 )
2894 });
2895 let uniform_q4tp = ts.iter().all(|t| {
2900 matches!(
2901 t,
2902 Self::Mapped {
2903 dtype: TensorDtype::Q4TiledP,
2904 ..
2905 }
2906 )
2907 }) && ts
2908 .iter()
2909 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2910 let Some(pool) = pool else {
2911 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2912 t.matvec(x, o, None);
2913 }
2914 return;
2915 };
2916 if total_rows < 256
2917 || !(uniform_q8
2918 || uniform_f32
2919 || uniform_q4
2920 || uniform_vbit
2921 || uniform_q1
2922 || uniform_q1t
2923 || uniform_q4tp)
2924 {
2925 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2926 t.matvec(x, o, Some(pool));
2927 }
2928 return;
2929 }
2930
2931 if uniform_q4tp {
2932 let cols = ts[0].cols();
2938 let gpr = cols / GROUP_SIZE;
2939 let views: [Q4tpView; N] =
2940 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2941 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2942 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2943 let locate = |flat: usize| -> (usize, usize) {
2945 let mut acc = 0;
2946 for (i, &r) in rows_of.iter().enumerate() {
2947 if flat < acc + r {
2948 return (i, flat - acc);
2949 }
2950 acc += r;
2951 }
2952 (rows_of.len() - 1, 0)
2953 };
2954 let (views, outs_addr) = (&views, &outs_addr);
2955 if a8w8_enabled() {
2956 let act = split_act(x);
2957 let act = &act;
2958 let run = |start: usize, end: usize| {
2959 with_krow(gpr, |sc| {
2960 for flat in start..end {
2961 let (t, r) = locate(flat);
2962 let v = &views[t];
2963 v.scales_into(r, gpr, sc);
2964 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2965 for &(j, xv) in &act.outliers {
2966 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2967 acc += w * s * xv;
2968 }
2969 unsafe { *outs_addr[t].at(r) = acc };
2971 }
2972 });
2973 };
2974 pool.run_rows(total_rows, &run);
2975 } else {
2976 let run = |start: usize, end: usize| {
2977 with_krow(gpr, |sc| {
2978 for flat in start..end {
2979 let (t, r) = locate(flat);
2980 let v = &views[t];
2981 v.scales_into(r, gpr, sc);
2982 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2984 }
2985 });
2986 };
2987 pool.run_rows(total_rows, &run);
2988 }
2989 return;
2990 }
2991
2992 if uniform_q1 {
2993 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2996 if a8w8_enabled() {
2997 let act = split_act(x);
2998 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2999 let (act, gsum) = (&act, &gsum);
3000 let closures: [_; N] = std::array::from_fn(|i| {
3001 let (w, gpr, out) = (
3002 ts[i].bin_w().expect("uniform_q1 checked"),
3003 ts[i].cols() / GROUP_SIZE,
3004 outs_addr[i],
3005 );
3006 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3007 });
3008 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3009 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3010 pool.run_many(&parts);
3011 } else {
3012 let closures: [_; N] = std::array::from_fn(|i| {
3013 let (w, gpr, out) = (
3014 ts[i].bin_w().expect("uniform_q1 checked"),
3015 ts[i].cols() / GROUP_SIZE,
3016 outs_addr[i],
3017 );
3018 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3019 });
3020 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3021 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3022 pool.run_many(&parts);
3023 }
3024 return;
3025 }
3026
3027 if uniform_q1t {
3028 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3032 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3033 if a8w8_enabled() {
3034 let act = split_act(x);
3035 let act = &act;
3036 let x_ref = x;
3037 let closures: [_; N] = std::array::from_fn(|i| {
3038 let bytes = ts[i].quant_bytes();
3039 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3040 let gpr = cols / GROUP_SIZE;
3041 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3042 let out = outs_addr[i];
3043 move |s: usize, e: usize| {
3044 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3045 }
3046 });
3047 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3048 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3049 pool.run_many(&parts);
3050 } else {
3051 let x_ref = x;
3052 let closures: [_; N] = std::array::from_fn(|i| {
3053 let bytes = ts[i].quant_bytes();
3054 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3055 let gpr = cols / GROUP_SIZE;
3056 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3057 let out = outs_addr[i];
3058 move |s: usize, e: usize| {
3059 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3060 }
3061 });
3062 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3063 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3064 pool.run_many(&parts);
3065 }
3066 return;
3067 }
3068
3069 if uniform_q4 || uniform_vbit {
3070 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3071 if a8w8_enabled() {
3073 let act = split_act(x);
3074 let act = &act;
3075 if uniform_q4 {
3076 let closures: [_; N] = std::array::from_fn(|i| {
3077 let (packed, scales) =
3078 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3079 let (gpr, cols, out) =
3080 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3081 move |s: usize, e: usize| {
3082 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3083 }
3084 });
3085 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3086 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3087 pool.run_many(&parts);
3088 } else {
3089 let closures: [_; N] = std::array::from_fn(|i| {
3090 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3091 unreachable!()
3092 };
3093 let (bytes, rows, cols, out) = (
3094 ts[i].quant_bytes(),
3095 ts[i].rows(),
3096 ts[i].cols(),
3097 outs_addr[i],
3098 );
3099 move |s: usize, e: usize| {
3100 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3101 }
3102 });
3103 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3104 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3105 pool.run_many(&parts);
3106 }
3107 return;
3108 }
3109 if uniform_q4 {
3110 let closures: [_; N] = std::array::from_fn(|i| {
3111 let (packed, scales) =
3112 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3113 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3114 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3115 });
3116 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3117 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3118 pool.run_many(&parts);
3119 } else {
3120 let closures: [_; N] = std::array::from_fn(|i| {
3121 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3122 unreachable!()
3123 };
3124 let (bytes, rows, cols, out) = (
3125 ts[i].quant_bytes(),
3126 ts[i].rows(),
3127 ts[i].cols(),
3128 outs_addr[i],
3129 );
3130 move |s: usize, e: usize| {
3131 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3132 }
3133 });
3134 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3135 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3136 pool.run_many(&parts);
3137 }
3138 return;
3139 }
3140
3141 if uniform_f32 {
3142 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3143 let closures: [_; N] = std::array::from_fn(|i| {
3144 let Self::F32 { data, cols, .. } = ts[i] else {
3145 unreachable!()
3146 };
3147 let out = outs_addr[i];
3148 move |start: usize, end: usize| {
3149 for o in start..end {
3150 let row = &data[o * cols..(o + 1) * cols];
3151 let mut sum = 0.0f32;
3152 for j in 0..*cols {
3153 sum += row[j] * x[j];
3154 }
3155 unsafe { *out.at(o) = sum };
3157 }
3158 }
3159 });
3160 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3161 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3162 pool.run_many(&parts);
3163 return;
3164 }
3165
3166 struct Ctx<'a> {
3169 bytes: &'a [u8],
3170 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3171 rep: &'a [u8],
3172 row_scale: &'a [f32],
3173 cols: usize,
3174 xs: std::borrow::Cow<'a, [f32]>,
3175 }
3176 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3177 let Self::Mapped {
3178 dtype,
3179 cols,
3180 row_scale,
3181 col_field,
3182 repack,
3183 ..
3184 } = ts[i]
3185 else {
3186 unreachable!()
3187 };
3188 Ctx {
3189 bytes: ts[i].quant_bytes(),
3190 rep: repack,
3191 row_scale,
3192 cols: *cols,
3193 xs: prescale(x, col_field, *dtype),
3194 }
3195 });
3196 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3197 #[cfg(target_arch = "aarch64")]
3198 if sdot_enabled() {
3199 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3200 let closures: [_; N] = std::array::from_fn(|i| {
3201 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3202 move |start: usize, end: usize| {
3203 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3204 }
3205 });
3206 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3207 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3208 pool.run_many(&parts);
3209 return;
3210 }
3211 #[cfg(target_arch = "x86_64")]
3212 if avx2_a8w8_enabled() {
3213 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3214 let closures: [_; N] = std::array::from_fn(|i| {
3215 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3216 move |start: usize, end: usize| {
3217 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3218 }
3219 });
3220 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3221 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3222 pool.run_many(&parts);
3223 return;
3224 }
3225 let closures: [_; N] = std::array::from_fn(|i| {
3226 let (c, out) = (&ctxs[i], outs_addr[i]);
3227 move |start: usize, end: usize| {
3228 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
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 }
3235}
3236
3237impl QTensor {
3238 #[allow(clippy::needless_range_loop)]
3243 pub fn matvec2_many<const N: usize>(
3244 ts: [&QTensor; N],
3245 x1: &[f32],
3246 x2: &[f32],
3247 mut o1s: [&mut [f32]; N],
3248 mut o2s: [&mut [f32]; N],
3249 pool: Option<&Pool>,
3250 ) {
3251 if Self::any_flint_exact_cpu_default(&ts) {
3252 return float_activations_scope(|| {
3253 Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3254 });
3255 }
3256 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3257 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3258 for i in 0..N {
3259 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3260 }
3261 return;
3262 }
3263 let uniform_q8 = ts.iter().all(|t| {
3264 matches!(
3265 t,
3266 Self::Mapped {
3267 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3268 ..
3269 }
3270 )
3271 });
3272 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3273 let uniform_q4 = ts.iter().all(|t| {
3274 matches!(
3275 t,
3276 Self::Mapped {
3277 dtype: TensorDtype::Q4Block,
3278 ..
3279 }
3280 )
3281 });
3282 let uniform_vbit = ts.iter().all(|t| {
3283 matches!(
3284 t,
3285 Self::Mapped {
3286 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3287 ..
3288 }
3289 )
3290 });
3291 let fusable = pool.is_some()
3292 && total_rows >= 256
3293 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3294 if !fusable {
3295 for i in 0..N {
3296 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3297 }
3298 return;
3299 }
3300 let pool = pool.unwrap();
3301
3302 if uniform_q4 || uniform_vbit {
3303 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3304 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3305 if a8w8_enabled() {
3307 let a1 = split_act(x1);
3308 let a2 = split_act(x2);
3309 let (a1, a2) = (&a1, &a2);
3310 if uniform_q4 {
3311 let closures: [_; N] = std::array::from_fn(|i| {
3312 let (packed, scales) =
3313 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3314 let (gpr, cols, o1, o2) =
3315 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3316 move |s: usize, e: usize| {
3317 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3318 }
3319 });
3320 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3321 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3322 pool.run_many(&parts);
3323 } else {
3324 let closures: [_; N] = std::array::from_fn(|i| {
3325 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3326 unreachable!()
3327 };
3328 let (bytes, rows, cols, o1, o2) = (
3329 ts[i].quant_bytes(),
3330 ts[i].rows(),
3331 ts[i].cols(),
3332 p1[i],
3333 p2[i],
3334 );
3335 move |s: usize, e: usize| {
3336 vbit_range2_a8w8(
3337 bytes,
3338 vbit_offsets,
3339 x1,
3340 x2,
3341 a1,
3342 a2,
3343 rows,
3344 cols,
3345 o1,
3346 o2,
3347 s,
3348 e,
3349 )
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 }
3356 return;
3357 }
3358 if uniform_q4 {
3359 let closures: [_; N] = std::array::from_fn(|i| {
3360 let (packed, scales) =
3361 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3362 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3363 move |s: usize, e: usize| {
3364 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3365 }
3366 });
3367 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3368 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3369 pool.run_many(&parts);
3370 } else {
3371 let closures: [_; N] = std::array::from_fn(|i| {
3372 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3373 unreachable!()
3374 };
3375 let (bytes, rows, cols, o1, o2) = (
3376 ts[i].quant_bytes(),
3377 ts[i].rows(),
3378 ts[i].cols(),
3379 p1[i],
3380 p2[i],
3381 );
3382 move |s: usize, e: usize| {
3383 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3384 }
3385 });
3386 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3387 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3388 pool.run_many(&parts);
3389 }
3390 return;
3391 }
3392
3393 if uniform_f32 {
3394 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3395 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3396 let closures: [_; N] = std::array::from_fn(|i| {
3397 let Self::F32 { data, cols, .. } = ts[i] else {
3398 unreachable!()
3399 };
3400 let (o1, o2) = (p1[i], p2[i]);
3401 move |start: usize, end: usize| {
3402 for o in start..end {
3403 let row = &data[o * cols..(o + 1) * cols];
3404 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3405 for j in 0..*cols {
3406 s1 += row[j] * x1[j];
3407 s2 += row[j] * x2[j];
3408 }
3409 unsafe {
3411 *o1.at(o) = s1;
3412 *o2.at(o) = s2;
3413 }
3414 }
3415 }
3416 });
3417 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3418 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3419 pool.run_many(&parts);
3420 return;
3421 }
3422
3423 struct Ctx<'a> {
3424 bytes: &'a [u8],
3425 row_scale: &'a [f32],
3426 cols: usize,
3427 xs1: std::borrow::Cow<'a, [f32]>,
3428 xs2: std::borrow::Cow<'a, [f32]>,
3429 }
3430 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3431 let Self::Mapped {
3432 dtype,
3433 cols,
3434 row_scale,
3435 col_field,
3436 ..
3437 } = ts[i]
3438 else {
3439 unreachable!()
3440 };
3441 Ctx {
3442 bytes: ts[i].quant_bytes(),
3443 row_scale,
3444 cols: *cols,
3445 xs1: prescale(x1, col_field, *dtype),
3446 xs2: prescale(x2, col_field, *dtype),
3447 }
3448 });
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 #[cfg(target_arch = "aarch64")]
3452 if sdot_enabled() {
3453 let acts: [(SplitAct, SplitAct); N] =
3454 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3455 let closures: [_; N] = std::array::from_fn(|i| {
3456 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3457 move |start: usize, end: usize| {
3458 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3459 }
3460 });
3461 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3462 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3463 pool.run_many(&parts);
3464 return;
3465 }
3466 #[cfg(target_arch = "x86_64")]
3467 if avx2_a8w8_enabled() {
3468 let acts: [(SplitAct, SplitAct); N] =
3469 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3470 let closures: [_; N] = std::array::from_fn(|i| {
3471 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3472 move |start: usize, end: usize| {
3473 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3474 }
3475 });
3476 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3477 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3478 pool.run_many(&parts);
3479 return;
3480 }
3481 let closures: [_; N] = std::array::from_fn(|i| {
3482 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3483 move |start: usize, end: usize| {
3484 q8_range2_f32(
3485 c.bytes,
3486 c.row_scale,
3487 &c.xs1,
3488 &c.xs2,
3489 c.cols,
3490 o1,
3491 o2,
3492 start,
3493 end,
3494 )
3495 }
3496 });
3497 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3498 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3499 pool.run_many(&parts);
3500 }
3501
3502 pub fn matvec_silu_mul(
3507 gate: &QTensor,
3508 up: &QTensor,
3509 x: &[f32],
3510 out: &mut [f32],
3511 pool: Option<&Pool>,
3512 ) -> bool {
3513 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3514 }
3515
3516 pub fn matvec_silu_mul_limited(
3522 gate: &QTensor,
3523 up: &QTensor,
3524 x: &[f32],
3525 out: &mut [f32],
3526 limit: f32,
3527 pool: Option<&Pool>,
3528 ) -> bool {
3529 if gate.flint_exact_cpu_default() || up.flint_exact_cpu_default() {
3530 return float_activations_scope(|| {
3531 Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3532 });
3533 }
3534 if gate.has_prism_contract()
3535 || up.has_prism_contract()
3536 || gate.per_tensor_only()
3537 || up.per_tensor_only()
3538 {
3539 return false;
3543 }
3544 let inter = gate.rows();
3545 debug_assert_eq!(up.rows(), inter);
3546 debug_assert_eq!(out.len(), inter);
3547 debug_assert_eq!(gate.cols(), up.cols());
3548 if !a8w8_enabled() {
3549 return false;
3550 }
3551 let act = split_act(x);
3552 let act = &act;
3553 let x_ref = x;
3554 let out_addr = SendMut(out.as_mut_ptr());
3555
3556 match (gate, up) {
3557 (
3559 Self::Mapped {
3560 dtype: TensorDtype::Q4Block,
3561 ..
3562 },
3563 Self::Mapped {
3564 dtype: TensorDtype::Q4Block,
3565 ..
3566 },
3567 ) => {
3568 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3569 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3570 let gpr = gate.cols() / GROUP_SIZE;
3571 let cols = gate.cols();
3572 let run = move |start: usize, end: usize| {
3573 for r in start..end {
3574 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3575 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3576 for &(j, xv) in &act.outliers {
3577 let flat = r * cols + j;
3578 let gb = gp[flat / 2];
3579 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3580 let gsc = f16_to_f32(u16::from_le_bytes([
3581 gs[(flat / GROUP_SIZE) * 2],
3582 gs[(flat / GROUP_SIZE) * 2 + 1],
3583 ]));
3584 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3585 let ub = up_p[flat / 2];
3586 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3587 let usc = f16_to_f32(u16::from_le_bytes([
3588 up_s[(flat / GROUP_SIZE) * 2],
3589 up_s[(flat / GROUP_SIZE) * 2 + 1],
3590 ]));
3591 uv += ((un as i32 - 8) as f32) * usc * xv;
3592 }
3593 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3595 }
3596 };
3597 dispatch_rows(pool, inter, &run);
3598 true
3599 }
3600 (
3604 Self::Mapped {
3605 dtype: TensorDtype::Q4Tiled,
3606 ..
3607 },
3608 Self::Mapped {
3609 dtype: TensorDtype::Q4Tiled,
3610 ..
3611 },
3612 ) => {
3613 let g_bytes = gate.quant_bytes();
3614 let u_bytes = up.quant_bytes();
3615 let gpr = gate.cols() / GROUP_SIZE;
3616 let run = move |start: usize, end: usize| {
3617 for r in start..end {
3618 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3619 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3620 for &(j, xv) in &act.outliers {
3621 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3622 gv += w * s * xv;
3623 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3624 uv += w * s * xv;
3625 }
3626 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3628 }
3629 };
3630 dispatch_rows(pool, inter, &run);
3631 true
3632 }
3633 (
3636 Self::Mapped {
3637 dtype: TensorDtype::Q4TiledP,
3638 ..
3639 },
3640 Self::Mapped {
3641 dtype: TensorDtype::Q4TiledP,
3642 ..
3643 },
3644 ) => {
3645 let cols = gate.cols();
3646 let gpr = cols / GROUP_SIZE;
3647 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3648 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3649 let run = |start: usize, end: usize| {
3650 with_krows(gpr, |gsc, usc| {
3651 for r in start..end {
3652 gv_view.scales_into(r, gpr, gsc);
3653 uv_view.scales_into(r, gpr, usc);
3654 let mut gv =
3655 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3656 let mut uv =
3657 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3658 for &(j, xv) in &act.outliers {
3659 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3660 gv += w * s * xv;
3661 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3662 uv += w * s * xv;
3663 }
3664 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3666 }
3667 });
3668 };
3669 dispatch_rows(pool, inter, &run);
3670 true
3671 }
3672 (
3679 Self::Mapped {
3680 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3681 ..
3682 },
3683 Self::Mapped {
3684 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3685 ..
3686 },
3687 ) => {
3688 let gw = gate.bin_w().expect("binary dtype matched");
3689 let uw = up.bin_w().expect("binary dtype matched");
3690 let gpr = gate.cols() / GROUP_SIZE;
3691 let gsum = q1_group_sums(&act.xq, gpr);
3692 let gsum = &gsum;
3693 let run = move |start: usize, end: usize| {
3694 for r in start..end {
3695 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3696 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3697 for &(j, xv) in &act.outliers {
3698 let (w, s) = gw.outlier(r, gpr, j);
3699 gv += w * s * xv;
3700 let (w, s) = uw.outlier(r, gpr, j);
3701 uv += w * s * xv;
3702 }
3703 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3705 }
3706 };
3707 dispatch_rows(pool, inter, &run);
3708 true
3709 }
3710 (
3714 Self::Mapped {
3715 dtype: TensorDtype::Q2TiledP,
3716 ..
3717 },
3718 Self::Mapped {
3719 dtype: TensorDtype::Q2TiledP,
3720 ..
3721 },
3722 ) => {
3723 let cols = gate.cols();
3724 let gpr = cols / GROUP_SIZE;
3725 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3726 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3727 let gsum = q1_group_sums(&act.xq, gpr);
3728 let gsum = &gsum;
3729 let run = move |start: usize, end: usize| {
3730 with_krows(gpr, |gsc, usc| {
3731 for r in start..end {
3732 gv_view.scales_into(r, gpr, gsc);
3733 uv_view.scales_into(r, gpr, usc);
3734 let mut gv =
3735 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3736 let mut uv =
3737 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3738 for &(j, xv) in &act.outliers {
3739 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3740 gv += w * s * xv;
3741 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3742 uv += w * s * xv;
3743 }
3744 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3746 }
3747 });
3748 };
3749 dispatch_rows(pool, inter, &run);
3750 true
3751 }
3752 (
3757 Self::Mapped {
3758 dtype: TensorDtype::Q8Row,
3759 row_scale: g_rs,
3760 ..
3761 },
3762 Self::Mapped {
3763 dtype: TensorDtype::Q8Row,
3764 row_scale: u_rs,
3765 ..
3766 },
3767 ) => {
3768 let g_bytes = gate.quant_bytes();
3769 let u_bytes = up.quant_bytes();
3770 let cols = gate.cols();
3771 let run = move |start: usize, end: usize| {
3772 for r in start..end {
3773 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3774 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3775 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3777 }
3778 };
3779 dispatch_rows(pool, inter, &run);
3780 true
3781 }
3782 (
3784 Self::Mapped {
3785 dtype: TensorDtype::Q1T,
3786 ..
3787 },
3788 Self::Mapped {
3789 dtype: TensorDtype::Q1T,
3790 ..
3791 },
3792 ) => {
3793 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3794 let g_bytes = gate.quant_bytes();
3795 let u_bytes = up.quant_bytes();
3796 let gpr = gate.cols() / GROUP_SIZE;
3797 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3798 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3799 let run = move |start: usize, end: usize| {
3800 for r in start..end {
3801 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3802 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3803 for &(j, xv) in &act.outliers {
3804 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3805 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3806 }
3807 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3808 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3809 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3811 }
3812 };
3813 dispatch_rows(pool, inter, &run);
3814 true
3815 }
3816 _ => false,
3817 }
3818 }
3819
3820 pub fn moe_gate_up_many(
3835 pairs: &[(&QTensor, &QTensor)],
3836 x: &[f32],
3837 outs: &mut [Vec<f32>],
3838 pool: Option<&Pool>,
3839 ) -> bool {
3840 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3841 }
3842
3843 pub fn moe_gate_up_many_limited(
3847 pairs: &[(&QTensor, &QTensor)],
3848 x: &[f32],
3849 outs: &mut [Vec<f32>],
3850 limit: f32,
3851 pool: Option<&Pool>,
3852 ) -> bool {
3853 if Self::any_flint_exact_cpu_default_pairs(pairs) {
3854 return float_activations_scope(|| {
3855 Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
3856 });
3857 }
3858 if pairs.is_empty() || pairs.len() != outs.len() {
3859 return false;
3860 }
3861 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3862 return false;
3863 }
3864 if !a8w8_enabled() {
3865 if limit > 0.0 {
3866 return false;
3869 }
3870 let groups = vec![vec![0]; pairs.len()];
3871 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3872 }
3873 let inter = pairs[0].0.rows();
3874 let cols = pairs[0].0.cols();
3875 if cols % GROUP_SIZE != 0 {
3876 return false;
3877 }
3878 let gpr = cols / GROUP_SIZE;
3879 let q2 = matches!(
3882 pairs[0].0,
3883 Self::Mapped {
3884 dtype: TensorDtype::Q2TiledP,
3885 ..
3886 }
3887 );
3888 let want = if q2 {
3889 TensorDtype::Q2TiledP
3890 } else {
3891 TensorDtype::Q4TiledP
3892 };
3893 let mut views = Vec::with_capacity(pairs.len() * 2);
3894 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3895 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3896 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3897 if !both
3898 || g.rows() != inter
3899 || u.rows() != inter
3900 || g.cols() != cols
3901 || u.cols() != cols
3902 || o.len() != inter
3903 {
3904 return false;
3905 }
3906 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3907 views.push(mk(g.quant_bytes(), inter, cols));
3908 views.push(mk(u.quant_bytes(), inter, cols));
3909 }
3910 let act = split_act(x);
3911 let gsum = if q2 {
3912 q1_group_sums(&act.xq, gpr)
3913 } else {
3914 Vec::new()
3915 };
3916 let (act, gsum) = (&act, &gsum);
3917 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3918 let (views, ptrs) = (&views, &ptrs);
3919 let run = |start: usize, end: usize| {
3920 with_krows(gpr, |gsc, usc| {
3921 for flat in start..end {
3922 let (e, r) = (flat / inter, flat % inter);
3923 let gv_view = &views[e * 2];
3924 let uv_view = &views[e * 2 + 1];
3925 gv_view.scales_into(r, gpr, gsc);
3926 uv_view.scales_into(r, gpr, usc);
3927 let (mut gv, mut uv) = if q2 {
3928 (
3929 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3930 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3931 )
3932 } else {
3933 (
3934 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3935 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3936 )
3937 };
3938 for &(j, xv) in &act.outliers {
3939 let (og, ou) = if q2 {
3940 (
3941 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3942 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3943 )
3944 } else {
3945 (
3946 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3947 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3948 )
3949 };
3950 gv += og.0 * og.1 * xv;
3951 uv += ou.0 * ou.1 * xv;
3952 }
3953 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3955 }
3956 });
3957 };
3958 dispatch_rows(pool, pairs.len() * inter, &run);
3959 true
3960 }
3961
3962 pub fn moe_down_many(
3971 downs: &[&QTensor],
3972 gs: &[Vec<f32>],
3973 weights: &[f32],
3974 out: &mut [f32],
3975 pool: Option<&Pool>,
3976 ) -> bool {
3977 if Self::any_flint_exact_cpu_default(downs) {
3978 return float_activations_scope(|| {
3979 Self::moe_down_many(downs, gs, weights, out, pool)
3980 });
3981 }
3982 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3983 return false;
3984 }
3985 if downs.iter().any(|d| d.per_tensor_only()) {
3986 return false;
3987 }
3988 if !a8w8_enabled() {
3989 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3990 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3991 return false;
3992 }
3993 out.fill(0.0);
3994 for (row, &w) in terms.iter().zip(weights) {
3995 for (o, &v) in out.iter_mut().zip(row) {
3996 *o += w * v;
3997 }
3998 }
3999 return true;
4000 }
4001 let rows = out.len();
4002 let cols = downs[0].cols();
4003 if cols % GROUP_SIZE != 0 {
4004 return false;
4005 }
4006 let gpr = cols / GROUP_SIZE;
4007 let mut views = Vec::with_capacity(downs.len());
4008 for (d, g) in downs.iter().zip(gs.iter()) {
4009 if !matches!(
4010 d,
4011 Self::Mapped {
4012 dtype: TensorDtype::Q4TiledP,
4013 ..
4014 }
4015 ) || d.rows() != rows
4016 || d.cols() != cols
4017 || g.len() != cols
4018 {
4019 return false;
4020 }
4021 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4022 }
4023 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4025 let out_addr = SendMut(out.as_mut_ptr());
4033 let (views, acts, weights) = (&views, &acts, &weights);
4034 let run = |start: usize, end: usize| {
4035 with_krow(gpr, |sc| {
4036 for r in start..end {
4037 let mut acc = 0f32;
4038 for (e, v) in views.iter().enumerate() {
4039 v.scales_into(r, gpr, sc);
4040 let a = &acts[e];
4041 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4042 for &(j, xv) in &a.outliers {
4043 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4044 d += w * s * xv;
4045 }
4046 acc += weights[e] * d;
4047 }
4048 unsafe { *out_addr.at(r) = acc };
4050 }
4051 });
4052 };
4053 dispatch_rows(pool, rows, &run);
4054 true
4055 }
4056
4057 pub fn moe_gate_up_rows(
4069 pairs: &[(&QTensor, &QTensor)],
4070 groups: &[Vec<usize>],
4071 xs: &[f32],
4072 outs: &mut [Vec<f32>],
4073 pool: Option<&Pool>,
4074 ) -> bool {
4075 if Self::any_flint_exact_cpu_default_pairs(pairs) {
4076 return float_activations_scope(|| {
4077 Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4078 });
4079 }
4080 if pairs.is_empty() || pairs.len() != groups.len() {
4081 return false;
4082 }
4083 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4084 return false;
4085 }
4086 let inter = pairs[0].0.rows();
4087 let cols = pairs[0].0.cols();
4088 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4089 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4090 return false;
4091 }
4092 let b = xs.len() / cols;
4093 let gpr = cols / GROUP_SIZE;
4094 let mut views = Vec::with_capacity(pairs.len() * 2);
4095 for (g, u) in pairs {
4096 let q4tp = |t: &QTensor| {
4097 matches!(
4098 t,
4099 Self::Mapped {
4100 dtype: TensorDtype::Q4TiledP,
4101 ..
4102 }
4103 )
4104 };
4105 if g.has_prism_contract()
4106 || u.has_prism_contract()
4107 || !q4tp(g)
4108 || !q4tp(u)
4109 || g.rows() != inter
4110 || u.rows() != inter
4111 || g.cols() != cols
4112 || u.cols() != cols
4113 {
4114 return false;
4115 }
4116 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4117 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4118 }
4119 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4120 return false;
4121 }
4122 let quantized = a8w8_enabled();
4123 let acts: Vec<SplitAct> = if quantized {
4124 (0..b)
4125 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4126 .collect()
4127 } else {
4128 Vec::new()
4129 };
4130 let mut offs = Vec::with_capacity(groups.len());
4131 let mut o = 0usize;
4132 for g in groups {
4133 offs.push(o);
4134 o += g.len();
4135 }
4136 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4137 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4138 let run = |start: usize, end: usize| {
4139 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4140 for flat in start..end {
4141 let (e, r) = (flat / inter, flat % inter);
4142 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4143 gv_view.scales_into(r, gpr, &mut gsc);
4144 uv_view.scales_into(r, gpr, &mut usc);
4145 for (k, &t) in groups[e].iter().enumerate() {
4146 if !quantized {
4147 let x = &xs[t * cols..(t + 1) * cols];
4148 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4149 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4150 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4151 continue;
4152 }
4153 let act = &acts[t];
4154 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4155 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4156 for &(j, xv) in &act.outliers {
4157 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4158 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4159 gv += og.0 * og.1 * xv;
4160 uv += ou.0 * ou.1 * xv;
4161 }
4162 let silu_g = gv / (1.0 + (-gv).exp());
4163 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4166 }
4167 }
4168 };
4169 dispatch_rows(pool, pairs.len() * inter, &run);
4170 true
4171 }
4172
4173 pub fn moe_down_rows(
4180 downs: &[&QTensor],
4181 group_lens: &[usize],
4182 gs: &[Vec<f32>],
4183 outs: &mut [Vec<f32>],
4184 pool: Option<&Pool>,
4185 ) -> bool {
4186 if Self::any_flint_exact_cpu_default(downs) {
4187 return float_activations_scope(|| {
4188 Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4189 });
4190 }
4191 if downs.is_empty() || downs.len() != group_lens.len() {
4192 return false;
4193 }
4194 if downs.iter().any(|d| d.per_tensor_only()) {
4195 return false;
4196 }
4197 let rows = downs[0].rows();
4198 let cols = downs[0].cols();
4199 let n_pairs: usize = group_lens.iter().sum();
4200 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4201 return false;
4202 }
4203 let gpr = cols / GROUP_SIZE;
4204 let mut views = Vec::with_capacity(downs.len());
4205 for d in downs {
4206 if d.has_prism_contract()
4207 || !matches!(
4208 d,
4209 Self::Mapped {
4210 dtype: TensorDtype::Q4TiledP,
4211 ..
4212 }
4213 )
4214 || d.rows() != rows
4215 || d.cols() != cols
4216 {
4217 return false;
4218 }
4219 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4220 }
4221 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4222 return false;
4223 }
4224 let quantized = a8w8_enabled();
4225 let acts: Vec<SplitAct> = if quantized {
4226 gs.iter().map(|g| split_act(g)).collect()
4227 } else {
4228 Vec::new()
4229 };
4230 let mut offs = Vec::with_capacity(group_lens.len());
4231 let mut o = 0usize;
4232 for &l in group_lens {
4233 offs.push(o);
4234 o += l;
4235 }
4236 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4237 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4238 let run = |start: usize, end: usize| {
4239 let mut sc = vec![0f32; gpr];
4240 for flat in start..end {
4241 let (e, r) = (flat / rows, flat % rows);
4242 let v = &views[e];
4243 v.scales_into(r, gpr, &mut sc);
4244 for k in 0..group_lens[e] {
4245 if !quantized {
4246 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4247 unsafe { *ptrs[offs[e] + k].at(r) = d };
4248 continue;
4249 }
4250 let a = &acts[offs[e] + k];
4251 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4252 for &(j, xv) in &a.outliers {
4253 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4254 d += w * s * xv;
4255 }
4256 unsafe { *ptrs[offs[e] + k].at(r) = d };
4258 }
4259 }
4260 };
4261 dispatch_rows(pool, downs.len() * rows, &run);
4262 true
4263 }
4264}
4265
4266#[cfg(target_os = "macos")]
4271mod accel_blas {
4272 #[link(name = "Accelerate", kind = "framework")]
4273 unsafe extern "C" {
4274 pub fn cblas_sgemm(
4275 order: i32,
4276 trans_a: i32,
4277 trans_b: i32,
4278 m: i32,
4279 n: i32,
4280 k: i32,
4281 alpha: f32,
4282 a: *const f32,
4283 lda: i32,
4284 b: *const f32,
4285 ldb: i32,
4286 beta: f32,
4287 c: *mut f32,
4288 ldc: i32,
4289 );
4290 }
4291}
4292
4293#[cfg(target_os = "macos")]
4294pub(crate) fn accel_gemm_enabled() -> bool {
4295 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4296 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4297}
4298
4299#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4302pub(crate) fn accel_gemm_enabled() -> bool {
4303 true
4304}
4305
4306#[cfg(target_arch = "aarch64")]
4313#[allow(clippy::too_many_arguments)]
4314pub(crate) fn neon_gemm_rm(
4315 m: usize,
4316 n: usize,
4317 k: usize,
4318 alpha: f32,
4319 a: &[f32],
4320 lda: usize,
4321 b_mat: &[f32],
4322 ldb: usize,
4323 b_rows_are_n: bool,
4324 c: &mut [f32],
4325 ldc: usize,
4326) {
4327 debug_assert!(a.len() >= (m - 1) * lda + k);
4328 debug_assert!(c.len() >= (m - 1) * ldc + n);
4329 unsafe {
4331 use core::arch::aarch64::*;
4332 let mut i = 0usize;
4333 while i < m {
4334 let mi = (m - i).min(4);
4335 let mut j = 0usize;
4336 while j < n {
4337 let nj = (n - j).min(8);
4338 if mi == 4 && nj == 8 {
4339 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4340 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4341 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4342 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4343 for p in 0..k {
4344 let (b0, b1) = if b_rows_are_n {
4345 let base = b_mat.as_ptr().add(j * ldb + p);
4348 let g = |o: usize| *base.add(o * ldb);
4349 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4350 } else {
4351 let base = b_mat.as_ptr().add(p * ldb + j);
4352 (
4353 [*base, *base.add(1), *base.add(2), *base.add(3)],
4354 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4355 )
4356 };
4357 let bv0 = vld1q_f32(b0.as_ptr());
4358 let bv1 = vld1q_f32(b1.as_ptr());
4359 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4360 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4361 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4362 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4363 c0a = vfmaq_f32(c0a, a0, bv0);
4364 c0b = vfmaq_f32(c0b, a0, bv1);
4365 c1a = vfmaq_f32(c1a, a1, bv0);
4366 c1b = vfmaq_f32(c1b, a1, bv1);
4367 c2a = vfmaq_f32(c2a, a2, bv0);
4368 c2b = vfmaq_f32(c2b, a2, bv1);
4369 c3a = vfmaq_f32(c3a, a3, bv0);
4370 c3b = vfmaq_f32(c3b, a3, bv1);
4371 }
4372 let al = vdupq_n_f32(alpha);
4373 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4374 .iter()
4375 .enumerate()
4376 {
4377 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4378 vst1q_f32(dst, vmulq_f32(*ca, al));
4379 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4380 }
4381 } else {
4382 for r in 0..mi {
4383 for q in 0..nj {
4384 let mut acc = 0f32;
4385 for p in 0..k {
4386 let bv = if b_rows_are_n {
4387 b_mat[(j + q) * ldb + p]
4388 } else {
4389 b_mat[p * ldb + j + q]
4390 };
4391 acc += a[(i + r) * lda + p] * bv;
4392 }
4393 c[(i + r) * ldc + j + q] = acc * alpha;
4394 }
4395 }
4396 }
4397 j += nj;
4398 }
4399 i += mi;
4400 }
4401 }
4402}
4403
4404#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4406#[allow(clippy::too_many_arguments)]
4407pub(crate) fn sgemm_rm(
4408 m: usize,
4409 n: usize,
4410 k: usize,
4411 alpha: f32,
4412 a: &[f32],
4413 lda: usize,
4414 b_mat: &[f32],
4415 ldb: usize,
4416 b_rows_are_n: bool,
4417 c: &mut [f32],
4418 ldc: usize,
4419) {
4420 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4421}
4422
4423#[allow(clippy::too_many_arguments)]
4427pub fn sgemm_public(
4428 m: usize,
4429 n: usize,
4430 k: usize,
4431 alpha: f32,
4432 a: &[f32],
4433 lda: usize,
4434 b_mat: &[f32],
4435 ldb: usize,
4436 b_rows_are_n: bool,
4437 c: &mut [f32],
4438 ldc: usize,
4439) {
4440 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4441 {
4442 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4443 }
4444 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4449 {
4450 if m == 0 || n == 0 {
4451 return;
4452 }
4453 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4454 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4455 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4456 unsafe {
4458 matrixmultiply::sgemm(
4459 m,
4460 k,
4461 n,
4462 alpha,
4463 a.as_ptr(),
4464 lda as isize,
4465 1,
4466 b_mat.as_ptr(),
4467 rsb,
4468 csb,
4469 0.0,
4470 c.as_mut_ptr(),
4471 ldc as isize,
4472 1,
4473 );
4474 }
4475 }
4476}
4477
4478#[cfg(target_os = "macos")]
4481#[allow(clippy::too_many_arguments)]
4482pub(crate) fn sgemm_rm(
4483 m: usize,
4484 n: usize,
4485 k: usize,
4486 alpha: f32,
4487 a: &[f32],
4488 lda: usize,
4489 b_mat: &[f32],
4490 ldb: usize,
4491 b_rows_are_n: bool,
4492 c: &mut [f32],
4493 ldc: usize,
4494) {
4495 debug_assert!(a.len() >= (m - 1) * lda + k);
4496 debug_assert!(c.len() >= (m - 1) * ldc + n);
4497 #[cfg(target_arch = "aarch64")]
4502 if std::env::var("CMF_FORCE_NEON_GEMM")
4503 .map(|v| v == "1")
4504 .unwrap_or(false)
4505 {
4506 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4507 }
4508 unsafe {
4509 accel_blas::cblas_sgemm(
4510 101, 111, if b_rows_are_n { 112 } else { 111 },
4513 m as i32,
4514 n as i32,
4515 k as i32,
4516 alpha,
4517 a.as_ptr(),
4518 lda as i32,
4519 b_mat.as_ptr(),
4520 ldb as i32,
4521 0.0,
4522 c.as_mut_ptr(),
4523 ldc as i32,
4524 );
4525 }
4526}
4527
4528#[cfg(target_os = "macos")]
4535fn qmatmat_accel(
4536 q: &[u8],
4537 row_scale: &[f32],
4538 pre: &[std::borrow::Cow<'_, [f32]>],
4539 rows: usize,
4540 cols: usize,
4541 out: &mut [f32],
4542 pool: Option<&Pool>,
4543) {
4544 const TR: usize = 2048;
4549 let b = pre.len();
4550 thread_local! {
4551 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4552 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4553 }
4554 XPANEL.with(|xp| {
4555 WTILE.with(|wt| {
4556 let mut xpanel = xp.borrow_mut();
4557 xpanel.clear();
4558 for x in pre {
4559 xpanel.extend_from_slice(x);
4560 }
4561 let mut wtile = wt.borrow_mut();
4562 wtile.resize(TR * cols, 0.0);
4563 let mut r0 = 0usize;
4564 while r0 < rows {
4565 let tr = TR.min(rows - r0);
4566 let wt_addr = SendMut(wtile.as_mut_ptr());
4568 let run = |start: usize, end: usize| {
4569 for r in start..end {
4570 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4571 let s = row_scale[r0 + r];
4572 let dst =
4574 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4575 for (d, &v) in dst.iter_mut().zip(row) {
4576 *d = (v as i8) as f32 * s;
4577 }
4578 }
4579 };
4580 dispatch_rows(pool, tr, &run);
4581 unsafe {
4583 accel_blas::cblas_sgemm(
4584 101, 111, 112, b as i32,
4588 tr as i32,
4589 cols as i32,
4590 1.0,
4591 xpanel.as_ptr(),
4592 cols as i32,
4593 wtile.as_ptr(),
4594 cols as i32,
4595 0.0,
4596 out.as_mut_ptr().add(r0),
4597 rows as i32,
4598 );
4599 }
4600 r0 += tr;
4601 }
4602 })
4603 });
4604}
4605
4606fn qmatmat(
4607 q: &[u8],
4608 row_scale: &[f32],
4609 pre: &[std::borrow::Cow<'_, [f32]>],
4610 rows: usize,
4611 cols: usize,
4612 out: &mut [f32],
4613 pool: Option<&Pool>,
4614) {
4615 let b = pre.len();
4616 debug_assert_eq!(out.len(), b * rows);
4617 #[cfg(target_os = "macos")]
4622 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4623 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4624 return;
4625 }
4626 #[cfg(target_arch = "aarch64")]
4627 if sdot_enabled() {
4628 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4629 let out_addr = SendMut(out.as_mut_ptr());
4630 let blocked_ok = blocked_enabled();
4633 let use_i8mm = i8mm_enabled();
4634 if blocked_ok {
4635 let run = |start: usize, end: usize| {
4636 let mut o = start;
4637 while o < end {
4638 if o + 2 <= end {
4639 let r0 = &q[o * cols..(o + 1) * cols];
4640 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4641 let mut bi = 0usize;
4642 while bi + 4 <= acts.len() {
4643 let xs = [
4644 acts[bi].xq.as_slice(),
4645 acts[bi + 1].xq.as_slice(),
4646 acts[bi + 2].xq.as_slice(),
4647 acts[bi + 3].xq.as_slice(),
4648 ];
4649 let d = if use_i8mm {
4650 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4651 } else {
4652 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4653 };
4654 for (r, row) in [r0, r1].into_iter().enumerate() {
4655 for k in 0..4 {
4656 let act = &acts[bi + k];
4657 let mut v = d[r][k] as f32 * act.sx;
4658 for &(j, xv) in &act.outliers {
4659 v += (row[j] as i8) as f32 * xv;
4660 }
4661 unsafe {
4662 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4663 };
4664 }
4665 }
4666 bi += 4;
4667 }
4668 while bi < acts.len() {
4669 for (r, row) in [r0, r1].into_iter().enumerate() {
4670 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4671 unsafe { *out_addr.at(bi * rows + o + r) = v };
4672 }
4673 bi += 1;
4674 }
4675 o += 2;
4676 } else {
4677 let row = &q[o * cols..(o + 1) * cols];
4678 for (bi, act) in acts.iter().enumerate() {
4679 let v = row_dot_sdot(row, act) * row_scale[o];
4680 unsafe { *out_addr.at(bi * rows + o) = v };
4681 }
4682 o += 1;
4683 }
4684 }
4685 };
4686 dispatch_rows(pool, rows, &run);
4687 return;
4688 }
4689 let run = |start: usize, end: usize| {
4690 for o in start..end {
4691 let row = &q[o * cols..(o + 1) * cols];
4692 for (bi, act) in acts.iter().enumerate() {
4693 let v = row_dot_sdot(row, act) * row_scale[o];
4694 unsafe { *out_addr.at(bi * rows + o) = v };
4695 }
4696 }
4697 };
4698 dispatch_rows(pool, rows, &run);
4699 return;
4700 }
4701 #[cfg(target_arch = "x86_64")]
4706 if avx2_a8w8_enabled() {
4707 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4708 let out_addr = SendMut(out.as_mut_ptr());
4709 let blocked_ok = blocked_enabled();
4712 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4713 let run = |start: usize, end: usize| {
4714 let mut o = start;
4715 while o < end {
4716 if o + 2 <= end {
4717 let r0 = &q[o * cols..(o + 1) * cols];
4718 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4719 let mut bi = 0usize;
4720 while bi + 4 <= acts.len() {
4721 let xs = [
4722 acts[bi].xq.as_slice(),
4723 acts[bi + 1].xq.as_slice(),
4724 acts[bi + 2].xq.as_slice(),
4725 acts[bi + 3].xq.as_slice(),
4726 ];
4727 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4728 for (r, row) in [r0, r1].into_iter().enumerate() {
4729 for k in 0..4 {
4730 let act = &acts[bi + k];
4731 let mut v = d[r][k] as f32 * act.sx;
4732 for &(j, xv) in &act.outliers {
4733 v += (row[j] as i8) as f32 * xv;
4734 }
4735 unsafe {
4736 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4737 };
4738 }
4739 }
4740 bi += 4;
4741 }
4742 while bi < acts.len() {
4743 for (r, row) in [r0, r1].into_iter().enumerate() {
4744 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4745 unsafe { *out_addr.at(bi * rows + o + r) = v };
4746 }
4747 bi += 1;
4748 }
4749 o += 2;
4750 } else {
4751 let row = &q[o * cols..(o + 1) * cols];
4752 for (bi, act) in acts.iter().enumerate() {
4753 let v = row_dot_avx2(row, act) * row_scale[o];
4754 unsafe { *out_addr.at(bi * rows + o) = v };
4755 }
4756 o += 1;
4757 }
4758 }
4759 };
4760 dispatch_rows(pool, rows, &run);
4761 return;
4762 }
4763 let run = |start: usize, end: usize| {
4764 for o in start..end {
4765 let row = &q[o * cols..(o + 1) * cols];
4766 for (bi, act) in acts.iter().enumerate() {
4767 let v = row_dot_avx2(row, act) * row_scale[o];
4768 unsafe { *out_addr.at(bi * rows + o) = v };
4769 }
4770 }
4771 };
4772 dispatch_rows(pool, rows, &run);
4773 return;
4774 }
4775 let out_addr = SendMut(out.as_mut_ptr());
4776 let run = |start: usize, end: usize| {
4777 for o in start..end {
4778 let row = &q[o * cols..(o + 1) * cols];
4779 for (bi, x) in pre.iter().enumerate() {
4780 let mut acc = 0f32;
4781 for j in 0..cols {
4782 acc += (row[j] as i8) as f32 * x[j];
4783 }
4784 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4785 }
4786 }
4787 };
4788 dispatch_rows(pool, rows, &run);
4789}
4790
4791fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4794 match pool {
4795 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4796 _ => run(0, rows),
4797 }
4798}
4799
4800fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4802 let groups = rows * cols / GROUP_SIZE;
4803 bytes.split_at(groups * 16)
4804}
4805
4806#[inline]
4811fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4812 #[cfg(target_arch = "aarch64")]
4813 unsafe {
4814 return vbit_fill4_neon(data, buf);
4815 }
4816 #[cfg(target_arch = "x86_64")]
4817 if avx2_enabled() {
4818 return unsafe { vbit_fill4_avx2(data, buf) };
4819 }
4820 #[allow(unreachable_code)]
4821 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4822 let u = unpack8::<4>(&data[blk * 4..]);
4823 for k in 0..8 {
4824 chunk[k] = (u[k] - 7) as i8 as u8;
4825 }
4826 }
4827}
4828
4829#[cfg(target_arch = "aarch64")]
4830#[target_feature(enable = "neon")]
4831unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4832 unsafe {
4835 use core::arch::aarch64::*;
4836 let n = buf.len();
4837 let mask = vdupq_n_u8(0x0F);
4838 let seven = vdupq_n_s8(7);
4839 let mut g = 0usize;
4840 while g * 32 + 32 <= n {
4841 let b = vld1q_u8(data.as_ptr().add(g * 16));
4842 let hi = vshrq_n_u8::<4>(b);
4843 let lo = vandq_u8(b, mask);
4844 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4845 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4846 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4847 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4848 g += 1;
4849 }
4850 }
4851}
4852
4853#[cfg(target_arch = "x86_64")]
4854#[target_feature(enable = "avx2")]
4855unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4856 unsafe {
4858 use core::arch::x86_64::*;
4859 let n = buf.len();
4860 let mask = _mm_set1_epi8(0x0F);
4861 let seven = _mm256_set1_epi8(7);
4862 let mut g = 0usize;
4863 while g * 32 + 32 <= n {
4864 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4865 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4866 let lo = _mm_and_si128(b, mask);
4867 let z = _mm256_sub_epi8(
4868 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4869 seven,
4870 );
4871 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4872 g += 1;
4873 }
4874 }
4875}
4876
4877#[inline(always)]
4882fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4883 let mut acc = 0u64;
4884 for i in 0..B {
4885 acc = (acc << 8) | data[i] as u64;
4886 }
4887 let mask = (1u64 << B) - 1;
4888 let mut out = [0i32; 8];
4889 for (k, o) in out.iter_mut().enumerate() {
4890 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4891 }
4892 out
4893}
4894
4895#[allow(clippy::too_many_arguments)]
4901fn vbitmatvec(
4902 bytes: &[u8],
4903 offsets: &[usize],
4904 x: &[f32],
4905 rows: usize,
4906 cols: usize,
4907 out: &mut [f32],
4908 pool: Option<&Pool>,
4909) {
4910 debug_assert_eq!(out.len(), rows);
4911 debug_assert_eq!(offsets.len(), rows + 1);
4912
4913 if a8w8_enabled() {
4917 let act = split_act(x);
4918 let out_addr = SendMut(out.as_mut_ptr());
4919 let run = move |start: usize, end: usize| {
4920 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4921 };
4922 dispatch_rows(pool, rows, &run);
4923 return;
4924 }
4925
4926 let out_addr = SendMut(out.as_mut_ptr());
4927 let run = move |start: usize, end: usize| {
4928 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4929 };
4930 dispatch_rows(pool, rows, &run);
4931}
4932
4933#[allow(clippy::too_many_arguments)]
4937fn vbit_range_a8w8(
4938 bytes: &[u8],
4939 offsets: &[usize],
4940 x: &[f32],
4941 act: &SplitAct,
4942 rows: usize,
4943 cols: usize,
4944 out: SendMut,
4945 start: usize,
4946 end: usize,
4947) {
4948 let ng = cols / GROUP_SIZE;
4949 let bits = &bytes[..rows];
4950 let sc_off = rows;
4951 let row_dot = |r: usize| -> f32 {
4952 let b = bits[r] as usize;
4953 let l = (1i32 << (b - 1)) - 1;
4954 let mask = (1u64 << b) - 1;
4955 let data = &bytes[offsets[r]..offsets[r + 1]];
4956 if b == 8 {
4957 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4959 let mut dot = 0f32;
4960 for g in 0..ng {
4961 let so = (r * ng + g) * 2;
4962 let sgf = f16_to_f32(u16::from_le_bytes([
4963 bytes[sc_off + so],
4964 bytes[sc_off + so + 1],
4965 ]));
4966 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4967 let mut gd = 0f32;
4968 for &xv in xg.iter() {
4969 if nbits < 8 {
4970 acc = (acc << 8) | data[idx] as u64;
4971 idx += 1;
4972 nbits += 8;
4973 }
4974 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4975 nbits -= 8;
4976 gd += (u - l) as f32 * xv;
4977 }
4978 dot += gd * sgf;
4979 }
4980 return dot;
4981 }
4982 thread_local! {
4986 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4987 const { std::cell::RefCell::new(Vec::new()) };
4988 }
4989 #[inline(always)]
4990 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4991 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4992 let u = unpack8::<B>(&data[blk * B..]);
4993 for k in 0..8 {
4994 chunk[k] = (u[k] - l) as i8 as u8;
4995 }
4996 }
4997 }
4998 let _ = mask;
4999 VBIT_SCRATCH.with(|scratch| {
5000 let mut buf = scratch.borrow_mut();
5001 buf.resize(cols, 0);
5002 match b {
5003 3 => fill::<3>(data, l, &mut buf),
5004 4 => vbit_fill4(data, &mut buf),
5005 5 => fill::<5>(data, l, &mut buf),
5006 6 => fill::<6>(data, l, &mut buf),
5007 _ => unreachable!(),
5008 }
5009 let mut dot = 0f32;
5010 for g in 0..ng {
5011 let so = (r * ng + g) * 2;
5012 let s = f16_to_f32(u16::from_le_bytes([
5013 bytes[sc_off + so],
5014 bytes[sc_off + so + 1],
5015 ]));
5016 let d = dot_i8_i8(
5017 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5018 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5019 ) as f32
5020 * act.sx;
5021 dot += d * s;
5022 }
5023 for &(j, xv) in &act.outliers {
5024 let so = (r * ng + j / GROUP_SIZE) * 2;
5025 let s = f16_to_f32(u16::from_le_bytes([
5026 bytes[sc_off + so],
5027 bytes[sc_off + so + 1],
5028 ]));
5029 dot += (buf[j] as i8) as f32 * s * xv;
5031 }
5032 dot
5033 })
5034 };
5035 for r in start..end {
5036 unsafe { *out.at(r) = row_dot(r) };
5038 }
5039}
5040
5041#[allow(clippy::too_many_arguments)]
5043fn vbit_range_f32(
5044 bytes: &[u8],
5045 offsets: &[usize],
5046 x: &[f32],
5047 rows: usize,
5048 cols: usize,
5049 out: SendMut,
5050 start: usize,
5051 end: usize,
5052) {
5053 let ng = cols / GROUP_SIZE;
5054 let bits = &bytes[..rows];
5055 let sc_off = rows;
5056 #[inline(always)]
5060 fn dot_row<const B: usize>(
5061 data: &[u8],
5062 bytes: &[u8],
5063 sc_off: usize,
5064 r: usize,
5065 ng: usize,
5066 x: &[f32],
5067 ) -> f32 {
5068 let l = ((1i32 << (B - 1)) - 1) as f32;
5069 let gbytes = GROUP_SIZE * B / 8;
5070 let mut dot = 0f32;
5071 for g in 0..ng {
5072 let so = (r * ng + g) * 2;
5073 let s = f16_to_f32(u16::from_le_bytes([
5074 bytes[sc_off + so],
5075 bytes[sc_off + so + 1],
5076 ]));
5077 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5078 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5079 let mut gd = 0f32;
5080 for blk in 0..GROUP_SIZE / 8 {
5081 let u = unpack8::<B>(&gd0[blk * B..]);
5082 let xb = &xg[blk * 8..blk * 8 + 8];
5083 for k in 0..8 {
5084 gd += (u[k] as f32 - l) * xb[k];
5085 }
5086 }
5087 dot += gd * s;
5088 }
5089 dot
5090 }
5091 for r in start..end {
5092 let data = &bytes[offsets[r]..offsets[r + 1]];
5093 let v = match bits[r] {
5094 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5095 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5096 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5097 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5098 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5099 b => unreachable!("vbit bit-width {b} (validated at load)"),
5100 };
5101 unsafe { *out.at(r) = v };
5103 }
5104}
5105
5106#[allow(clippy::too_many_arguments)]
5111fn vbitmatvec2(
5112 bytes: &[u8],
5113 offsets: &[usize],
5114 x1: &[f32],
5115 x2: &[f32],
5116 rows: usize,
5117 cols: usize,
5118 o1: &mut [f32],
5119 o2: &mut [f32],
5120 pool: Option<&Pool>,
5121) {
5122 debug_assert_eq!(o1.len(), rows);
5123 debug_assert_eq!(o2.len(), rows);
5124
5125 if a8w8_enabled() {
5126 let a1 = split_act(x1);
5127 let a2 = split_act(x2);
5128 let p1 = SendMut(o1.as_mut_ptr());
5129 let p2 = SendMut(o2.as_mut_ptr());
5130 let run = move |start: usize, end: usize| {
5131 vbit_range2_a8w8(
5132 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5133 )
5134 };
5135 dispatch_rows(pool, rows, &run);
5136 return;
5137 }
5138
5139 let p1 = SendMut(o1.as_mut_ptr());
5140 let p2 = SendMut(o2.as_mut_ptr());
5141 let run = move |start: usize, end: usize| {
5142 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5143 };
5144 dispatch_rows(pool, rows, &run);
5145}
5146
5147#[allow(clippy::too_many_arguments)]
5151fn vbit_range2_a8w8(
5152 bytes: &[u8],
5153 offsets: &[usize],
5154 x1: &[f32],
5155 x2: &[f32],
5156 a1: &SplitAct,
5157 a2: &SplitAct,
5158 rows: usize,
5159 cols: usize,
5160 p1: SendMut,
5161 p2: SendMut,
5162 start: usize,
5163 end: usize,
5164) {
5165 let ng = cols / GROUP_SIZE;
5166 let bits = &bytes[..rows];
5167 let sc_off = rows;
5168 let row_dots = |r: usize| -> (f32, f32) {
5169 let b = bits[r] as usize;
5170 let l = (1i32 << (b - 1)) - 1;
5171 let data = &bytes[offsets[r]..offsets[r + 1]];
5172 if b == 8 {
5173 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5176 let (mut d1, mut d2) = (0f32, 0f32);
5177 for g in 0..ng {
5178 let so = (r * ng + g) * 2;
5179 let sgf = f16_to_f32(u16::from_le_bytes([
5180 bytes[sc_off + so],
5181 bytes[sc_off + so + 1],
5182 ]));
5183 let (mut g1, mut g2) = (0f32, 0f32);
5184 for k in 0..GROUP_SIZE {
5185 if nbits < 8 {
5186 acc = (acc << 8) | data[idx] as u64;
5187 idx += 1;
5188 nbits += 8;
5189 }
5190 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5191 nbits -= 8;
5192 let w = (u - l) as f32;
5193 g1 += w * x1[g * GROUP_SIZE + k];
5194 g2 += w * x2[g * GROUP_SIZE + k];
5195 }
5196 d1 += g1 * sgf;
5197 d2 += g2 * sgf;
5198 }
5199 return (d1, d2);
5200 }
5201 thread_local! {
5202 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5203 const { std::cell::RefCell::new(Vec::new()) };
5204 }
5205 #[inline(always)]
5206 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5207 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5208 let u = unpack8::<B>(&data[blk * B..]);
5209 for k in 0..8 {
5210 chunk[k] = (u[k] - l) as i8 as u8;
5211 }
5212 }
5213 }
5214 VBIT_SCRATCH2.with(|scratch| {
5215 let mut buf = scratch.borrow_mut();
5216 buf.resize(cols, 0);
5217 match b {
5218 3 => fill::<3>(data, l, &mut buf),
5219 4 => vbit_fill4(data, &mut buf),
5220 5 => fill::<5>(data, l, &mut buf),
5221 6 => fill::<6>(data, l, &mut buf),
5222 _ => unreachable!(),
5223 }
5224 let (mut d1, mut d2) = (0f32, 0f32);
5225 for g in 0..ng {
5226 let so = (r * ng + g) * 2;
5227 let s = f16_to_f32(u16::from_le_bytes([
5228 bytes[sc_off + so],
5229 bytes[sc_off + so + 1],
5230 ]));
5231 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5232 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5233 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5234 d1 += v1 * s;
5235 d2 += v2 * s;
5236 }
5237 for &(j, xv) in &a1.outliers {
5238 let so = (r * ng + j / GROUP_SIZE) * 2;
5239 let s = f16_to_f32(u16::from_le_bytes([
5240 bytes[sc_off + so],
5241 bytes[sc_off + so + 1],
5242 ]));
5243 d1 += (buf[j] as i8) as f32 * s * xv;
5244 }
5245 for &(j, xv) in &a2.outliers {
5246 let so = (r * ng + j / GROUP_SIZE) * 2;
5247 let s = f16_to_f32(u16::from_le_bytes([
5248 bytes[sc_off + so],
5249 bytes[sc_off + so + 1],
5250 ]));
5251 d2 += (buf[j] as i8) as f32 * s * xv;
5252 }
5253 (d1, d2)
5254 })
5255 };
5256 for r in start..end {
5257 let (v1, v2) = row_dots(r);
5258 unsafe {
5260 *p1.at(r) = v1;
5261 *p2.at(r) = v2;
5262 }
5263 }
5264}
5265
5266#[allow(clippy::too_many_arguments)]
5270fn vbit_range2_f32(
5271 bytes: &[u8],
5272 offsets: &[usize],
5273 x1: &[f32],
5274 x2: &[f32],
5275 rows: usize,
5276 cols: usize,
5277 p1: SendMut,
5278 p2: SendMut,
5279 start: usize,
5280 end: usize,
5281) {
5282 let ng = cols / GROUP_SIZE;
5283 let bits = &bytes[..rows];
5284 let sc_off = rows;
5285 #[inline(always)]
5286 #[allow(clippy::too_many_arguments)]
5287 fn dot_row2<const B: usize>(
5288 data: &[u8],
5289 bytes: &[u8],
5290 sc_off: usize,
5291 r: usize,
5292 ng: usize,
5293 x1: &[f32],
5294 x2: &[f32],
5295 ) -> (f32, f32) {
5296 let l = ((1i32 << (B - 1)) - 1) as f32;
5297 let gbytes = GROUP_SIZE * B / 8;
5298 let (mut d1, mut d2) = (0f32, 0f32);
5299 for g in 0..ng {
5300 let so = (r * ng + g) * 2;
5301 let s = f16_to_f32(u16::from_le_bytes([
5302 bytes[sc_off + so],
5303 bytes[sc_off + so + 1],
5304 ]));
5305 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5306 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5307 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5308 let (mut g1, mut g2) = (0f32, 0f32);
5309 for blk in 0..GROUP_SIZE / 8 {
5310 let u = unpack8::<B>(&gd0[blk * B..]);
5311 for k in 0..8 {
5312 let w = u[k] as f32 - l;
5313 g1 += w * x1g[blk * 8 + k];
5314 g2 += w * x2g[blk * 8 + k];
5315 }
5316 }
5317 d1 += g1 * s;
5318 d2 += g2 * s;
5319 }
5320 (d1, d2)
5321 }
5322 for r in start..end {
5323 let data = &bytes[offsets[r]..offsets[r + 1]];
5324 let (v1, v2) = match bits[r] {
5325 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5326 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5327 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5328 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5329 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5330 b => unreachable!("vbit bit-width {b} (validated at load)"),
5331 };
5332 unsafe {
5334 *p1.at(r) = v1;
5335 *p2.at(r) = v2;
5336 }
5337 }
5338}
5339
5340#[inline]
5347#[allow(unreachable_code)]
5348fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5349 #[cfg(target_arch = "aarch64")]
5350 unsafe {
5351 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5352 }
5353 #[cfg(target_arch = "x86_64")]
5354 unsafe {
5355 if vnni_tiles_enabled() {
5356 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5357 }
5358 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5359 }
5360 let mut acc = 0f32;
5361 for gi in 0..gpr {
5362 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5363 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5364 let mut d = 0i32;
5365 for (k, &b) in tile[2..].iter().enumerate() {
5366 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5367 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5368 }
5369 acc += d as f32 * s;
5370 }
5371 acc
5372}
5373
5374#[cfg(target_arch = "aarch64")]
5375#[target_feature(enable = "neon,dotprod")]
5376unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5377 unsafe {
5380 use core::arch::aarch64::*;
5381 use core::arch::asm;
5382 let lomask = vdupq_n_u8(0x0F);
5383 let eight = vdupq_n_s8(8);
5384 let mut acc = 0f32;
5385 for gi in 0..gpr {
5386 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5387 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5388 let b = vld1q_u8(t.add(2));
5389 let lo = vandq_u8(b, lomask);
5390 let hi = vshrq_n_u8::<4>(b);
5391 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5392 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5393 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5394 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5395 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5396 asm!(
5397 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5398 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5399 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5400 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5401 options(pure, nomem, nostack),
5402 );
5403 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5404 }
5405 acc
5406 }
5407}
5408
5409#[cfg(target_arch = "x86_64")]
5410#[target_feature(enable = "avx2")]
5411unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5412 unsafe {
5414 use core::arch::x86_64::*;
5415 let lomask = _mm_set1_epi8(0x0F);
5416 let eight = _mm256_set1_epi8(8);
5417 let ones = _mm256_set1_epi16(1);
5418 let mut acc = 0f32;
5419 for gi in 0..gpr {
5420 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5421 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5422 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5423 let lo = _mm_and_si128(b, lomask);
5424 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5425 let w = _mm256_sub_epi8(
5426 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5427 eight,
5428 );
5429 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5430 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5431 let d = _mm256_madd_epi16(p16, ones);
5432 let hi128 = _mm256_extracti128_si256::<1>(d);
5433 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5434 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5435 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5436 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5437 }
5438 acc
5439 }
5440}
5441
5442#[cfg(target_arch = "x86_64")]
5446#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5447unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5448 unsafe {
5450 use core::arch::x86_64::*;
5451 let lomask = _mm_set1_epi8(0x0F);
5452 let eight = _mm256_set1_epi8(8);
5453 let mut acc = 0f32;
5454 for gi in 0..gpr {
5455 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5456 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5457 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5458 let lo = _mm_and_si128(b, lomask);
5459 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5460 let w = _mm256_sub_epi8(
5461 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5462 eight,
5463 );
5464 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5465 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5466 acc += d as f32 * s;
5467 }
5468 acc
5469 }
5470}
5471
5472#[cfg(target_arch = "x86_64")]
5477#[target_feature(enable = "avx2,fma")]
5482unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5483 unsafe {
5485 use core::arch::x86_64::*;
5486 let lomask = _mm_set1_epi8(0x0F);
5487 let eight = _mm256_set1_epi8(8);
5488 let ones = _mm256_set1_epi16(1);
5489 let mut f0 = _mm256_setzero_ps();
5502 let mut f1 = _mm256_setzero_ps();
5503 let mut f2 = _mm256_setzero_ps();
5504 let mut f3 = _mm256_setzero_ps();
5505 for gi in 0..gpr {
5506 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5507 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5508 let sv = _mm256_set1_ps(s);
5509 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5510 let lo = _mm_and_si128(bb, lomask);
5511 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5512 let w = _mm256_sub_epi8(
5513 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5514 eight,
5515 );
5516 let aw = _mm256_abs_epi8(w);
5517 let off = gi * GROUP_SIZE;
5518 let dot = |xq: &[i8]| {
5519 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5520 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5521 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5522 };
5523 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5524 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5525 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5526 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5527 }
5528 [
5529 hsum256_ps(f0),
5530 hsum256_ps(f1),
5531 hsum256_ps(f2),
5532 hsum256_ps(f3),
5533 ]
5534 }
5535}
5536
5537#[cfg(target_arch = "x86_64")]
5540#[target_feature(enable = "avx2")]
5541#[inline]
5542unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5543 unsafe {
5545 use core::arch::x86_64::*;
5546 let hi = _mm256_extractf128_ps::<1>(v);
5547 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5548 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5549 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5550 _mm_cvtss_f32(s)
5551 }
5552}
5553
5554#[cfg(target_arch = "x86_64")]
5556#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5557unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5558 unsafe {
5560 use core::arch::x86_64::*;
5561 let lomask = _mm_set1_epi8(0x0F);
5562 let eight = _mm256_set1_epi8(8);
5563 let mut f0 = _mm256_setzero_ps();
5566 let mut f1 = _mm256_setzero_ps();
5567 let mut f2 = _mm256_setzero_ps();
5568 let mut f3 = _mm256_setzero_ps();
5569 for gi in 0..gpr {
5570 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5571 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5572 let sv = _mm256_set1_ps(s);
5573 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5574 let lo = _mm_and_si128(bb, lomask);
5575 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5576 let w = _mm256_sub_epi8(
5577 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5578 eight,
5579 );
5580 let aw = _mm256_abs_epi8(w);
5581 let off = gi * GROUP_SIZE;
5582 let dot = |xq: &[i8]| {
5583 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5584 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5585 _mm256_setzero_si256(),
5586 aw,
5587 _mm256_sign_epi8(x, w),
5588 ))
5589 };
5590 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5591 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5592 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5593 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5594 }
5595 let acc = [
5596 hsum256_ps(f0),
5597 hsum256_ps(f1),
5598 hsum256_ps(f2),
5599 hsum256_ps(f3),
5600 ];
5601 acc
5602 }
5603}
5604
5605#[cfg(target_arch = "aarch64")]
5610#[target_feature(enable = "neon,dotprod")]
5611unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5612 unsafe {
5614 use core::arch::aarch64::*;
5615 use core::arch::asm;
5616 let lomask = vdupq_n_u8(0x0F);
5617 let eight = vdupq_n_s8(8);
5618 let mut acc = [0f32; 4];
5619 for gi in 0..gpr {
5620 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5621 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5622 let b = vld1q_u8(t.add(2));
5623 let lo = vandq_u8(b, lomask);
5624 let hi = vshrq_n_u8::<4>(b);
5625 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5626 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5627 for (k, xq) in xs.iter().enumerate() {
5628 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5629 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5630 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5631 asm!(
5632 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5633 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5634 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5635 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5636 options(pure, nomem, nostack),
5637 );
5638 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5639 }
5640 }
5641 acc
5642 }
5643}
5644
5645#[inline]
5647fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5648 let gi = j / GROUP_SIZE;
5649 let k = j % GROUP_SIZE;
5650 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5651 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5652 let byte = tile[2 + k / 2];
5653 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5654 ((nib as i32 - 8) as f32, s)
5655}
5656
5657#[inline]
5660fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5661 let mut acc = 0f32;
5662 for gi in 0..gpr {
5663 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5664 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5665 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5666 let mut ga = 0f32;
5667 for (k, &b) in tile[2..].iter().enumerate() {
5668 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5669 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5670 }
5671 acc += ga * s;
5672 }
5673 acc
5674}
5675
5676struct Q4tpView<'a> {
5680 nib: &'a [u8],
5681 params: &'a [u8],
5682 codes: &'a [u8],
5683 stride: usize,
5684 zero_rung: bool,
5686}
5687
5688impl<'a> Q4tpView<'a> {
5689 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5690 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5691 Self {
5692 nib: &bytes[..params_off],
5693 params: &bytes[params_off..codes_off],
5694 codes: &bytes[codes_off..],
5695 stride,
5696 zero_rung: false,
5697 }
5698 }
5699
5700 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5702 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5703 Self {
5704 nib: &bytes[..params_off],
5705 params: &bytes[params_off..codes_off],
5706 codes: &bytes[codes_off..],
5707 stride,
5708 zero_rung: true,
5709 }
5710 }
5711
5712 #[inline]
5728 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5729 let tab = if self.zero_rung {
5730 q2tp_ladder(self.params, r)
5731 } else {
5732 q4tp_ladder(self.params, r)
5733 };
5734 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5735 let out = &mut out[..gpr];
5736 let mut chunks = out.chunks_exact_mut(8);
5737 let mut ci = 0usize;
5738 for c in &mut chunks {
5739 let w = u64::from(codes[ci])
5740 | u64::from(codes[ci + 1]) << 8
5741 | u64::from(codes[ci + 2]) << 16
5742 | u64::from(codes[ci + 3]) << 24
5743 | u64::from(codes[ci + 4]) << 32;
5744 for (k, o) in c.iter_mut().enumerate() {
5745 *o = tab[((w >> (5 * k)) & 31) as usize];
5746 }
5747 ci += 5;
5748 }
5749 let tail = &codes[ci..];
5752 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5753 *o = tab[q4tp_code(tail, k)];
5754 }
5755 }
5756}
5757
5758#[inline]
5759fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5760 #[cfg(target_arch = "aarch64")]
5761 unsafe {
5762 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5763 }
5764 #[cfg(target_arch = "x86_64")]
5765 unsafe {
5766 if vnni_tiles_enabled() {
5767 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5768 }
5769 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5770 }
5771 #[allow(unreachable_code)]
5772 {
5773 let mut acc = 0f32;
5774 for gi in 0..gpr {
5775 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5776 let s = scales[gi];
5777 let mut d = 0i32;
5778 for (k, &b) in tile.iter().enumerate() {
5779 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5780 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5781 }
5782 acc += d as f32 * s;
5783 }
5784 acc
5785 }
5786}
5787
5788#[cfg(target_arch = "aarch64")]
5791#[target_feature(enable = "neon,dotprod")]
5792unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5793 unsafe {
5796 use core::arch::aarch64::*;
5797 use core::arch::asm;
5798 let lomask = vdupq_n_u8(0x0F);
5799 let eight = vdupq_n_s8(8);
5800 let mut acc = 0f32;
5801 for gi in 0..gpr {
5802 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5803 let s = *scales.get_unchecked(gi);
5804 let b = vld1q_u8(t);
5805 let lo = vandq_u8(b, lomask);
5806 let hi = vshrq_n_u8::<4>(b);
5807 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5808 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5809 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5810 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5811 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5812 asm!(
5813 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5814 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5815 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5816 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5817 options(pure, nomem, nostack),
5818 );
5819 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5820 }
5821 acc
5822 }
5823}
5824
5825#[cfg(target_arch = "x86_64")]
5826#[target_feature(enable = "avx2")]
5827unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5828 unsafe {
5830 use core::arch::x86_64::*;
5831 let lomask = _mm_set1_epi8(0x0F);
5832 let eight = _mm256_set1_epi8(8);
5833 let ones = _mm256_set1_epi16(1);
5834 let mut acc = 0f32;
5835 for gi in 0..gpr {
5836 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5837 let s = *scales.get_unchecked(gi);
5838 let b = _mm_loadu_si128(t as *const __m128i);
5839 let lo = _mm_and_si128(b, lomask);
5840 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5841 let w = _mm256_sub_epi8(
5842 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5843 eight,
5844 );
5845 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5846 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5847 let d = _mm256_madd_epi16(p16, ones);
5848 let hi128 = _mm256_extracti128_si256::<1>(d);
5849 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5850 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5851 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5852 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5853 }
5854 acc
5855 }
5856}
5857
5858#[cfg(target_arch = "x86_64")]
5861#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5862unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5863 unsafe {
5865 use core::arch::x86_64::*;
5866 let lomask = _mm_set1_epi8(0x0F);
5867 let eight = _mm256_set1_epi8(8);
5868 let mut acc = 0f32;
5869 for gi in 0..gpr {
5870 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5871 let s = *scales.get_unchecked(gi);
5872 let b = _mm_loadu_si128(t as *const __m128i);
5873 let lo = _mm_and_si128(b, lomask);
5874 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5875 let w = _mm256_sub_epi8(
5876 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5877 eight,
5878 );
5879 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5880 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5881 }
5882 acc
5883 }
5884}
5885
5886#[inline]
5889fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5890 #[cfg(target_arch = "x86_64")]
5891 if avx2_enabled() {
5892 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5894 }
5895 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5896}
5897
5898#[inline]
5899fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5900 let mut acc = 0f32;
5901 for gi in 0..gpr {
5902 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5903 let s = scales[gi];
5904 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5905 let mut ga = 0f32;
5906 for (k, &b) in tile.iter().enumerate() {
5907 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5908 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5909 }
5910 acc += ga * s;
5911 }
5912 acc
5913}
5914
5915#[cfg(target_arch = "x86_64")]
5919#[target_feature(enable = "avx2")]
5920unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5921 unsafe {
5924 use core::arch::x86_64::*;
5925 let mask = _mm_set1_epi8(15);
5926 let eight = _mm_set1_epi8(8);
5927 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5928 let mut acc = 0.0f32;
5929 for gi in 0..gpr {
5930 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5931 let lo = _mm_and_si128(packed, mask);
5932 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5933 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5934 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5935 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5936 let a = _mm256_mul_ps(
5937 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5938 _mm256_loadu_ps(xp),
5939 );
5940 let b = _mm256_mul_ps(
5941 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5942 _mm256_loadu_ps(xp.add(8)),
5943 );
5944 let c = _mm256_mul_ps(
5945 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5946 _mm256_loadu_ps(xp.add(16)),
5947 );
5948 let d = _mm256_mul_ps(
5949 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5950 _mm256_loadu_ps(xp.add(24)),
5951 );
5952 let mut pairs = [0.0f32; 16];
5953 _mm256_storeu_ps(
5954 pairs.as_mut_ptr(),
5955 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5956 );
5957 _mm256_storeu_ps(
5958 pairs.as_mut_ptr().add(8),
5959 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5960 );
5961 let mut ga = 0.0f32;
5962 for v in pairs {
5963 ga += v;
5964 }
5965 acc += ga * scales[gi];
5966 }
5967 acc
5968 }
5969}
5970
5971#[inline]
5974fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5975 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5976 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5977 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5978 ((n as i32 - 8) as f32, scales[gi])
5979}
5980
5981fn q4tp_matvec(
5983 bytes: &[u8],
5984 x: &[f32],
5985 rows: usize,
5986 cols: usize,
5987 out: &mut [f32],
5988 pool: Option<&Pool>,
5989) {
5990 debug_assert_eq!(out.len(), rows);
5991 let gpr = cols / GROUP_SIZE;
5992 let v = Q4tpView::new(bytes, rows, cols);
5993 let out_addr = SendMut(out.as_mut_ptr());
5994 if a8w8_enabled() {
5995 let act = split_act(x);
5996 let run = |start: usize, end: usize| {
5997 with_krow(gpr, |sc| {
5999 for r in start..end {
6000 v.scales_into(r, gpr, sc);
6001 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
6002 for &(j, xv) in &act.outliers {
6003 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6004 acc += w * s * xv;
6005 }
6006 unsafe { *out_addr.at(r) = acc };
6008 }
6009 })
6010 };
6011 dispatch_rows(pool, rows, &run);
6012 return;
6013 }
6014 let run = |start: usize, end: usize| {
6015 with_krow(gpr, |sc| {
6016 for r in start..end {
6017 v.scales_into(r, gpr, sc);
6018 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
6020 }
6021 })
6022 };
6023 dispatch_rows(pool, rows, &run);
6024}
6025
6026#[allow(clippy::too_many_arguments)]
6029fn q4tp_matvec2(
6030 bytes: &[u8],
6031 x1: &[f32],
6032 x2: &[f32],
6033 rows: usize,
6034 cols: usize,
6035 o1: &mut [f32],
6036 o2: &mut [f32],
6037 pool: Option<&Pool>,
6038) {
6039 let gpr = cols / GROUP_SIZE;
6040 let v = Q4tpView::new(bytes, rows, cols);
6041 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6042 let run = |start: usize, end: usize| {
6043 let mut sc = vec![0f32; gpr];
6044 for r in start..end {
6045 v.scales_into(r, gpr, &mut sc);
6046 unsafe {
6048 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6049 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6050 }
6051 }
6052 };
6053 dispatch_rows(pool, rows, &run);
6054}
6055
6056#[inline]
6059fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6060 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6061 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6062 let c = (byte >> (2 * (k % 4))) & 3;
6063 (c as f32 - 1.5, scales[gi])
6064}
6065
6066#[cfg(target_arch = "x86_64")]
6067const Q2TP_DECODE_U32: [u32; 256] = {
6068 let mut tab = [0u32; 256];
6069 let mut b = 0usize;
6070 while b < 256 {
6071 tab[b] = ((b as u32) & 3)
6072 | ((((b as u32) >> 2) & 3) << 8)
6073 | ((((b as u32) >> 4) & 3) << 16)
6074 | ((((b as u32) >> 6) & 3) << 24);
6075 b += 1;
6076 }
6077 tab
6078};
6079
6080#[cfg(target_arch = "x86_64")]
6084#[target_feature(enable = "avx2")]
6085unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6086 use core::arch::x86_64::*;
6087 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6088 let codes = _mm256_setr_epi32(
6089 Q2TP_DECODE_U32[ch[0] as usize] as i32,
6090 Q2TP_DECODE_U32[ch[1] as usize] as i32,
6091 Q2TP_DECODE_U32[ch[2] as usize] as i32,
6092 Q2TP_DECODE_U32[ch[3] as usize] as i32,
6093 Q2TP_DECODE_U32[ch[4] as usize] as i32,
6094 Q2TP_DECODE_U32[ch[5] as usize] as i32,
6095 Q2TP_DECODE_U32[ch[6] as usize] as i32,
6096 Q2TP_DECODE_U32[ch[7] as usize] as i32,
6097 );
6098 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6099 let pair = _mm256_maddubs_epi16(codes, xv);
6100 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6101 let sum128 = _mm_add_epi32(
6102 _mm256_castsi256_si128(quad),
6103 _mm256_extracti128_si256(quad, 1),
6104 );
6105 let sum64 = _mm_hadd_epi32(sum128, sum128);
6106 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6107}
6108
6109#[cfg(target_arch = "x86_64")]
6110#[inline]
6111fn q2tp_avx2_enabled() -> bool {
6112 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6113 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6114}
6115
6116#[inline]
6123fn dot_q2tp_row_i8(
6124 chunks: &[u8],
6125 r: usize,
6126 gpr: usize,
6127 xq: &[i8],
6128 gsum: &[i32],
6129 scales: &[f32],
6130) -> f32 {
6131 let mut acc = 0f32;
6132 let base = r * gpr * Q2TP_CHUNK;
6133 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6134 let mut codes = [0i8; GROUP_SIZE];
6135 #[cfg(target_arch = "x86_64")]
6136 let avx2 = q2tp_avx2_enabled();
6141 for gi in 0..gpr {
6142 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
6143 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6144 #[cfg(target_arch = "aarch64")]
6145 let dot = unsafe {
6151 use core::arch::aarch64::*;
6152 let b = vld1_u8(ch.as_ptr());
6153 let three = vdup_n_u8(3);
6154 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
6155 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
6156 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
6157 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
6158 let x4 = vld4_s8(xg.as_ptr());
6159 let mut acc4 = vdupq_n_s32(0);
6160 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
6161 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
6162 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
6163 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
6164 vaddvq_s32(acc4)
6165 };
6166 #[cfg(target_arch = "x86_64")]
6167 let dot: i32 = if avx2 {
6168 unsafe { q2tp_code_dot_avx2(ch, xg) }
6171 } else {
6172 ch.iter()
6173 .enumerate()
6174 .map(|(k, &b)| {
6175 ((b & 3) as i32) * xg[k * 4] as i32
6176 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
6177 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
6178 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
6179 })
6180 .sum()
6181 };
6182 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6183 let dot: i32 = {
6184 for (k, &b) in ch.iter().enumerate() {
6185 codes[k * 4] = (b & 3) as i8;
6186 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6187 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6188 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6189 }
6190 codes
6191 .iter()
6192 .zip(xg)
6193 .map(|(&c, &x)| c as i32 * x as i32)
6194 .sum()
6195 };
6196 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6197 }
6198 acc
6199}
6200
6201fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6205 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6206}
6207
6208#[inline]
6212fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6213 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6214}
6215
6216#[inline]
6217fn q2tp_row_exact_center(
6218 chunks: &[u8],
6219 r: usize,
6220 gpr: usize,
6221 x: &[f32],
6222 scales: &[f32],
6223 center: f32,
6224) -> f32 {
6225 let mut acc = 0f32;
6226 for gi in 0..gpr {
6227 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6228 let s = scales[gi];
6229 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6230 let mut g = 0f32;
6231 for (k, &b) in ch.iter().enumerate() {
6232 g += ((b & 3) as f32 - center) * xb[k * 4]
6233 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6234 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6235 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6236 }
6237 acc += s * g;
6238 }
6239 acc
6240}
6241
6242fn q2tp_matvec(
6243 bytes: &[u8],
6244 x: &[f32],
6245 rows: usize,
6246 cols: usize,
6247 out: &mut [f32],
6248 pool: Option<&Pool>,
6249) {
6250 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6251}
6252
6253fn q2tp_affine_matvec(
6254 bytes: &[u8],
6255 x: &[f32],
6256 rows: usize,
6257 cols: usize,
6258 out: &mut [f32],
6259 pool: Option<&Pool>,
6260) {
6261 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6262}
6263
6264fn q2tp_matvec_mode(
6265 bytes: &[u8],
6266 x: &[f32],
6267 rows: usize,
6268 cols: usize,
6269 out: &mut [f32],
6270 pool: Option<&Pool>,
6271 affine: bool,
6272) {
6273 debug_assert_eq!(out.len(), rows);
6274 let gpr = cols / GROUP_SIZE;
6275 let v = Q4tpView::new_q2(bytes, rows, cols);
6276 let out_addr = SendMut(out.as_mut_ptr());
6277 if !affine && a8w8_enabled() {
6282 let act = split_act(x);
6283 let gsum = q1_group_sums(&act.xq, gpr);
6284 let (act, gsum) = (&act, &gsum);
6285 let run = move |start: usize, end: usize| {
6286 with_krow(gpr, |sc| {
6287 for r in start..end {
6288 v.scales_into(r, gpr, sc);
6289 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6290 for &(j, xv) in &act.outliers {
6291 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6292 acc += w * s * xv;
6293 }
6294 unsafe { *out_addr.at(r) = acc };
6296 }
6297 })
6298 };
6299 dispatch_rows(pool, rows, &run);
6300 return;
6301 }
6302 let run = |start: usize, end: usize| {
6303 with_krow(gpr, |sc| {
6304 for r in start..end {
6305 v.scales_into(r, gpr, sc);
6306 unsafe {
6308 *out_addr.at(r) = if affine {
6309 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6310 } else {
6311 q2tp_row_exact(v.nib, r, gpr, x, sc)
6312 }
6313 };
6314 }
6315 })
6316 };
6317 dispatch_rows(pool, rows, &run);
6318}
6319
6320#[allow(clippy::too_many_arguments)]
6322fn q2tp_matvec2(
6323 bytes: &[u8],
6324 x1: &[f32],
6325 x2: &[f32],
6326 rows: usize,
6327 cols: usize,
6328 o1: &mut [f32],
6329 o2: &mut [f32],
6330 pool: Option<&Pool>,
6331) {
6332 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6333}
6334
6335#[allow(clippy::too_many_arguments)]
6336fn q2tp_affine_matvec2(
6337 bytes: &[u8],
6338 x1: &[f32],
6339 x2: &[f32],
6340 rows: usize,
6341 cols: usize,
6342 o1: &mut [f32],
6343 o2: &mut [f32],
6344 pool: Option<&Pool>,
6345) {
6346 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6347}
6348
6349#[allow(clippy::too_many_arguments)]
6350fn q2tp_matvec2_mode(
6351 bytes: &[u8],
6352 x1: &[f32],
6353 x2: &[f32],
6354 rows: usize,
6355 cols: usize,
6356 o1: &mut [f32],
6357 o2: &mut [f32],
6358 pool: Option<&Pool>,
6359 affine: bool,
6360) {
6361 let gpr = cols / GROUP_SIZE;
6362 let v = Q4tpView::new_q2(bytes, rows, cols);
6363 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6364 let run = |start: usize, end: usize| {
6365 let mut sc = vec![0f32; gpr];
6366 for r in start..end {
6367 v.scales_into(r, gpr, &mut sc);
6368 unsafe {
6370 *p1.at(r) = if affine {
6371 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6372 } else {
6373 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6374 };
6375 *p2.at(r) = if affine {
6376 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6377 } else {
6378 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6379 };
6380 }
6381 }
6382 };
6383 dispatch_rows(pool, rows, &run);
6384}
6385
6386pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6393 let gpr = cols / GROUP_SIZE;
6398 let v = Q4tpView::new_q2(bytes, rows, cols);
6399 with_krow(gpr, |sc| {
6400 for r in 0..rows {
6401 v.scales_into(r, gpr, sc);
6402 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6403 }
6404 });
6405}
6406
6407pub fn q2tp_affine_matvec_for_test(
6410 bytes: &[u8],
6411 x: &[f32],
6412 rows: usize,
6413 cols: usize,
6414 out: &mut [f32],
6415) {
6416 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6417}
6418
6419pub fn q2tp_matmat_for_test(
6420 bytes: &[u8],
6421 xs_all: &[f32],
6422 b: usize,
6423 rows: usize,
6424 cols: usize,
6425 out: &mut [f32],
6426) {
6427 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6428}
6429
6430fn q2tp_matmat(
6431 bytes: &[u8],
6432 xs_all: &[f32],
6433 b: usize,
6434 rows: usize,
6435 cols: usize,
6436 out: &mut [f32],
6437 pool: Option<&Pool>,
6438) {
6439 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6440}
6441
6442fn q2tp_affine_matmat(
6443 bytes: &[u8],
6444 xs_all: &[f32],
6445 b: usize,
6446 rows: usize,
6447 cols: usize,
6448 out: &mut [f32],
6449 pool: Option<&Pool>,
6450) {
6451 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6452}
6453
6454fn q2tp_matmat_mode(
6455 bytes: &[u8],
6456 xs_all: &[f32],
6457 b: usize,
6458 rows: usize,
6459 cols: usize,
6460 out: &mut [f32],
6461 pool: Option<&Pool>,
6462 affine: bool,
6463) {
6464 debug_assert_eq!(out.len(), b * rows);
6465 let gpr = cols / GROUP_SIZE;
6466 let v = Q4tpView::new_q2(bytes, rows, cols);
6467 let out_addr = SendMut(out.as_mut_ptr());
6468 let run = |start: usize, end: usize| {
6469 let mut sc = vec![0f32; gpr];
6470 for r in start..end {
6471 v.scales_into(r, gpr, &mut sc);
6472 for bi in 0..b {
6473 let x = &xs_all[bi * cols..(bi + 1) * cols];
6474 unsafe {
6476 *out_addr.at(bi * rows + r) = if affine {
6477 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6478 } else {
6479 q2tp_row_exact(v.nib, r, gpr, x, &sc)
6480 }
6481 };
6482 }
6483 }
6484 };
6485 dispatch_rows(pool, rows, &run);
6486}
6487
6488#[cfg(target_arch = "aarch64")]
6496#[target_feature(enable = "neon,dotprod")]
6497unsafe fn dot_q4tp_row_1x4_sdot_v1(
6498 nib: &[u8],
6499 r: usize,
6500 gpr: usize,
6501 xs: [&[i8]; 4],
6502 scales: &[f32],
6503) -> [f32; 4] {
6504 unsafe {
6505 use core::arch::aarch64::*;
6506 use core::arch::asm;
6507 let lomask = vdupq_n_u8(0x0F);
6508 let eight = vdupq_n_s8(8);
6509 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6510 for gi in 0..gpr {
6511 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6512 let s = *scales.get_unchecked(gi);
6513 let bb = vld1q_u8(t);
6514 let lo = vandq_u8(bb, lomask);
6515 let hi = vshrq_n_u8::<4>(bb);
6516 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6517 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6518 let mut d = [0f32; 4];
6519 for (k, dk) in d.iter_mut().enumerate() {
6520 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6521 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6522 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6523 asm!(
6524 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6525 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6526 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6527 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6528 options(pure, nomem, nostack),
6529 );
6530 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6531 }
6532 f0 += d[0];
6533 f1 += d[1];
6534 f2 += d[2];
6535 f3 += d[3];
6536 }
6537 [f0, f1, f2, f3]
6538 }
6539}
6540
6541#[allow(dead_code)]
6548static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6549
6550#[cfg(test)]
6553static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6554
6555#[cfg(target_arch = "x86_64")]
6561fn q4tp_blocked_x86() -> bool {
6562 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6563 1 => false,
6564 2 => avx512vnni_enabled(),
6569 _ => blocked_enabled() && avx512vnni_enabled(),
6573 }
6574}
6575
6576#[cfg(target_arch = "aarch64")]
6578#[allow(dead_code)]
6579fn q4tp_v1() -> bool {
6580 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6581 1 => true,
6582 2 => false,
6583 _ => {
6584 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6585 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6586 }
6587 }
6588}
6589
6590#[cfg(target_arch = "x86_64")]
6602#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6603unsafe fn dot_q4tp_2x8_avx512(
6604 nib: &[u8],
6605 r0: usize,
6606 gpr: usize,
6607 xs: [&[i8]; 8],
6608 sc0: &[f32],
6609 sc1: &[f32],
6610) -> [[f32; 8]; 2] {
6611 unsafe {
6614 use core::arch::x86_64::*;
6615 let lomask = _mm256_set1_epi8(0x0F);
6616 let eight = _mm256_set1_epi8(8);
6617 let zero = _mm512_setzero_si512();
6618 let mut v0 = [_mm512_setzero_ps(); 8];
6619 let mut v1 = [_mm512_setzero_ps(); 8];
6620 let pairs = gpr / 2;
6621 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6622 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6623 let bb = _mm256_loadu_si256(t as *const __m256i);
6624 let lo = _mm256_and_si256(bb, lomask);
6625 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6626 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6627 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6628 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6629 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6630 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6631 };
6632 for gp in 0..pairs {
6633 let gi = gp * 2;
6634 let (wa0, neg0) = unpack(r0, gi);
6635 let (wa1, neg1) = unpack(r0 + 1, gi);
6636 let off = gi * GROUP_SIZE;
6637 let sv = |sc: &[f32]| {
6638 _mm512_insertf32x8::<1>(
6639 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6640 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6641 )
6642 };
6643 let s0 = sv(sc0);
6644 let s1 = sv(sc1);
6645 for k in 0..8 {
6646 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6647 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6648 zero,
6649 wa0,
6650 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6651 ));
6652 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6653 zero,
6654 wa1,
6655 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6656 ));
6657 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6658 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6659 }
6660 }
6661 let mut acc = [[0f32; 8]; 2];
6662 for k in 0..8 {
6663 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6664 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6665 }
6666 if gpr % 2 == 1 {
6667 let off = (gpr - 1) * GROUP_SIZE;
6668 for j in off..off + GROUP_SIZE {
6669 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6670 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6671 for k in 0..8 {
6672 let x = *xs[k].get_unchecked(j) as f32;
6673 acc[0][k] += w0 * sa * x;
6674 acc[1][k] += w1 * sb * x;
6675 }
6676 }
6677 }
6678 acc
6679 }
6680}
6681
6682#[cfg(target_arch = "x86_64")]
6687#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6688unsafe fn dot_q4tp_row_1x8_avx512(
6689 nib: &[u8],
6690 r: usize,
6691 gpr: usize,
6692 xs: [&[i8]; 8],
6693 scales: &[f32],
6694) -> [f32; 8] {
6695 unsafe {
6697 use core::arch::x86_64::*;
6698 let lomask = _mm256_set1_epi8(0x0F);
6699 let eight = _mm256_set1_epi8(8);
6700 let zero = _mm512_setzero_si512();
6701 let (mut v0, mut v1, mut v2, mut v3) = (
6702 _mm512_setzero_ps(),
6703 _mm512_setzero_ps(),
6704 _mm512_setzero_ps(),
6705 _mm512_setzero_ps(),
6706 );
6707 let (mut v4, mut v5, mut v6, mut v7) = (
6708 _mm512_setzero_ps(),
6709 _mm512_setzero_ps(),
6710 _mm512_setzero_ps(),
6711 _mm512_setzero_ps(),
6712 );
6713 let pairs = gpr / 2;
6714 for gp in 0..pairs {
6715 let gi = gp * 2;
6716 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6717 let bb = _mm256_loadu_si256(t as *const __m256i);
6718 let lo = _mm256_and_si256(bb, lomask);
6719 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6720 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6725 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6726 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6727 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6728 let wabs = _mm512_abs_epi8(w);
6729 let neg = _mm512_movepi8_mask(w);
6730 let off = gi * GROUP_SIZE;
6731 let sv = _mm512_insertf32x8::<1>(
6732 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6733 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6734 );
6735 let dot = |x: &[i8]| -> __m512 {
6736 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6737 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6738 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6739 };
6740 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6741 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6742 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6743 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6744 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6745 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6746 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6747 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6748 }
6749 let mut acc = [
6750 _mm512_reduce_add_ps(v0),
6751 _mm512_reduce_add_ps(v1),
6752 _mm512_reduce_add_ps(v2),
6753 _mm512_reduce_add_ps(v3),
6754 _mm512_reduce_add_ps(v4),
6755 _mm512_reduce_add_ps(v5),
6756 _mm512_reduce_add_ps(v6),
6757 _mm512_reduce_add_ps(v7),
6758 ];
6759 if gpr % 2 == 1 {
6762 let off = (gpr - 1) * GROUP_SIZE;
6763 for j in off..off + GROUP_SIZE {
6764 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6765 let ws = w * s;
6766 for k in 0..8 {
6767 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6768 }
6769 }
6770 }
6771 acc
6772 }
6773}
6774
6775#[cfg(target_arch = "x86_64")]
6788#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6789unsafe fn dot_q4tp_row_1x4_avx512(
6790 nib: &[u8],
6791 r: usize,
6792 gpr: usize,
6793 xs: [&[i8]; 4],
6794 scales: &[f32],
6795) -> [f32; 4] {
6796 unsafe {
6798 use core::arch::x86_64::*;
6799 let lomask = _mm256_set1_epi8(0x0F);
6800 let eight = _mm256_set1_epi8(8);
6801 let zero = _mm512_setzero_si512();
6802 let (mut v0, mut v1, mut v2, mut v3) = (
6803 _mm512_setzero_ps(),
6804 _mm512_setzero_ps(),
6805 _mm512_setzero_ps(),
6806 _mm512_setzero_ps(),
6807 );
6808 let pairs = gpr / 2;
6809 for gp in 0..pairs {
6810 let gi = gp * 2;
6811 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6812 let bb = _mm256_loadu_si256(t as *const __m256i);
6813 let lo = _mm256_and_si256(bb, lomask);
6814 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6815 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6820 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6821 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6822 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6823 let wabs = _mm512_abs_epi8(w);
6824 let neg = _mm512_movepi8_mask(w);
6825 let off = gi * GROUP_SIZE;
6826 let sv = _mm512_insertf32x8::<1>(
6827 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6828 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6829 );
6830 let dot = |x: &[i8]| -> __m512 {
6831 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6832 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6833 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6834 };
6835 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6836 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6837 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6838 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6839 }
6840 let mut acc = [
6841 _mm512_reduce_add_ps(v0),
6842 _mm512_reduce_add_ps(v1),
6843 _mm512_reduce_add_ps(v2),
6844 _mm512_reduce_add_ps(v3),
6845 ];
6846 if gpr % 2 == 1 {
6849 let off = (gpr - 1) * GROUP_SIZE;
6850 for j in off..off + GROUP_SIZE {
6851 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6852 let ws = w * s;
6853 for k in 0..4 {
6854 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6855 }
6856 }
6857 }
6858 acc
6859 }
6860}
6861
6862#[cfg(target_arch = "aarch64")]
6866#[target_feature(enable = "neon,dotprod")]
6867unsafe fn dot_q4tp_row_1x4_sdot(
6868 nib: &[u8],
6869 r: usize,
6870 gpr: usize,
6871 xs: [&[i8]; 4],
6872 scales: &[f32],
6873) -> [f32; 4] {
6874 unsafe {
6876 use core::arch::aarch64::*;
6877 use core::arch::asm;
6878 let lomask = vdupq_n_u8(0x0F);
6879 let eight = vdupq_n_s8(8);
6880 let (mut v0, mut v1, mut v2, mut v3) = (
6896 vdupq_n_f32(0.0),
6897 vdupq_n_f32(0.0),
6898 vdupq_n_f32(0.0),
6899 vdupq_n_f32(0.0),
6900 );
6901 for gi in 0..gpr {
6902 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6903 let s = *scales.get_unchecked(gi);
6904 let bb = vld1q_u8(t);
6905 let lo = vandq_u8(bb, lomask);
6906 let hi = vshrq_n_u8::<4>(bb);
6907 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6908 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6909 let off = gi * GROUP_SIZE;
6910 let dot4 = |x: &[i8]| -> int32x4_t {
6911 let x0 = vld1q_s8(x.as_ptr().add(off));
6912 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6913 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6914 asm!(
6915 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6916 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6917 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6918 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6919 options(pure, nomem, nostack),
6920 );
6921 vaddq_s32(a0, a1)
6922 };
6923 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6924 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6925 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6926 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6927 }
6928 [
6929 vaddvq_f32(v0),
6930 vaddvq_f32(v1),
6931 vaddvq_f32(v2),
6932 vaddvq_f32(v3),
6933 ]
6934 }
6935}
6936
6937fn q4tp_matmat(
6945 bytes: &[u8],
6946 xs_all: &[f32],
6947 b: usize,
6948 rows: usize,
6949 cols: usize,
6950 out: &mut [f32],
6951 pool: Option<&Pool>,
6952) {
6953 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6954}
6955
6956#[allow(clippy::too_many_arguments)]
6960fn q4tp_matmat_with(
6961 bytes: &[u8],
6962 xs_all: &[f32],
6963 b: usize,
6964 rows: usize,
6965 cols: usize,
6966 out: &mut [f32],
6967 pool: Option<&Pool>,
6968 exact: bool,
6969) {
6970 debug_assert_eq!(out.len(), b * rows);
6971 let gpr = cols / GROUP_SIZE;
6972 let v = Q4tpView::new(bytes, rows, cols);
6973
6974 #[cfg(target_os = "macos")]
6978 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6979 dequant_matmat_accel(
6980 &|r, dst| {
6981 let mut sc = [0f32; 32];
6982 let mut scv;
6983 let s: &[f32] = if gpr <= 32 {
6984 v.scales_into(r, gpr, &mut sc);
6985 &sc[..gpr]
6986 } else {
6987 scv = vec![0f32; gpr];
6988 v.scales_into(r, gpr, &mut scv);
6989 &scv
6990 };
6991 for gi in 0..gpr {
6992 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6993 for (k, &bb) in tile.iter().enumerate() {
6994 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6995 dst[gi * GROUP_SIZE + k * 2 + 1] =
6996 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6997 }
6998 }
6999 },
7000 xs_all,
7001 b,
7002 rows,
7003 cols,
7004 out,
7005 pool,
7006 );
7007 return;
7008 }
7009
7010 let out_addr = SendMut(out.as_mut_ptr());
7011 if a8w8_enabled() {
7012 let acts: Vec<SplitAct> = (0..b)
7013 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7014 .collect();
7015 let acts = &acts;
7016 #[cfg(target_arch = "aarch64")]
7023 let blocked_ok = sdot_enabled() && blocked_enabled();
7024 #[cfg(target_arch = "x86_64")]
7031 let blocked_ok = q4tp_blocked_x86() && !exact;
7032 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7033 let blocked_ok = {
7034 let _ = exact;
7035 false
7036 };
7037 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7047 .ok()
7048 .and_then(|v| v.parse().ok())
7049 .filter(|v| *v > 0)
7050 .unwrap_or(256);
7051 let run = |start: usize, end: usize| {
7052 for abase in (0..acts.len()).step_by(panel_cols) {
7053 let alen = (acts.len() - abase).min(panel_cols);
7054 let mut sc = vec![0f32; gpr];
7055 #[cfg(target_arch = "x86_64")]
7056 let mut r_lo = start;
7057 #[cfg(target_arch = "x86_64")]
7058 if blocked_ok && alen >= 8 {
7059 let mut sc1 = vec![0f32; gpr];
7060 while r_lo + 2 <= end {
7061 v.scales_into(r_lo, gpr, &mut sc);
7062 v.scales_into(r_lo + 1, gpr, &mut sc1);
7063 let mut bi = 0usize;
7064 while bi + 8 <= alen {
7065 let xs = [
7066 acts[abase + bi].xq.as_slice(),
7067 acts[abase + bi + 1].xq.as_slice(),
7068 acts[abase + bi + 2].xq.as_slice(),
7069 acts[abase + bi + 3].xq.as_slice(),
7070 acts[abase + bi + 4].xq.as_slice(),
7071 acts[abase + bi + 5].xq.as_slice(),
7072 acts[abase + bi + 6].xq.as_slice(),
7073 acts[abase + bi + 7].xq.as_slice(),
7074 ];
7075 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7076 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7077 for k in 0..8 {
7078 let act = &acts[abase + bi + k];
7079 let mut acc = dr[k] * act.sx;
7080 for &(j, xv) in &act.outliers {
7081 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7082 acc += w * s * xv;
7083 }
7084 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7086 }
7087 }
7088 bi += 8;
7089 }
7090 for row in [r_lo, r_lo + 1] {
7093 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7094 for b2 in bi..alen {
7095 let act = &acts[abase + b2];
7096 let xs4 = [
7097 act.xq.as_slice(),
7098 act.xq.as_slice(),
7099 act.xq.as_slice(),
7100 act.xq.as_slice(),
7101 ];
7102 let d =
7103 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7104 let mut acc = d[0] * act.sx;
7105 for &(j, xv) in &act.outliers {
7106 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7107 acc += w * s * xv;
7108 }
7109 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7111 }
7112 }
7113 r_lo += 2;
7114 }
7115 }
7116 #[cfg(target_arch = "x86_64")]
7117 let row_start = r_lo;
7118 #[cfg(not(target_arch = "x86_64"))]
7119 let row_start = start;
7120 for r in row_start..end {
7121 v.scales_into(r, gpr, &mut sc);
7122 let mut bi = 0usize;
7123 #[cfg(target_arch = "x86_64")]
7124 if blocked_ok {
7125 while bi + 8 <= alen {
7126 let xs = [
7127 acts[abase + bi].xq.as_slice(),
7128 acts[abase + bi + 1].xq.as_slice(),
7129 acts[abase + bi + 2].xq.as_slice(),
7130 acts[abase + bi + 3].xq.as_slice(),
7131 acts[abase + bi + 4].xq.as_slice(),
7132 acts[abase + bi + 5].xq.as_slice(),
7133 acts[abase + bi + 6].xq.as_slice(),
7134 acts[abase + bi + 7].xq.as_slice(),
7135 ];
7136 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
7137 for k in 0..8 {
7138 let act = &acts[abase + bi + k];
7139 let mut acc = d[k] * act.sx;
7140 for &(j, xv) in &act.outliers {
7141 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7142 acc += w * s * xv;
7143 }
7144 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7146 }
7147 bi += 8;
7148 }
7149 while bi + 4 <= alen {
7150 let xs = [
7151 acts[abase + bi].xq.as_slice(),
7152 acts[abase + bi + 1].xq.as_slice(),
7153 acts[abase + bi + 2].xq.as_slice(),
7154 acts[abase + bi + 3].xq.as_slice(),
7155 ];
7156 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
7157 for k in 0..4 {
7158 let act = &acts[abase + bi + k];
7159 let mut acc = d[k] * act.sx;
7160 for &(j, xv) in &act.outliers {
7161 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7162 acc += w * s * xv;
7163 }
7164 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7166 }
7167 bi += 4;
7168 }
7169 }
7170 #[cfg(target_arch = "aarch64")]
7171 if blocked_ok {
7172 while bi + 4 <= alen {
7173 let xs = [
7174 acts[abase + bi].xq.as_slice(),
7175 acts[abase + bi + 1].xq.as_slice(),
7176 acts[abase + bi + 2].xq.as_slice(),
7177 acts[abase + bi + 3].xq.as_slice(),
7178 ];
7179 let d = unsafe {
7180 if exact || q4tp_v1() {
7181 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
7182 } else {
7183 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7184 }
7185 };
7186 for k in 0..4 {
7187 let act = &acts[abase + bi + k];
7188 let mut acc = d[k] * act.sx;
7189 for &(j, xv) in &act.outliers {
7190 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7191 acc += w * s * xv;
7192 }
7193 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7195 }
7196 bi += 4;
7197 }
7198 }
7199 let _ = blocked_ok;
7200 while bi < alen {
7201 let act = &acts[abase + bi];
7202 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7203 for &(j, xv) in &act.outliers {
7204 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7205 acc += w * s * xv;
7206 }
7207 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7209 bi += 1;
7210 }
7211 }
7212 }
7213 };
7214 dispatch_rows(pool, rows, &run);
7215 return;
7216 }
7217
7218 let run = |start: usize, end: usize| {
7219 let mut sc = vec![0f32; gpr];
7220 for r in start..end {
7221 v.scales_into(r, gpr, &mut sc);
7222 for bi in 0..b {
7223 let x = &xs_all[bi * cols..(bi + 1) * cols];
7224 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7226 }
7227 }
7228 };
7229 dispatch_rows(pool, rows, &run);
7230}
7231
7232fn q4t_matvec(
7234 bytes: &[u8],
7235 x: &[f32],
7236 rows: usize,
7237 cols: usize,
7238 out: &mut [f32],
7239 pool: Option<&Pool>,
7240) {
7241 debug_assert_eq!(out.len(), rows);
7242 let gpr = cols / GROUP_SIZE;
7243 let out_addr = SendMut(out.as_mut_ptr());
7244 if a8w8_enabled() {
7245 let act = split_act(x);
7246 let run = move |start: usize, end: usize| {
7247 for r in start..end {
7248 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7249 for &(j, xv) in &act.outliers {
7250 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7251 acc += w * s * xv;
7252 }
7253 unsafe { *out_addr.at(r) = acc };
7255 }
7256 };
7257 dispatch_rows(pool, rows, &run);
7258 return;
7259 }
7260 let run = move |start: usize, end: usize| {
7261 for r in start..end {
7262 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7264 }
7265 };
7266 dispatch_rows(pool, rows, &run);
7267}
7268
7269#[allow(clippy::too_many_arguments)]
7271fn q4t_matvec2(
7272 bytes: &[u8],
7273 x1: &[f32],
7274 x2: &[f32],
7275 rows: usize,
7276 cols: usize,
7277 o1: &mut [f32],
7278 o2: &mut [f32],
7279 pool: Option<&Pool>,
7280) {
7281 let gpr = cols / GROUP_SIZE;
7282 let p1 = SendMut(o1.as_mut_ptr());
7283 let p2 = SendMut(o2.as_mut_ptr());
7284 if a8w8_enabled() {
7285 let a1 = split_act(x1);
7286 let a2 = split_act(x2);
7287 let run = move |start: usize, end: usize| {
7288 for r in start..end {
7289 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7290 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7291 for &(j, xv) in &a1.outliers {
7292 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7293 v1 += w * s * xv;
7294 }
7295 for &(j, xv) in &a2.outliers {
7296 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7297 v2 += w * s * xv;
7298 }
7299 unsafe {
7301 *p1.at(r) = v1;
7302 *p2.at(r) = v2;
7303 }
7304 }
7305 };
7306 dispatch_rows(pool, rows, &run);
7307 return;
7308 }
7309 let run = move |start: usize, end: usize| {
7310 for r in start..end {
7311 unsafe {
7313 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7314 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7315 }
7316 }
7317 };
7318 dispatch_rows(pool, rows, &run);
7319}
7320
7321#[allow(clippy::too_many_arguments)]
7323#[cfg(target_os = "macos")]
7329fn dequant_matmat_accel(
7330 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7331 xs_all: &[f32],
7332 b: usize,
7333 rows: usize,
7334 cols: usize,
7335 out: &mut [f32],
7336 pool: Option<&Pool>,
7337) {
7338 const TR: usize = 2048;
7339 thread_local! {
7340 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7341 }
7342 WTILE.with(|wt| {
7343 let mut wtile = wt.borrow_mut();
7344 wtile.resize(TR * cols, 0.0);
7345 let mut r0 = 0usize;
7346 while r0 < rows {
7347 let tr = TR.min(rows - r0);
7348 let wt_addr = SendMut(wtile.as_mut_ptr());
7349 let run = |start: usize, end: usize| {
7350 for r in start..end {
7351 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7353 dequant_row(r0 + r, dst);
7354 }
7355 };
7356 dispatch_rows(pool, tr, &run);
7357 unsafe {
7358 accel_blas::cblas_sgemm(
7359 101, 111, 112, b as i32,
7363 tr as i32,
7364 cols as i32,
7365 1.0,
7366 xs_all.as_ptr(),
7367 cols as i32,
7368 wtile.as_ptr(),
7369 cols as i32,
7370 0.0,
7371 out.as_mut_ptr().add(r0),
7372 rows as i32,
7373 );
7374 }
7375 r0 += tr;
7376 }
7377 });
7378}
7379
7380fn q4t_matmat(
7381 bytes: &[u8],
7382 xs_all: &[f32],
7383 b: usize,
7384 rows: usize,
7385 cols: usize,
7386 out: &mut [f32],
7387 pool: Option<&Pool>,
7388) {
7389 debug_assert_eq!(out.len(), b * rows);
7390 let gpr = cols / GROUP_SIZE;
7391 #[cfg(target_os = "macos")]
7395 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7396 dequant_matmat_accel(
7397 &|r, dst| {
7398 for gi in 0..gpr {
7399 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7400 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7401 for (k, &bb) in tile[2..].iter().enumerate() {
7402 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7403 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7404 }
7405 }
7406 },
7407 xs_all,
7408 b,
7409 rows,
7410 cols,
7411 out,
7412 pool,
7413 );
7414 return;
7415 }
7416 let out_addr = SendMut(out.as_mut_ptr());
7417 if a8w8_enabled() {
7418 let acts: Vec<SplitAct> = (0..b)
7419 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7420 .collect();
7421 let acts = &acts;
7422 #[cfg(target_arch = "x86_64")]
7423 let blocked_ok = avx2_enabled() && blocked_enabled();
7424 #[cfg(target_arch = "aarch64")]
7425 let blocked_ok = sdot_enabled() && blocked_enabled();
7426 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7427 let blocked_ok = false;
7428 let run = move |start: usize, end: usize| {
7429 for r in start..end {
7430 let mut bi = 0usize;
7431 #[cfg(target_arch = "aarch64")]
7432 if blocked_ok {
7433 while bi + 4 <= acts.len() {
7434 let xs = [
7435 acts[bi].xq.as_slice(),
7436 acts[bi + 1].xq.as_slice(),
7437 acts[bi + 2].xq.as_slice(),
7438 acts[bi + 3].xq.as_slice(),
7439 ];
7440 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7441 for k in 0..4 {
7442 let act = &acts[bi + k];
7443 let mut acc = d[k] * act.sx;
7444 for &(j, xv) in &act.outliers {
7445 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7446 acc += w * sc * xv;
7447 }
7448 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7450 }
7451 bi += 4;
7452 }
7453 }
7454 #[cfg(target_arch = "x86_64")]
7455 if blocked_ok {
7456 while bi + 4 <= acts.len() {
7457 let xs = [
7458 acts[bi].xq.as_slice(),
7459 acts[bi + 1].xq.as_slice(),
7460 acts[bi + 2].xq.as_slice(),
7461 acts[bi + 3].xq.as_slice(),
7462 ];
7463 let d = unsafe {
7464 if vnni_tiles_enabled() {
7465 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7466 } else {
7467 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7468 }
7469 };
7470 for k in 0..4 {
7471 let act = &acts[bi + k];
7472 let mut acc = d[k] * act.sx;
7473 for &(j, xv) in &act.outliers {
7474 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7475 acc += w * sc * xv;
7476 }
7477 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7479 }
7480 bi += 4;
7481 }
7482 }
7483 let _ = blocked_ok;
7484 while bi < acts.len() {
7485 let act = &acts[bi];
7486 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7487 for &(j, xv) in &act.outliers {
7488 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7489 acc += w * s * xv;
7490 }
7491 unsafe { *out_addr.at(bi * rows + r) = acc };
7493 bi += 1;
7494 }
7495 }
7496 };
7497 dispatch_rows(pool, rows, &run);
7498 return;
7499 }
7500 let run = move |start: usize, end: usize| {
7501 for r in start..end {
7502 for bi in 0..b {
7503 let x = &xs_all[bi * cols..(bi + 1) * cols];
7504 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7506 }
7507 }
7508 };
7509 dispatch_rows(pool, rows, &run);
7510}
7511
7512fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7521 (0..gpr)
7522 .map(|gi| {
7523 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7524 .iter()
7525 .map(|&v| v as i32)
7526 .sum()
7527 })
7528 .collect()
7529}
7530
7531#[inline]
7535#[allow(unreachable_code)]
7536#[cfg(target_arch = "x86_64")]
7541#[target_feature(enable = "avx2")]
7542unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7543 unsafe {
7545 use core::arch::x86_64::*;
7546 let expand = _mm256_setr_epi8(
7548 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,
7549 3, 3, 3,
7550 );
7551 let bitsel = _mm256_setr_epi8(
7552 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7553 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7554 );
7555 let ones8 = _mm256_set1_epi8(1);
7556 let ones16 = _mm256_set1_epi16(1);
7557 let mut acc = 0f32;
7558 for gi in 0..gpr {
7559 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7560 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7561 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7562 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7563 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7564 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7565 let sel = _mm256_and_si256(x, mask);
7566 let p16 = _mm256_maddubs_epi16(ones8, sel);
7568 let d32 = _mm256_madd_epi16(p16, ones16);
7569 let hi128 = _mm256_extracti128_si256::<1>(d32);
7570 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7571 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7572 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7573 let msum = _mm_cvtsi128_si32(s32);
7574 let d = 2 * msum - gsum[gi];
7577 acc += d as f32 * s;
7578 }
7579 acc
7580 }
7581}
7582
7583#[cfg(target_arch = "x86_64")]
7586#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7587unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7588 unsafe {
7590 use core::arch::x86_64::*;
7591 let expand = _mm256_setr_epi8(
7592 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,
7593 3, 3, 3,
7594 );
7595 let bitsel = _mm256_setr_epi8(
7596 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7597 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7598 );
7599 let ones8 = _mm256_set1_epi8(1);
7600 let mut acc = 0f32;
7601 for gi in 0..gpr {
7602 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7603 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7604 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7605 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7606 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7607 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7608 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7609 let d = 2 * msum - gsum[gi];
7610 acc += d as f32 * s;
7611 }
7612 acc
7613 }
7614}
7615
7616#[cfg(target_arch = "x86_64")]
7618#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7619unsafe fn dot_q1_row_1x4_vnni(
7620 bytes: &[u8],
7621 r: usize,
7622 gpr: usize,
7623 xs: [&[i8]; 4],
7624 gsums: [&[i32]; 4],
7625) -> [f32; 4] {
7626 unsafe {
7628 use core::arch::x86_64::*;
7629 let expand = _mm256_setr_epi8(
7630 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,
7631 3, 3, 3,
7632 );
7633 let bitsel = _mm256_setr_epi8(
7634 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7635 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7636 );
7637 let ones8 = _mm256_set1_epi8(1);
7638 let mut acc = [0f32; 4];
7639 for gi in 0..gpr {
7640 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7641 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7642 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7643 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7644 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7645 for (k, xq) in xs.iter().enumerate() {
7646 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7647 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7648 let d = 2 * msum - gsums[k][gi];
7649 acc[k] += d as f32 * s;
7650 }
7651 }
7652 acc
7653 }
7654}
7655
7656#[cfg(target_arch = "x86_64")]
7659#[target_feature(enable = "avx2")]
7660unsafe fn dot_q1_row_1x4_avx2(
7661 bytes: &[u8],
7662 r: usize,
7663 gpr: usize,
7664 xs: [&[i8]; 4],
7665 gsums: [&[i32]; 4],
7666) -> [f32; 4] {
7667 unsafe {
7669 use core::arch::x86_64::*;
7670 let expand = _mm256_setr_epi8(
7671 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,
7672 3, 3, 3,
7673 );
7674 let bitsel = _mm256_setr_epi8(
7675 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7676 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7677 );
7678 let ones8 = _mm256_set1_epi8(1);
7679 let ones16 = _mm256_set1_epi16(1);
7680 let mut acc = [0f32; 4];
7681 for gi in 0..gpr {
7682 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7683 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7684 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7685 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7686 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7687 for (k, xq) in xs.iter().enumerate() {
7688 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7689 let sel = _mm256_and_si256(x, mask);
7690 let p16 = _mm256_maddubs_epi16(ones8, sel);
7691 let d32 = _mm256_madd_epi16(p16, ones16);
7692 let hi128 = _mm256_extracti128_si256::<1>(d32);
7693 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7694 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7695 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7696 let msum = _mm_cvtsi128_si32(s32);
7697 let d = 2 * msum - gsums[k][gi];
7698 acc[k] += d as f32 * s;
7699 }
7700 }
7701 acc
7702 }
7703}
7704
7705#[allow(unreachable_code)]
7706fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7707 #[cfg(target_arch = "aarch64")]
7708 unsafe {
7709 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7710 }
7711 #[cfg(target_arch = "x86_64")]
7712 if avx2_enabled() {
7713 unsafe {
7714 if vnni_tiles_enabled() {
7715 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7716 }
7717 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7718 }
7719 }
7720 let _ = gsum;
7721 let mut acc = 0f32;
7722 for gi in 0..gpr {
7723 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7724 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7725 let mut d = 0i32;
7726 for (j, &b) in tile[2..].iter().enumerate() {
7727 for k in 0..8 {
7728 let w = ((b >> k) & 1) as i32 * 2 - 1;
7729 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7730 }
7731 }
7732 acc += d as f32 * s;
7733 }
7734 acc
7735}
7736
7737#[cfg(target_arch = "aarch64")]
7746#[target_feature(enable = "neon,dotprod")]
7747unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7748 unsafe {
7751 use core::arch::aarch64::*;
7752 use core::arch::asm;
7753 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7754 let m = vld1q_u8(MASKS.as_ptr());
7755 macro_rules! tile_dot {
7757 ($t:expr, $x:expr) => {{
7758 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7759 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7760 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7761 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7762 let x0 = vld1q_s8($x);
7763 let x1 = vld1q_s8($x.add(16));
7764 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7765 asm!(
7766 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7767 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7768 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7769 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7770 options(pure, nomem, nostack),
7771 );
7772 vaddq_s32(a0, a1)
7773 }};
7774 }
7775 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7784 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7785 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7786 const IW11: [u8; 16] = [
7787 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7788 ];
7789 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7790 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7791 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7792 let isc = vld1_u8(ISC.as_ptr());
7793 macro_rules! tile_dot_tbl {
7795 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7796 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7797 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7798 let x0 = vld1q_s8($x);
7799 let x1 = vld1q_s8($x.add(16));
7800 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7801 asm!(
7802 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7803 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7804 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7805 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7806 options(pure, nomem, nostack),
7807 );
7808 vaddq_s32(a0, a1)
7809 }};
7810 }
7811 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7812 let row_base = r * gpr * Q1_TILE;
7813 let abs_end = bytes.len();
7814 let xp = xq.as_ptr();
7815 let gp = gsum.as_ptr();
7816 let mut accv = vdupq_n_f32(0.0);
7817 let mut gi = 0;
7818 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7821 let t0 = base.add(gi * Q1_TILE);
7822 let ld_a = vld1q_u8(t0);
7823 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7824 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7825 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7826 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7827 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7828 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7830 let g = vld1q_s32(gp.add(gi));
7831 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7832 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7833 let scf: float32x4_t;
7834 asm!(
7835 "fcvtl {o:v}.4s, {i:v}.4h",
7836 o = out(vreg) scf, i = in(vreg) sc16,
7837 options(pure, nomem, nostack),
7838 );
7839 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7840 gi += 4;
7841 }
7842 let mut acc = vaddvq_f32(accv);
7843 while gi < gpr {
7844 let t = base.add(gi * Q1_TILE);
7845 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7846 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7847 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7848 gi += 1;
7849 }
7850 acc
7851 }
7852}
7853
7854#[cfg(target_arch = "aarch64")]
7859#[target_feature(enable = "neon,dotprod")]
7860unsafe fn dot_q1_row_1x4_sdot(
7861 bytes: &[u8],
7862 r: usize,
7863 gpr: usize,
7864 xs: [&[i8]; 4],
7865 gs: [&[i32]; 4],
7866) -> [f32; 4] {
7867 unsafe {
7869 use core::arch::aarch64::*;
7870 use core::arch::asm;
7871 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7872 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7873 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7874 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7875 const IW11: [u8; 16] = [
7876 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7877 ];
7878 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7879 let m = vld1q_u8(MASKS.as_ptr());
7880 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7881 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7882 let isc = vld1_u8(ISC.as_ptr());
7883 macro_rules! sdot2 {
7884 ($w0:expr, $w1:expr, $x:expr) => {{
7885 let x0 = vld1q_s8($x);
7886 let x1 = vld1q_s8($x.add(16));
7887 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7888 asm!(
7889 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7890 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7891 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7892 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7893 options(pure, nomem, nostack),
7894 );
7895 vaddq_s32(a0, a1)
7896 }};
7897 }
7898 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7899 let row_base = r * gpr * Q1_TILE;
7900 let abs_end = bytes.len();
7901 let mut accv = [vdupq_n_f32(0.0); 4];
7902 let mut gi = 0;
7903 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7904 let t0 = base.add(gi * Q1_TILE);
7905 let ld_a = vld1q_u8(t0);
7906 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7907 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7909 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7910 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7911 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7912 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7913 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7914 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7915 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7916 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7917 let scf: float32x4_t;
7918 asm!(
7919 "fcvtl {o:v}.4s, {i:v}.4h",
7920 o = out(vreg) scf, i = in(vreg) sc16,
7921 options(pure, nomem, nostack),
7922 );
7923 for k in 0..4 {
7924 let xp = xs[k].as_ptr();
7925 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7926 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7927 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7928 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7929 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7930 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7931 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7932 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7933 }
7934 gi += 4;
7935 }
7936 let mut acc = [
7937 vaddvq_f32(accv[0]),
7938 vaddvq_f32(accv[1]),
7939 vaddvq_f32(accv[2]),
7940 vaddvq_f32(accv[3]),
7941 ];
7942 while gi < gpr {
7943 let t = base.add(gi * Q1_TILE);
7944 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7945 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7946 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7947 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7948 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7949 for k in 0..4 {
7950 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7951 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7952 }
7953 gi += 1;
7954 }
7955 acc
7956 }
7957}
7958
7959#[inline]
7961fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7962 let gi = j / GROUP_SIZE;
7963 let k = j % GROUP_SIZE;
7964 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7965 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7966 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7967 ((bit as i32 * 2 - 1) as f32, s)
7968}
7969
7970#[inline]
7972fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7973 let mut acc = 0f32;
7974 for gi in 0..gpr {
7975 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7976 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7977 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7978 let mut ga = 0f32;
7979 for (j, &b) in tile[2..].iter().enumerate() {
7980 for k in 0..8 {
7981 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7982 }
7983 }
7984 acc += ga * s;
7985 }
7986 acc
7987}
7988
7989#[allow(clippy::too_many_arguments)]
7992fn bin_range_a8w8(
7993 w: BinW,
7994 gpr: usize,
7995 act: &SplitAct,
7996 gsum: &[i32],
7997 out: SendMut,
7998 start: usize,
7999 end: usize,
8000) {
8001 for r in start..end {
8002 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
8003 for &(j, xv) in &act.outliers {
8004 let (wv, s) = w.outlier(r, gpr, j);
8005 acc += wv * s * xv;
8006 }
8007 unsafe { *out.at(r) = acc };
8009 }
8010}
8011
8012fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
8014 for r in start..end {
8015 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
8017 }
8018}
8019
8020#[derive(Clone, Copy)]
8028pub(crate) enum BinW<'a> {
8029 Q1(&'a [u8]),
8030 Q1g {
8031 bytes: &'a [u8],
8032 sc_off: usize,
8034 rows: usize,
8037 },
8038}
8039
8040impl<'a> BinW<'a> {
8041 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8042 Self::Q1g {
8043 bytes,
8044 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8045 rows,
8046 }
8047 }
8048
8049 #[inline]
8050 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8051 match self {
8052 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8053 Self::Q1g {
8054 bytes,
8055 sc_off,
8056 rows,
8057 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8058 }
8059 }
8060
8061 #[inline]
8062 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8063 match self {
8064 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8065 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8066 }
8067 }
8068
8069 #[inline]
8070 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8071 match self {
8072 Self::Q1(b) => q1_outlier(b, r, gpr, j),
8073 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8074 }
8075 }
8076
8077 #[cfg(target_arch = "aarch64")]
8082 #[inline]
8083 unsafe fn dot_1x4_sdot(
8084 self,
8085 r: usize,
8086 gpr: usize,
8087 xs: [&[i8]; 4],
8088 gs: [&[i32]; 4],
8089 ) -> [f32; 4] {
8090 unsafe {
8092 match self {
8093 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8094 Self::Q1g {
8095 bytes,
8096 sc_off,
8097 rows,
8098 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8099 }
8100 }
8101 }
8102
8103 #[cfg(target_arch = "x86_64")]
8108 #[inline]
8109 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
8110 unsafe {
8112 match (self, vnni_tiles_enabled()) {
8113 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
8114 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
8115 (Self::Q1g { bytes, sc_off, .. }, true) => {
8116 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
8117 }
8118 (Self::Q1g { bytes, sc_off, .. }, false) => {
8119 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
8120 }
8121 }
8122 }
8123 }
8124}
8125
8126#[inline]
8136fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
8137 u32::from_le_bytes([
8138 bytes[t * 4],
8139 bytes[t * 4 + 1],
8140 bytes[t * 4 + 2],
8141 bytes[t * 4 + 3],
8142 ])
8143}
8144
8145#[inline]
8147fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
8148 let o = sc_off + (t / 4) * 2;
8149 u16::from_le_bytes([bytes[o], bytes[o + 1]])
8150}
8151
8152#[allow(unreachable_code)]
8154#[allow(clippy::too_many_arguments)]
8155fn dot_q1g_row_i8(
8156 bytes: &[u8],
8157 sc_off: usize,
8158 rows: usize,
8159 r: usize,
8160 gpr: usize,
8161 xq: &[i8],
8162 gsum: &[i32],
8163) -> f32 {
8164 #[cfg(target_arch = "aarch64")]
8165 unsafe {
8166 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
8167 }
8168 #[cfg(target_arch = "x86_64")]
8169 if avx2_enabled() {
8170 unsafe {
8171 if vnni_tiles_enabled() {
8172 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
8173 }
8174 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
8175 }
8176 }
8177 let _ = (gsum, rows);
8178 let mut acc = 0f32;
8179 for gi in 0..gpr {
8180 let t = r * gpr + gi;
8181 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8182 let bits = q1g_bits(bytes, t);
8183 let mut d = 0i32;
8184 for k in 0..GROUP_SIZE {
8185 let w = ((bits >> k) & 1) as i32 * 2 - 1;
8186 d += w * xq[gi * GROUP_SIZE + k] as i32;
8187 }
8188 acc += d as f32 * s;
8189 }
8190 acc
8191}
8192
8193#[cfg(target_arch = "x86_64")]
8198#[target_feature(enable = "avx2")]
8199#[inline]
8200unsafe fn q1g_hsum4(
8201 d0: core::arch::x86_64::__m256i,
8202 d1: core::arch::x86_64::__m256i,
8203 d2: core::arch::x86_64::__m256i,
8204 d3: core::arch::x86_64::__m256i,
8205) -> core::arch::x86_64::__m128i {
8206 unsafe {
8208 use core::arch::x86_64::*;
8209 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8210 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8211 }
8212}
8213
8214#[cfg(target_arch = "x86_64")]
8218#[target_feature(enable = "avx2")]
8219#[inline]
8220unsafe fn q1g_group_acc(
8221 acc: f32,
8222 msum4: core::arch::x86_64::__m128i,
8223 gsum4: &[i32],
8224 s: f32,
8225) -> f32 {
8226 unsafe {
8228 use core::arch::x86_64::*;
8229 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8230 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8231 let mut f = [0f32; 4];
8232 _mm_storeu_ps(f.as_mut_ptr(), d);
8233 let mut acc = acc;
8234 for v in f {
8235 acc += v * s;
8236 }
8237 acc
8238 }
8239}
8240
8241#[cfg(target_arch = "x86_64")]
8246#[target_feature(enable = "avx2")]
8247unsafe fn dot_q1g_row_avx2(
8248 bytes: &[u8],
8249 sc_off: usize,
8250 r: usize,
8251 gpr: usize,
8252 xq: &[i8],
8253 gsum: &[i32],
8254) -> f32 {
8255 unsafe {
8258 use core::arch::x86_64::*;
8259 let expand = _mm256_setr_epi8(
8260 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,
8261 3, 3, 3,
8262 );
8263 let bitsel = _mm256_setr_epi8(
8264 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8265 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8266 );
8267 let ones8 = _mm256_set1_epi8(1);
8268 let ones16 = _mm256_set1_epi16(1);
8269 let bp = bytes.as_ptr().add(r * gpr * 4);
8270 let mut acc = 0f32;
8271 let mut gi = 0;
8272 while gi < gpr {
8273 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8274 let mut d = [_mm256_setzero_si256(); 4];
8275 for (k, dk) in d.iter_mut().enumerate() {
8276 let t = gi + k;
8277 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8278 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8279 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8280 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8281 let sel = _mm256_and_si256(x, mask);
8282 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8283 }
8284 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8285 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8286 gi += 4;
8287 }
8288 acc
8289 }
8290}
8291
8292#[cfg(target_arch = "x86_64")]
8295#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8296unsafe fn dot_q1g_row_vnni(
8297 bytes: &[u8],
8298 sc_off: usize,
8299 r: usize,
8300 gpr: usize,
8301 xq: &[i8],
8302 gsum: &[i32],
8303) -> f32 {
8304 unsafe {
8306 use core::arch::x86_64::*;
8307 let expand = _mm256_setr_epi8(
8308 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,
8309 3, 3, 3,
8310 );
8311 let bitsel = _mm256_setr_epi8(
8312 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8313 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8314 );
8315 let ones8 = _mm256_set1_epi8(1);
8316 let bp = bytes.as_ptr().add(r * gpr * 4);
8317 let mut acc = 0f32;
8318 let mut gi = 0;
8319 while gi < gpr {
8320 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8321 let mut d = [_mm256_setzero_si256(); 4];
8322 for (k, dk) in d.iter_mut().enumerate() {
8323 let t = gi + k;
8324 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8325 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8326 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8327 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8328 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8329 }
8330 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8331 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8332 gi += 4;
8333 }
8334 acc
8335 }
8336}
8337
8338#[cfg(target_arch = "x86_64")]
8342#[target_feature(enable = "avx2")]
8343unsafe fn dot_q1g_row_1x4_avx2(
8344 bytes: &[u8],
8345 sc_off: usize,
8346 r: usize,
8347 gpr: usize,
8348 xs: [&[i8]; 4],
8349 gsums: [&[i32]; 4],
8350) -> [f32; 4] {
8351 unsafe {
8353 use core::arch::x86_64::*;
8354 let expand = _mm256_setr_epi8(
8355 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,
8356 3, 3, 3,
8357 );
8358 let bitsel = _mm256_setr_epi8(
8359 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8360 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8361 );
8362 let ones8 = _mm256_set1_epi8(1);
8363 let ones16 = _mm256_set1_epi16(1);
8364 let bp = bytes.as_ptr().add(r * gpr * 4);
8365 let mut acc = [0f32; 4];
8366 let mut gi = 0;
8367 while gi < gpr {
8368 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8369 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8371 for t in 0..4 {
8372 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8373 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8374 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8375 for (k, xq) in xs.iter().enumerate() {
8376 let x = _mm256_loadu_si256(
8377 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8378 );
8379 let sel = _mm256_and_si256(x, mask);
8380 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8381 }
8382 }
8383 for k in 0..4 {
8384 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8385 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8386 }
8387 gi += 4;
8388 }
8389 acc
8390 }
8391}
8392
8393#[cfg(target_arch = "x86_64")]
8395#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8396unsafe fn dot_q1g_row_1x4_vnni(
8397 bytes: &[u8],
8398 sc_off: usize,
8399 r: usize,
8400 gpr: usize,
8401 xs: [&[i8]; 4],
8402 gsums: [&[i32]; 4],
8403) -> [f32; 4] {
8404 unsafe {
8406 use core::arch::x86_64::*;
8407 let expand = _mm256_setr_epi8(
8408 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,
8409 3, 3, 3,
8410 );
8411 let bitsel = _mm256_setr_epi8(
8412 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8413 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8414 );
8415 let ones8 = _mm256_set1_epi8(1);
8416 let bp = bytes.as_ptr().add(r * gpr * 4);
8417 let mut acc = [0f32; 4];
8418 let mut gi = 0;
8419 while gi < gpr {
8420 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8421 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8422 for t in 0..4 {
8423 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8424 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8425 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8426 for (k, xq) in xs.iter().enumerate() {
8427 let x = _mm256_loadu_si256(
8428 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8429 );
8430 d[k][t] = _mm256_dpbusd_epi32(
8431 _mm256_setzero_si256(),
8432 ones8,
8433 _mm256_and_si256(x, mask),
8434 );
8435 }
8436 }
8437 for k in 0..4 {
8438 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8439 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8440 }
8441 gi += 4;
8442 }
8443 acc
8444 }
8445}
8446
8447#[cfg(target_arch = "aarch64")]
8458#[target_feature(enable = "neon,dotprod")]
8459#[allow(clippy::too_many_arguments)]
8460unsafe fn dot_q1g_row_sdot(
8461 bytes: &[u8],
8462 sc_off: usize,
8463 rows: usize,
8464 r: usize,
8465 gpr: usize,
8466 xq: &[i8],
8467 gsum: &[i32],
8468) -> f32 {
8469 unsafe {
8472 use core::arch::aarch64::*;
8473 use core::arch::asm;
8474 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8475 const IW: [[u8; 16]; 8] = [
8477 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8478 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8479 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8480 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8481 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8482 [
8483 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8484 ],
8485 [
8486 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8487 ],
8488 [
8489 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8490 ],
8491 ];
8492 let m = vld1q_u8(MASKS.as_ptr());
8493 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8494 macro_rules! tile_dot {
8495 ($w0:expr, $w1:expr, $x:expr) => {{
8496 let x0 = vld1q_s8($x);
8497 let x1 = vld1q_s8($x.add(16));
8498 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8499 asm!(
8500 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8501 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8502 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8503 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8504 options(pure, nomem, nostack),
8505 );
8506 vaddq_s32(a0, a1)
8507 }};
8508 }
8509 macro_rules! tbl_mask {
8510 ($ld:expr, $i:expr) => {
8511 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8512 };
8513 }
8514 let bp = bytes.as_ptr().add(r * gpr * 4);
8515 let xp = xq.as_ptr();
8516 let gp = gsum.as_ptr();
8517 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8518 let mut accv = vdupq_n_f32(0.0);
8519 let mut gi = 0;
8520 while gi + 4 <= vec_end {
8521 let ld = vld1q_u8(bp.add(gi * 4));
8522 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8523 let d1 = tile_dot!(
8524 tbl_mask!(ld, 2),
8525 tbl_mask!(ld, 3),
8526 xp.add((gi + 1) * GROUP_SIZE)
8527 );
8528 let d2 = tile_dot!(
8529 tbl_mask!(ld, 4),
8530 tbl_mask!(ld, 5),
8531 xp.add((gi + 2) * GROUP_SIZE)
8532 );
8533 let d3 = tile_dot!(
8534 tbl_mask!(ld, 6),
8535 tbl_mask!(ld, 7),
8536 xp.add((gi + 3) * GROUP_SIZE)
8537 );
8538 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8539 let g = vld1q_s32(gp.add(gi));
8540 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8541 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8542 let scf: float32x4_t;
8543 asm!(
8544 "fcvtl {o:v}.4s, {i:v}.4h",
8545 o = out(vreg) scf, i = in(vreg) h,
8546 options(pure, nomem, nostack),
8547 );
8548 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8549 gi += 4;
8550 }
8551 let mut acc = vaddvq_f32(accv);
8552 while gi < gpr {
8553 let t = bp.add(gi * 4);
8554 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8555 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8556 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8557 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8558 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8559 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8560 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8561 gi += 1;
8562 }
8563 acc
8564 }
8565}
8566
8567#[cfg(target_arch = "aarch64")]
8571#[target_feature(enable = "neon,dotprod")]
8572#[allow(clippy::too_many_arguments)]
8573unsafe fn dot_q1g_row_1x4_sdot(
8574 bytes: &[u8],
8575 sc_off: usize,
8576 rows: usize,
8577 r: usize,
8578 gpr: usize,
8579 xs: [&[i8]; 4],
8580 gs: [&[i32]; 4],
8581) -> [f32; 4] {
8582 unsafe {
8584 use core::arch::aarch64::*;
8585 use core::arch::asm;
8586 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8587 const IW: [[u8; 16]; 8] = [
8588 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8589 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8590 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8591 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8592 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8593 [
8594 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8595 ],
8596 [
8597 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8598 ],
8599 [
8600 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8601 ],
8602 ];
8603 let m = vld1q_u8(MASKS.as_ptr());
8604 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8605 macro_rules! sdot2 {
8606 ($w0:expr, $w1:expr, $x:expr) => {{
8607 let x0 = vld1q_s8($x);
8608 let x1 = vld1q_s8($x.add(16));
8609 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8610 asm!(
8611 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8612 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8613 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8614 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8615 options(pure, nomem, nostack),
8616 );
8617 vaddq_s32(a0, a1)
8618 }};
8619 }
8620 let bp = bytes.as_ptr().add(r * gpr * 4);
8621 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8622 let mut accv = [vdupq_n_f32(0.0); 4];
8623 let mut gi = 0;
8624 while gi + 4 <= vec_end {
8625 let ld = vld1q_u8(bp.add(gi * 4));
8626 let w: [int8x16_t; 8] =
8628 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8629 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8630 let scf: float32x4_t;
8631 asm!(
8632 "fcvtl {o:v}.4s, {i:v}.4h",
8633 o = out(vreg) scf, i = in(vreg) h,
8634 options(pure, nomem, nostack),
8635 );
8636 for k in 0..4 {
8637 let xp = xs[k].as_ptr();
8638 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8639 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8640 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8641 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8642 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8643 let g = vld1q_s32(gs[k].as_ptr().add(gi));
8644 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8645 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8646 }
8647 gi += 4;
8648 }
8649 let mut acc = [
8650 vaddvq_f32(accv[0]),
8651 vaddvq_f32(accv[1]),
8652 vaddvq_f32(accv[2]),
8653 vaddvq_f32(accv[3]),
8654 ];
8655 while gi < gpr {
8656 let t = bp.add(gi * 4);
8657 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8658 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8659 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8660 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8661 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8662 for k in 0..4 {
8663 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8664 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8665 }
8666 gi += 1;
8667 }
8668 acc
8669 }
8670}
8671
8672#[inline]
8674fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8675 let t = r * gpr + j / GROUP_SIZE;
8676 let k = j % GROUP_SIZE;
8677 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8678 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8679 ((bit as i32 * 2 - 1) as f32, s)
8680}
8681
8682#[inline]
8684fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8685 let mut acc = 0f32;
8686 for gi in 0..gpr {
8687 let t = r * gpr + gi;
8688 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8689 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8690 let mut ga = 0f32;
8691 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8692 for k in 0..8 {
8693 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8694 }
8695 }
8696 acc += ga * s;
8697 }
8698 acc
8699}
8700
8701fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8706 let entries = base_len + (rows + 1) * 4;
8707 (base_len, entries, entries <= bytes.len())
8708}
8709
8710#[inline]
8712fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8713 let o = rp_off + r * 4;
8714 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8715}
8716
8717const SIGN5: [[f32; 5]; 256] = {
8721 let mut lut = [[0.0f32; 5]; 256];
8722 let pow3 = [1u16, 3, 9, 27, 81];
8723 let mut byte = 0usize;
8724 while byte < 256 {
8725 let mut i = 0usize;
8726 while i < 5 {
8727 let code = (byte as u16 / pow3[i]) % 3;
8728 lut[byte][i] = if code == 1 {
8729 1.0
8730 } else if code == 2 {
8731 -1.0
8732 } else {
8733 0.0
8734 };
8735 i += 1;
8736 }
8737 byte += 1;
8738 }
8739 lut
8740};
8741
8742const SIGN5_I8: [[i8; 5]; 256] = {
8744 let mut lut = [[0i8; 5]; 256];
8745 let pow3 = [1u16, 3, 9, 27, 81];
8746 let mut byte = 0usize;
8747 while byte < 256 {
8748 let mut i = 0usize;
8749 while i < 5 {
8750 let code = (byte as u16 / pow3[i]) % 3;
8751 lut[byte][i] = if code == 1 {
8752 1
8753 } else if code == 2 {
8754 -1
8755 } else {
8756 0
8757 };
8758 i += 1;
8759 }
8760 byte += 1;
8761 }
8762 lut
8763};
8764
8765const SIGN5_U64: [u64; 256] = {
8771 let mut lut = [0u64; 256];
8772 let pow3 = [1u16, 3, 9, 27, 81];
8773 let mut byte = 0usize;
8774 while byte < 256 {
8775 let mut v = 0u64;
8776 let mut i = 0usize;
8777 while i < 5 {
8778 let code = (byte as u16 / pow3[i]) % 3;
8779 let s: u8 = if code == 1 {
8780 1
8781 } else if code == 2 {
8782 0xFF
8783 } else {
8784 0
8785 };
8786 v |= (s as u64) << (i * 8);
8787 i += 1;
8788 }
8789 lut[byte] = v;
8790 byte += 1;
8791 }
8792 lut
8793};
8794
8795#[inline]
8800fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8801 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8802 let off = (r * gpr + j / GROUP_SIZE) * TILE;
8803 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8804 let within = j % GROUP_SIZE;
8805 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8806}
8807
8808#[cfg(target_arch = "aarch64")]
8811#[target_feature(enable = "neon,dotprod")]
8812#[inline]
8813unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8814 unsafe {
8816 use core::arch::aarch64::*;
8817 use core::arch::asm;
8818 let w0 = vld1q_s8(w);
8819 let w1 = vld1q_s8(w.add(16));
8820 let x0 = vld1q_s8(x);
8821 let x1 = vld1q_s8(x.add(16));
8822 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8823 asm!(
8824 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8825 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8826 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8827 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8828 options(pure, nomem, nostack),
8829 );
8830 vaddvq_s32(vaddq_s32(a0, a1))
8831 }
8832}
8833
8834#[cfg(target_arch = "x86_64")]
8837#[target_feature(enable = "avx2")]
8838#[inline]
8839unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8840 unsafe {
8842 use core::arch::x86_64::*;
8843 let wv = _mm256_loadu_si256(w as *const __m256i);
8844 let xv = _mm256_loadu_si256(x as *const __m256i);
8845 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8846 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8847 let hi128 = _mm256_extracti128_si256::<1>(d);
8848 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8849 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8850 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8851 _mm_cvtsi128_si32(s32)
8852 }
8853}
8854
8855#[inline]
8860fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8861 debug_assert!(dst.len() >= 40);
8862 unsafe {
8865 let p = dst.as_mut_ptr();
8866 for bi in 0..7 {
8867 core::ptr::write_unaligned(
8868 p.add(bi * 5) as *mut u64,
8869 SIGN5_U64[*codes.add(bi) as usize],
8870 );
8871 }
8872 }
8873}
8874
8875#[inline]
8880fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8881 #[cfg(target_arch = "aarch64")]
8882 unsafe {
8883 return sdot32_i8(w, x);
8884 }
8885 #[cfg(target_arch = "x86_64")]
8886 unsafe {
8887 return i8dot32_avx2(w, x);
8888 }
8889 #[allow(unreachable_code)]
8890 unsafe {
8891 let mut s = 0i32;
8892 for k in 0..GROUP_SIZE {
8893 s += *w.add(k) as i32 * *x.add(k) as i32;
8894 }
8895 s
8896 }
8897}
8898
8899#[inline]
8900unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8901 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8902 (
8903 SIGN5_U64[*codes as usize],
8904 SIGN5_U64[*codes.add(1) as usize],
8905 SIGN5_U64[*codes.add(2) as usize],
8906 SIGN5_U64[*codes.add(3) as usize],
8907 SIGN5_U64[*codes.add(4) as usize],
8908 SIGN5_U64[*codes.add(5) as usize],
8909 SIGN5_U64[*codes.add(6) as usize],
8910 )
8911 };
8912
8913 let u0 = s0 | (s1 << 40);
8914 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8915 let u2 = (s3 >> 8) | (s4 << 32);
8916 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8917
8918 (u0, u1, u2, u3)
8919}
8920
8921#[cfg(target_arch = "aarch64")]
8925#[target_feature(enable = "neon,dotprod")]
8926unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8927 use core::arch::aarch64::*;
8928 use core::arch::asm;
8929 unsafe {
8930 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8931 let mut acc = 0f32;
8932 let bytes_ptr = bytes.as_ptr();
8933 let xq_ptr = xq.as_ptr();
8934 let row_off = r * gpr * TILE;
8935
8936 let gpr2 = gpr & !1;
8937 let mut gi = 0;
8938 while gi < gpr2 {
8939 let off0 = row_off + gi * TILE;
8940 let off1 = off0 + TILE;
8941 let s0 = f16_to_f32(u16::from_le_bytes([
8942 *bytes_ptr.add(off0),
8943 *bytes_ptr.add(off0 + 1),
8944 ]));
8945 let s1 = f16_to_f32(u16::from_le_bytes([
8946 *bytes_ptr.add(off1),
8947 *bytes_ptr.add(off1 + 1),
8948 ]));
8949
8950 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8951 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8952
8953 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8954 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8955 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8956 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8957
8958 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8959 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8960 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8961 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8962
8963 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8964 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8965 asm!(
8966 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8967 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8968 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8969 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8970 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8971 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8972 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8973 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8974 options(pure, nomem, nostack),
8975 );
8976 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8977 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8978 acc += d0 as f32 * s0 + d1 as f32 * s1;
8979 gi += 2;
8980 }
8981
8982 if gi < gpr {
8983 let off = row_off + gi * TILE;
8984 let s = f16_to_f32(u16::from_le_bytes([
8985 *bytes_ptr.add(off),
8986 *bytes_ptr.add(off + 1),
8987 ]));
8988 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8989 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8990 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8991 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8992 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8993 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8994 asm!(
8995 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8996 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8997 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8998 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8999 options(pure, nomem, nostack),
9000 );
9001 let d = vaddvq_s32(vaddq_s32(a0, a1));
9002 acc += d as f32 * s;
9003 }
9004 acc
9005 }
9006}
9007
9008#[cfg(target_arch = "x86_64")]
9010#[target_feature(enable = "avx2")]
9011unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9012 use core::arch::x86_64::*;
9013 unsafe {
9014 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9015 let mut acc = 0f32;
9016 let bytes_ptr = bytes.as_ptr();
9017 let xq_ptr = xq.as_ptr();
9018 let row_off = r * gpr * TILE;
9019
9020 let ones = _mm256_set1_epi16(1);
9021 for gi in 0..gpr {
9022 let off = row_off + gi * TILE;
9023 let s = f16_to_f32(u16::from_le_bytes([
9024 *bytes_ptr.add(off),
9025 *bytes_ptr.add(off + 1),
9026 ]));
9027 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9028 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9029 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9030 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9031 let d256 = _mm256_madd_epi16(p16, ones);
9032 let d128 = _mm_add_epi32(
9033 _mm256_castsi256_si128(d256),
9034 _mm256_extracti128_si256(d256, 1),
9035 );
9036 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
9037 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
9038 acc += d32 as f32 * s;
9039 }
9040 acc
9041 }
9042}
9043
9044#[cfg(target_arch = "x86_64")]
9046#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9047unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9048 use core::arch::x86_64::*;
9049 unsafe {
9051 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9052 let mut acc = 0f32;
9053 let bytes_ptr = bytes.as_ptr();
9054 let xq_ptr = xq.as_ptr();
9055 let row_off = r * gpr * TILE;
9056 for gi in 0..gpr {
9057 let off = row_off + gi * TILE;
9058 let s = f16_to_f32(u16::from_le_bytes([
9059 *bytes_ptr.add(off),
9060 *bytes_ptr.add(off + 1),
9061 ]));
9062 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9063 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9064 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9065 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9066 acc += d as f32 * s;
9067 }
9068 acc
9069 }
9070}
9071
9072#[inline]
9076fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9077 #[cfg(target_arch = "aarch64")]
9078 unsafe {
9079 return q1t_dot_row_sdot(bytes, r, gpr, xq);
9080 }
9081 #[cfg(target_arch = "x86_64")]
9082 unsafe {
9083 if vnni_tiles_enabled() {
9084 return q1t_dot_row_vnni(bytes, r, gpr, xq);
9085 }
9086 return q1t_dot_row_avx2(bytes, r, gpr, xq);
9087 }
9088 #[allow(unreachable_code)]
9089 {
9090 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9091 let mut acc = 0f32;
9092 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
9094 let off = (r * gpr + gi) * TILE;
9095 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9096 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9097 let mut d = 0i32;
9098 for k in 0..GROUP_SIZE {
9099 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9100 }
9101 acc += d as f32 * s;
9102 }
9103 acc
9104 }
9105}
9106
9107fn q1t_row_outlier_correction(
9114 bytes: &[u8],
9115 r: usize,
9116 rp_off: usize,
9117 entries_off: usize,
9118 has_ov: bool,
9119 x: &[f32],
9120) -> f32 {
9121 if !has_ov {
9122 return 0.0;
9123 }
9124 let (c0, c1) = (
9125 q1t_rowptr(bytes, rp_off, r),
9126 q1t_rowptr(bytes, rp_off, r + 1),
9127 );
9128 let mut corr = 0f32;
9129 for p in c0..c1 {
9130 let e = entries_off + p * 4;
9131 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9132 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9133 corr += val * x[col];
9134 }
9135 corr
9136}
9137
9138fn q1t_dequant_row(
9142 bytes: &[u8],
9143 r: usize,
9144 gpr: usize,
9145 rp_off: usize,
9146 entries_off: usize,
9147 has_ov: bool,
9148 buf: &mut [f32],
9149) {
9150 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9151 for g in 0..gpr {
9152 let off = (r * gpr + g) * TILE;
9153 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9154 let codes = &bytes[off + 2..off + TILE];
9155 let bc = g * GROUP_SIZE;
9156 for bi in 0..6 {
9158 let lut = &SIGN5[codes[bi] as usize];
9159 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
9160 for i in 0..5 {
9161 d[i] = lut[i] * s;
9162 }
9163 }
9164 let lut = &SIGN5[codes[6] as usize];
9165 buf[bc + 30] = lut[0] * s;
9166 buf[bc + 31] = lut[1] * s;
9167 }
9168 if !has_ov {
9169 return;
9170 }
9171 let (c0, c1) = (
9172 q1t_rowptr(bytes, rp_off, r),
9173 q1t_rowptr(bytes, rp_off, r + 1),
9174 );
9175 for p in c0..c1 {
9176 let e = entries_off + p * 4;
9177 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9178 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9179 }
9180}
9181
9182fn q1t_add_overlay(
9186 bytes: &[u8],
9187 x: &[f32],
9188 rows: usize,
9189 cols: usize,
9190 out: &mut [f32],
9191 pool: Option<&Pool>,
9192) {
9193 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9194 let gpr = cols / GROUP_SIZE;
9195 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9196 if !has_ov {
9197 return;
9198 }
9199 let out_addr = SendMut(out.as_mut_ptr());
9200 let run = move |start: usize, end: usize| {
9201 for r in start..end {
9202 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9203 unsafe { *out_addr.at(r) += corr };
9205 }
9206 };
9207 dispatch_rows(pool, rows, &run);
9208}
9209
9210#[allow(clippy::too_many_arguments)]
9213fn q1t_range_a8w8(
9214 bytes: &[u8],
9215 gpr: usize,
9216 rp_off: usize,
9217 ent_off: usize,
9218 has_ov: bool,
9219 act: &SplitAct,
9220 x: &[f32],
9221 out: SendMut,
9222 start: usize,
9223 end: usize,
9224) {
9225 for r in start..end {
9226 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9227 for &(j, xv) in &act.outliers {
9228 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9229 }
9230 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9231 unsafe { *out.at(r) = acc };
9233 }
9234}
9235
9236#[allow(clippy::too_many_arguments)]
9239fn q1t_range_f32_batch(
9240 bytes: &[u8],
9241 gpr: usize,
9242 rp_off: usize,
9243 ent_off: usize,
9244 has_ov: bool,
9245 x: &[f32],
9246 out: SendMut,
9247 start: usize,
9248 end: usize,
9249) {
9250 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9251 let mut sg = [0f32; GROUP_SIZE];
9252 for r in start..end {
9253 let mut acc = 0f32;
9254 for g in 0..gpr {
9255 let off = (r * gpr + g) * TILE;
9256 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9257 let codes = &bytes[off + 2..off + TILE];
9258 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9259 for bi in 0..6 {
9260 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9261 }
9262 let lut = &SIGN5[codes[6] as usize];
9263 sg[30] = lut[0];
9264 sg[31] = lut[1];
9265 let mut gsum = 0f32;
9266 for k in 0..GROUP_SIZE {
9267 gsum += sg[k] * xg[k];
9268 }
9269 acc += s * gsum;
9270 }
9271 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9272 unsafe { *out.at(r) = acc };
9274 }
9275}
9276
9277fn q1t_matvec(
9281 bytes: &[u8],
9282 x: &[f32],
9283 rows: usize,
9284 cols: usize,
9285 out: &mut [f32],
9286 pool: Option<&Pool>,
9287) {
9288 debug_assert_eq!(out.len(), rows);
9289 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9290 let gpr = cols / GROUP_SIZE;
9291 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9292 let out_addr = SendMut(out.as_mut_ptr());
9293 if a8w8_enabled() {
9297 let act = split_act(x);
9298 let act = &act;
9299 let run = move |start: usize, end: usize| {
9300 for r in start..end {
9301 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9302 for &(j, xv) in &act.outliers {
9303 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9304 }
9305 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9306 unsafe { *out_addr.at(r) = acc };
9308 }
9309 };
9310 dispatch_rows(pool, rows, &run);
9311 return;
9312 }
9313 let run = move |start: usize, end: usize| {
9314 let mut sg = [0f32; GROUP_SIZE];
9318 for r in start..end {
9319 let mut acc = 0f32;
9320 for g in 0..gpr {
9321 let off = (r * gpr + g) * TILE;
9322 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9323 let codes = &bytes[off + 2..off + TILE];
9324 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9325 for bi in 0..6 {
9326 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9327 }
9328 let lut = &SIGN5[codes[6] as usize];
9329 sg[30] = lut[0];
9330 sg[31] = lut[1];
9331 let mut gsum = 0f32;
9332 for k in 0..GROUP_SIZE {
9333 gsum += sg[k] * xg[k];
9334 }
9335 acc += s * gsum;
9336 }
9337 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9338 unsafe { *out_addr.at(r) = acc };
9339 }
9340 };
9341 dispatch_rows(pool, rows, &run);
9342}
9343
9344#[cfg(target_arch = "aarch64")]
9350#[target_feature(enable = "neon,dotprod")]
9351unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9352 use core::arch::aarch64::*;
9353 use core::arch::asm;
9354 unsafe {
9356 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9357 let bytes_ptr = bytes.as_ptr();
9358 let row_off = r * gpr * TILE;
9359 let xp = [xa.as_ptr(), xb.as_ptr()];
9360 let mut acc = [0f32; 2];
9361 macro_rules! sdot2 {
9362 ($w0:expr, $w1:expr, $x:expr) => {{
9363 let x0 = vld1q_s8($x);
9364 let x1 = vld1q_s8($x.add(16));
9365 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9366 asm!(
9367 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9368 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9369 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9370 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9371 options(pure, nomem, nostack),
9372 );
9373 vaddvq_s32(vaddq_s32(a0, a1))
9374 }};
9375 }
9376 let gpr2 = gpr & !1;
9377 let mut gi = 0;
9378 while gi < gpr2 {
9379 let off0 = row_off + gi * TILE;
9380 let off1 = off0 + TILE;
9381 let s0 = f16_to_f32(u16::from_le_bytes([
9382 *bytes_ptr.add(off0),
9383 *bytes_ptr.add(off0 + 1),
9384 ]));
9385 let s1 = f16_to_f32(u16::from_le_bytes([
9386 *bytes_ptr.add(off1),
9387 *bytes_ptr.add(off1 + 1),
9388 ]));
9389 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9390 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9391 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9392 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9393 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9394 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9395 for k in 0..2 {
9396 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9397 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9398 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9399 }
9400 gi += 2;
9401 }
9402 if gi < gpr {
9403 let off = row_off + gi * TILE;
9404 let s = f16_to_f32(u16::from_le_bytes([
9405 *bytes_ptr.add(off),
9406 *bytes_ptr.add(off + 1),
9407 ]));
9408 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9409 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9410 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9411 for k in 0..2 {
9412 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9413 acc[k] += d as f32 * s;
9414 }
9415 }
9416 acc
9417 }
9418}
9419
9420fn q1t_matvec2(
9426 bytes: &[u8],
9427 x1: &[f32],
9428 x2: &[f32],
9429 rows: usize,
9430 cols: usize,
9431 o1: &mut [f32],
9432 o2: &mut [f32],
9433 pool: Option<&Pool>,
9434) {
9435 debug_assert_eq!(o1.len(), rows);
9436 debug_assert_eq!(o2.len(), rows);
9437 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9438 let gpr = cols / GROUP_SIZE;
9439 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9440 let out1 = SendMut(o1.as_mut_ptr());
9441 let out2 = SendMut(o2.as_mut_ptr());
9442 if a8w8_enabled() {
9443 let a1 = split_act(x1);
9444 let a2 = split_act(x2);
9445 let (a1, a2) = (&a1, &a2);
9446 let run = move |start: usize, end: usize| {
9447 for r in start..end {
9448 #[cfg(target_arch = "aarch64")]
9449 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9452 #[cfg(not(target_arch = "aarch64"))]
9453 let ds = [
9454 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9455 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9456 ];
9457 let mut acc1 = ds[0] * a1.sx;
9458 for &(j, xv) in &a1.outliers {
9459 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9460 }
9461 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9462 let mut acc2 = ds[1] * a2.sx;
9463 for &(j, xv) in &a2.outliers {
9464 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9465 }
9466 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9467 unsafe {
9469 *out1.at(r) = acc1;
9470 *out2.at(r) = acc2;
9471 }
9472 }
9473 };
9474 dispatch_rows(pool, rows, &run);
9475 return;
9476 }
9477 let run = move |start: usize, end: usize| {
9478 let mut sg = [0f32; GROUP_SIZE];
9481 for r in start..end {
9482 let mut acc1 = 0f32;
9483 let mut acc2 = 0f32;
9484 for g in 0..gpr {
9485 let off = (r * gpr + g) * TILE;
9486 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9487 let codes = &bytes[off + 2..off + TILE];
9488 for bi in 0..6 {
9489 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9490 }
9491 let lut = &SIGN5[codes[6] as usize];
9492 sg[30] = lut[0];
9493 sg[31] = lut[1];
9494 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9495 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9496 let mut gsum1 = 0f32;
9497 for k in 0..GROUP_SIZE {
9498 gsum1 += sg[k] * xg1[k];
9499 }
9500 acc1 += s * gsum1;
9501 let mut gsum2 = 0f32;
9502 for k in 0..GROUP_SIZE {
9503 gsum2 += sg[k] * xg2[k];
9504 }
9505 acc2 += s * gsum2;
9506 }
9507 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9508 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9509 unsafe {
9511 *out1.at(r) = acc1;
9512 *out2.at(r) = acc2;
9513 }
9514 }
9515 };
9516 dispatch_rows(pool, rows, &run);
9517}
9518
9519fn q1t_matmat(
9522 bytes: &[u8],
9523 xs: &[f32],
9524 b: usize,
9525 rows: usize,
9526 cols: usize,
9527 out: &mut [f32],
9528 pool: Option<&Pool>,
9529) {
9530 debug_assert_eq!(out.len(), b * rows);
9531 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9532 let gpr = cols / GROUP_SIZE;
9533 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9534 let out_addr = SendMut(out.as_mut_ptr());
9535 if a8w8_enabled() {
9539 let acts: Vec<SplitAct> = (0..b)
9540 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9541 .collect();
9542 let acts = &acts;
9543 let run = move |start: usize, end: usize| {
9544 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
9548 for g in 0..gpr {
9549 let off = (r * gpr + g) * TILE;
9550 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9551 q1t_unpack_group_i8(
9552 bytes.as_ptr().wrapping_add(off + 2),
9553 &mut sg[g * GROUP_SIZE..],
9554 );
9555 }
9556 for bi in 0..b {
9557 let act = &acts[bi];
9558 let mut isum = 0f32;
9559 for g in 0..gpr {
9560 let d = q1t_i8dot32(
9561 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9562 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9563 );
9564 isum += d as f32 * sc[g];
9565 }
9566 let mut acc = isum * act.sx;
9567 for &(j, xv) in &act.outliers {
9568 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9569 }
9570 accs[bi] = acc;
9571 }
9572 if has_ov {
9576 let (c0, c1) = (
9577 q1t_rowptr(bytes, rp_off, r),
9578 q1t_rowptr(bytes, rp_off, r + 1),
9579 );
9580 for p in c0..c1 {
9581 let e = ent_off + p * 4;
9582 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9583 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9584 for bi in 0..b {
9585 accs[bi] += val * xs[bi * cols + col];
9586 }
9587 }
9588 }
9589 for bi in 0..b {
9590 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9591 }
9592 }
9593 };
9594 dispatch_rows(pool, rows, &run);
9595 return;
9596 }
9597 let run = move |start: usize, end: usize| {
9598 let mut buf = vec![0f32; cols];
9599 for r in start..end {
9600 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9601 for bi in 0..b {
9602 let xr = &xs[bi * cols..(bi + 1) * cols];
9603 let mut acc = 0f32;
9604 for j in 0..cols {
9605 acc += buf[j] * xr[j];
9606 }
9607 unsafe { *out_addr.at(bi * rows + r) = acc };
9608 }
9609 }
9610 };
9611 dispatch_rows(pool, rows, &run);
9612}
9613
9614#[cfg(test)]
9617fn q1_matvec(
9618 bytes: &[u8],
9619 x: &[f32],
9620 rows: usize,
9621 cols: usize,
9622 out: &mut [f32],
9623 pool: Option<&Pool>,
9624) {
9625 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9626}
9627
9628fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9631 debug_assert_eq!(out.len(), rows);
9632 let gpr = cols / GROUP_SIZE;
9633 let out_addr = SendMut(out.as_mut_ptr());
9634 if a8w8_enabled() {
9635 let act = split_act(x);
9636 let gsum = q1_group_sums(&act.xq, gpr);
9637 let (act, gsum) = (&act, &gsum);
9638 let run =
9639 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9640 dispatch_rows(pool, rows, &run);
9641 return;
9642 }
9643 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9644 dispatch_rows(pool, rows, &run);
9645}
9646
9647#[cfg(test)]
9649#[allow(clippy::too_many_arguments)]
9650fn q1_matvec2(
9651 bytes: &[u8],
9652 x1: &[f32],
9653 x2: &[f32],
9654 rows: usize,
9655 cols: usize,
9656 o1: &mut [f32],
9657 o2: &mut [f32],
9658 pool: Option<&Pool>,
9659) {
9660 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9661}
9662
9663#[allow(clippy::too_many_arguments)]
9665fn bin_matvec2(
9666 w: BinW,
9667 x1: &[f32],
9668 x2: &[f32],
9669 rows: usize,
9670 cols: usize,
9671 o1: &mut [f32],
9672 o2: &mut [f32],
9673 pool: Option<&Pool>,
9674) {
9675 let gpr = cols / GROUP_SIZE;
9676 let p1 = SendMut(o1.as_mut_ptr());
9677 let p2 = SendMut(o2.as_mut_ptr());
9678 if a8w8_enabled() {
9679 let a1 = split_act(x1);
9680 let a2 = split_act(x2);
9681 let g1 = q1_group_sums(&a1.xq, gpr);
9682 let g2 = q1_group_sums(&a2.xq, gpr);
9683 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9684 let run = move |start: usize, end: usize| {
9685 for r in start..end {
9686 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9687 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9688 for &(j, xv) in &a1.outliers {
9689 let (wv, s) = w.outlier(r, gpr, j);
9690 v1 += wv * s * xv;
9691 }
9692 for &(j, xv) in &a2.outliers {
9693 let (wv, s) = w.outlier(r, gpr, j);
9694 v2 += wv * s * xv;
9695 }
9696 unsafe {
9698 *p1.at(r) = v1;
9699 *p2.at(r) = v2;
9700 }
9701 }
9702 };
9703 dispatch_rows(pool, rows, &run);
9704 return;
9705 }
9706 let run = move |start: usize, end: usize| {
9707 for r in start..end {
9708 unsafe {
9710 *p1.at(r) = w.row_exact(r, gpr, x1);
9711 *p2.at(r) = w.row_exact(r, gpr, x2);
9712 }
9713 }
9714 };
9715 dispatch_rows(pool, rows, &run);
9716}
9717
9718#[cfg(test)]
9720#[allow(clippy::too_many_arguments)]
9721fn q1_matmat(
9722 bytes: &[u8],
9723 xs_all: &[f32],
9724 b: usize,
9725 rows: usize,
9726 cols: usize,
9727 out: &mut [f32],
9728 pool: Option<&Pool>,
9729) {
9730 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9731}
9732
9733#[allow(clippy::too_many_arguments)]
9735fn bin_matmat(
9736 w: BinW,
9737 xs_all: &[f32],
9738 b: usize,
9739 rows: usize,
9740 cols: usize,
9741 out: &mut [f32],
9742 pool: Option<&Pool>,
9743) {
9744 debug_assert_eq!(out.len(), b * rows);
9745 let gpr = cols / GROUP_SIZE;
9746 let out_addr = SendMut(out.as_mut_ptr());
9747 if a8w8_enabled() {
9748 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9749 .map(|bi| {
9750 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9751 let gsum = q1_group_sums(&act.xq, gpr);
9752 (act, gsum)
9753 })
9754 .collect();
9755 let acts = &acts;
9756 #[cfg(target_arch = "x86_64")]
9757 let blocked_ok = avx2_enabled() && blocked_enabled();
9758 #[cfg(target_arch = "aarch64")]
9759 let blocked_ok = sdot_enabled() && blocked_enabled();
9760 let run = move |start: usize, end: usize| {
9761 for r in start..end {
9762 let mut bi = 0usize;
9763 #[cfg(target_arch = "aarch64")]
9766 if blocked_ok {
9767 while bi + 4 <= acts.len() {
9768 let xs = [
9769 acts[bi].0.xq.as_slice(),
9770 acts[bi + 1].0.xq.as_slice(),
9771 acts[bi + 2].0.xq.as_slice(),
9772 acts[bi + 3].0.xq.as_slice(),
9773 ];
9774 let gs = [
9775 acts[bi].1.as_slice(),
9776 acts[bi + 1].1.as_slice(),
9777 acts[bi + 2].1.as_slice(),
9778 acts[bi + 3].1.as_slice(),
9779 ];
9780 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9782 for k in 0..4 {
9783 let (act, _) = &acts[bi + k];
9784 let mut acc = d[k] * act.sx;
9785 for &(j, xv) in &act.outliers {
9786 let (wv, sc) = w.outlier(r, gpr, j);
9787 acc += wv * sc * xv;
9788 }
9789 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9791 }
9792 bi += 4;
9793 }
9794 }
9795 #[cfg(target_arch = "x86_64")]
9796 if blocked_ok {
9797 while bi + 4 <= acts.len() {
9798 let xs = [
9799 acts[bi].0.xq.as_slice(),
9800 acts[bi + 1].0.xq.as_slice(),
9801 acts[bi + 2].0.xq.as_slice(),
9802 acts[bi + 3].0.xq.as_slice(),
9803 ];
9804 let gs = [
9805 acts[bi].1.as_slice(),
9806 acts[bi + 1].1.as_slice(),
9807 acts[bi + 2].1.as_slice(),
9808 acts[bi + 3].1.as_slice(),
9809 ];
9810 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9812 for k in 0..4 {
9813 let (act, _) = &acts[bi + k];
9814 let mut acc = d[k] * act.sx;
9815 for &(j, xv) in &act.outliers {
9816 let (wv, sc) = w.outlier(r, gpr, j);
9817 acc += wv * sc * xv;
9818 }
9819 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9821 }
9822 bi += 4;
9823 }
9824 }
9825 while bi < acts.len() {
9826 let (act, gsum) = &acts[bi];
9827 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9828 for &(j, xv) in &act.outliers {
9829 let (wv, s) = w.outlier(r, gpr, j);
9830 acc += wv * s * xv;
9831 }
9832 unsafe { *out_addr.at(bi * rows + r) = acc };
9834 bi += 1;
9835 }
9836 }
9837 };
9838 dispatch_rows(pool, rows, &run);
9839 return;
9840 }
9841 let run = move |start: usize, end: usize| {
9842 for r in start..end {
9843 for bi in 0..b {
9844 let x = &xs_all[bi * cols..(bi + 1) * cols];
9845 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9847 }
9848 }
9849 };
9850 dispatch_rows(pool, rows, &run);
9851}
9852
9853fn q4matvec(
9859 bytes: &[u8],
9860 x: &[f32],
9861 rows: usize,
9862 cols: usize,
9863 out: &mut [f32],
9864 pool: Option<&Pool>,
9865) {
9866 debug_assert_eq!(out.len(), rows);
9867 let (packed, scales) = q4_split(bytes, rows, cols);
9868 let gpr = cols / GROUP_SIZE;
9869 let out_addr = SendMut(out.as_mut_ptr());
9870
9871 if a8w8_enabled() {
9872 let act = split_act(x);
9873 let run = move |start: usize, end: usize| {
9874 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9875 };
9876 dispatch_rows(pool, rows, &run);
9877 return;
9878 }
9879
9880 let run =
9881 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9882 dispatch_rows(pool, rows, &run);
9883}
9884
9885#[inline]
9888#[allow(unreachable_code)]
9889#[cfg(target_arch = "x86_64")]
9894#[target_feature(enable = "avx2")]
9895unsafe fn dot_q4b_row_1x4_avx2(
9896 buf: &[u8],
9897 scales: &[u8],
9898 g0: usize,
9899 gpr: usize,
9900 xs: [&[i8]; 4],
9901) -> [f32; 4] {
9902 unsafe {
9904 use core::arch::x86_64::*;
9905 let ones = _mm256_set1_epi16(1);
9906 let mut acc = [0f32; 4];
9907 for gi in 0..gpr {
9908 let s = f16_to_f32(u16::from_le_bytes([
9909 scales[(g0 + gi) * 2],
9910 scales[(g0 + gi) * 2 + 1],
9911 ]));
9912 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9913 let aw = _mm256_abs_epi8(w);
9914 for (k, xq) in xs.iter().enumerate() {
9915 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9916 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9917 let d = _mm256_madd_epi16(p16, ones);
9918 let hi128 = _mm256_extracti128_si256::<1>(d);
9919 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9920 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9921 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9922 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9923 }
9924 }
9925 acc
9926 }
9927}
9928
9929#[cfg(target_arch = "x86_64")]
9931#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9932unsafe fn dot_q4b_row_1x4_vnni(
9933 buf: &[u8],
9934 scales: &[u8],
9935 g0: usize,
9936 gpr: usize,
9937 xs: [&[i8]; 4],
9938) -> [f32; 4] {
9939 unsafe {
9941 use core::arch::x86_64::*;
9942 let mut acc = [0f32; 4];
9943 for gi in 0..gpr {
9944 let s = f16_to_f32(u16::from_le_bytes([
9945 scales[(g0 + gi) * 2],
9946 scales[(g0 + gi) * 2 + 1],
9947 ]));
9948 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9949 let aw = _mm256_abs_epi8(w);
9950 for (k, xq) in xs.iter().enumerate() {
9951 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9952 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9953 acc[k] += d as f32 * s;
9954 }
9955 }
9956 acc
9957 }
9958}
9959
9960#[cfg(target_arch = "x86_64")]
9966#[target_feature(enable = "avx2")]
9967unsafe fn dot_q4b_row_1x4_sx_avx2(
9968 buf: &[u8],
9969 scales: &[u8],
9970 g0: usize,
9971 gpr: usize,
9972 xs: [&[i8]; 4],
9973 sxs: [f32; 4],
9974) -> [f32; 4] {
9975 unsafe {
9977 use core::arch::x86_64::*;
9978 let ones = _mm256_set1_epi16(1);
9979 let mut acc = [0f32; 4];
9980 for gi in 0..gpr {
9981 let s = f16_to_f32(u16::from_le_bytes([
9982 scales[(g0 + gi) * 2],
9983 scales[(g0 + gi) * 2 + 1],
9984 ]));
9985 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9986 let aw = _mm256_abs_epi8(w);
9987 for (k, xq) in xs.iter().enumerate() {
9988 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9989 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9990 let d = _mm256_madd_epi16(p16, ones);
9991 let hi128 = _mm256_extracti128_si256::<1>(d);
9992 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9993 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9994 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9995 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9996 }
9997 }
9998 acc
9999 }
10000}
10001
10002#[cfg(target_arch = "x86_64")]
10005#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10006unsafe fn dot_q4b_row_1x4_sx_vnni(
10007 buf: &[u8],
10008 scales: &[u8],
10009 g0: usize,
10010 gpr: usize,
10011 xs: [&[i8]; 4],
10012 sxs: [f32; 4],
10013) -> [f32; 4] {
10014 unsafe {
10016 use core::arch::x86_64::*;
10017 let mut acc = [0f32; 4];
10018 for gi in 0..gpr {
10019 let s = f16_to_f32(u16::from_le_bytes([
10020 scales[(g0 + gi) * 2],
10021 scales[(g0 + gi) * 2 + 1],
10022 ]));
10023 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10024 let aw = _mm256_abs_epi8(w);
10025 for (k, xq) in xs.iter().enumerate() {
10026 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10027 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10028 acc[k] += (d as f32 * sxs[k]) * s;
10029 }
10030 }
10031 acc
10032 }
10033}
10034
10035#[allow(unreachable_code)]
10036fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10037 #[cfg(target_arch = "aarch64")]
10038 unsafe {
10039 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
10040 }
10041 #[cfg(target_arch = "x86_64")]
10042 unsafe {
10043 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10044 }
10045 let mut acc = 0f32;
10046 for gi in 0..gpr {
10047 let g = g0 + gi;
10048 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10049 let mut d = 0i32;
10050 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10051 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10052 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10053 }
10054 acc += d as f32 * s;
10055 }
10056 acc
10057}
10058
10059#[inline]
10061#[allow(unreachable_code)]
10062fn dot_q4_row_i8_2(
10063 packed: &[u8],
10064 scales: &[u8],
10065 g0: usize,
10066 gpr: usize,
10067 xq1: &[i8],
10068 xq2: &[i8],
10069) -> (f32, f32) {
10070 #[cfg(target_arch = "aarch64")]
10071 unsafe {
10072 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10073 }
10074 #[cfg(target_arch = "x86_64")]
10075 unsafe {
10076 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10077 }
10078 (
10079 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10080 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10081 )
10082}
10083
10084#[allow(clippy::too_many_arguments)]
10087fn q4_range_a8w8(
10088 packed: &[u8],
10089 scales: &[u8],
10090 gpr: usize,
10091 cols: usize,
10092 act: &SplitAct,
10093 out: SendMut,
10094 start: usize,
10095 end: usize,
10096) {
10097 for r in start..end {
10098 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10099 for &(j, xv) in &act.outliers {
10101 let flat = r * cols + j;
10102 let byte = packed[flat / 2];
10103 let nib = if flat & 1 == 0 {
10104 byte & 0x0F
10105 } else {
10106 byte >> 4
10107 };
10108 let s = f16_to_f32(u16::from_le_bytes([
10109 scales[(flat / GROUP_SIZE) * 2],
10110 scales[(flat / GROUP_SIZE) * 2 + 1],
10111 ]));
10112 acc += ((nib as i32 - 8) as f32) * s * xv;
10113 }
10114 unsafe { *out.at(r) = acc };
10116 }
10117}
10118
10119#[allow(clippy::too_many_arguments)]
10122fn q4_range2_a8w8(
10123 packed: &[u8],
10124 scales: &[u8],
10125 gpr: usize,
10126 cols: usize,
10127 a1: &SplitAct,
10128 a2: &SplitAct,
10129 p1: SendMut,
10130 p2: SendMut,
10131 start: usize,
10132 end: usize,
10133) {
10134 for r in start..end {
10135 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
10136 let mut acc1 = s1 * a1.sx;
10137 let mut acc2 = s2 * a2.sx;
10138 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
10140 for &(j, xv) in outliers {
10141 let flat = r * cols + j;
10142 let byte = packed[flat / 2];
10143 let nib = if flat & 1 == 0 {
10144 byte & 0x0F
10145 } else {
10146 byte >> 4
10147 };
10148 let s = f16_to_f32(u16::from_le_bytes([
10149 scales[(flat / GROUP_SIZE) * 2],
10150 scales[(flat / GROUP_SIZE) * 2 + 1],
10151 ]));
10152 *acc += ((nib as i32 - 8) as f32) * s * xv;
10153 }
10154 };
10155 fix(&a1.outliers, &mut acc1);
10156 fix(&a2.outliers, &mut acc2);
10157 unsafe {
10159 *p1.at(r) = acc1;
10160 *p2.at(r) = acc2;
10161 }
10162 }
10163}
10164
10165fn q4_range_f32(
10167 packed: &[u8],
10168 scales: &[u8],
10169 gpr: usize,
10170 x: &[f32],
10171 out: SendMut,
10172 start: usize,
10173 end: usize,
10174) {
10175 for r in start..end {
10176 let mut acc = 0f32;
10177 for gi in 0..gpr {
10178 let g = r * gpr + gi;
10179 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10180 let pk = &packed[g * 16..(g + 1) * 16];
10181 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10182 let mut ga = 0f32;
10183 for (k, &b) in pk.iter().enumerate() {
10184 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10185 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10186 }
10187 acc += ga * s;
10188 }
10189 unsafe { *out.at(r) = acc };
10191 }
10192}
10193
10194#[allow(clippy::too_many_arguments)]
10198fn q4matvec2(
10199 bytes: &[u8],
10200 x1: &[f32],
10201 x2: &[f32],
10202 rows: usize,
10203 cols: usize,
10204 o1: &mut [f32],
10205 o2: &mut [f32],
10206 pool: Option<&Pool>,
10207) {
10208 debug_assert_eq!(o1.len(), rows);
10209 debug_assert_eq!(o2.len(), rows);
10210 let (packed, scales) = q4_split(bytes, rows, cols);
10211 let gpr = cols / GROUP_SIZE;
10212
10213 if a8w8_enabled() {
10214 let a1 = split_act(x1);
10215 let a2 = split_act(x2);
10216 let p1 = SendMut(o1.as_mut_ptr());
10217 let p2 = SendMut(o2.as_mut_ptr());
10218 let run = move |start: usize, end: usize| {
10219 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10220 };
10221 dispatch_rows(pool, rows, &run);
10222 return;
10223 }
10224
10225 let p1 = SendMut(o1.as_mut_ptr());
10226 let p2 = SendMut(o2.as_mut_ptr());
10227 let run = move |start: usize, end: usize| {
10228 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10229 };
10230 dispatch_rows(pool, rows, &run);
10231}
10232
10233#[allow(clippy::too_many_arguments)]
10235fn q4_range2_f32(
10236 packed: &[u8],
10237 scales: &[u8],
10238 gpr: usize,
10239 x1: &[f32],
10240 x2: &[f32],
10241 p1: SendMut,
10242 p2: SendMut,
10243 start: usize,
10244 end: usize,
10245) {
10246 for r in start..end {
10247 let (mut acc1, mut acc2) = (0f32, 0f32);
10248 for gi in 0..gpr {
10249 let g = r * gpr + gi;
10250 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10251 let pk = &packed[g * 16..(g + 1) * 16];
10252 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10253 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10254 let (mut g1, mut g2) = (0f32, 0f32);
10255 for (k, &b) in pk.iter().enumerate() {
10256 let wl = (b & 0x0F) as f32 - 8.0;
10257 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10258 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10259 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10260 }
10261 acc1 += g1 * s;
10262 acc2 += g2 * s;
10263 }
10264 unsafe {
10266 *p1.at(r) = acc1;
10267 *p2.at(r) = acc2;
10268 }
10269 }
10270}
10271
10272thread_local! {
10273 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10276 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10277}
10278
10279#[allow(clippy::too_many_arguments)]
10285fn q4matmat(
10286 bytes: &[u8],
10287 xs_all: &[f32],
10288 b: usize,
10289 rows: usize,
10290 cols: usize,
10291 out: &mut [f32],
10292 pool: Option<&Pool>,
10293) {
10294 debug_assert_eq!(xs_all.len(), b * cols);
10295 debug_assert_eq!(out.len(), b * rows);
10296 let (packed, scales) = q4_split(bytes, rows, cols);
10297 let gpr = cols / GROUP_SIZE;
10298 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10299
10300 if a8w8_enabled() {
10301 let acts: Vec<SplitAct> = (0..b)
10302 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10303 .collect();
10304 let acts = &acts;
10305 let out_addr = SendMut(out.as_mut_ptr());
10306 let run = move |start: usize, end: usize| {
10307 ROW_I8.with(|rb| {
10308 let mut buf = rb.borrow_mut();
10309 buf.resize(cols, 0);
10310 for r in start..end {
10311 for gi in 0..gpr {
10315 let g = r * gpr + gi;
10316 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10317 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10318 buf[gi * GROUP_SIZE + k * 2 + 1] =
10319 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10320 }
10321 }
10322 let mut bi = 0usize;
10323 #[cfg(target_arch = "x86_64")]
10324 if avx2_enabled() && blocked_enabled() {
10325 while bi + 4 <= acts.len() {
10326 let xs = [
10327 acts[bi].xq.as_slice(),
10328 acts[bi + 1].xq.as_slice(),
10329 acts[bi + 2].xq.as_slice(),
10330 acts[bi + 3].xq.as_slice(),
10331 ];
10332 let d = unsafe {
10333 if vnni_tiles_enabled() {
10334 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10335 } else {
10336 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10337 }
10338 };
10339 for k in 0..4 {
10340 let act = &acts[bi + k];
10341 let mut acc = d[k] * act.sx;
10342 for &(j, xv) in &act.outliers {
10343 acc += (buf[j] as i8) as f32
10344 * gscale((r * cols + j) / GROUP_SIZE)
10345 * xv;
10346 }
10347 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10349 }
10350 bi += 4;
10351 }
10352 }
10353 while bi < acts.len() {
10354 let act = &acts[bi];
10355 let mut acc = 0f32;
10356 for gi in 0..gpr {
10357 let d = dot_i8_i8(
10358 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10359 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10360 );
10361 acc += d as f32 * gscale(r * gpr + gi);
10362 }
10363 acc *= act.sx;
10364 for &(j, xv) in &act.outliers {
10366 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10367 }
10368 unsafe { *out_addr.at(bi * rows + r) = acc };
10370 bi += 1;
10371 }
10372 }
10373 })
10374 };
10375 dispatch_rows(pool, rows, &run);
10376 return;
10377 }
10378
10379 let out_addr = SendMut(out.as_mut_ptr());
10380 let run = move |start: usize, end: usize| {
10381 ROW_F32.with(|rb| {
10382 let mut buf = rb.borrow_mut();
10383 buf.resize(cols, 0.0);
10384 for r in start..end {
10385 for gi in 0..gpr {
10388 let g = r * gpr + gi;
10389 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10390 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10391 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10392 }
10393 }
10394 for bi in 0..b {
10395 let x = &xs_all[bi * cols..(bi + 1) * cols];
10396 let mut acc = 0f32;
10397 for gi in 0..gpr {
10398 let mut ga = 0f32;
10399 for k in 0..GROUP_SIZE / 2 {
10404 let e = gi * GROUP_SIZE + k * 2;
10405 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10406 }
10407 acc += ga * gscale(r * gpr + gi);
10408 }
10409 unsafe { *out_addr.at(bi * rows + r) = acc };
10411 }
10412 }
10413 })
10414 };
10415 dispatch_rows(pool, rows, &run);
10416}
10417
10418#[allow(clippy::too_many_arguments)]
10423fn vbitmatmat(
10424 bytes: &[u8],
10425 offsets: &[usize],
10426 xs_all: &[f32],
10427 b: usize,
10428 rows: usize,
10429 cols: usize,
10430 out: &mut [f32],
10431 pool: Option<&Pool>,
10432) {
10433 debug_assert_eq!(xs_all.len(), b * cols);
10434 debug_assert_eq!(out.len(), b * rows);
10435 debug_assert_eq!(offsets.len(), rows + 1);
10436 let ng = cols / GROUP_SIZE;
10437 let bits = &bytes[..rows];
10438 let sc_off = rows;
10439 let gscale = |r: usize, g: usize| {
10440 let so = (r * ng + g) * 2;
10441 f16_to_f32(u16::from_le_bytes([
10442 bytes[sc_off + so],
10443 bytes[sc_off + so + 1],
10444 ]))
10445 };
10446
10447 let decode_f32 = |r: usize, dst: &mut [f32]| {
10449 let bw = bits[r] as usize;
10450 let l = ((1i32 << (bw - 1)) - 1) as f32;
10451 let data = &bytes[offsets[r]..offsets[r + 1]];
10452 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10453 for d in dst.iter_mut() {
10454 while nbits < bw {
10455 acc = (acc << 8) | data[idx] as u64;
10456 idx += 1;
10457 nbits += 8;
10458 }
10459 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10460 nbits -= bw;
10461 *d = u - l;
10462 }
10463 };
10464
10465 if a8w8_enabled() {
10466 let acts: Vec<SplitAct> = (0..b)
10467 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10468 .collect();
10469 let acts = &acts;
10470 let out_addr = SendMut(out.as_mut_ptr());
10471 let run = move |start: usize, end: usize| {
10472 for r in start..end {
10473 let bw = bits[r] as usize;
10474 if bw == 8 {
10475 ROW_F32.with(|rb| {
10478 let mut buf = rb.borrow_mut();
10479 buf.resize(cols, 0.0);
10480 decode_f32(r, &mut buf);
10481 for bi in 0..b {
10482 let x = &xs_all[bi * cols..(bi + 1) * cols];
10483 let mut dot = 0f32;
10484 for g in 0..ng {
10485 let mut gd = 0f32;
10486 for k in 0..GROUP_SIZE {
10487 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10488 }
10489 dot += gd * gscale(r, g);
10490 }
10491 unsafe { *out_addr.at(bi * rows + r) = dot };
10493 }
10494 });
10495 continue;
10496 }
10497 let l = (1i32 << (bw - 1)) - 1;
10498 let data = &bytes[offsets[r]..offsets[r + 1]];
10499 ROW_I8.with(|rb| {
10500 let mut buf = rb.borrow_mut();
10501 buf.resize(cols, 0);
10502 #[inline(always)]
10503 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10504 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10505 let u = unpack8::<B>(&data[blk * B..]);
10506 for k in 0..8 {
10507 chunk[k] = (u[k] - l) as i8 as u8;
10508 }
10509 }
10510 }
10511 match bw {
10512 3 => fill::<3>(data, l, &mut buf),
10513 4 => vbit_fill4(data, &mut buf),
10514 5 => fill::<5>(data, l, &mut buf),
10515 6 => fill::<6>(data, l, &mut buf),
10516 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10517 }
10518 let mut bi = 0usize;
10519 #[cfg(target_arch = "x86_64")]
10523 if avx2_enabled() && blocked_enabled() {
10524 while bi + 4 <= acts.len() {
10525 let xs = [
10526 acts[bi].xq.as_slice(),
10527 acts[bi + 1].xq.as_slice(),
10528 acts[bi + 2].xq.as_slice(),
10529 acts[bi + 3].xq.as_slice(),
10530 ];
10531 let sxs = [
10532 acts[bi].sx,
10533 acts[bi + 1].sx,
10534 acts[bi + 2].sx,
10535 acts[bi + 3].sx,
10536 ];
10537 let d = unsafe {
10538 if vnni_tiles_enabled() {
10539 dot_q4b_row_1x4_sx_vnni(
10540 &buf,
10541 &bytes[sc_off..],
10542 r * ng,
10543 ng,
10544 xs,
10545 sxs,
10546 )
10547 } else {
10548 dot_q4b_row_1x4_sx_avx2(
10549 &buf,
10550 &bytes[sc_off..],
10551 r * ng,
10552 ng,
10553 xs,
10554 sxs,
10555 )
10556 }
10557 };
10558 for k in 0..4 {
10559 let act = &acts[bi + k];
10560 let mut dot = d[k];
10561 for &(j, xv) in &act.outliers {
10562 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10563 }
10564 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10566 }
10567 bi += 4;
10568 }
10569 }
10570 while bi < acts.len() {
10571 let act = &acts[bi];
10572 let mut dot = 0f32;
10573 for g in 0..ng {
10574 let d = dot_i8_i8(
10575 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10576 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10577 ) as f32
10578 * act.sx;
10579 dot += d * gscale(r, g);
10580 }
10581 for &(j, xv) in &act.outliers {
10582 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10583 }
10584 unsafe { *out_addr.at(bi * rows + r) = dot };
10586 bi += 1;
10587 }
10588 });
10589 }
10590 };
10591 dispatch_rows(pool, rows, &run);
10592 return;
10593 }
10594
10595 let out_addr = SendMut(out.as_mut_ptr());
10596 let run = move |start: usize, end: usize| {
10597 ROW_F32.with(|rb| {
10598 let mut buf = rb.borrow_mut();
10599 buf.resize(cols, 0.0);
10600 for r in start..end {
10601 decode_f32(r, &mut buf);
10602 for bi in 0..b {
10603 let x = &xs_all[bi * cols..(bi + 1) * cols];
10604 let mut dot = 0f32;
10605 for g in 0..ng {
10606 let mut gd = 0f32;
10607 for k in 0..GROUP_SIZE {
10608 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10609 }
10610 dot += gd * gscale(r, g);
10611 }
10612 unsafe { *out_addr.at(bi * rows + r) = dot };
10614 }
10615 }
10616 })
10617 };
10618 dispatch_rows(pool, rows, &run);
10619}
10620
10621pub(crate) fn gpu_batch_job<'a>(
10625 t: &'a QTensor,
10626 x: &[f32],
10627) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10628 match t {
10629 QTensor::Mapped {
10630 model,
10631 idx,
10632 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10633 rows,
10634 cols,
10635 row_scale,
10636 col_field,
10637 ..
10638 } => Some((
10639 model.clone(),
10640 crate::gpu::BatchJob {
10641 idx: *idx,
10642 rows: *rows,
10643 cols: *cols,
10644 row_scale,
10645 xs: prescale(x, col_field, *dt).into_owned(),
10646 layout: crate::gpu::BatchLayout::Q8,
10647 },
10648 )),
10649 QTensor::Mapped {
10651 model,
10652 idx,
10653 dtype: TensorDtype::Q1,
10654 rows,
10655 cols,
10656 ..
10657 } => Some((
10658 model.clone(),
10659 crate::gpu::BatchJob {
10660 idx: *idx,
10661 rows: *rows,
10662 cols: *cols,
10663 row_scale: &[],
10664 xs: x.to_vec(),
10665 layout: crate::gpu::BatchLayout::Q1,
10666 },
10667 )),
10668 QTensor::Mapped {
10670 model,
10671 idx,
10672 dtype: TensorDtype::Q1G,
10673 rows,
10674 cols,
10675 ..
10676 } => Some((
10677 model.clone(),
10678 crate::gpu::BatchJob {
10679 idx: *idx,
10680 rows: *rows,
10681 cols: *cols,
10682 row_scale: &[],
10683 xs: x.to_vec(),
10684 layout: crate::gpu::BatchLayout::Q1g,
10685 },
10686 )),
10687 QTensor::Mapped {
10692 model,
10693 idx,
10694 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10695 rows,
10696 cols,
10697 ..
10698 } => Some((
10699 model.clone(),
10700 crate::gpu::BatchJob {
10701 idx: *idx,
10702 rows: *rows,
10703 cols: *cols,
10704 row_scale: &[],
10705 xs: x.to_vec(),
10706 layout: if *dt == TensorDtype::Q4Tiled {
10707 crate::gpu::BatchLayout::Q4t
10708 } else {
10709 crate::gpu::BatchLayout::Q4tp
10710 },
10711 },
10712 )),
10713 _ => None,
10714 }
10715}
10716
10717thread_local! {
10718 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10719 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10720}
10721
10722pub(crate) fn prescale<'a>(
10723 x: &'a [f32],
10724 col_field: &[f32],
10725 dtype: TensorDtype,
10726) -> std::borrow::Cow<'a, [f32]> {
10727 if dtype == TensorDtype::Q8_2f {
10728 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10729 } else {
10730 std::borrow::Cow::Borrowed(x)
10731 }
10732}
10733
10734pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10737 x: &[f32],
10738 col_field: &[f32],
10739 dtype: TensorDtype,
10740 buf_id: u8,
10741 f: F,
10742) -> R {
10743 if dtype == TensorDtype::Q8_2f {
10744 if buf_id == 1 {
10745 PRESCALE_BUF1.with(|b| {
10746 let mut buf = b.borrow_mut();
10747 buf.clear();
10748 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10749 f(&buf)
10750 })
10751 } else {
10752 PRESCALE_BUF2.with(|b| {
10753 let mut buf = b.borrow_mut();
10754 buf.clear();
10755 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10756 f(&buf)
10757 })
10758 }
10759 } else {
10760 f(x)
10761 }
10762}
10763
10764#[cfg(target_arch = "x86_64")]
10769pub(crate) fn avx2_enabled() -> bool {
10770 use std::sync::OnceLock;
10771 static ON: OnceLock<bool> = OnceLock::new();
10772 *ON.get_or_init(|| {
10773 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10774 && std::arch::is_x86_feature_detected!("avx2")
10775 && std::arch::is_x86_feature_detected!("fma")
10776 })
10777}
10778
10779#[cfg(target_arch = "x86_64")]
10784fn avx2_a8w8_enabled() -> bool {
10785 if FLOAT_ACTIVATIONS.get() {
10786 return false;
10787 }
10788 use std::sync::OnceLock;
10789 static ON: OnceLock<bool> = OnceLock::new();
10790 *ON.get_or_init(|| {
10791 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10792 })
10793}
10794
10795thread_local! {
10796 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10797}
10798
10799pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10802 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10803 impl Drop for Restore {
10804 fn drop(&mut self) {
10805 FULL_GPU_Q8.set(self.0);
10806 }
10807 }
10808 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10809}
10810
10811thread_local! {
10815 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10816}
10817
10818fn flint_exact_cpu_default_enabled() -> bool {
10825 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10826 *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
10827}
10828
10829impl QTensor {
10830 #[inline]
10831 fn flint_exact_cpu_default(&self) -> bool {
10832 !FLOAT_ACTIVATIONS.get()
10833 && flint_exact_cpu_default_enabled()
10834 && matches!(
10835 self,
10836 Self::Mapped { model, .. } if model.header.arch.arch_name == "flint"
10837 )
10838 }
10839
10840 #[inline]
10841 fn any_flint_exact_cpu_default(ts: &[&QTensor]) -> bool {
10842 ts.iter().any(|t| t.flint_exact_cpu_default())
10843 }
10844
10845 #[inline]
10846 fn any_flint_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
10847 pairs
10848 .iter()
10849 .any(|(gate, up)| gate.flint_exact_cpu_default() || up.flint_exact_cpu_default())
10850 }
10851}
10852
10853pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
10858 float_activations_scope(f)
10859}
10860
10861#[inline]
10865pub(crate) fn float_activations_active() -> bool {
10866 FLOAT_ACTIVATIONS.get()
10867}
10868
10869pub(crate) struct FloatActivationsScope {
10873 previous: bool,
10874 _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
10875}
10876
10877impl Drop for FloatActivationsScope {
10878 fn drop(&mut self) {
10879 FLOAT_ACTIVATIONS.set(self.previous);
10880 }
10881}
10882
10883#[inline]
10884pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
10885 FloatActivationsScope {
10886 previous: FLOAT_ACTIVATIONS.replace(true),
10887 _not_send: std::marker::PhantomData,
10888 }
10889}
10890
10891#[doc(hidden)]
10897pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10898 let _scope = enter_float_activations_scope();
10899 f()
10900}
10901
10902static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10914
10915pub(crate) fn row_exact() -> bool {
10916 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10917}
10918
10919fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10920 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10921 impl Drop for Restore<'_> {
10922 fn drop(&mut self) {
10923 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10924 }
10925 }
10926 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10927 let _restore = Restore(active);
10928 f()
10929}
10930
10931pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10933 counted_row_exact_scope(&ROW_EXACT, f)
10934}
10935
10936#[inline]
10940pub(crate) fn a8w8_enabled() -> bool {
10941 #[cfg(target_arch = "aarch64")]
10942 {
10943 sdot_enabled()
10944 }
10945 #[cfg(target_arch = "x86_64")]
10946 {
10947 avx2_a8w8_enabled()
10948 }
10949 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10950 {
10951 false
10952 }
10953}
10954
10955#[inline]
10958#[allow(unreachable_code)]
10959fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10960 #[cfg(target_arch = "aarch64")]
10961 unsafe {
10962 return dot_i8_sdot(w, xq);
10963 }
10964 #[cfg(target_arch = "x86_64")]
10965 unsafe {
10966 if avx512vnni_enabled() {
10967 return dot_i8_i8_vnni(w, xq);
10968 }
10969 return dot_i8_i8_avx2(w, xq);
10970 }
10971 w.iter()
10972 .zip(xq)
10973 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10974 .sum()
10975}
10976
10977#[cfg(target_arch = "x86_64")]
10981fn avx512vnni_enabled() -> bool {
10982 use std::sync::OnceLock;
10983 static ON: OnceLock<bool> = OnceLock::new();
10984 *ON.get_or_init(|| {
10985 std::env::var("CMF_AVX512")
10986 .map(|v| v != "0")
10987 .unwrap_or(true)
10988 && std::arch::is_x86_feature_detected!("avx512f")
10989 && std::arch::is_x86_feature_detected!("avx512bw")
10990 && std::arch::is_x86_feature_detected!("avx512vl")
10991 && std::arch::is_x86_feature_detected!("avx512vnni")
10992 })
10993}
10994
10995#[cfg(target_arch = "x86_64")]
11003fn vnni_tiles_enabled() -> bool {
11004 use std::sync::OnceLock;
11005 static ON: OnceLock<bool> = OnceLock::new();
11006 *ON.get_or_init(|| {
11007 std::env::var("CMF_VNNI_TILES")
11008 .map(|v| v != "0")
11009 .unwrap_or(true)
11010 && avx512vnni_enabled()
11011 })
11012}
11013
11014#[cfg(target_arch = "x86_64")]
11019#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11020#[inline]
11021unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
11022 unsafe {
11024 use core::arch::x86_64::*;
11025 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
11026 let hi128 = _mm256_extracti128_si256::<1>(d);
11027 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11028 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11029 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11030 _mm_cvtsi128_si32(s32)
11031 }
11032}
11033
11034#[cfg(target_arch = "x86_64")]
11039#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11040unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11041 unsafe {
11043 use core::arch::x86_64::*;
11044 let n = w.len();
11045 let mut j = 0usize;
11046 let mut total: i32;
11047 {
11052 #[inline(always)]
11053 unsafe fn step(
11054 w: *const u8,
11055 x: *const i8,
11056 acc: core::arch::x86_64::__m512i,
11057 ) -> core::arch::x86_64::__m512i {
11058 unsafe {
11059 use core::arch::x86_64::*;
11060 let wv = _mm512_loadu_si512(w as *const _);
11061 let xv = _mm512_loadu_si512(x as *const _);
11062 let aw = _mm512_abs_epi8(wv);
11063 let neg = _mm512_movepi8_mask(wv);
11064 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
11065 _mm512_dpbusd_epi32(acc, aw, sx)
11066 }
11067 }
11068 let (mut a0, mut a1, mut a2, mut a3) = (
11069 _mm512_setzero_si512(),
11070 _mm512_setzero_si512(),
11071 _mm512_setzero_si512(),
11072 _mm512_setzero_si512(),
11073 );
11074 while j + 256 <= n {
11075 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11076 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
11077 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
11078 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
11079 j += 256;
11080 }
11081 while j + 64 <= n {
11082 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11083 j += 64;
11084 }
11085 let s01 = _mm512_add_epi32(a0, a1);
11086 let s23 = _mm512_add_epi32(a2, a3);
11087 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
11088 }
11089 if j + 32 <= n {
11091 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11092 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11093 let d = _mm256_dpbusd_epi32(
11094 _mm256_setzero_si256(),
11095 _mm256_abs_epi8(wv),
11096 _mm256_sign_epi8(xv, wv),
11097 );
11098 let hi128 = _mm256_extracti128_si256::<1>(d);
11099 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11100 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11101 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11102 total += _mm_cvtsi128_si32(s32);
11103 j += 32;
11104 }
11105 while j < n {
11106 total += (w[j] as i8) as i32 * xq[j] as i32;
11107 j += 1;
11108 }
11109 total
11110 }
11111}
11112
11113#[cfg(target_arch = "x86_64")]
11115#[target_feature(enable = "avx2,fma")]
11116unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
11117 unsafe {
11119 use core::arch::x86_64::*;
11120 let n = x.len();
11121 let wp = w.as_ptr();
11122 let xp = x.as_ptr();
11123 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
11124 let mut j = 0usize;
11125 while j + 16 <= n {
11126 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
11127 let lo = _mm256_cvtepi8_epi32(wb);
11128 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
11129 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
11130 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
11131 j += 16;
11132 }
11133 let acc = _mm256_add_ps(a0, a1);
11134 let hi128 = _mm256_extractf128_ps::<1>(acc);
11135 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
11136 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
11137 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
11138 let mut sum = _mm_cvtss_f32(s32);
11139 while j < n {
11140 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
11141 j += 1;
11142 }
11143 sum
11144 }
11145}
11146
11147#[cfg(target_arch = "x86_64")]
11152#[target_feature(enable = "avx2")]
11153unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
11154 unsafe {
11156 use core::arch::x86_64::*;
11157 let n = w.len();
11158 let ones = _mm256_set1_epi16(1);
11159 let mut acc = _mm256_setzero_si256();
11160 let mut j = 0usize;
11161 while j + 32 <= n {
11162 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11163 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11164 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
11165 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
11166 j += 32;
11167 }
11168 let hi128 = _mm256_extracti128_si256::<1>(acc);
11169 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
11170 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11171 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11172 let mut s = _mm_cvtsi128_si32(s32);
11173 while j < n {
11174 s += (w[j] as i8) as i32 * xq[j] as i32;
11175 j += 1;
11176 }
11177 s
11178 }
11179}
11180
11181#[cfg(target_arch = "aarch64")]
11185#[target_feature(enable = "neon,i8mm")]
11186unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11187 unsafe {
11189 use core::arch::aarch64::*;
11190 use core::arch::asm;
11191 let n = w0.len();
11192 let w0p = w0.as_ptr() as *const i8;
11193 let w1p = w1.as_ptr() as *const i8;
11194 let mut acc01 = vdupq_n_s32(0);
11197 let mut acc23 = vdupq_n_s32(0);
11198 let mut i = 0usize;
11199 while i + 8 <= n {
11200 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
11201 let xb01 = vcombine_s8(
11202 vld1_s8(xs[0].as_ptr().add(i)),
11203 vld1_s8(xs[1].as_ptr().add(i)),
11204 );
11205 let xb23 = vcombine_s8(
11206 vld1_s8(xs[2].as_ptr().add(i)),
11207 vld1_s8(xs[3].as_ptr().add(i)),
11208 );
11209 asm!(
11210 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
11211 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
11212 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
11213 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
11214 options(pure, nomem, nostack),
11215 );
11216 i += 8;
11217 }
11218 let mut out = [[0i32; 4]; 2];
11219 let a01: [i32; 4] = core::mem::transmute(acc01);
11220 let a23: [i32; 4] = core::mem::transmute(acc23);
11221 out[0][0] = a01[0];
11222 out[0][1] = a01[1];
11223 out[1][0] = a01[2];
11224 out[1][1] = a01[3];
11225 out[0][2] = a23[0];
11226 out[0][3] = a23[1];
11227 out[1][2] = a23[2];
11228 out[1][3] = a23[3];
11229 if i < n {
11230 for (k, x) in xs.iter().enumerate() {
11231 for j in i..n {
11232 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11233 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11234 }
11235 }
11236 }
11237 out
11238 }
11239}
11240
11241#[cfg(target_arch = "aarch64")]
11245#[target_feature(enable = "neon,dotprod")]
11246unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11247 unsafe {
11249 use core::arch::aarch64::*;
11250 use core::arch::asm;
11251 let n = w0.len();
11252 let w0p = w0.as_ptr() as *const i8;
11253 let w1p = w1.as_ptr() as *const i8;
11254 let mut acc = [[vdupq_n_s32(0); 4]; 2];
11255 let mut i = 0usize;
11256 while i + 16 <= n {
11257 let wv0 = vld1q_s8(w0p.add(i));
11258 let wv1 = vld1q_s8(w1p.add(i));
11259 for (k, x) in xs.iter().enumerate() {
11260 let xv = vld1q_s8(x.as_ptr().add(i));
11261 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11262 asm!(
11263 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11264 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11265 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11266 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11267 options(pure, nomem, nostack),
11268 );
11269 acc[0][k] = a0;
11270 acc[1][k] = a1;
11271 }
11272 i += 16;
11273 }
11274 let mut out = [[0i32; 4]; 2];
11275 for r in 0..2 {
11276 for k in 0..4 {
11277 out[r][k] = vaddvq_s32(acc[r][k]);
11278 }
11279 }
11280 if i < n {
11281 for (k, x) in xs.iter().enumerate() {
11282 for j in i..n {
11283 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11284 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11285 }
11286 }
11287 }
11288 out
11289 }
11290}
11291
11292#[cfg(target_arch = "x86_64")]
11298#[target_feature(enable = "avx2")]
11299unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11300 unsafe {
11302 use core::arch::x86_64::*;
11303 let n = w0.len();
11304 let ones = _mm256_set1_epi16(1);
11305 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11306 let mut j = 0usize;
11307 while j + 32 <= n {
11308 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11309 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11310 let aw0 = _mm256_abs_epi8(wv0);
11311 let aw1 = _mm256_abs_epi8(wv1);
11312 for (k, x) in xs.iter().enumerate() {
11313 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11314 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11315 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11316 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11317 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11318 }
11319 j += 32;
11320 }
11321 let mut out = [[0i32; 4]; 2];
11322 for r in 0..2 {
11323 for k in 0..4 {
11324 let a = acc[r][k];
11325 let hi128 = _mm256_extracti128_si256::<1>(a);
11326 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11327 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11328 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11329 out[r][k] = _mm_cvtsi128_si32(s32);
11330 }
11331 }
11332 if j < n {
11333 for (k, x) in xs.iter().enumerate() {
11334 for i in j..n {
11335 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11336 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11337 }
11338 }
11339 }
11340 out
11341 }
11342}
11343
11344#[cfg(target_arch = "x86_64")]
11349#[inline]
11350fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11351 let dot = if avx512vnni_enabled() && row.len() >= 64 {
11352 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11353 } else {
11354 unsafe { dot_i8_i8_avx2(row, &act.xq) }
11355 };
11356 let mut acc = dot as f32 * act.sx;
11357 for &(j, xv) in &act.outliers {
11358 acc += (row[j] as i8) as f32 * xv;
11359 }
11360 acc
11361}
11362
11363#[cfg(target_arch = "x86_64")]
11367#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11368unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11369 unsafe {
11371 use core::arch::x86_64::*;
11372 let n = w.len();
11373 let flip = _mm512_set1_epi8(-128); #[inline(always)]
11375 unsafe fn step(
11376 w: *const u8,
11377 x: *const i8,
11378 flip: core::arch::x86_64::__m512i,
11379 acc: core::arch::x86_64::__m512i,
11380 ) -> core::arch::x86_64::__m512i {
11381 unsafe {
11382 use core::arch::x86_64::*;
11383 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11384 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11385 }
11386 }
11387 let (mut a0, mut a1, mut a2, mut a3) = (
11388 _mm512_setzero_si512(),
11389 _mm512_setzero_si512(),
11390 _mm512_setzero_si512(),
11391 _mm512_setzero_si512(),
11392 );
11393 let mut j = 0usize;
11394 while j + 256 <= n {
11395 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11396 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11397 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11398 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11399 j += 256;
11400 }
11401 while j + 64 <= n {
11402 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11403 j += 64;
11404 }
11405 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11406 _mm512_add_epi32(a0, a1),
11407 _mm512_add_epi32(a2, a3),
11408 ));
11409 while j < n {
11411 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11412 j += 1;
11413 }
11414 total
11415 }
11416}
11417
11418#[cfg(target_arch = "x86_64")]
11424#[target_feature(enable = "avx2")]
11425unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11426 unsafe {
11429 use core::arch::x86_64::*;
11430 let lomask = _mm_set1_epi8(0x0F);
11431 let eight = _mm256_set1_epi8(8);
11432 let ones = _mm256_set1_epi16(1);
11433 let mut acc = 0f32;
11434 for gi in 0..gpr {
11435 let g = g0 + gi;
11436 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11437 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11438 let lo = _mm_and_si128(b, lomask);
11439 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11440 let w = _mm256_sub_epi8(
11441 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11442 eight,
11443 );
11444 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11445 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11446 let d = _mm256_madd_epi16(p16, ones);
11447 let hi128 = _mm256_extracti128_si256::<1>(d);
11448 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11449 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11450 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11451 acc += _mm_cvtsi128_si32(s32) as f32 * s;
11452 }
11453 acc
11454 }
11455}
11456
11457#[cfg(target_arch = "x86_64")]
11460#[target_feature(enable = "avx2")]
11461unsafe fn dot_q4_row_avx2_2(
11462 packed: &[u8],
11463 scales: &[u8],
11464 g0: usize,
11465 gpr: usize,
11466 xq1: &[i8],
11467 xq2: &[i8],
11468) -> (f32, f32) {
11469 unsafe {
11471 use core::arch::x86_64::*;
11472 let lomask = _mm_set1_epi8(0x0F);
11473 let eight = _mm256_set1_epi8(8);
11474 let ones = _mm256_set1_epi16(1);
11475 let (mut acc1, mut acc2) = (0f32, 0f32);
11476 #[inline(always)]
11477 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11478 unsafe {
11479 use core::arch::x86_64::*;
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 _mm_cvtsi128_si32(s32)
11485 }
11486 }
11487 for gi in 0..gpr {
11488 let g = g0 + gi;
11489 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11490 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11491 let lo = _mm_and_si128(b, lomask);
11492 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11493 let w = _mm256_sub_epi8(
11494 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11495 eight,
11496 );
11497 let aw = _mm256_abs_epi8(w);
11498 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11499 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11500 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11501 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11502 acc1 += hsum(d1) as f32 * s;
11503 acc2 += hsum(d2) as f32 * s;
11504 }
11505 (acc1, acc2)
11506 }
11507}
11508
11509#[cfg(target_arch = "x86_64")]
11511fn q8_range_avx2(
11512 q: &[u8],
11513 row_scale: &[f32],
11514 act: &SplitAct,
11515 cols: usize,
11516 out_addr: SendMut,
11517 start: usize,
11518 end: usize,
11519) {
11520 for o in start..end {
11521 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11522 unsafe { *out_addr.at(o) = v };
11524 }
11525}
11526
11527#[cfg(target_arch = "x86_64")]
11529#[allow(clippy::too_many_arguments)]
11530fn q8_range2_avx2(
11531 q: &[u8],
11532 row_scale: &[f32],
11533 a1: &SplitAct,
11534 a2: &SplitAct,
11535 cols: usize,
11536 p1: SendMut,
11537 p2: SendMut,
11538 start: usize,
11539 end: usize,
11540) {
11541 for o in start..end {
11542 let row = &q[o * cols..(o + 1) * cols];
11543 unsafe {
11545 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11546 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11547 }
11548 }
11549}
11550
11551#[cfg(target_arch = "aarch64")]
11562fn i8mm_enabled() -> bool {
11563 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11564 *ON.get_or_init(|| {
11565 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11566 && std::arch::is_aarch64_feature_detected!("i8mm")
11567 })
11568}
11569
11570#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11574fn sdot_enabled() -> bool {
11575 if FLOAT_ACTIVATIONS.get() {
11576 return false;
11577 }
11578 use std::sync::OnceLock;
11579 static ON: OnceLock<bool> = OnceLock::new();
11580 *ON.get_or_init(|| {
11581 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11582 if !want {
11583 return false;
11584 }
11585
11586 #[cfg(target_arch = "aarch64")]
11587 {
11588 if std::arch::is_aarch64_feature_detected!("dotprod") {
11589 return true;
11590 }
11591 #[cfg(target_os = "android")]
11592 {
11593 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11594 if cpuinfo.lines().any(|l| {
11595 (l.starts_with("Features") || l.starts_with("features"))
11596 && l.contains("asimddp")
11597 }) {
11598 return true;
11599 }
11600 }
11601 }
11602 false
11603 }
11604 #[cfg(not(target_arch = "aarch64"))]
11605 {
11606 false
11607 }
11608 })
11609}
11610
11611struct SplitAct {
11616 xq: Vec<i8>,
11617 sx: f32,
11618 outliers: Vec<(usize, f32)>,
11619 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11622 xsum: i32,
11623}
11624
11625thread_local! {
11626 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11629 const { std::cell::RefCell::new(Vec::new()) };
11630}
11631
11632impl Drop for SplitAct {
11633 fn drop(&mut self) {
11634 let buf = std::mem::take(&mut self.xq);
11635 if buf.capacity() > 0 {
11636 XQ_FREE.with(|f| {
11637 let mut f = f.borrow_mut();
11638 if f.len() < 16 {
11639 f.push(buf);
11640 }
11641 });
11642 }
11643 }
11644}
11645
11646thread_local! {
11647 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11654 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11659 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11660}
11661
11662#[inline]
11665fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11666 KROW.with(|s| {
11667 let b = unsafe { &mut *s.get() };
11670 if b.len() < n {
11671 b.resize(n, 0.0);
11672 }
11673 f(&mut b[..n])
11674 })
11675}
11676
11677#[inline(always)]
11688fn q8_round(t: f32) -> i8 {
11689 let t = t.clamp(-127.0, 127.0);
11690 let i = t as i32;
11691 let f = t - i as f32;
11692 let r = if f >= 0.5 {
11693 i + 1
11694 } else if f <= -0.5 {
11695 i - 1
11696 } else {
11697 i
11698 };
11699 r as i8
11700}
11701
11702#[inline]
11703fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11704 KROWS.with(|s| {
11705 let b = unsafe { &mut *s.get() };
11708 for row in b.iter_mut() {
11709 if row.len() < n {
11710 row.resize(n, 0.0);
11711 }
11712 }
11713 let (a, b) = b.split_at_mut(1);
11714 f(&mut a[0][..n], &mut b[0][..n])
11715 })
11716}
11717
11718#[inline]
11719fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11720 if limit > 0.0 {
11721 up = up.clamp(-limit, limit);
11722 gate = gate.min(limit);
11723 }
11724 gate / (1.0 + (-gate).exp()) * up
11725}
11726
11727fn split_act(x: &[f32]) -> SplitAct {
11728 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11729 let n = x.len();
11730 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11731 let thr = 8.0 * rms;
11732 let mut outliers: Vec<(usize, f32)> = Vec::new();
11736 let mut amax = 0f32;
11737 for (j, &v) in x.iter().enumerate() {
11738 let a = v.abs();
11739 if a > thr {
11740 outliers.push((j, v));
11741 } else if a > amax {
11742 amax = a;
11743 }
11744 }
11745 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11746 let inv = 1.0 / sx;
11747 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11748 xq.clear();
11749 xq.reserve(n);
11750 if outliers.is_empty() {
11751 xq.extend(
11752 x.iter()
11753 .map(|&v| q8_round(v * inv)),
11754 );
11755 } else {
11756 xq.extend(x.iter().map(|&v| {
11758 if v.abs() > thr {
11759 0
11760 } else {
11761 q8_round(v * inv)
11762 }
11763 }));
11764 }
11765 let xsum = xq.iter().map(|&v| v as i32).sum();
11766 SplitAct {
11767 xq,
11768 sx,
11769 outliers,
11770 xsum,
11771 }
11772}
11773
11774fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11775 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11776 let n = x.len();
11777 let rms = (x
11778 .iter()
11779 .zip(col)
11780 .map(|(&a, &c)| {
11781 let v = a * c;
11782 (v * v) as f64
11783 })
11784 .sum::<f64>()
11785 / n.max(1) as f64)
11786 .sqrt() as f32;
11787 let thr = 8.0 * rms;
11788
11789 let mut outliers = Vec::new();
11790 let mut amax = 0f32;
11791 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11792 let v = a * c;
11793 let s = v.abs();
11794 if s > thr {
11795 outliers.push((j, v));
11796 } else if s > amax {
11797 amax = s;
11798 }
11799 }
11800
11801 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11802 let inv = 1.0 / sx;
11803 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11804 xq.clear();
11805 xq.reserve(n);
11806 if outliers.is_empty() {
11807 xq.extend(
11808 x.iter()
11809 .zip(col)
11810 .map(|(&a, &c)| q8_round((a * c) * inv)),
11811 );
11812 } else {
11813 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11814 let v = a * c;
11815 if v.abs() > thr {
11816 0
11817 } else {
11818 q8_round(v * inv)
11819 }
11820 }));
11821 }
11822 let xsum = xq.iter().map(|&v| v as i32).sum();
11823 SplitAct {
11824 xq,
11825 sx,
11826 outliers,
11827 xsum,
11828 }
11829}
11830
11831#[cfg(target_arch = "aarch64")]
11834#[target_feature(enable = "neon,dotprod")]
11835unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11836 unsafe {
11838 use core::arch::aarch64::*;
11839 use core::arch::asm;
11840 let wp = w.as_ptr() as *const i8;
11841 let n = w.len();
11842 let (mut a0, mut a1, mut a2, mut a3) = (
11843 vdupq_n_s32(0),
11844 vdupq_n_s32(0),
11845 vdupq_n_s32(0),
11846 vdupq_n_s32(0),
11847 );
11848 let mut i = 0;
11849 while i + 64 <= n {
11850 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11851 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11852 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11853 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11854 asm!(
11855 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11856 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11857 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11858 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11859 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11860 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11861 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11862 options(pure, nomem, nostack),
11863 );
11864 i += 64;
11865 }
11866 while i + 16 <= n {
11867 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11868 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11869 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11870 i += 16;
11871 }
11872 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11873 while i < n {
11874 s += (*wp.add(i)) as i32 * xq[i] as i32;
11875 i += 1;
11876 }
11877 s
11878 }
11879}
11880
11881#[cfg(target_arch = "aarch64")]
11885#[target_feature(enable = "neon,dotprod")]
11886unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11887 unsafe {
11889 use core::arch::aarch64::*;
11890 use core::arch::asm;
11891 let n = xq.len();
11892 let px = xq.as_ptr();
11893 let (p0, p1, p2, p3) = (
11894 w0.as_ptr() as *const i8,
11895 w1.as_ptr() as *const i8,
11896 w2.as_ptr() as *const i8,
11897 w3.as_ptr() as *const i8,
11898 );
11899 let (mut a0, mut a1, mut a2, mut a3) = (
11900 vdupq_n_s32(0),
11901 vdupq_n_s32(0),
11902 vdupq_n_s32(0),
11903 vdupq_n_s32(0),
11904 );
11905 let mut i = 0;
11906 while i + 16 <= n {
11907 let x = vld1q_s8(px.add(i));
11908 let v0 = vld1q_s8(p0.add(i));
11909 let v1 = vld1q_s8(p1.add(i));
11910 let v2 = vld1q_s8(p2.add(i));
11911 let v3 = vld1q_s8(p3.add(i));
11912 asm!(
11913 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11914 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11915 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11916 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11917 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11918 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11919 options(pure, nomem, nostack),
11920 );
11921 i += 16;
11922 }
11923 let mut r = [
11924 vaddvq_s32(a0),
11925 vaddvq_s32(a1),
11926 vaddvq_s32(a2),
11927 vaddvq_s32(a3),
11928 ];
11929 while i < n {
11930 let xi = *px.add(i) as i32;
11931 r[0] += (*p0.add(i)) as i32 * xi;
11932 r[1] += (*p1.add(i)) as i32 * xi;
11933 r[2] += (*p2.add(i)) as i32 * xi;
11934 r[3] += (*p3.add(i)) as i32 * xi;
11935 i += 1;
11936 }
11937 r
11938 }
11939}
11940
11941#[cfg(target_arch = "aarch64")]
11948#[target_feature(enable = "neon,dotprod")]
11949unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11950 unsafe {
11953 use core::arch::aarch64::*;
11954 use core::arch::asm;
11955 let n = xq.len();
11956 let px = xq.as_ptr();
11957 let pg = g.as_ptr() as *const i8;
11958 let (mut a0, mut a1, mut a2, mut a3) = (
11959 vdupq_n_s32(0),
11960 vdupq_n_s32(0),
11961 vdupq_n_s32(0),
11962 vdupq_n_s32(0),
11963 );
11964 let mut i = 0;
11965 while i + 16 <= n {
11966 let x = vld1q_s8(px.add(i));
11967 let base = pg.add(4 * i);
11968 let v0 = vld1q_s8(base);
11969 let v1 = vld1q_s8(base.add(16));
11970 let v2 = vld1q_s8(base.add(32));
11971 let v3 = vld1q_s8(base.add(48));
11972 asm!(
11973 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11974 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11975 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11976 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11977 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11978 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11979 options(pure, nomem, nostack),
11980 );
11981 i += 16;
11982 }
11983 [
11984 vaddvq_s32(a0),
11985 vaddvq_s32(a1),
11986 vaddvq_s32(a2),
11987 vaddvq_s32(a3),
11988 ]
11989 }
11990}
11991
11992#[cfg(target_arch = "aarch64")]
11998fn q8_range_sdot(
11999 q: &[u8],
12000 rep: &[u8],
12001 row_scale: &[f32],
12002 act: &SplitAct,
12003 cols: usize,
12004 out_addr: SendMut,
12005 start: usize,
12006 end: usize,
12007) {
12008 let mut o = start;
12009 if !rep.is_empty() {
12012 while o < end && o % 4 != 0 {
12013 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12014 unsafe { *out_addr.at(o) = v };
12015 o += 1;
12016 }
12017 }
12018 while o + 4 <= end {
12019 let r = if rep.is_empty() {
12020 unsafe {
12021 dot_i8_sdot_4rows(
12022 &q[o * cols..(o + 1) * cols],
12023 &q[(o + 1) * cols..(o + 2) * cols],
12024 &q[(o + 2) * cols..(o + 3) * cols],
12025 &q[(o + 3) * cols..(o + 4) * cols],
12026 &act.xq,
12027 )
12028 }
12029 } else {
12030 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
12031 };
12032 for k in 0..4 {
12033 let mut acc = r[k] as f32 * act.sx;
12034 for &(j, xv) in &act.outliers {
12035 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
12036 }
12037 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
12039 }
12040 o += 4;
12041 }
12042 while o < end {
12043 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12044 unsafe { *out_addr.at(o) = v };
12045 o += 1;
12046 }
12047}
12048
12049#[cfg(target_arch = "aarch64")]
12052#[allow(clippy::too_many_arguments)]
12053fn q8_range2_sdot(
12054 q: &[u8],
12055 row_scale: &[f32],
12056 a1: &SplitAct,
12057 a2: &SplitAct,
12058 cols: usize,
12059 p1: SendMut,
12060 p2: SendMut,
12061 start: usize,
12062 end: usize,
12063) {
12064 for o in start..end {
12065 let row = &q[o * cols..(o + 1) * cols];
12066 unsafe {
12068 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
12069 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
12070 }
12071 }
12072}
12073
12074#[allow(clippy::too_many_arguments)]
12076fn q8_range2_f32(
12077 q: &[u8],
12078 row_scale: &[f32],
12079 x1: &[f32],
12080 x2: &[f32],
12081 cols: usize,
12082 p1: SendMut,
12083 p2: SendMut,
12084 start: usize,
12085 end: usize,
12086) {
12087 for o in start..end {
12088 let row = &q[o * cols..(o + 1) * cols];
12089 unsafe {
12091 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
12092 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
12093 }
12094 }
12095}
12096
12097fn q8_range_f32(
12099 q: &[u8],
12100 row_scale: &[f32],
12101 xs: &[f32],
12102 cols: usize,
12103 out_addr: SendMut,
12104 start: usize,
12105 end: usize,
12106) {
12107 for o in start..end {
12108 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12109 unsafe { *out_addr.at(o) = v };
12111 }
12112}
12113
12114#[inline]
12118fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
12119 #[cfg(target_arch = "aarch64")]
12120 return row_dot_sdot(row, act);
12121 #[cfg(target_arch = "x86_64")]
12122 return row_dot_avx2(row, act);
12123 #[allow(unreachable_code)]
12124 q8_row_dot_scalar(row, act)
12125}
12126
12127#[allow(dead_code)]
12128fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
12129 let mut acc = 0i32;
12130 for (k, &b) in row.iter().enumerate() {
12131 acc += (b as i8) as i32 * act.xq[k] as i32;
12132 }
12133 let mut acc = acc as f32 * act.sx;
12134 for &(j, xv) in &act.outliers {
12135 acc += (row[j] as i8) as f32 * xv;
12136 }
12137 acc
12138}
12139
12140#[cfg(target_arch = "aarch64")]
12143#[inline]
12144fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
12145 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
12146 for &(j, xv) in &act.outliers {
12147 acc += (row[j] as i8) as f32 * xv;
12148 }
12149 acc
12150}
12151
12152#[cfg(target_arch = "aarch64")]
12160#[target_feature(enable = "neon,dotprod")]
12161unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12162 unsafe {
12165 use core::arch::aarch64::*;
12166 use core::arch::asm;
12167 let lomask = vdupq_n_u8(0x0F);
12168 let eight = vdupq_n_s8(8);
12169 let mut acc = 0f32;
12170 for gi in 0..gpr {
12171 let g = g0 + gi;
12172 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12173 let b = vld1q_u8(packed.as_ptr().add(g * 16));
12174 let lo = vandq_u8(b, lomask);
12175 let hi = vshrq_n_u8::<4>(b);
12176 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12177 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12178 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
12179 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
12180 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
12181 asm!(
12182 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
12183 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
12184 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12185 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
12186 options(pure, nomem, nostack),
12187 );
12188 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12189 }
12190 acc
12191 }
12192}
12193
12194#[cfg(target_arch = "aarch64")]
12199#[target_feature(enable = "neon,dotprod")]
12200unsafe fn dot_q4_row_sdot2(
12201 packed: &[u8],
12202 scales: &[u8],
12203 g0: usize,
12204 gpr: usize,
12205 xq1: &[i8],
12206 xq2: &[i8],
12207) -> (f32, f32) {
12208 unsafe {
12211 use core::arch::aarch64::*;
12212 use core::arch::asm;
12213 let lomask = vdupq_n_u8(0x0F);
12214 let eight = vdupq_n_s8(8);
12215 let (mut acc1, mut acc2) = (0f32, 0f32);
12216 for gi in 0..gpr {
12217 let g = g0 + gi;
12218 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12219 let b = vld1q_u8(packed.as_ptr().add(g * 16));
12220 let lo = vandq_u8(b, lomask);
12221 let hi = vshrq_n_u8::<4>(b);
12222 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12223 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12224 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
12225 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
12226 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
12227 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
12228 let (mut a0, mut a1, mut b0, mut b1) = (
12229 vdupq_n_s32(0),
12230 vdupq_n_s32(0),
12231 vdupq_n_s32(0),
12232 vdupq_n_s32(0),
12233 );
12234 asm!(
12235 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
12236 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
12237 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
12238 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
12239 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12240 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
12241 e0 = in(vreg) e0, e1 = in(vreg) e1,
12242 x10 = in(vreg) x10, x11 = in(vreg) x11,
12243 x20 = in(vreg) x20, x21 = in(vreg) x21,
12244 options(pure, nomem, nostack),
12245 );
12246 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12247 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
12248 }
12249 (acc1, acc2)
12250 }
12251}
12252
12253#[inline]
12258pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12259 #[cfg(target_arch = "aarch64")]
12260 unsafe {
12261 return axpy_i8_f32_neon(acc, row, w);
12262 }
12263 #[cfg(target_arch = "x86_64")]
12264 if avx2_enabled() {
12265 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12266 }
12267 #[allow(unreachable_code)]
12268 {
12269 for (a, &b) in acc.iter_mut().zip(row) {
12270 *a += w * b as f32;
12271 }
12272 }
12273}
12274
12275#[cfg(target_arch = "x86_64")]
12277#[target_feature(enable = "avx2,fma")]
12278unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12279 unsafe {
12281 use core::arch::x86_64::*;
12282 let n = acc.len().min(row.len());
12283 let ap = acc.as_mut_ptr();
12284 let rp = row.as_ptr();
12285 let wv = _mm256_set1_ps(w);
12286 let mut j = 0usize;
12287 while j + 16 <= n {
12288 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12289 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12290 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12291 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12292 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12293 _mm256_storeu_ps(ap.add(j), v0);
12294 _mm256_storeu_ps(ap.add(j + 8), v1);
12295 j += 16;
12296 }
12297 while j < n {
12298 *ap.add(j) += w * (*rp.add(j)) as f32;
12299 j += 1;
12300 }
12301 }
12302}
12303
12304#[cfg(target_arch = "aarch64")]
12305#[target_feature(enable = "neon")]
12306unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12307 unsafe {
12309 use core::arch::aarch64::*;
12310 let n = acc.len().min(row.len());
12311 let ap = acc.as_mut_ptr();
12312 let rp = row.as_ptr();
12313 let wv = vdupq_n_f32(w);
12314 let mut j = 0usize;
12315 while j + 16 <= n {
12316 let rb = vld1q_s8(rp.add(j));
12317 let lo = vmovl_s8(vget_low_s8(rb));
12318 let hi = vmovl_s8(vget_high_s8(rb));
12319 for (off, half) in [(0, lo), (8, hi)] {
12320 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12321 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12322 let o = j + off;
12323 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12324 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12325 }
12326 j += 16;
12327 }
12328 while j < n {
12329 *ap.add(j) += w * (*rp.add(j)) as f32;
12330 j += 1;
12331 }
12332 }
12333}
12334
12335#[inline]
12338pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12339 #[cfg(target_arch = "aarch64")]
12340 unsafe {
12341 return dot_i8_f32_neon(w, x);
12342 }
12343 #[cfg(target_arch = "x86_64")]
12344 if avx2_enabled() {
12345 return unsafe { dot_i8_f32_avx2(w, x) };
12346 }
12347 #[allow(unreachable_code)]
12348 {
12349 let mut sum = 0.0f32;
12350 for (j, &b) in w.iter().enumerate() {
12351 sum += (b as i8) as f32 * x[j];
12352 }
12353 sum
12354 }
12355}
12356
12357#[inline]
12361fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12362 #[cfg(target_arch = "aarch64")]
12363 unsafe {
12364 return dot_i8_col_f32_neon(w, x, col);
12365 }
12366 #[allow(unreachable_code)]
12367 {
12368 let mut sum = 0.0f32;
12369 for (j, &b) in w.iter().enumerate() {
12370 sum += (b as i8) as f32 * x[j] * col[j];
12371 }
12372 sum
12373 }
12374}
12375
12376#[cfg(target_arch = "aarch64")]
12377#[target_feature(enable = "neon")]
12378unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12379 unsafe {
12381 use core::arch::aarch64::*;
12382 let n = x.len();
12383 let wp = w.as_ptr() as *const i8;
12384 let xp = x.as_ptr();
12385 let cp = col.as_ptr();
12386 let (mut a0, mut a1, mut a2, mut a3) = (
12387 vdupq_n_f32(0.0),
12388 vdupq_n_f32(0.0),
12389 vdupq_n_f32(0.0),
12390 vdupq_n_f32(0.0),
12391 );
12392 let mut j = 0usize;
12393 while j + 16 <= n {
12394 let wb = vld1q_s8(wp.add(j));
12395 let lo = vmovl_s8(vget_low_s8(wb));
12396 let hi = vmovl_s8(vget_high_s8(wb));
12397 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12398 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12399 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12400 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12401 a0 = vfmaq_f32(
12402 a0,
12403 w0,
12404 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12405 );
12406 a1 = vfmaq_f32(
12407 a1,
12408 w1,
12409 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12410 );
12411 a2 = vfmaq_f32(
12412 a2,
12413 w2,
12414 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12415 );
12416 a3 = vfmaq_f32(
12417 a3,
12418 w3,
12419 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12420 );
12421 j += 16;
12422 }
12423 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12424 while j < n {
12425 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12426 j += 1;
12427 }
12428 sum
12429 }
12430}
12431
12432#[cfg(target_arch = "aarch64")]
12433#[target_feature(enable = "neon")]
12434unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12435 unsafe {
12437 use core::arch::aarch64::*;
12438 let n = x.len();
12439 let wp = w.as_ptr() as *const i8;
12440 let xp = x.as_ptr();
12441 let (mut a0, mut a1, mut a2, mut a3) = (
12442 vdupq_n_f32(0.0),
12443 vdupq_n_f32(0.0),
12444 vdupq_n_f32(0.0),
12445 vdupq_n_f32(0.0),
12446 );
12447 let mut j = 0usize;
12448 while j + 16 <= n {
12449 let wb = vld1q_s8(wp.add(j));
12450 let lo = vmovl_s8(vget_low_s8(wb));
12451 let hi = vmovl_s8(vget_high_s8(wb));
12452 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12453 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12454 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12455 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12456 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12457 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12458 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12459 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12460 j += 16;
12461 }
12462 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12463 while j < n {
12464 sum += (*wp.add(j)) as f32 * *xp.add(j);
12465 j += 1;
12466 }
12467 sum
12468 }
12469}
12470
12471#[allow(clippy::too_many_arguments)]
12472fn qmatvec(
12473 q: &[u8],
12474 rep: &[u8],
12475 row_scale: &[f32],
12476 x: &[f32],
12477 col_field: &[f32],
12478 dtype: TensorDtype,
12479 rows: usize,
12480 cols: usize,
12481 out: &mut [f32],
12482 pool: Option<&Pool>,
12483) {
12484 debug_assert_eq!(out.len(), rows);
12485 #[cfg(not(target_arch = "aarch64"))]
12486 let _ = rep;
12487
12488 #[cfg(target_arch = "aarch64")]
12489 if sdot_enabled() {
12490 let act = if dtype == TensorDtype::Q8_2f {
12491 split_act_q8_2f(x, col_field)
12492 } else {
12493 split_act(x)
12494 };
12495 let out_addr = SendMut(out.as_mut_ptr());
12496 let run_range = |start: usize, end: usize| {
12497 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12498 };
12499 match pool {
12500 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12501 _ => run_range(0, rows),
12502 }
12503 return;
12504 }
12505 #[cfg(target_arch = "x86_64")]
12508 if avx2_a8w8_enabled() {
12509 let act = if dtype == TensorDtype::Q8_2f {
12510 split_act_q8_2f(x, col_field)
12511 } else {
12512 split_act(x)
12513 };
12514 let out_addr = SendMut(out.as_mut_ptr());
12515 let run_range = |start: usize, end: usize| {
12516 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12517 };
12518 match pool {
12519 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12520 _ => run_range(0, rows),
12521 }
12522 return;
12523 }
12524
12525 prescale_with(x, col_field, dtype, 1, |xs| {
12526 let out_addr = SendMut(out.as_mut_ptr());
12527 let run_range = move |start: usize, end: usize| {
12528 for o in start..end {
12529 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12530 unsafe { *out_addr.at(o) = v };
12532 }
12533 };
12534 match pool {
12535 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12536 _ => run_range(0, rows),
12537 }
12538 });
12539}
12540
12541#[allow(clippy::too_many_arguments)]
12542fn qmatvec2(
12543 q: &[u8],
12544 row_scale: &[f32],
12545 x1: &[f32],
12546 x2: &[f32],
12547 col_field: &[f32],
12548 dtype: TensorDtype,
12549 rows: usize,
12550 cols: usize,
12551 o1: &mut [f32],
12552 o2: &mut [f32],
12553 pool: Option<&Pool>,
12554) {
12555 #[cfg(target_arch = "aarch64")]
12556 if sdot_enabled() {
12557 let a1s = if dtype == TensorDtype::Q8_2f {
12558 split_act_q8_2f(x1, col_field)
12559 } else {
12560 split_act(x1)
12561 };
12562 let a2s = if dtype == TensorDtype::Q8_2f {
12563 split_act_q8_2f(x2, col_field)
12564 } else {
12565 split_act(x2)
12566 };
12567 let p1 = SendMut(o1.as_mut_ptr());
12568 let p2 = SendMut(o2.as_mut_ptr());
12569 let run_range = |start: usize, end: usize| {
12570 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12571 };
12572 match pool {
12573 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12574 _ => run_range(0, rows),
12575 }
12576 return;
12577 }
12578 #[cfg(target_arch = "x86_64")]
12579 if avx2_a8w8_enabled() {
12580 let a1s = if dtype == TensorDtype::Q8_2f {
12581 split_act_q8_2f(x1, col_field)
12582 } else {
12583 split_act(x1)
12584 };
12585 let a2s = if dtype == TensorDtype::Q8_2f {
12586 split_act_q8_2f(x2, col_field)
12587 } else {
12588 split_act(x2)
12589 };
12590 let p1 = SendMut(o1.as_mut_ptr());
12591 let p2 = SendMut(o2.as_mut_ptr());
12592 let run_range = |start: usize, end: usize| {
12593 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12594 };
12595 match pool {
12596 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12597 _ => run_range(0, rows),
12598 }
12599 return;
12600 }
12601
12602 prescale_with(x1, col_field, dtype, 1, |x1s| {
12603 prescale_with(x2, col_field, dtype, 2, |x2s| {
12604 let p1 = SendMut(o1.as_mut_ptr());
12605 let p2 = SendMut(o2.as_mut_ptr());
12606 let run_range = move |start: usize, end: usize| {
12607 for o in start..end {
12608 let row = &q[o * cols..(o + 1) * cols];
12609 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12610 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12611 unsafe {
12613 *p1.at(o) = s1;
12614 *p2.at(o) = s2;
12615 }
12616 }
12617 };
12618 match pool {
12619 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12620 _ => run_range(0, rows),
12621 }
12622 });
12623 });
12624}
12625
12626#[derive(Clone, Copy)]
12627struct SendMut(*mut f32);
12628unsafe impl Send for SendMut {}
12629unsafe impl Sync for SendMut {}
12630
12631impl SendMut {
12632 #[inline]
12633 fn at(self, i: usize) -> *mut f32 {
12634 unsafe { self.0.add(i) }
12635 }
12636}
12637
12638#[cfg(test)]
12639mod tests {
12640 #[test]
12644 fn q8_round_is_round_clamp() {
12645 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12646 let mut probe = vec![
12647 0.0f32,
12648 -0.0,
12649 f32::NAN,
12650 f32::INFINITY,
12651 f32::NEG_INFINITY,
12652 f32::MAX,
12653 f32::MIN,
12654 1e30,
12655 -1e30,
12656 f32::MIN_POSITIVE,
12657 -f32::MIN_POSITIVE,
12658 ];
12659 for k in -300i32..=300 {
12660 let h = k as f32 * 0.5;
12661 let up = f32::from_bits(h.to_bits() + 1);
12662 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12663 for t in [h, up, down] {
12664 probe.push(t);
12665 probe.push(-t);
12666 }
12667 }
12668 let mut t = -140.0f32;
12669 while t < 140.0 {
12670 probe.push(t);
12671 t += 0.000_731;
12672 }
12673 for t in probe {
12674 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12675 }
12676 }
12677
12678 use super::*;
12679
12680 #[test]
12684 fn f32_matmat_equals_matvec_bitwise() {
12685 let (rows, cols, b) = (5usize, 37usize, 7usize);
12686 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12687 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12688 let t = QTensor::from_f32(w, rows, cols);
12689 let mut out = vec![0f32; b * rows];
12690 t.matmat(&xs, b, &mut out, None);
12691 for bi in 0..b {
12692 let mut one = vec![0f32; rows];
12693 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12694 for o in 0..rows {
12695 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12696 }
12697 }
12698 }
12699
12700 #[test]
12701 fn q2tp_i8_dot_matches_exact_on_grid() {
12702 let (rows, cols) = (5, 64);
12706 let gpr = cols / GROUP_SIZE;
12707 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12711 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12712 .collect();
12713 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12714 let x: Vec<f32> = (0..cols)
12715 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12716 .collect();
12717 let act = split_act(&x);
12718 assert!(
12719 act.outliers.is_empty(),
12720 "on-grid input must have no outliers"
12721 );
12722 let gsum = q1_group_sums(&act.xq, gpr);
12723 for r in 0..rows {
12724 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12725 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12726 assert!(
12727 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12728 "row {r}: exact {exact} vs i8 {fast}"
12729 );
12730 }
12731 }
12732
12733 #[cfg(target_arch = "x86_64")]
12734 #[test]
12735 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12736 if !std::arch::is_x86_feature_detected!("avx2") {
12741 return;
12742 }
12743 let mut seed = 0x9e3779b9u32;
12744 let mut next = || {
12745 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12746 seed
12747 };
12748 for _ in 0..20_000 {
12749 let mut ch = [0u8; Q2TP_CHUNK];
12750 let mut x = [0i8; GROUP_SIZE];
12751 for b in &mut ch {
12752 *b = next() as u8;
12753 }
12754 for v in &mut x {
12755 *v = (next() >> 24) as i8;
12756 }
12757 let mut reference = 0i32;
12758 for (k, &b) in ch.iter().enumerate() {
12759 reference += (b & 3) as i32 * x[k * 4] as i32;
12760 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12761 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12762 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12763 }
12764 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12767 assert_eq!(got, reference, "packed q2 lane mismatch");
12768 }
12769 }
12770
12771 #[test]
12772 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12773 let (rows, cols) = (1usize, GROUP_SIZE);
12774 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12775 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12779 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12780 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
12782 let mut raw = vec![0.0f32; rows];
12783 let mut affine = vec![0.0f32; rows];
12784 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12785 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12786 assert_eq!(raw, vec![-24.0]);
12787 assert_eq!(affine, vec![-8.0]);
12788 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12789 }
12790
12791 #[test]
12792 fn q8_row_dot_fast_matches_scalar() {
12793 let cols = 96;
12796 let row: Vec<u8> = (0..cols)
12797 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12798 .collect();
12799 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12800 let act = split_act(&x);
12801 let fast = q8_row_dot(&row, &act);
12802 let scalar = q8_row_dot_scalar(&row, &act);
12803 assert!(
12804 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12805 "fast {fast} vs scalar {scalar}"
12806 );
12807 }
12808
12809 #[test]
12810 fn f32_matvec_matches_matvec_rows_bitexact() {
12811 let (rows, cols) = (300, 40);
12812 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12813 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12814 let qt = QTensor::from_f32(w.clone(), rows, cols);
12815
12816 let mut a = vec![0.0f32; rows];
12817 matvec_rows(None, &w, &x, &mut a);
12818 let mut b = vec![0.0f32; rows];
12819 qt.matvec(&x, &mut b, None);
12820 assert_eq!(a, b);
12821 }
12822
12823 #[test]
12824 fn sdot_kernel_exact_on_grid() {
12825 eprintln!("sdot_enabled = {}", sdot_enabled());
12830 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
12832 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12833 .collect();
12834 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12835 let x: Vec<f32> = (0..cols)
12836 .map(|i| match i % 3 {
12837 0 => 1.0,
12838 1 => -1.0,
12839 _ => 0.0,
12840 })
12841 .collect();
12842 let mut a = vec![0.0f32; rows];
12843 qmatvec(
12844 &w,
12845 &[],
12846 &scales,
12847 &x,
12848 &[],
12849 TensorDtype::Q8Row,
12850 rows,
12851 cols,
12852 &mut a,
12853 None,
12854 );
12855 for o in 0..rows {
12856 let mut acc = 0.0f32;
12857 for j in 0..cols {
12858 acc += (w[o * cols + j] as i8) as f32 * x[j];
12859 }
12860 let expect = acc * scales[o];
12861 assert!(
12862 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12863 "row {o}: {} vs {expect}",
12864 a[o]
12865 );
12866 }
12867 }
12868
12869 #[test]
12870 fn q1_tbl_fast_path_matches_reference() {
12871 let (rows, cols) = (5, 256);
12876 let gpr = cols / GROUP_SIZE;
12877 let mut bytes = Vec::new();
12878 for t in 0..rows * gpr {
12879 let s = 0.007 + (t % 11) as f32 * 0.004;
12880 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12881 for j in 0..4 {
12882 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12883 }
12884 }
12885 let x: Vec<f32> = (0..cols)
12886 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12887 .collect();
12888 let mut w = vec![0.0f32; rows * cols];
12889 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12890 let mut got = vec![0.0f32; rows];
12891 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12892 for o in 0..rows {
12893 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12894 assert!(
12895 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12896 "row {o}: {} vs {expect}",
12897 got[o]
12898 );
12899 }
12900 let b = 5usize;
12903 let mut xs_all = Vec::new();
12904 for bi in 0..b {
12905 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12906 }
12907 let mut mm = vec![0.0f32; b * rows];
12908 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12909 for bi in 0..b {
12910 let mut single = vec![0.0f32; rows];
12911 q1_matvec(
12912 &bytes,
12913 &xs_all[bi * cols..(bi + 1) * cols],
12914 rows,
12915 cols,
12916 &mut single,
12917 None,
12918 );
12919 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12920 }
12921 }
12922
12923 #[test]
12924 fn q1_kernels_match_exact_reference() {
12925 let (rows, cols) = (7, 96);
12927 let gpr = cols / GROUP_SIZE;
12928 let mut bytes = Vec::new();
12929 for t in 0..rows * gpr {
12930 let s = 0.01 + (t % 13) as f32 * 0.003;
12931 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12932 for j in 0..4 {
12933 bytes.push(((t * 31 + j * 97) % 251) as u8);
12934 }
12935 }
12936 let x: Vec<f32> = (0..cols)
12938 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12939 .collect();
12940 let mut w = vec![0.0f32; rows * cols];
12942 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12943 let mut expect = vec![0.0f32; rows];
12944 for o in 0..rows {
12945 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12946 }
12947 let mut got = vec![0.0f32; rows];
12948 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12949 for o in 0..rows {
12950 assert!(
12951 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12952 "row {o}: {} vs {}",
12953 got[o],
12954 expect[o]
12955 );
12956 }
12957 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12959 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12960 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12961 assert_eq!(a1, got);
12962 let mut xs = x.clone();
12963 xs.extend_from_slice(&x2);
12964 let mut mm = vec![0.0f32; 2 * rows];
12965 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12966 assert_eq!(&mm[..rows], got.as_slice());
12967 assert_eq!(&mm[rows..], a2.as_slice());
12968 }
12969
12970 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12973 let tiles = rows * cols / GROUP_SIZE;
12974 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12975 for t in 0..tiles {
12976 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12977 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12978 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12979 q1.extend_from_slice(&bits.to_le_bytes());
12980 }
12981 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12982 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12983 (q1, rp.data)
12984 }
12985
12986 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12989 let mut x: Vec<f32> = (0..cols)
12990 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12991 .collect();
12992 x[3 % cols] = 41.0;
12993 x[(cols / 2 + 5) % cols] = -37.5;
12994 x
12995 }
12996
12997 #[test]
13004 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
13005 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
13006 let gpr = cols / GROUP_SIZE;
13007 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
13008 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13009 let x = q1g_test_x(cols, rows);
13010 let act = split_act(&x);
13011 assert!(!act.outliers.is_empty(), "test wants the outlier path");
13012 let gsum = q1_group_sums(&act.xq, gpr);
13013 for r in 0..rows {
13014 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
13015 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
13016 let (a, b) = (
13017 wa.dot_i8(r, gpr, &act.xq, &gsum),
13018 wb.dot_i8(r, gpr, &act.xq, &gsum),
13019 );
13020 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
13021 for j in [0, 31, 32, 127, cols - 1] {
13022 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
13023 assert_eq!(
13024 (a.0.to_bits(), a.1.to_bits()),
13025 (b.0.to_bits(), b.1.to_bits())
13026 );
13027 }
13028 #[cfg(target_arch = "x86_64")]
13029 if avx2_enabled() {
13030 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
13031 unsafe {
13033 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
13034 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
13035 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
13036 }
13037 if std::is_x86_feature_detected!("avx512vnni")
13038 && std::is_x86_feature_detected!("avx512bw")
13039 && std::is_x86_feature_detected!("avx512vl")
13040 {
13041 unsafe {
13043 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
13044 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
13045 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
13046 }
13047 }
13048 }
13049 }
13050 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
13052 bin_matvec(wa, &x, rows, cols, &mut ga, None);
13053 bin_matvec(wb, &x, rows, cols, &mut gb, None);
13054 assert_eq!(ga, gb, "{rows}x{cols} matvec");
13055 let x2 = q1g_test_x(cols, rows + 1);
13056 let (mut a1, mut a2, mut b1, mut b2) = (
13057 vec![0f32; rows],
13058 vec![0f32; rows],
13059 vec![0f32; rows],
13060 vec![0f32; rows],
13061 );
13062 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
13063 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
13064 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
13065 let b = 5usize;
13066 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
13067 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
13068 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
13069 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
13070 assert_eq!(ma, mb, "{rows}x{cols} matmat");
13071 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
13074 .map(|k| {
13075 let a = split_act(&xs[k * cols..(k + 1) * cols]);
13076 let g = q1_group_sums(&a.xq, gpr);
13077 (a, g)
13078 })
13079 .collect();
13080 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
13081 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
13082 for r in 0..rows {
13083 #[cfg(target_arch = "x86_64")]
13084 if avx2_enabled() {
13085 let (a, b) = unsafe {
13087 (
13088 wa.dot_1x4_x86(r, gpr, xq4, gs4),
13089 wb.dot_1x4_x86(r, gpr, xq4, gs4),
13090 )
13091 };
13092 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13093 }
13094 #[cfg(target_arch = "aarch64")]
13095 if sdot_enabled() {
13096 let (a, b) = unsafe {
13098 (
13099 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
13100 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
13101 )
13102 };
13103 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13104 }
13105 let _ = (r, &xq4, &gs4);
13106 }
13107 }
13108 }
13109
13110 #[test]
13116 #[ignore]
13117 fn q1g_cpu_matvec_speed() {
13118 use crate::pool::Pool;
13119 let threads = std::env::var("CMF_THREADS")
13120 .ok()
13121 .and_then(|v| v.parse().ok())
13122 .unwrap_or(8usize);
13123 let pool = Pool::new(threads);
13124 for &(name, rows, cols, reps) in &[
13125 ("gate/up_proj", 6144usize, 2048usize, 41usize),
13126 ("down_proj", 2048, 6144, 41),
13127 ("lm_head", 151_669, 2048, 9),
13128 ] {
13129 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
13130 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13131 let x = q1g_test_x(cols, 1);
13132 let mut out = vec![0f32; rows];
13133 let mut time = |w: BinW| {
13134 let t0 = std::time::Instant::now();
13135 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
13136 t0.elapsed().as_secs_f64()
13137 };
13138 let _ = (time(wa), time(wb));
13139 let (mut ta, mut tb) = (Vec::new(), Vec::new());
13140 for r in 0..reps {
13141 if r % 2 == 0 {
13142 ta.push(time(wa));
13143 tb.push(time(wb));
13144 } else {
13145 tb.push(time(wb));
13146 ta.push(time(wa));
13147 }
13148 }
13149 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
13150 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
13151 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
13152 eprintln!(
13153 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
13154 a8w8_enabled(),
13155 ma * 1e6,
13156 mb * 1e6,
13157 mb / ma
13158 );
13159 }
13160 }
13161
13162 #[test]
13167 fn q1g_mapped_tensor_paths_equal_q1() {
13168 use crate::pool::Pool;
13169 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13170 let (rows, cols) = (300usize, 256usize);
13171 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
13172 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
13173 let header: CmfHeader = serde_json::from_value(serde_json::json!({
13174 "format": "cmf", "version": cortiq_core::CMF_VERSION,
13175 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
13176 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13177 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
13178 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13179 "max_position_embeddings": 8},
13180 "quant_type": "VBIT"
13181 }))
13182 .unwrap();
13183 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
13184 name: name.into(),
13185 dtype,
13186 shape: vec![rows, cols],
13187 data: data.clone(),
13188 };
13189 let specs = [
13190 spec("g1", TensorDtype::Q1, &gq1),
13191 spec("u1", TensorDtype::Q1, &uq1),
13192 spec("gg", TensorDtype::Q1G, &gq1g),
13193 spec("ug", TensorDtype::Q1G, &uq1g),
13194 ];
13195 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
13196 std::fs::create_dir_all(&dir).unwrap();
13197 let path = dir.join("m.cmf");
13198 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13199 let model = Arc::new(CmfModel::open(&path).unwrap());
13200 assert!(model.verify().is_empty(), "{:?}", model.verify());
13201 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13202 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
13203 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
13204 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
13205 let pool = Pool::new(3);
13206 let x = q1g_test_x(cols, 4);
13207 let run = |a: &QTensor| {
13208 let mut o = vec![0f32; rows];
13209 a.matvec(&x, &mut o, Some(&pool));
13210 o
13211 };
13212 assert_eq!(run(&g1), run(&gg), "matvec");
13213 let many = |a: &QTensor, b: &QTensor| {
13214 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
13215 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
13216 (oa, ob)
13217 };
13218 let want = many(&g1, &u1);
13219 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
13220 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
13221 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
13222 let fused = |g: &QTensor, u: &QTensor| {
13223 let mut o = vec![0f32; rows];
13224 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
13225 (ok, o)
13226 };
13227 let (ok_a, fa) = fused(&g1, &u1);
13228 let (ok_b, fb) = fused(&gg, &ug);
13229 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
13230 assert_eq!(fa, fb, "fused gate/up");
13231 let b = 6usize;
13232 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
13233 let mm = |a: &QTensor| {
13234 let mut o = vec![0f32; b * rows];
13235 a.matmat(&xs, b, &mut o, Some(&pool));
13236 o
13237 };
13238 assert_eq!(mm(&g1), mm(&gg), "matmat");
13239 let x2 = q1g_test_x(cols, 77);
13240 let pair = |a: &QTensor| {
13241 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13242 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13243 (o1, o2)
13244 };
13245 assert_eq!(pair(&g1), pair(&gg), "matvec2");
13246 for r in [0, 1, rows / 2, rows - 1] {
13247 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
13248 g1.row_f32(r, &mut ra);
13249 gg.row_f32(r, &mut rb);
13250 assert_eq!(ra, rb, "embedding row {r}");
13251 }
13252 drop((g1, u1, gg, ug, model));
13253 let _ = std::fs::remove_dir_all(&dir);
13254 }
13255
13256 #[test]
13263 fn blr_mapped_tensor_paths_equal_dense() {
13264 use crate::pool::Pool;
13265 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13266 let (rows, cols, r) = (96usize, 128usize, 16usize);
13267 let mut rng = crate::blr::tests::Rng(41);
13268 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13269 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13270 let mut q8: Vec<u8> = (0..rows * cols)
13272 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13273 .collect();
13274 for o in 0..rows {
13275 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13276 }
13277 let header: CmfHeader = serde_json::from_value(serde_json::json!({
13278 "format": "cmf", "version": cortiq_core::CMF_VERSION,
13279 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13280 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13281 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13282 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13283 "max_position_embeddings": 8},
13284 "quant_type": "Q8_2F"
13285 }))
13286 .unwrap();
13287 let specs = [
13288 TensorSpec {
13289 name: "g".into(),
13290 dtype: TensorDtype::BinaryLowRank,
13291 shape: vec![rows, cols],
13292 data: gate.clone(),
13293 },
13294 TensorSpec {
13295 name: "u".into(),
13296 dtype: TensorDtype::BinaryLowRank,
13297 shape: vec![rows, cols],
13298 data: up.clone(),
13299 },
13300 TensorSpec {
13301 name: "q".into(),
13302 dtype: TensorDtype::Q8Row,
13303 shape: vec![rows, cols],
13304 data: q8,
13305 },
13306 ];
13307 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13308 std::fs::create_dir_all(&dir).unwrap();
13309 let path = dir.join("m.cmf");
13310 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13311 let model = Arc::new(CmfModel::open(&path).unwrap());
13312 assert!(model.verify().is_empty(), "{:?}", model.verify());
13313 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13314 let (g, u, q) = (t("g"), t("u"), t("q"));
13315 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13316 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13317 assert!(g.as_f32().is_none() && g.model_idx().is_none());
13318 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13319 assert!(!g.sparse_col_ok());
13320 assert_eq!((g.rows(), g.cols()), (rows, cols));
13321 let deq = |bytes: &[u8]| {
13322 let mut w = vec![0f32; rows * cols];
13323 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13324 w
13325 };
13326 let (wg, wu) = (deq(&gate), deq(&up));
13327 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13328 (0..rows)
13329 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13330 .collect()
13331 };
13332 let close = |a: &[f32], b: &[f32], what: &str| {
13333 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13334 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13335 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13336 };
13337 let pool = Pool::new(3);
13338 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13339 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13340 let mut o = vec![0f32; rows];
13341 g.matvec(&x, &mut o, Some(&pool));
13342 close(&o, &mv(&wg, &x), "matvec");
13343 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13344 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13345 close(&o1, &mv(&wg, &x), "matvec2 a");
13346 close(&o2, &mv(&wg, &x2), "matvec2 b");
13347 for b in [2usize, 5] {
13348 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13349 let mut om = vec![0f32; b * rows];
13350 g.matmat(&xs, b, &mut om, Some(&pool));
13351 for t in 0..b {
13352 close(
13353 &om[t * rows..(t + 1) * rows],
13354 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13355 "matmat",
13356 );
13357 }
13358 }
13359 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13362 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13363 close(&a, &mv(&wg, &x), "matvec_many blr gate");
13364 close(&c, &mv(&wu, &x), "matvec_many blr up");
13365 let mut qo = vec![0f32; rows];
13366 q.matvec(&x, &mut qo, Some(&pool));
13367 assert_eq!(bq, qo, "matvec_many q8_row sibling");
13368 let mut f = vec![0f32; rows];
13369 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13370 let mut row = vec![0f32; cols];
13371 g.row_f32(rows / 3, &mut row);
13372 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13373 let mut scratch = vec![0f32; cols];
13374 let d = g.row_dot(7, &x, &mut scratch);
13375 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13376 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13377 let mut col = vec![0f32; rows];
13378 g.add_col_scaled(cols / 2, 1.5, &mut col);
13379 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13380 close(&col, &want, "add_col_scaled");
13381 drop((g, u, q, model));
13382 let _ = std::fs::remove_dir_all(&dir);
13383 }
13384
13385 #[test]
13386 fn repack_is_bit_identical() {
13387 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
13393 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13394 .collect();
13395 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13396 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13397 let rep = q8_repack_layout(&w, rows, cols);
13398 for g in 0..rows / 4 {
13400 for c in 0..cols / 16 {
13401 for lane in 0..4 {
13402 assert_eq!(
13403 &rep[g * 4 * cols + c * 64 + lane * 16
13404 ..g * 4 * cols + c * 64 + lane * 16 + 16],
13405 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13406 );
13407 }
13408 }
13409 }
13410 let mut a = vec![0.0f32; rows];
13411 qmatvec(
13412 &w,
13413 &[],
13414 &scales,
13415 &x,
13416 &[],
13417 TensorDtype::Q8Row,
13418 rows,
13419 cols,
13420 &mut a,
13421 None,
13422 );
13423 let mut b = vec![0.0f32; rows];
13424 qmatvec(
13425 &w,
13426 &rep,
13427 &scales,
13428 &x,
13429 &[],
13430 TensorDtype::Q8Row,
13431 rows,
13432 cols,
13433 &mut b,
13434 None,
13435 );
13436 assert_eq!(a, b, "full-range repack output diverged");
13437
13438 #[cfg(target_arch = "aarch64")]
13439 if sdot_enabled() {
13440 let act = split_act(&x);
13442 let mut c1 = vec![0.0f32; rows];
13443 let mut c2 = vec![0.0f32; rows];
13444 q8_range_sdot(
13445 &w,
13446 &[],
13447 &scales,
13448 &act,
13449 cols,
13450 SendMut(c1.as_mut_ptr()),
13451 3,
13452 rows - 2,
13453 );
13454 q8_range_sdot(
13455 &w,
13456 &rep,
13457 &scales,
13458 &act,
13459 cols,
13460 SendMut(c2.as_mut_ptr()),
13461 3,
13462 rows - 2,
13463 );
13464 assert_eq!(c1, c2, "unaligned-range repack output diverged");
13465 }
13466 }
13467
13468 #[test]
13469 fn sdot_a8w8_noise_is_bounded() {
13470 let (rows, cols) = (16, 512);
13474 let w: Vec<u8> = (0..rows * cols)
13475 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13476 .collect();
13477 let scales = vec![0.01f32; rows];
13478 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13479 let mut a = vec![0.0f32; rows];
13480 qmatvec(
13481 &w,
13482 &[],
13483 &scales,
13484 &x,
13485 &[],
13486 TensorDtype::Q8Row,
13487 rows,
13488 cols,
13489 &mut a,
13490 None,
13491 );
13492 let (mut num, mut den) = (0f64, 0f64);
13493 for o in 0..rows {
13494 let mut acc = 0.0f32;
13495 for j in 0..cols {
13496 acc += (w[o * cols + j] as i8) as f32 * x[j];
13497 }
13498 let expect = acc * scales[o];
13499 num += ((a[o] - expect) as f64).powi(2);
13500 den += (expect as f64).powi(2);
13501 }
13502 let rel = (num / den.max(1e-12)).sqrt();
13503 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13504 }
13505
13506 #[test]
13507 fn i8_dot_neon_matches_scalar() {
13508 let n = 100;
13509 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13510 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13511 let mut scalar = 0.0f32;
13512 for j in 0..n {
13513 scalar += (w[j] as i8) as f32 * x[j];
13514 }
13515 let fast = dot_i8_f32(&w, &x);
13516 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13517 }
13518
13519 #[test]
13521 fn vbitmatvec_matches_full_dequant() {
13522 let (rows, cols) = (6, 64);
13523 let ng = cols / GROUP_SIZE;
13524 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13526 let mut bytes = bits.clone();
13527 for g in 0..rows * ng {
13528 let s = 0.02 + 0.001 * g as f32;
13529 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13530 }
13531 for r in 0..rows {
13532 let b = bits[r] as usize;
13533 let (mut acc, mut nb) = (0u64, 0usize);
13534 let mut rowbytes = Vec::new();
13535 for i in 0..cols {
13536 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13537 acc = (acc << b) | v;
13538 nb += b;
13539 while nb >= 8 {
13540 nb -= 8;
13541 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13542 }
13543 }
13544 if nb > 0 {
13545 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13546 }
13547 bytes.extend_from_slice(&rowbytes);
13548 }
13549 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13550
13551 let mut reference = vec![0f32; rows * cols];
13552 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13553 let mut expect = vec![0f32; rows];
13554 for r in 0..rows {
13555 expect[r] = reference[r * cols..(r + 1) * cols]
13556 .iter()
13557 .zip(&x)
13558 .map(|(w, xv)| w * xv)
13559 .sum();
13560 }
13561 let mut got = vec![0f32; rows];
13562 let offsets = vbit_row_offsets(&bytes, rows, cols);
13563 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13564 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13568 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13569 for r in 0..rows {
13570 assert!(
13571 (got[r] - expect[r]).abs() < tol * scale,
13572 "row {r}: {} vs {}",
13573 got[r],
13574 expect[r]
13575 );
13576 }
13577 }
13578
13579 #[test]
13584 #[cfg(target_arch = "x86_64")]
13585 fn vbit_matmat_blocked_matches_per_row() {
13586 let (rows, cols, b) = (64usize, 128usize, 9usize);
13587 let ng = cols / GROUP_SIZE;
13588 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13589 let mut bytes = bits.clone();
13590 for g in 0..rows * ng {
13591 let sc = 0.02 + 0.0005 * g as f32;
13592 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13593 }
13594 for r in 0..rows {
13595 let bw = bits[r] as usize;
13596 let (mut acc, mut nb) = (0u64, 0usize);
13597 let mut rowbytes = Vec::new();
13598 for i in 0..cols {
13599 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13600 acc = (acc << bw) | v;
13601 nb += bw;
13602 while nb >= 8 {
13603 nb -= 8;
13604 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13605 }
13606 }
13607 if nb > 0 {
13608 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13609 }
13610 bytes.extend_from_slice(&rowbytes);
13611 }
13612 let x: Vec<f32> = (0..b * cols)
13613 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13614 .collect();
13615 let offsets = vbit_row_offsets(&bytes, rows, cols);
13616 let mut y_a = vec![0f32; b * rows];
13617 let mut y_b = vec![0f32; b * rows];
13618 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13619 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13620 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13621 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13622 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13623 let max_d = y_a
13624 .iter()
13625 .zip(&y_b)
13626 .map(|(p, q)| (p - q).abs())
13627 .fold(0.0f32, f32::max);
13628 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13629 }
13630
13631 #[test]
13639 fn q4t_matmat_blocked_matches_per_row() {
13640 let (rows, cols, b) = (16usize, 64usize, 9usize);
13641 let gpr = cols / GROUP_SIZE;
13642 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13643 for r in 0..rows {
13644 for g in 0..gpr {
13645 let t = (r * gpr + g) * Q4_TILE;
13646 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13647 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13648 for k in 0..16 {
13649 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13650 }
13651 }
13652 }
13653 let x: Vec<f32> = (0..b * cols)
13654 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13655 .collect();
13656 let mut y_blk = vec![0f32; b * rows];
13657 let mut y_row = vec![0f32; b * rows];
13658 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13659 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13660 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13661 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13662 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13663 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13664 }
13665
13666 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13673 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13674 let gpr = cols / GROUP_SIZE;
13675 let stride = q4tp_code_stride(gpr);
13676 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13677 let mut b = vec![0u8; codes_off + rows * stride];
13678 for r in 0..rows {
13679 for g in 0..gpr {
13680 let t = (r * gpr + g) * Q4TP_NIB;
13681 for k in 0..16 {
13682 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13683 }
13684 }
13685 let lo = -6.0 - 0.03 * (r % 17) as f32;
13686 let step = 0.01 + 0.004 * (r % 11) as f32;
13687 let p = params_off + r * 4;
13688 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13689 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13690 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13691 for g in 0..gpr {
13692 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13693 }
13694 }
13695 b
13696 }
13697
13698 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13702 let gpr = cols / GROUP_SIZE;
13703 let v = Q4tpView::new(bytes, rows, cols);
13704 let mut out = vec![0u8; rows * gpr * Q4_TILE];
13705 let mut sc = vec![0f32; gpr];
13706 for r in 0..rows {
13707 v.scales_into(r, gpr, &mut sc);
13708 for g in 0..gpr {
13709 let t = (r * gpr + g) * Q4_TILE;
13710 let s = sc[g];
13711 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13712 let src = (r * gpr + g) * Q4TP_NIB;
13713 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13714 }
13715 }
13716 out
13717 }
13718
13719 #[test]
13725 fn q4tp_exact_path_matches_dequant_reference() {
13726 let (rows, cols) = (256usize, 512usize);
13727 let gpr = cols / GROUP_SIZE;
13728 let bytes = synth_q4tp(rows, cols);
13729 let mut w = vec![0f32; rows * cols];
13730 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13731
13732 let x: Vec<f32> = (0..cols)
13733 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13734 .collect();
13735 let v = Q4tpView::new(&bytes, rows, cols);
13736 let mut sc = vec![0f32; gpr];
13737 for r in 0..rows {
13738 v.scales_into(r, gpr, &mut sc);
13739 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13740 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13741 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13745 assert!(
13746 (got - want).abs() <= 1e-5 * mag,
13747 "row {r}: kernel {got} vs dequant {want}"
13748 );
13749 }
13750 }
13751
13752 #[test]
13758 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13759 let (rows, cols) = (256usize, 512usize);
13760 let bytes = synth_q4tp(rows, cols);
13761 let twin = q4tp_as_q4t(&bytes, rows, cols);
13762 let x: Vec<f32> = (0..cols)
13763 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13764 .collect();
13765
13766 let mut got = vec![0f32; rows];
13767 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13768 let mut want = vec![0f32; rows];
13769 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13770
13771 let mut w = vec![0f32; rows * cols];
13774 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13775 for r in 0..rows {
13776 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13777 assert!(
13778 (got[r] - want[r]).abs() <= 1e-3 * mag,
13779 "row {r}: q4tp {} vs q4t {}",
13780 got[r],
13781 want[r]
13782 );
13783 }
13784 }
13785
13786 #[test]
13791 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13792 let (rows, cols, b) = (256usize, 512usize, 5usize);
13793 let bytes = synth_q4tp(rows, cols);
13794 let twin = q4tp_as_q4t(&bytes, rows, cols);
13795 let xs: Vec<f32> = (0..b * cols)
13796 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13797 .collect();
13798
13799 let mut got = vec![0f32; b * rows];
13800 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13801 let mut want = vec![0f32; b * rows];
13802 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13803
13804 let mut w = vec![0f32; rows * cols];
13805 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13806 for t in 0..b {
13807 for r in 0..rows {
13808 let mag: f32 = (0..cols)
13809 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13810 .sum();
13811 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13812 assert!(
13813 (g - wa).abs() <= 1e-3 * mag,
13814 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13815 );
13816 }
13817 }
13818 }
13819
13820 #[test]
13821 fn q4tp_matvec2_matches_the_single_stream_kernel() {
13822 let (rows, cols) = (128usize, 256usize);
13823 let gpr = cols / GROUP_SIZE;
13824 let bytes = synth_q4tp(rows, cols);
13825 let xs: Vec<f32> = (0..2 * cols)
13826 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13827 .collect();
13828
13829 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13830 q4tp_matvec2(
13831 &bytes,
13832 &xs[..cols],
13833 &xs[cols..],
13834 rows,
13835 cols,
13836 &mut o1,
13837 &mut o2,
13838 None,
13839 );
13840
13841 let v = Q4tpView::new(&bytes, rows, cols);
13844 let mut sc = vec![0f32; gpr];
13845 for r in 0..rows {
13846 v.scales_into(r, gpr, &mut sc);
13847 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13848 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13849 }
13850 }
13851
13852 #[test]
13859 fn q4tp_matvec_keeps_pace_with_q4t() {
13860 let (rows, cols) = (4096usize, 3072usize);
13861 let bytes = synth_q4tp(rows, cols);
13862 let twin = q4tp_as_q4t(&bytes, rows, cols);
13863 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13864 let mut o = vec![0f32; rows];
13865 let n = 12;
13866 let mut best = (f64::MAX, f64::MAX);
13867 for _ in 0..3 {
13870 let t0 = std::time::Instant::now();
13871 for _ in 0..n {
13872 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13873 }
13874 best.0 = best.0.min(t0.elapsed().as_secs_f64());
13875 let t0 = std::time::Instant::now();
13876 for _ in 0..n {
13877 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13878 }
13879 best.1 = best.1.min(t0.elapsed().as_secs_f64());
13880 }
13881 let ratio = best.1 / best.0;
13882 println!(
13883 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13884 best.0 * 1e3 / n as f64,
13885 best.1 * 1e3 / n as f64
13886 );
13887 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13888 }
13889
13890 #[cfg(target_os = "macos")]
13891 #[test]
13892 fn q4t_matmat_accel_matches_dequant_reference() {
13893 if !accel_gemm_enabled() {
13894 return; }
13896 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
13898 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13899 for r in 0..rows {
13900 for g in 0..gpr {
13901 let t = (r * gpr + g) * Q4_TILE;
13902 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13903 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13904 for k in 0..16 {
13905 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13906 }
13907 }
13908 }
13909 let x: Vec<f32> = (0..b * cols)
13910 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13911 .collect();
13912 let mut got = vec![0f32; b * rows];
13913 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13914 let mut w = vec![0f32; rows * cols];
13916 for r in 0..rows {
13917 for g in 0..gpr {
13918 let t = (r * gpr + g) * Q4_TILE;
13919 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13920 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13921 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13922 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13923 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13924 }
13925 }
13926 }
13927 for bi in 0..b {
13928 for r in 0..rows {
13929 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13930 let d = (got[bi * rows + r] - want).abs();
13931 assert!(
13932 d <= want.abs().max(1.0) * 1e-4,
13933 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13934 got[bi * rows + r]
13935 );
13936 }
13937 }
13938 }
13939
13940 #[test]
13941 fn q4matvec_matches_full_dequant() {
13942 let (rows, cols) = (8, 64);
13943 let groups = rows * cols / GROUP_SIZE;
13944 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13946 for i in 0..groups * 16 {
13947 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13948 }
13949 for g in 0..groups {
13950 let s = 0.01 + 0.003 * g as f32;
13951 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13952 }
13953 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13954
13955 let mut reference = vec![0.0f32; rows * cols];
13956 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13957 let mut expect = vec![0.0f32; rows];
13958 for r in 0..rows {
13959 expect[r] = reference[r * cols..(r + 1) * cols]
13960 .iter()
13961 .zip(&x)
13962 .map(|(w, xv)| w * xv)
13963 .sum();
13964 }
13965
13966 let mut got = vec![0.0f32; rows];
13967 q4matvec(&bytes, &x, rows, cols, &mut got, None);
13968 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13972 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13973 for r in 0..rows {
13974 assert!(
13975 (got[r] - expect[r]).abs() < tol * scale,
13976 "row {r}: {} vs {}",
13977 got[r],
13978 expect[r]
13979 );
13980 }
13981 }
13982
13983 #[test]
13986 fn vbitmatvec2_equals_two_singles() {
13987 let (rows, cols) = (6, 64);
13988 let ng = cols / GROUP_SIZE;
13989 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13990 let mut bytes = bits.clone();
13991 for g in 0..rows * ng {
13992 let s = 0.02 + 0.001 * g as f32;
13993 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13994 }
13995 for r in 0..rows {
13996 let b = bits[r] as usize;
13997 let (mut acc, mut nb) = (0u64, 0usize);
13998 let mut rowbytes = Vec::new();
13999 for i in 0..cols {
14000 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14001 acc = (acc << b) | v;
14002 nb += b;
14003 while nb >= 8 {
14004 nb -= 8;
14005 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14006 }
14007 }
14008 if nb > 0 {
14009 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14010 }
14011 bytes.extend_from_slice(&rowbytes);
14012 }
14013 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14014 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
14015 let offsets = vbit_row_offsets(&bytes, rows, cols);
14016
14017 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14018 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
14019 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
14020 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14021 vbitmatvec2(
14022 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
14023 );
14024 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
14025 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
14026 }
14027
14028 #[test]
14030 fn q4matvec2_equals_two_singles() {
14031 let (rows, cols) = (8, 128);
14032 let groups = rows * cols / GROUP_SIZE;
14033 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14034 for i in 0..groups * 16 {
14035 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14036 }
14037 for g in 0..groups {
14038 let s = 0.01 + 0.003 * g as f32;
14039 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14040 }
14041 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14044 x1[9] = 250.0;
14045 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14046
14047 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14048 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
14049 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
14050 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14051 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
14052 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
14053 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
14054 }
14055
14056 #[test]
14059 fn matvec_many_equals_separate_matvecs() {
14060 use crate::pool::Pool;
14061 let (r1, r2, cols) = (300, 200, 64);
14062 let mk = |salt: usize, rows: usize| {
14063 QTensor::from_f32(
14064 (0..rows * cols)
14065 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
14066 .collect(),
14067 rows,
14068 cols,
14069 )
14070 };
14071 let (a, b) = (mk(1, r1), mk(5, r2));
14072 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
14073 let pool = Pool::new(3);
14074
14075 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
14076 a.matvec(&x, &mut ea, Some(&pool));
14077 b.matvec(&x, &mut eb, Some(&pool));
14078 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
14079 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14080 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
14081 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
14082 }
14083
14084 #[test]
14089 fn q4tp_matvec_many_equals_separate_matvecs() {
14090 use crate::pool::Pool;
14091 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14092
14093 let (r1, r2, cols) = (300usize, 200usize, 64usize);
14094 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14095 "arch_name": "tiny-q4tp",
14096 "hidden_size": cols,
14097 "intermediate_size": cols * 2,
14098 "num_layers": 1,
14099 "num_attention_heads": 2,
14100 "num_kv_heads": 1,
14101 "head_dim": 32,
14102 "vocab_size": r1,
14103 "layer_types": ["FullAttention"],
14104 "rms_norm_eps": 1e-6,
14105 "max_position_embeddings": 8,
14106 "linear_conv_kernel_dim": 0,
14107 "linear_num_key_heads": 0,
14108 "linear_num_value_heads": 0
14109 }))
14110 .unwrap();
14111 let header = CmfHeader {
14112 format: "cmf".into(),
14113 version: CMF_VERSION,
14114 arch,
14115 quant_type: QuantType::Q4Block,
14116 provenance: None,
14117 tokenizer_config: None,
14118 section_hashes: None,
14119 skills: Vec::new(),
14120 shard: None,
14121 calibration: None,
14122 routing: None,
14123 genome: None,
14124 lineage: Vec::new(),
14125 router: None,
14126 segments: Vec::new(),
14127 };
14128 let specs = [
14129 TensorSpec {
14130 name: "q".into(),
14131 dtype: TensorDtype::Q4TiledP,
14132 shape: vec![r1, cols],
14133 data: synth_q4tp(r1, cols),
14134 },
14135 TensorSpec {
14136 name: "kv".into(),
14137 dtype: TensorDtype::Q4TiledP,
14138 shape: vec![r2, cols],
14139 data: synth_q4tp(r2, cols),
14140 },
14141 ];
14142 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
14143 std::fs::create_dir_all(&dir).unwrap();
14144 let path = dir.join("m.cmf");
14145 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14146 let model = Arc::new(CmfModel::open(&path).unwrap());
14147 let (a, b) = (
14148 QTensor::from_model(&model, "q").unwrap(),
14149 QTensor::from_model(&model, "kv").unwrap(),
14150 );
14151 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
14152 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
14153 let x: Vec<f32> = (0..cols)
14154 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
14155 .collect();
14156 let pool = Pool::new(3);
14157 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14158 a.matvec(&x, &mut ea, Some(&pool));
14159 b.matvec(&x, &mut eb, Some(&pool));
14160 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14161 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14162 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
14163 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
14164 let _ = std::fs::remove_dir_all(&dir);
14165 }
14166
14167 #[test]
14173 fn multi_token_moe_rows_equal_single_token_decode() {
14174 use crate::pool::Pool;
14175 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14176
14177 let (h, inter, ne) = (64usize, 128usize, 3usize);
14178 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14179 "arch_name": "tiny-q4tp-moe",
14180 "hidden_size": h,
14181 "intermediate_size": inter,
14182 "num_layers": 1,
14183 "num_attention_heads": 2,
14184 "num_kv_heads": 1,
14185 "head_dim": 32,
14186 "vocab_size": 8,
14187 "layer_types": ["FullAttention"],
14188 "rms_norm_eps": 1e-6,
14189 "max_position_embeddings": 8,
14190 "linear_conv_kernel_dim": 0,
14191 "linear_num_key_heads": 0,
14192 "linear_num_value_heads": 0
14193 }))
14194 .unwrap();
14195 let header = CmfHeader {
14196 format: "cmf".into(),
14197 version: CMF_VERSION,
14198 arch,
14199 quant_type: QuantType::Q4Block,
14200 provenance: None,
14201 tokenizer_config: None,
14202 section_hashes: None,
14203 skills: Vec::new(),
14204 shard: None,
14205 calibration: None,
14206 routing: None,
14207 genome: None,
14208 lineage: Vec::new(),
14209 router: None,
14210 segments: Vec::new(),
14211 };
14212 let mut specs = Vec::new();
14213 for e in 0..ne {
14214 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
14215 .into_iter()
14216 .enumerate()
14217 {
14218 let mut data = synth_q4tp(r, c);
14221 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
14222 .iter_mut()
14223 .enumerate()
14224 {
14225 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
14226 }
14227 specs.push(TensorSpec {
14228 name: format!("{n}{e}"),
14229 dtype: TensorDtype::Q4TiledP,
14230 shape: vec![r, c],
14231 data,
14232 });
14233 }
14234 }
14235 let dir = std::env::temp_dir().join(format!(
14236 "cmf-moe-rows-{}-{}",
14237 std::process::id(),
14238 FLOAT_ACTIVATIONS.get()
14239 ));
14240 std::fs::create_dir_all(&dir).unwrap();
14241 let path = dir.join("m.cmf");
14242 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14243 let model = Arc::new(CmfModel::open(&path).unwrap());
14244 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
14245 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
14246 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
14247 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
14248 let b = 4usize;
14249 let mut xs: Vec<f32> = (0..b * h)
14250 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
14251 .collect();
14252 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
14255 (vec![2, 0], vec![0.6, 0.4]),
14256 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14257 (vec![1], vec![1.0]),
14258 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14259 ];
14260 let pool = Pool::new(3);
14261 let mut want = vec![0f32; b * h];
14263 for (tk, (idx, w)) in routes.iter().enumerate() {
14264 let x = &xs[tk * h..(tk + 1) * h];
14265 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14266 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14267 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14268 if FLOAT_ACTIVATIONS.get() {
14269 for (slot, &e) in idx.iter().enumerate() {
14270 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14271 g[e].matvec(x, &mut gate, Some(&pool));
14272 u[e].matvec(x, &mut up, Some(&pool));
14273 for (v, u) in gate.iter_mut().zip(up) {
14274 *v = (*v / (1.0 + (-*v).exp())) * u;
14275 }
14276 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14277 }
14278 }
14279 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14280 assert!(QTensor::moe_down_many(
14281 &downs,
14282 &gs,
14283 w,
14284 &mut want[tk * h..(tk + 1) * h],
14285 Some(&pool)
14286 ));
14287 }
14288 if FLOAT_ACTIVATIONS.get() {
14289 for (tk, (idx, w)) in routes.iter().enumerate() {
14290 let mut scalar = vec![0.0; h];
14291 for (&e, &weight) in idx.iter().zip(w) {
14292 let (mut gate, mut up, mut down) =
14293 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14294 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14295 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14296 for (v, u) in gate.iter_mut().zip(up) {
14297 *v = (*v / (1.0 + (-*v).exp())) * u;
14298 }
14299 d[e].matvec(&gate, &mut down, Some(&pool));
14300 for (v, d) in scalar.iter_mut().zip(down) {
14301 *v += weight * d;
14302 }
14303 }
14304 assert_eq!(
14305 &want[tk * h..(tk + 1) * h],
14306 scalar,
14307 "float many equals scalar experts"
14308 );
14309 }
14310 }
14311 let mut experts: Vec<usize> = Vec::new();
14313 let mut groups: Vec<Vec<usize>> = Vec::new();
14314 for (tk, (idx, _)) in routes.iter().enumerate() {
14315 for &e in idx {
14316 match experts.iter().position(|&x| x == e) {
14317 Some(k) => groups[k].push(tk),
14318 None => {
14319 experts.push(e);
14320 groups.push(vec![tk]);
14321 }
14322 }
14323 }
14324 }
14325 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14326 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14327 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14328 assert!(QTensor::moe_gate_up_rows(
14329 &pairs,
14330 &groups,
14331 &xs,
14332 &mut gs,
14333 Some(&pool)
14334 ));
14335 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14336 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14337 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14338 assert!(QTensor::moe_down_rows(
14339 &downs,
14340 &lens,
14341 &gs,
14342 &mut ds,
14343 Some(&pool)
14344 ));
14345 let slot = |tk: usize, e: usize| {
14346 let k = experts.iter().position(|&x| x == e).unwrap();
14347 groups[..k].iter().map(|g| g.len()).sum::<usize>()
14348 + groups[k].iter().position(|&x| x == tk).unwrap()
14349 };
14350 let mut got = vec![0f32; b * h];
14351 for (tk, (idx, w)) in routes.iter().enumerate() {
14352 for i in 0..h {
14353 let mut acc = 0f32;
14354 for (&e, &we) in idx.iter().zip(w) {
14355 acc += we * ds[slot(tk, e)][i];
14356 }
14357 got[tk * h + i] = acc;
14358 }
14359 }
14360 assert!(want.iter().any(|v| *v != 0.0));
14361 assert_eq!(
14362 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14363 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14364 "multi-token MoE must equal decode bit for bit"
14365 );
14366
14367 let b5 = 5usize;
14370 let x5: Vec<f32> = (0..b5 * h)
14371 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14372 .collect();
14373 let mut mm = vec![0f32; b5 * inter];
14374 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14375 for tk in 0..b5 {
14376 let mut mv = vec![0f32; inter];
14377 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14378 assert_eq!(
14379 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14380 mm[tk * inter..(tk + 1) * inter]
14381 .iter()
14382 .map(|v| v.to_bits())
14383 .collect::<Vec<_>>(),
14384 "row-exact matmat token {tk}"
14385 );
14386 }
14387 let _ = std::fs::remove_dir_all(&dir);
14390 }
14391
14392 #[test]
14400 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14401 use crate::pool::Pool;
14402 use std::sync::atomic::Ordering::Relaxed;
14403 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14404 Q4TP_ALT.store(2, Relaxed);
14407 let pool = Pool::new(3);
14408 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14412 let bytes = synth_q4tp(rows, cols);
14413 let mut xs: Vec<f32> = (0..b * cols)
14414 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14415 .collect();
14416 xs[3] = 7.5; let run = |exact: bool| {
14418 let mut out = vec![0f32; b * rows];
14419 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14420 out
14421 };
14422 let (fast, exact) = (run(false), run(true));
14423 #[cfg(not(target_arch = "aarch64"))]
14424 let _ = fast;
14425 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14426 let mut matvecs = vec![0f32; b * rows];
14427 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14428 q4tp_matvec(
14429 &bytes,
14430 &xs[bi * cols..(bi + 1) * cols],
14431 rows,
14432 cols,
14433 o,
14434 Some(&pool),
14435 );
14436 }
14437 assert!(matvecs.iter().any(|v| *v != 0.0));
14438 assert_eq!(
14439 bits(&exact),
14440 bits(&matvecs),
14441 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14442 );
14443 #[cfg(target_arch = "aarch64")]
14444 {
14445 let gpr = cols / GROUP_SIZE;
14447 let v = Q4tpView::new(&bytes, rows, cols);
14448 let mut old = vec![0f32; b * rows];
14449 let mut sc = vec![0f32; gpr];
14450 let a8w8 = a8w8_enabled();
14451 let blocked = sdot_enabled() && blocked_enabled();
14452 let acts: Vec<SplitAct> = (0..b)
14453 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14454 .collect();
14455 for r in 0..rows {
14456 v.scales_into(r, gpr, &mut sc);
14457 if !a8w8 {
14458 for bi in 0..b {
14459 let x = &xs[bi * cols..(bi + 1) * cols];
14460 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14461 }
14462 continue;
14463 }
14464 let finish = |d: f32, act: &SplitAct| {
14465 let mut acc = d * act.sx;
14466 for &(j, xv) in &act.outliers {
14467 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14468 acc += w * s * xv;
14469 }
14470 acc
14471 };
14472 let mut bi = 0usize;
14473 while blocked && bi + 4 <= b {
14474 let xs4 = [
14475 acts[bi].xq.as_slice(),
14476 acts[bi + 1].xq.as_slice(),
14477 acts[bi + 2].xq.as_slice(),
14478 acts[bi + 3].xq.as_slice(),
14479 ];
14480 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14481 for k in 0..4 {
14482 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14483 }
14484 bi += 4;
14485 }
14486 for (bi, act) in acts.iter().enumerate().skip(bi) {
14487 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14488 old[bi * rows + r] = finish(d, act);
14489 }
14490 }
14491 assert_eq!(
14492 bits(&fast),
14493 bits(&old),
14494 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14495 );
14496 if blocked {
14499 assert_ne!(
14500 bits(&fast),
14501 bits(&matvecs),
14502 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14503 );
14504 }
14505 }
14506 }
14507 Q4TP_ALT.store(0, Relaxed);
14508 }
14509
14510 #[test]
14511 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14512 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14513 let active = AtomicUsize::new(0);
14516 counted_row_exact_scope(&active, || {
14517 assert_eq!(active.load(Ordering::Acquire), 1);
14518 counted_row_exact_scope(&active, || {
14519 assert_eq!(active.load(Ordering::Acquire), 2);
14520 });
14521 assert_eq!(active.load(Ordering::Acquire), 1);
14522 });
14523 assert_eq!(active.load(Ordering::Acquire), 0);
14524
14525 let both_entered = Barrier::new(2);
14526 let release_last = Barrier::new(2);
14527 std::thread::scope(|s| {
14528 let first = s.spawn(|| counted_row_exact_scope(&active, || {
14529 both_entered.wait();
14530 }));
14531 let last = s.spawn(|| counted_row_exact_scope(&active, || {
14532 both_entered.wait();
14533 release_last.wait();
14534 }));
14535 first.join().unwrap();
14536 let after_first = active.load(Ordering::Acquire);
14537 release_last.wait();
14538 last.join().unwrap();
14539 assert_eq!(after_first, 1, "second request must remain exact");
14540 });
14541 assert_eq!(active.load(Ordering::Acquire), 0);
14542 let panic = std::panic::catch_unwind(|| {
14543 counted_row_exact_scope(&active, || panic!("scope unwind"));
14544 });
14545 assert!(panic.is_err());
14546 assert_eq!(active.load(Ordering::Acquire), 0);
14547 }
14548
14549 #[test]
14550 #[cfg(target_arch = "x86_64")]
14551 fn q4tp_float_avx2_is_bitwise_scalar() {
14552 if !avx2_enabled() {
14553 return;
14554 }
14555 for cols in [32, 64, 96, 2048, 4096] {
14556 let rows = 9;
14557 let bytes = synth_q4tp(rows, cols);
14558 let v = Q4tpView::new(&bytes, rows, cols);
14559 let gpr = cols / GROUP_SIZE;
14560 let mut sc = vec![0.0; gpr];
14561 for seed in 1..=5 {
14562 let xs: Vec<f32> = (0..cols)
14563 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14564 .collect();
14565 for r in 0..rows {
14566 v.scales_into(r, gpr, &mut sc);
14567 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14568 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14569 assert_eq!(
14570 scalar.to_bits(),
14571 vector.to_bits(),
14572 "cols={cols} row={r} seed={seed}"
14573 );
14574 }
14575 }
14576 }
14577 }
14578
14579 #[test]
14580 fn multi_token_moe_rows_float_equal_single_token_decode() {
14581 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14582 }
14583
14584 #[test]
14585 fn full_gpu_q8_scope_is_nested_and_thread_local() {
14586 assert!(!FULL_GPU_Q8.get());
14587 let before = gpu_split_frac();
14588 {
14589 let _guard = enter_full_gpu_q8_scope();
14590 assert_eq!(gpu_split_frac(), 1.0);
14591 {
14592 let _nested = enter_full_gpu_q8_scope();
14593 }
14594 assert_eq!(gpu_split_frac(), 1.0);
14595 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14596 }
14597 assert!(!FULL_GPU_Q8.get());
14598 assert_eq!(gpu_split_frac(), before);
14599 }
14600
14601 #[test]
14602 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14603 assert!(!FLOAT_ACTIVATIONS.get());
14604 let before = a8w8_enabled();
14605 float_activations_scope(|| {
14606 assert!(!a8w8_enabled());
14607 float_activations_scope(|| assert!(!a8w8_enabled()));
14608 assert!(FLOAT_ACTIVATIONS.get());
14609 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14610 .join()
14611 .unwrap();
14612 });
14613 assert!(!FLOAT_ACTIVATIONS.get());
14614 assert_eq!(a8w8_enabled(), before);
14615 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14616 assert!(!FLOAT_ACTIVATIONS.get());
14617 }
14618
14619 #[test]
14622 fn batched_matmat_equals_per_position_matvec() {
14623 let (rows, cols, b) = (8, 64, 5);
14624 let groups = rows * cols / GROUP_SIZE;
14626 let mut q4 = Vec::new();
14627 for i in 0..groups * 16 {
14628 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14629 }
14630 for g in 0..groups {
14631 q4.extend_from_slice(
14632 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14633 );
14634 }
14635 let ng = cols / GROUP_SIZE;
14637 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14638 let mut vb = bits.clone();
14639 for g in 0..rows * ng {
14640 vb.extend_from_slice(
14641 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14642 );
14643 }
14644 for r in 0..rows {
14645 let bw = bits[r] as usize;
14646 let (mut acc, mut nb) = (0u64, 0usize);
14647 let mut rowbytes = Vec::new();
14648 for i in 0..cols {
14649 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14650 acc = (acc << bw) | v;
14651 nb += bw;
14652 while nb >= 8 {
14653 nb -= 8;
14654 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14655 }
14656 }
14657 if nb > 0 {
14658 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14659 }
14660 vb.extend_from_slice(&rowbytes);
14661 }
14662 let offsets = vbit_row_offsets(&vb, rows, cols);
14663
14664 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14665
14666 let mut got = vec![0f32; b * rows];
14668 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14669 for bi in 0..b {
14670 let mut expect = vec![0f32; rows];
14671 q4matvec(
14672 &q4,
14673 &xs[bi * cols..(bi + 1) * cols],
14674 rows,
14675 cols,
14676 &mut expect,
14677 None,
14678 );
14679 assert_eq!(
14680 &got[bi * rows..(bi + 1) * rows],
14681 &expect[..],
14682 "q4 batch pos {bi}"
14683 );
14684 }
14685
14686 let mut got = vec![0f32; b * rows];
14688 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14689 for bi in 0..b {
14690 let mut expect = vec![0f32; rows];
14691 vbitmatvec(
14692 &vb,
14693 &offsets,
14694 &xs[bi * cols..(bi + 1) * cols],
14695 rows,
14696 cols,
14697 &mut expect,
14698 None,
14699 );
14700 assert_eq!(
14701 &got[bi * rows..(bi + 1) * rows],
14702 &expect[..],
14703 "vbit batch pos {bi}"
14704 );
14705 }
14706 }
14707
14708 #[test]
14712 fn q4_tiled_matches_q4_block_bitexact() {
14713 let (rows, cols, b) = (8usize, 128usize, 3usize);
14714 let groups = rows * cols / GROUP_SIZE;
14715 let mut split = Vec::with_capacity(groups * 18);
14716 for i in 0..groups * 16 {
14717 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14718 }
14719 for g in 0..groups {
14720 split.extend_from_slice(
14721 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14722 );
14723 }
14724 let (packed, scales) = split.split_at(groups * 16);
14726 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14727 for g in 0..groups {
14728 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14729 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14730 }
14731
14732 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14733 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14735
14736 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14737 q4matvec(&split, &x1, rows, cols, &mut a, None);
14738 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14739 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14740
14741 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14742 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14743 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14744 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14745 assert_eq!(a1, t1);
14746 assert_eq!(a2, t2);
14747
14748 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14749 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14750 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14751 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14752 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14753 }
14754
14755 #[test]
14762 fn q4matvec_sdot_outlier_exact() {
14763 let (rows, cols) = (4, 128);
14764 let groups = rows * cols / GROUP_SIZE;
14765 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14766 for i in 0..groups * 16 {
14767 bytes.push(((i * 11 + 5) % 256) as u8);
14768 }
14769 for g in 0..groups {
14770 let s = 0.02 + 0.002 * g as f32;
14771 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14772 }
14773 let mut x: Vec<f32> = (0..cols)
14774 .map(|i| match i % 3 {
14775 0 => 1.0,
14776 1 => -1.0,
14777 _ => 0.0,
14778 })
14779 .collect();
14780 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
14783 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14784 let mut expect = vec![0.0f32; rows];
14785 for r in 0..rows {
14786 expect[r] = reference[r * cols..(r + 1) * cols]
14787 .iter()
14788 .zip(&x)
14789 .map(|(w, xv)| w * xv)
14790 .sum();
14791 }
14792 let mut got = vec![0.0f32; rows];
14793 q4matvec(&bytes, &x, rows, cols, &mut got, None);
14794 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14795 for r in 0..rows {
14796 assert!(
14797 (got[r] - expect[r]).abs() < 2e-3 * scale,
14798 "row {r}: {} vs {} (outlier term must be exact)",
14799 got[r],
14800 expect[r]
14801 );
14802 }
14803 }
14804
14805 #[test]
14809 fn q1t_matvec_matches_reference() {
14810 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14811 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
14813 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14814 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14816 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14817 let mut bytes = Vec::new();
14818 for r in 0..rows {
14819 for g in 0..gpr {
14820 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14821 let mut c = [0u8; 7];
14822 for k in 0..GROUP_SIZE {
14823 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14825 0
14826 } else {
14827 ((k + r * 3 + g) % 3) as u8 };
14829 cortiq_core::quant::q1t_pack(&mut c, k, code);
14830 }
14831 bytes.extend_from_slice(&c);
14832 }
14833 }
14834 let mut row_ptr = vec![0u32; rows + 1];
14837 for &(idx, _) in &outliers {
14838 row_ptr[idx as usize / cols + 1] += 1;
14839 }
14840 for r in 0..rows {
14841 row_ptr[r + 1] += row_ptr[r];
14842 }
14843 for &p in &row_ptr {
14844 bytes.extend_from_slice(&p.to_le_bytes());
14845 }
14846 for &(idx, v) in &outliers {
14847 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14848 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14849 }
14850
14851 let mut refw = vec![0f32; rows * cols];
14852 dequant_q1t(&bytes, rows, cols, &mut refw);
14853 let x: Vec<f32> = (0..cols)
14856 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14857 .collect();
14858 let mut expect = vec![0f32; rows];
14859 for r in 0..rows {
14860 let mut a = 0.0f32;
14861 for j in 0..cols {
14862 a += refw[r * cols + j] * x[j];
14863 }
14864 expect[r] = a;
14865 }
14866 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14867 let mut got = vec![0f32; rows];
14868 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14869 for r in 0..rows {
14870 assert!(
14871 (got[r] - expect[r]).abs() < tol(expect[r]),
14872 "row {r}: {} vs {}",
14873 got[r],
14874 expect[r]
14875 );
14876 }
14877 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14879 let mut gm = vec![0f32; 2 * rows];
14880 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14881 for r in 0..rows {
14882 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14883 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14884 }
14885 let xb: Vec<f32> = (0..cols)
14889 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14890 .collect();
14891 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14892 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14893 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14894 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14895 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14896 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14897 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14898 }
14899
14900 #[test]
14903 fn q1t_matvec2_odd_gpr_matches_singles() {
14904 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14905 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
14907 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14908 for r in 0..rows {
14909 for g in 0..gpr {
14910 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14911 let mut c = [0u8; 7];
14912 for k in 0..GROUP_SIZE {
14913 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14914 }
14915 bytes.extend_from_slice(&c);
14916 }
14917 }
14918 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14919 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14920 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14921 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14922 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14923 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14924 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14925 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14926 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14927 }
14928
14929 #[test]
14933 #[ignore]
14934 fn q1t_matvec2_speed() {
14935 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14936 use std::time::Instant;
14937 let (rows, cols) = (8192usize, 4096usize);
14938 let gpr = cols / GROUP_SIZE;
14939 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14940 for r in 0..rows {
14941 for g in 0..gpr {
14942 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14943 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14944 let mut c = [0u8; 7];
14945 for k in 0..GROUP_SIZE {
14946 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14947 }
14948 bytes.extend_from_slice(&c);
14949 }
14950 }
14951 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14952 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14953 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14954 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14955 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14957 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14958 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14959 for _ in 0..8 {
14960 let t0 = Instant::now();
14961 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14962 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14963 let t1 = Instant::now();
14964 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14965 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14966 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14967 }
14968 assert_eq!(p1, s1);
14969 assert_eq!(p2, s2);
14970 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14971 }
14972
14973 #[test]
14977 #[ignore]
14978 fn q1t_matvec_speed() {
14979 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14980 use std::time::Instant;
14981 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
14983 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14984 for r in 0..rows {
14985 for g in 0..gpr {
14986 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14987 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14988 let mut c = [0u8; 7];
14989 for k in 0..GROUP_SIZE {
14990 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14991 }
14992 bytes.extend_from_slice(&c);
14993 }
14994 }
14995 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
14997 let mut idx = 0usize;
14998 while idx < n {
14999 row_ptr[idx / cols + 1] += 1;
15000 idx += stride;
15001 }
15002 for r in 0..rows {
15003 row_ptr[r + 1] += row_ptr[r];
15004 }
15005 for &p in &row_ptr {
15006 bytes.extend_from_slice(&p.to_le_bytes());
15007 }
15008 let mut idx = 0usize;
15009 while idx < n {
15010 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
15011 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
15012 idx += stride;
15013 }
15014 let x: Vec<f32> = (0..cols)
15017 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15018 .collect();
15019 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
15020
15021 let slow = |out: &mut [f32]| {
15023 let mut buf = vec![0f32; cols];
15024 for r in 0..rows {
15025 for g in 0..gpr {
15026 let off = (r * gpr + g) * Q1T_TILE;
15027 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
15028 let codes = &bytes[off + 2..off + Q1T_TILE];
15029 for k in 0..GROUP_SIZE {
15030 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
15031 1 => s,
15032 2 => -s,
15033 _ => 0.0,
15034 };
15035 }
15036 }
15037 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
15038 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
15039 }
15040 };
15041 let iters = 5;
15042 let mut a = vec![0f32; rows];
15043 slow(&mut a); let t = Instant::now();
15045 for _ in 0..iters {
15046 slow(&mut a);
15047 }
15048 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15049
15050 let mut b = vec![0f32; rows];
15051 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
15053 for _ in 0..iters {
15054 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
15055 }
15056 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15057
15058 for r in 0..rows {
15059 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
15060 }
15061 println!(
15062 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
15063 slow_ms / fast_ms
15064 );
15065 }
15066}
15067
15068#[cfg(test)]
15069mod gemm_bench {
15070 #[test]
15081 #[ignore]
15082 fn q4tp_matmat_throughput() {
15083 let _alt = super::Q4TP_ALT_TEST_LOCK
15084 .lock()
15085 .unwrap_or_else(|e| e.into_inner());
15086 let b: usize = std::env::var("CMF_BENCH_B")
15090 .ok()
15091 .and_then(|v| v.parse().ok())
15092 .unwrap_or(296);
15093 let (rows, cols) = (9216usize, 2304usize);
15094 let (_, _, _) = (rows, cols, b);
15095 let total =
15096 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
15097 .unwrap();
15098 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15103 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15104 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15105 let step = cortiq_core::quant::f32_to_f16(0.1);
15106 for r in 0..rows {
15107 let o = params_off + r * 4;
15108 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15109 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15110 }
15111 let _ = codes_off;
15112 let xs: Vec<f32> = (0..b * cols)
15113 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15114 .collect();
15115 let mut out = vec![0f32; b * rows];
15116 let pool = crate::pool::Pool::from_env();
15117 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15123 let reps: usize = std::env::var("CMF_BENCH_REPS")
15124 .ok()
15125 .and_then(|v| v.parse().ok())
15126 .unwrap_or(10);
15127 let mut best = [f64::MAX; 2];
15128 let mut sums = [0f32; 2];
15129 for _ in 0..reps {
15130 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
15131 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
15132 let t = std::time::Instant::now();
15133 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15134 best[k] = best[k].min(t.elapsed().as_secs_f64());
15135 sums[k] = out.iter().take(64).sum::<f32>();
15136 }
15137 }
15138 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15139 for (k, name) in ["previous", "tuned "].iter().enumerate() {
15140 println!(
15141 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15142 best[k] * 1e3,
15143 flops / best[k] / 1e9,
15144 sums[k]
15145 );
15146 }
15147 assert!(
15148 (sums[0] - sums[1]).abs() < 1e-2,
15149 "the tuned kernel changed the result: {} vs {}",
15150 sums[0],
15151 sums[1]
15152 );
15153 }
15154
15155 #[test]
15161 fn q4tp_matmat_blocked_matches_scalar() {
15162 use std::sync::atomic::Ordering::Relaxed;
15163 let _alt = super::Q4TP_ALT_TEST_LOCK
15164 .lock()
15165 .unwrap_or_else(|e| e.into_inner());
15166 for &(rows, cols, b) in &[
15174 (64usize, 128usize, 7usize),
15175 (33, 96, 4),
15176 (16, 256, 9),
15177 (192, 2304, 37),
15178 ] {
15179 let total = cortiq_core::quant::expected_nbytes(
15180 cortiq_core::TensorDtype::Q4TiledP,
15181 &[rows, cols],
15182 )
15183 .unwrap();
15184 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15185 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
15186 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15187 let step = cortiq_core::quant::f32_to_f16(0.1);
15188 for r in 0..rows {
15189 let o = params_off + r * 4;
15190 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15191 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15192 }
15193 let xs: Vec<f32> = (0..b * cols)
15194 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
15195 .collect();
15196 let mut got = vec![0f32; b * rows];
15197 let mut want = vec![0f32; b * rows];
15198 let gpr = cols / 32;
15199 let view = super::Q4tpView::new(&bytes, rows, cols);
15200 let pool = crate::pool::Pool::from_env();
15201 super::Q4TP_ALT.store(2, Relaxed);
15202 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
15203 super::Q4TP_ALT.store(1, Relaxed);
15204 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
15205 super::Q4TP_ALT.store(0, Relaxed);
15206 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15213 let (mut worst, mut at) = (0f32, 0usize);
15214 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
15215 if (g - w).abs() > worst {
15216 worst = (g - w).abs();
15217 at = i;
15218 }
15219 }
15220 assert!(
15221 worst <= 1e-4 * scale,
15222 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
15223 (scale {scale:.3e}) at cell {at}: {} vs {}",
15224 got[at],
15225 want[at]
15226 );
15227
15228 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
15235 for bi in 0..b {
15236 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
15237 for r in 0..rows {
15238 let mut sc = vec![0f32; gpr];
15239 view.scales_into(r, gpr, &mut sc);
15240 let mut exact = 0f64;
15241 for j in 0..cols {
15242 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15243 exact += w as f64 * sq as f64 * act.xq[j] as f64;
15244 }
15245 exact *= act.sx as f64;
15246 for &(j, xv) in &act.outliers {
15247 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15248 exact += w as f64 * sq as f64 * xv as f64;
15249 }
15250 let i = bi * rows + r;
15251 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
15252 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
15253 }
15254 }
15255 println!(
15256 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15257 per-column {e_scalar:.3e}"
15258 );
15259 assert!(
15264 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15265 "{rows}x{cols} b={b}: error against f64 too large — blocked \
15266 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15267 );
15268 }
15269 }
15270
15271 #[test]
15277 #[ignore]
15278 fn q4tp_matmat_row_exact_cost() {
15279 let _alt = super::Q4TP_ALT_TEST_LOCK
15280 .lock()
15281 .unwrap_or_else(|e| e.into_inner());
15282 let pool = crate::pool::Pool::from_env();
15283 let reps: usize = std::env::var("CMF_BENCH_REPS")
15284 .ok()
15285 .and_then(|v| v.parse().ok())
15286 .unwrap_or(30);
15287 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15288 let total = cortiq_core::quant::expected_nbytes(
15289 cortiq_core::TensorDtype::Q4TiledP,
15290 &[rows, cols],
15291 )
15292 .unwrap();
15293 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15294 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15295 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15296 let step = cortiq_core::quant::f32_to_f16(0.1);
15297 for r in 0..rows {
15298 let o = params_off + r * 4;
15299 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15300 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15301 }
15302 for &b in &[2usize, 4, 5, 8] {
15303 let xs: Vec<f32> = (0..b * cols)
15304 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15305 .collect();
15306 let mut out = vec![0f32; b * rows];
15307 let mut best = [f64::MAX; 3];
15308 for _ in 0..reps {
15309 for (k, best_k) in best.iter_mut().enumerate() {
15310 let t = std::time::Instant::now();
15311 match k {
15312 0 | 1 => super::q4tp_matmat_with(
15313 &bytes,
15314 &xs,
15315 b,
15316 rows,
15317 cols,
15318 &mut out,
15319 pool.as_deref(),
15320 k == 1,
15321 ),
15322 _ => {
15323 for (bi, o) in out.chunks_mut(rows).enumerate() {
15324 super::q4tp_matvec(
15325 &bytes,
15326 &xs[bi * cols..(bi + 1) * cols],
15327 rows,
15328 cols,
15329 o,
15330 pool.as_deref(),
15331 );
15332 }
15333 }
15334 }
15335 *best_k = best_k.min(t.elapsed().as_secs_f64());
15336 }
15337 }
15338 println!(
15339 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15340 {b} matvecs {:.3} ms",
15341 best[0] * 1e3,
15342 best[1] * 1e3,
15343 best[2] * 1e3
15344 );
15345 }
15346 }
15347 }
15348
15349 #[test]
15353 #[ignore]
15354 fn q4t_matmat_throughput() {
15355 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15356 let total =
15357 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15358 .unwrap();
15359 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15362 let sc = cortiq_core::quant::f32_to_f16(0.02);
15363 for t in bytes.chunks_mut(super::Q4_TILE) {
15364 t[..2].copy_from_slice(&sc.to_le_bytes());
15365 }
15366 let xs: Vec<f32> = (0..b * cols)
15367 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15368 .collect();
15369 let mut out = vec![0f32; b * rows];
15370 let pool = crate::pool::Pool::from_env();
15371 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15372 let reps: usize = std::env::var("CMF_BENCH_REPS")
15373 .ok()
15374 .and_then(|v| v.parse().ok())
15375 .unwrap_or(10);
15376 let mut best = f64::MAX;
15377 for _ in 0..reps {
15378 let t = std::time::Instant::now();
15379 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15380 best = best.min(t.elapsed().as_secs_f64());
15381 }
15382 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15383 println!(
15384 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15385 best * 1e3,
15386 flops / best / 1e9,
15387 out.iter().take(64).sum::<f32>()
15388 );
15389 }
15390}