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 match self {
1357 Self::F32 { data, .. } => {
1361 if !crate::f32_backend::matvec(data, x, out) {
1362 matvec_rows(pool, data, x, out);
1363 }
1364 }
1365 Self::Blr(t) => {
1366 if crate::gpu::enabled_here()
1367 && crate::blr::gpu_wanted()
1368 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1369 {
1370 return;
1371 }
1372 t.matvec(x, out, pool)
1373 }
1374 Self::Mapped {
1375 model,
1376 idx,
1377 dtype,
1378 rows,
1379 cols,
1380 row_scale,
1381 col_field,
1382 vbit_offsets,
1383 repack,
1384 } => {
1385 let _ = (model, idx);
1386 assert!(
1396 out.len() >= *rows && x.len() >= *cols,
1397 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1398 out.len(),
1399 x.len(),
1400 );
1401 let prism_forward =
1402 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1403 if *dtype == TensorDtype::Q2TiledP
1404 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1405 {
1406 use std::sync::atomic::{AtomicUsize, Ordering};
1407 static N: AtomicUsize = AtomicUsize::new(0);
1408 let n = N.fetch_add(1, Ordering::Relaxed);
1409 if n < 128 {
1410 eprintln!(
1411 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1412 model.tensors[*idx].name,
1413 prism_forward,
1414 rows,
1415 cols,
1416 crate::gpu::enabled_here(),
1417 crate::gpu::q2tp_gpu_opt_in(),
1418 crate::gpu::cur_layer(),
1419 );
1420 }
1421 }
1422 if prism_forward {
1427 let transformed = crate::prism::forward(model, &x[..*cols]);
1428 match dtype {
1429 TensorDtype::Q4Block => {
1430 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1431 }
1432 TensorDtype::Q4Tiled => {
1433 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1434 }
1435 TensorDtype::Q4TiledP => {
1436 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1437 }
1438 TensorDtype::Q2TiledP => {
1439 let affine =
1440 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1441 if *rows * *cols >= 8_388_608
1442 && crate::gpu::enabled_here()
1443 && crate::gpu::q2tp_gpu_opt_in()
1444 {
1445 let gpu_ok = if affine {
1446 crate::gpu::q2tp_affine_matvec(
1447 model,
1448 *idx,
1449 &transformed,
1450 *rows,
1451 *cols,
1452 out,
1453 )
1454 } else {
1455 crate::gpu::q2tp_matvec(
1456 model,
1457 *idx,
1458 &transformed,
1459 *rows,
1460 *cols,
1461 out,
1462 )
1463 };
1464 if gpu_ok {
1465 return;
1466 }
1467 }
1468 if affine {
1469 q2tp_affine_matvec(
1470 self.quant_bytes(),
1471 &transformed,
1472 *rows,
1473 *cols,
1474 out,
1475 pool,
1476 )
1477 } else {
1478 q2tp_matvec(
1479 self.quant_bytes(),
1480 &transformed,
1481 *rows,
1482 *cols,
1483 out,
1484 pool,
1485 )
1486 }
1487 }
1488 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1489 self.bin_w().expect("binary dtype"),
1490 &transformed,
1491 *rows,
1492 *cols,
1493 out,
1494 pool,
1495 ),
1496 TensorDtype::Q1T => {
1497 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1498 }
1499 TensorDtype::Mxfp4 => {
1500 mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1501 }
1502 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1503 self.quant_bytes(),
1504 vbit_offsets,
1505 &transformed,
1506 *rows,
1507 *cols,
1508 out,
1509 pool,
1510 ),
1511 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1512 self.quant_bytes(),
1513 repack,
1514 row_scale,
1515 &transformed,
1516 col_field,
1517 *dtype,
1518 *rows,
1519 *cols,
1520 out,
1521 pool,
1522 ),
1523 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1524 }
1525 return;
1526 }
1527 if *dtype == TensorDtype::Q4Block {
1528 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1532 let t0 = std::time::Instant::now();
1533 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1534 crate::gpu::ProbeArm::Gpu => {
1535 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1536 crate::gpu::probe_record(
1537 crate::gpu::OpClass::Matvec,
1538 true,
1539 t0.elapsed(),
1540 );
1541 return;
1542 }
1543 }
1544 crate::gpu::ProbeArm::CpuTimed => {
1545 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1546 crate::gpu::probe_record(
1547 crate::gpu::OpClass::Matvec,
1548 false,
1549 t0.elapsed(),
1550 );
1551 return;
1552 }
1553 crate::gpu::ProbeArm::Cpu => {}
1554 }
1555 }
1556 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1557 return;
1558 }
1559 if *dtype == TensorDtype::Q4Tiled {
1560 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1565 let t0 = std::time::Instant::now();
1566 let cls = crate::gpu::matvec_class(*rows, *cols);
1567 match crate::gpu::probe_arm(cls) {
1568 crate::gpu::ProbeArm::Gpu => {
1569 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1570 crate::gpu::probe_record(cls, true, t0.elapsed());
1571 return;
1572 }
1573 }
1574 crate::gpu::ProbeArm::CpuTimed => {
1575 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1576 crate::gpu::probe_record(cls, false, t0.elapsed());
1577 return;
1578 }
1579 crate::gpu::ProbeArm::Cpu => {}
1580 }
1581 }
1582 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1583 return;
1584 }
1585 if *dtype == TensorDtype::Q4TiledP {
1586 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1592 let t0 = std::time::Instant::now();
1593 let cls = crate::gpu::matvec_class(*rows, *cols);
1594 match crate::gpu::probe_arm(cls) {
1595 crate::gpu::ProbeArm::Gpu => {
1596 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1597 crate::gpu::probe_record(cls, true, t0.elapsed());
1598 return;
1599 }
1600 }
1601 crate::gpu::ProbeArm::CpuTimed => {
1602 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1603 crate::gpu::probe_record(cls, false, t0.elapsed());
1604 return;
1605 }
1606 crate::gpu::ProbeArm::Cpu => {}
1607 }
1608 }
1609 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1610 return;
1611 }
1612 if *dtype == TensorDtype::Q2TiledP {
1613 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1614 return;
1615 }
1616 if *dtype == TensorDtype::Mxfp4 {
1617 self.mxfp4_matvec(x, out, pool);
1618 return;
1619 }
1620 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1621 let w = self.bin_w().expect("binary dtype");
1626 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1627 let t0 = std::time::Instant::now();
1628 let arm = if crate::gpu::q1_force() {
1629 crate::gpu::ProbeArm::Gpu
1630 } else {
1631 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1632 };
1633 match arm {
1634 crate::gpu::ProbeArm::Gpu => {
1635 let ok = if *dtype == TensorDtype::Q1G {
1636 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1637 } else {
1638 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1639 };
1640 if ok {
1641 crate::gpu::probe_record(
1642 crate::gpu::OpClass::Matvec,
1643 true,
1644 t0.elapsed(),
1645 );
1646 return;
1647 }
1648 }
1649 crate::gpu::ProbeArm::CpuTimed => {
1650 bin_matvec(w, x, *rows, *cols, out, pool);
1651 crate::gpu::probe_record(
1652 crate::gpu::OpClass::Matvec,
1653 false,
1654 t0.elapsed(),
1655 );
1656 return;
1657 }
1658 crate::gpu::ProbeArm::Cpu => {}
1659 }
1660 }
1661 bin_matvec(w, x, *rows, *cols, out, pool);
1662 return;
1663 }
1664 if *dtype == TensorDtype::Q1T {
1665 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1669 let t0 = std::time::Instant::now();
1670 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1671 crate::gpu::ProbeArm::Gpu => {
1672 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1673 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1674 crate::gpu::probe_record(
1675 crate::gpu::OpClass::Matvec,
1676 true,
1677 t0.elapsed(),
1678 );
1679 return;
1680 }
1681 }
1682 crate::gpu::ProbeArm::CpuTimed => {
1683 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1684 crate::gpu::probe_record(
1685 crate::gpu::OpClass::Matvec,
1686 false,
1687 t0.elapsed(),
1688 );
1689 return;
1690 }
1691 crate::gpu::ProbeArm::Cpu => {}
1692 }
1693 }
1694 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1695 return;
1696 }
1697 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1698 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1699 return;
1700 }
1701 let xs = prescale(x, col_field, *dtype);
1702 if *rows >= crate::gpu::min_rows()
1707 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1708 && gpu_lmhead_enabled()
1709 && crate::gpu::enabled_here()
1710 {
1711 let t0 = std::time::Instant::now();
1714 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1715 crate::gpu::ProbeArm::Gpu => {}
1716 crate::gpu::ProbeArm::CpuTimed => {
1717 qmatvec(
1718 self.quant_bytes(),
1719 repack,
1720 row_scale,
1721 x,
1722 col_field,
1723 *dtype,
1724 *rows,
1725 *cols,
1726 out,
1727 pool,
1728 );
1729 crate::gpu::probe_record(
1730 crate::gpu::OpClass::Matvec,
1731 false,
1732 t0.elapsed(),
1733 );
1734 return;
1735 }
1736 crate::gpu::ProbeArm::Cpu => {
1737 qmatvec(
1738 self.quant_bytes(),
1739 repack,
1740 row_scale,
1741 x,
1742 col_field,
1743 *dtype,
1744 *rows,
1745 *cols,
1746 out,
1747 pool,
1748 );
1749 return;
1750 }
1751 }
1752 let frac = gpu_split_frac();
1753 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1754 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1755 let bytes = self.quant_bytes();
1756 let ok = std::thread::scope(|sc| {
1757 let g = sc.spawn(|| {
1758 crate::gpu::q8_matvec_range(
1759 model,
1760 *idx,
1761 cpu_rows,
1762 &row_scale[cpu_rows..],
1763 &xs,
1764 *rows - cpu_rows,
1765 *cols,
1766 out_gpu,
1767 )
1768 });
1769 if cpu_rows > 0 {
1770 let rep_cpu = if repack.is_empty() {
1773 &[][..]
1774 } else {
1775 &repack[..(cpu_rows / 4) * 4 * *cols]
1776 };
1777 qmatvec(
1778 &bytes[..cpu_rows * *cols],
1779 rep_cpu,
1780 &row_scale[..cpu_rows],
1781 x,
1782 col_field,
1783 *dtype,
1784 cpu_rows,
1785 *cols,
1786 out_cpu,
1787 pool,
1788 );
1789 }
1790 g.join().unwrap_or(false)
1791 });
1792 if ok {
1793 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1794 return;
1795 }
1796 qmatvec(
1799 &bytes[cpu_rows * *cols..(*rows) * *cols],
1800 &[],
1801 &row_scale[cpu_rows..],
1802 x,
1803 col_field,
1804 *dtype,
1805 *rows - cpu_rows,
1806 *cols,
1807 out_gpu,
1808 pool,
1809 );
1810 return;
1811 }
1812 qmatvec(
1813 self.quant_bytes(),
1814 repack,
1815 row_scale,
1816 x,
1817 col_field,
1818 *dtype,
1819 *rows,
1820 *cols,
1821 out,
1822 pool,
1823 );
1824 }
1825 }
1826 }
1827
1828 pub fn matvec2(
1830 &self,
1831 x1: &[f32],
1832 x2: &[f32],
1833 o1: &mut [f32],
1834 o2: &mut [f32],
1835 pool: Option<&Pool>,
1836 ) {
1837 match self {
1838 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1839 Self::Blr(_) => {
1840 self.matvec(x1, o1, pool);
1841 self.matvec(x2, o2, pool);
1842 }
1843 Self::Mapped {
1844 model,
1845 idx,
1846 dtype,
1847 rows,
1848 cols,
1849 row_scale,
1850 col_field,
1851 vbit_offsets,
1852 ..
1853 } => {
1854 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1855 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1856 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1857 match dtype {
1858 TensorDtype::Q4Block => {
1859 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1860 }
1861 TensorDtype::Q4Tiled => {
1862 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1863 }
1864 TensorDtype::Q4TiledP => {
1865 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1866 }
1867 TensorDtype::Q2TiledP => {
1868 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1869 q2tp_affine_matvec2(
1870 self.quant_bytes(),
1871 &tx1,
1872 &tx2,
1873 *rows,
1874 *cols,
1875 o1,
1876 o2,
1877 pool,
1878 )
1879 } else {
1880 q2tp_matvec2(
1881 self.quant_bytes(),
1882 &tx1,
1883 &tx2,
1884 *rows,
1885 *cols,
1886 o1,
1887 o2,
1888 pool,
1889 )
1890 }
1891 }
1892 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1893 self.bin_w().expect("binary dtype"),
1894 &tx1,
1895 &tx2,
1896 *rows,
1897 *cols,
1898 o1,
1899 o2,
1900 pool,
1901 ),
1902 TensorDtype::Q1T => {
1903 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1904 }
1905 TensorDtype::Mxfp4 => mxfp4::matvec2(
1906 self.quant_bytes(),
1907 &tx1,
1908 &tx2,
1909 *rows,
1910 *cols,
1911 o1,
1912 o2,
1913 pool,
1914 ),
1915 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1916 self.quant_bytes(),
1917 vbit_offsets,
1918 &tx1,
1919 &tx2,
1920 *rows,
1921 *cols,
1922 o1,
1923 o2,
1924 pool,
1925 ),
1926 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1927 self.quant_bytes(),
1928 row_scale,
1929 &tx1,
1930 &tx2,
1931 col_field,
1932 *dtype,
1933 *rows,
1934 *cols,
1935 o1,
1936 o2,
1937 pool,
1938 ),
1939 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1940 }
1941 return;
1942 }
1943 if *dtype == TensorDtype::Q4Block {
1944 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1945 return;
1946 }
1947 if *dtype == TensorDtype::Q4Tiled {
1948 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1949 return;
1950 }
1951 if *dtype == TensorDtype::Q4TiledP {
1952 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1953 return;
1954 }
1955 if *dtype == TensorDtype::Q2TiledP {
1956 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1957 return;
1958 }
1959 if *dtype == TensorDtype::Mxfp4 {
1960 mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1961 return;
1962 }
1963 if let Some(w) = self.bin_w() {
1964 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1965 return;
1966 }
1967 if *dtype == TensorDtype::Q1T {
1968 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1974 return;
1975 }
1976 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1977 vbitmatvec2(
1978 self.quant_bytes(),
1979 vbit_offsets,
1980 x1,
1981 x2,
1982 *rows,
1983 *cols,
1984 o1,
1985 o2,
1986 pool,
1987 );
1988 return;
1989 }
1990 qmatvec2(
1991 self.quant_bytes(),
1992 row_scale,
1993 x1,
1994 x2,
1995 col_field,
1996 *dtype,
1997 *rows,
1998 *cols,
1999 o1,
2000 o2,
2001 pool,
2002 );
2003 }
2004 }
2005 }
2006}
2007
2008impl QTensor {
2009 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2017 if self.has_prism_contract() {
2018 return None;
2019 }
2020 match self {
2021 Self::Mapped {
2022 model, idx, dtype, ..
2023 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2024 _ => None,
2025 }
2026 }
2027
2028 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2029 let cols = self.cols();
2030 let rows = self.rows();
2031 debug_assert_eq!(xs_all.len(), b * cols);
2032 debug_assert_eq!(out.len(), b * rows);
2033 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2034 if crate::gptq_capture::capturing() {
2038 if let Self::Mapped { model, idx, .. } = self {
2039 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2040 }
2041 if let Self::Blr(t) = self {
2042 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2043 }
2044 }
2045 match self {
2046 Self::Blr(t) => {
2047 if crate::gpu::enabled_here()
2048 && crate::blr::gpu_wanted()
2049 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2050 {
2051 return;
2052 }
2053 t.matmat(xs_all, b, out, pool)
2054 }
2055 Self::F32 { data, .. } => {
2056 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2057 return;
2058 }
2059 let out_addr = SendMut(out.as_mut_ptr());
2060 let run = |start: usize, end: usize| {
2067 let mut bi = start;
2068 while bi < end {
2069 let n = (end - bi).min(4);
2070 for o in 0..rows {
2071 let row = &data[o * cols..(o + 1) * cols];
2072 if n == 4 {
2073 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2074 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2075 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2076 for j in 0..cols {
2077 let w = row[j];
2078 a0 += w * x0[j];
2079 a1 += w * x1[j];
2080 a2 += w * x2[j];
2081 a3 += w * x3[j];
2082 }
2083 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2084 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2085 }
2086 } else {
2087 for k in 0..n {
2088 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2089 let mut acc = 0f32;
2090 for j in 0..cols {
2091 acc += row[j] * x[j];
2092 }
2093 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2094 }
2095 }
2096 }
2097 bi += n;
2098 }
2099 };
2100 dispatch_rows(pool, b, &run);
2101 }
2102 Self::Mapped {
2103 model,
2104 idx,
2105 dtype,
2106 row_scale,
2107 col_field,
2108 vbit_offsets,
2109 ..
2110 } => {
2111 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2112 let mut transformed = Vec::with_capacity(xs_all.len());
2113 for bi in 0..b {
2114 transformed.extend_from_slice(&crate::prism::forward(
2115 model,
2116 &xs_all[bi * cols..(bi + 1) * cols],
2117 ));
2118 }
2119 match dtype {
2120 TensorDtype::Q4Block => {
2121 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2122 }
2123 TensorDtype::Q4Tiled => {
2124 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2125 }
2126 TensorDtype::Q4TiledP => {
2127 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2128 }
2129 TensorDtype::Q2TiledP => {
2130 let affine =
2131 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2132 let gpu_batch_ok = if affine {
2138 b >= 2
2139 } else {
2140 b >= 32 && b * rows * cols >= 128_000_000
2141 };
2142 if gpu_batch_ok
2143 && cols % 32 == 0
2144 && crate::gpu::enabled_here()
2145 && crate::gpu::q2tp_gpu_opt_in()
2146 {
2147 let gpu_ok = if affine {
2148 crate::gpu::q2tp_affine_matmat(
2149 model,
2150 *idx,
2151 &transformed,
2152 b,
2153 rows,
2154 cols,
2155 out,
2156 )
2157 } else {
2158 crate::gpu::q2tp_matmat(
2159 model,
2160 *idx,
2161 &transformed,
2162 b,
2163 rows,
2164 cols,
2165 out,
2166 )
2167 };
2168 if gpu_ok {
2169 return;
2170 }
2171 }
2172 if affine
2176 && b == 1
2177 && cols % 32 == 0
2178 && crate::gpu::enabled_here()
2179 && crate::gpu::q2tp_gpu_opt_in()
2180 && crate::gpu::q2tp_affine_matvec(
2181 model,
2182 *idx,
2183 &transformed[..cols],
2184 rows,
2185 cols,
2186 &mut out[..rows],
2187 )
2188 {
2189 return;
2190 }
2191 if affine {
2192 q2tp_affine_matmat(
2193 self.quant_bytes(),
2194 &transformed,
2195 b,
2196 rows,
2197 cols,
2198 out,
2199 pool,
2200 )
2201 } else {
2202 q2tp_matmat(
2203 self.quant_bytes(),
2204 &transformed,
2205 b,
2206 rows,
2207 cols,
2208 out,
2209 pool,
2210 )
2211 }
2212 }
2213 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2214 self.bin_w().expect("binary dtype"),
2215 &transformed,
2216 b,
2217 rows,
2218 cols,
2219 out,
2220 pool,
2221 ),
2222 TensorDtype::Q1T => {
2223 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2224 }
2225 TensorDtype::Mxfp4 => {
2226 mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2227 }
2228 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2229 self.quant_bytes(),
2230 vbit_offsets,
2231 &transformed,
2232 b,
2233 rows,
2234 cols,
2235 out,
2236 pool,
2237 ),
2238 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2239 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2240 .map(|bi| {
2241 prescale(
2242 &transformed[bi * cols..(bi + 1) * cols],
2243 col_field,
2244 *dtype,
2245 )
2246 })
2247 .collect();
2248 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2249 }
2250 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2251 }
2252 return;
2253 }
2254 if *dtype == TensorDtype::Q4Block {
2255 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2256 return;
2257 }
2258 if *dtype == TensorDtype::Mxfp4 {
2259 self.mxfp4_matmat(xs_all, b, out, pool);
2260 return;
2261 }
2262 if *dtype == TensorDtype::Q4TiledP {
2263 if b >= 32
2277 && b * rows * cols >= 128_000_000
2278 && cols % 32 == 0
2279 && !row_exact()
2280 && !crate::gpu::mm_killed()
2281 && crate::gpu::enabled_here()
2282 {
2283 let class = if b >= 128 {
2284 crate::gpu::OpClass::MatmatWide
2285 } else {
2286 crate::gpu::OpClass::Matmat
2287 };
2288 if let Self::Mapped { model, idx, .. } = self {
2289 if crate::mm_ab::on() {
2301 let mut g = vec![0f32; b * rows];
2302 let t = std::time::Instant::now();
2303 let took = crate::gpu::q4tp_matmat(
2304 model, *idx, xs_all, b, rows, cols, &mut g,
2305 );
2306 let dg = t.elapsed();
2307 let t = std::time::Instant::now();
2308 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2309 let dc = t.elapsed();
2310 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2311 return;
2312 }
2313 let t0 = std::time::Instant::now();
2314 let resident = crate::gpu::weight_is_resident(model, *idx);
2318 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2319 crate::gpu::ProbeArm::Gpu => {
2320 if crate::gpu::q4tp_matmat(
2321 model, *idx, xs_all, b, rows, cols, out,
2322 ) {
2323 let el = t0.elapsed();
2324 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2334 let budget = std::time::Duration::from_secs_f64(
2335 flops / 1.5e12 * 8.0 + 0.020,
2336 );
2337 crate::gpu::mm_budget_check(
2338 "q4tp matmat",
2339 el,
2340 budget,
2341 crate::gpu::probe_was_cold() || !resident,
2342 );
2343 crate::gpu::probe_record(class, true, el);
2344 return;
2345 }
2346 }
2347 crate::gpu::ProbeArm::CpuTimed => {
2348 q4tp_matmat(
2349 self.quant_bytes(),
2350 xs_all,
2351 b,
2352 rows,
2353 cols,
2354 out,
2355 pool,
2356 );
2357 crate::gpu::probe_record(class, false, t0.elapsed());
2358 return;
2359 }
2360 crate::gpu::ProbeArm::Cpu => {}
2361 }
2362 }
2363 }
2364 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2365 return;
2366 }
2367 if *dtype == TensorDtype::Q2TiledP {
2368 if b >= 32
2374 && b * rows * cols >= 128_000_000
2375 && cols % 32 == 0
2376 && !crate::gpu::mm_killed()
2377 && crate::gpu::enabled_here()
2378 {
2379 let class = if b >= 128 {
2380 crate::gpu::OpClass::MatmatWide
2381 } else {
2382 crate::gpu::OpClass::Matmat
2383 };
2384 if let Self::Mapped { model, idx, .. } = self {
2385 let t0 = std::time::Instant::now();
2386 match crate::gpu::probe_arm(class) {
2387 crate::gpu::ProbeArm::Gpu => {
2388 if crate::gpu::q2tp_matmat(
2389 model, *idx, xs_all, b, rows, cols, out,
2390 ) {
2391 crate::gpu::probe_record(class, true, t0.elapsed());
2392 return;
2393 }
2394 }
2395 crate::gpu::ProbeArm::CpuTimed => {
2396 q2tp_matmat(
2397 self.quant_bytes(),
2398 xs_all,
2399 b,
2400 rows,
2401 cols,
2402 out,
2403 pool,
2404 );
2405 crate::gpu::probe_record(class, false, t0.elapsed());
2406 return;
2407 }
2408 crate::gpu::ProbeArm::Cpu => {}
2409 }
2410 }
2411 }
2412 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2417 return;
2418 }
2419 if *dtype == TensorDtype::Q4Tiled {
2420 if b >= 32
2432 && b * rows * cols >= 128_000_000
2433 && cols % 32 == 0
2434 && !crate::gpu::mm_killed()
2435 && crate::gpu::enabled_here()
2436 {
2437 let class = if b >= 128 {
2438 crate::gpu::OpClass::MatmatWide
2439 } else {
2440 crate::gpu::OpClass::Matmat
2441 };
2442 if let Self::Mapped { model, idx, .. } = self {
2443 let t0 = std::time::Instant::now();
2444 match crate::gpu::probe_arm(class) {
2445 crate::gpu::ProbeArm::Gpu => {
2446 if crate::gpu::q4t_matmat(
2447 model, *idx, xs_all, b, rows, cols, out,
2448 ) {
2449 let el = t0.elapsed();
2450 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2460 let budget = std::time::Duration::from_secs_f64(
2461 flops / 1.5e12 * 8.0 + 0.020,
2462 );
2463 crate::gpu::mm_budget_check(
2464 "q4t matmat",
2465 el,
2466 budget,
2467 crate::gpu::probe_was_cold(),
2468 );
2469 crate::gpu::probe_record(class, true, el);
2470 return;
2471 }
2472 }
2473 crate::gpu::ProbeArm::CpuTimed => {
2474 q4t_matmat(
2475 self.quant_bytes(),
2476 xs_all,
2477 b,
2478 rows,
2479 cols,
2480 out,
2481 pool,
2482 );
2483 crate::gpu::probe_record(class, false, t0.elapsed());
2484 return;
2485 }
2486 crate::gpu::ProbeArm::Cpu => {}
2487 }
2488 }
2489 }
2490 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2491 return;
2492 }
2493 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2494 let w = self.bin_w().expect("binary dtype");
2498 if b >= 32
2499 && b * rows * cols >= 128_000_000
2500 && cols % 64 == 0
2501 && crate::gpu::enabled_here()
2502 {
2503 if let Self::Mapped { model, idx, .. } = self {
2504 let t0 = std::time::Instant::now();
2505 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2506 crate::gpu::ProbeArm::Gpu => {
2507 let ok = if *dtype == TensorDtype::Q1G {
2508 crate::gpu::q1g_matmat(
2509 model, *idx, xs_all, b, rows, cols, out,
2510 )
2511 } else {
2512 crate::gpu::q1_matmat(
2513 model, *idx, xs_all, b, rows, cols, out,
2514 )
2515 };
2516 if ok {
2517 crate::gpu::probe_record(
2518 crate::gpu::OpClass::Matmat,
2519 true,
2520 t0.elapsed(),
2521 );
2522 return;
2523 }
2524 }
2525 crate::gpu::ProbeArm::CpuTimed => {
2526 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2527 crate::gpu::probe_record(
2528 crate::gpu::OpClass::Matmat,
2529 false,
2530 t0.elapsed(),
2531 );
2532 return;
2533 }
2534 crate::gpu::ProbeArm::Cpu => {}
2535 }
2536 }
2537 }
2538 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2539 return;
2540 }
2541 if *dtype == TensorDtype::Q1T {
2542 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2545 if let Self::Mapped { model, idx, .. } = self {
2546 let t0 = std::time::Instant::now();
2547 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2548 crate::gpu::ProbeArm::Gpu => {
2549 if crate::gpu::q1t_matmat(
2550 model, *idx, xs_all, b, rows, cols, out,
2551 ) {
2552 crate::gpu::probe_record(
2553 crate::gpu::OpClass::Matmat,
2554 true,
2555 t0.elapsed(),
2556 );
2557 return;
2558 }
2559 }
2560 crate::gpu::ProbeArm::CpuTimed => {
2561 q1t_matmat(
2562 self.quant_bytes(),
2563 xs_all,
2564 b,
2565 rows,
2566 cols,
2567 out,
2568 pool,
2569 );
2570 crate::gpu::probe_record(
2571 crate::gpu::OpClass::Matmat,
2572 false,
2573 t0.elapsed(),
2574 );
2575 return;
2576 }
2577 crate::gpu::ProbeArm::Cpu => {}
2578 }
2579 }
2580 }
2581 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2582 return;
2583 }
2584 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2585 vbitmatmat(
2586 self.quant_bytes(),
2587 vbit_offsets,
2588 xs_all,
2589 b,
2590 rows,
2591 cols,
2592 out,
2593 pool,
2594 );
2595 return;
2596 }
2597 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2598 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2599 .collect();
2600 if row_exact()
2606 && (1..=4).contains(&b)
2607 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2608 && crate::gpu::enabled_here()
2609 && crate::gpu::wgpu_active()
2610 {
2611 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2612 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2613 return;
2614 }
2615 }
2616 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2622 if let Self::Mapped { model, idx, .. } = self {
2623 let t0 = std::time::Instant::now();
2624 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2625 crate::gpu::ProbeArm::Gpu
2626 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2627 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2628 {
2629 let q = self.quant_bytes();
2633 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2634 return;
2635 }
2636 crate::gpu::ProbeArm::Gpu => {
2637 if *dtype == TensorDtype::Q8_2f
2645 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2646 && crate::gpu::q8_matmat_2f(
2647 model, *idx, row_scale, col_field, xs_all, b, rows,
2648 cols, out,
2649 )
2650 {
2651 crate::gpu::probe_record(
2652 crate::gpu::OpClass::Matmat,
2653 true,
2654 t0.elapsed(),
2655 );
2656 return;
2657 }
2658 let flat: Vec<f32> =
2659 pre.iter().flat_map(|v| v.iter().copied()).collect();
2660 if crate::gpu::q8_matmat(
2661 model, *idx, row_scale, &flat, b, rows, cols, out,
2662 ) {
2663 crate::gpu::probe_record(
2664 crate::gpu::OpClass::Matmat,
2665 true,
2666 t0.elapsed(),
2667 );
2668 return;
2669 }
2670 }
2671 crate::gpu::ProbeArm::CpuTimed => {
2672 let q = self.quant_bytes();
2673 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2674 crate::gpu::probe_record(
2675 crate::gpu::OpClass::Matmat,
2676 false,
2677 t0.elapsed(),
2678 );
2679 return;
2680 }
2681 crate::gpu::ProbeArm::Cpu => {}
2682 }
2683 }
2684 }
2685 let q = self.quant_bytes();
2686 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2687 }
2688 }
2689 }
2690}
2691
2692impl QTensor {
2693 fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2697 let (rows, cols) = (self.rows(), self.cols());
2698 if let Self::Mapped { model, idx, .. } = self {
2699 if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2700 let t0 = std::time::Instant::now();
2701 let cls = crate::gpu::matvec_class(rows, cols);
2702 match crate::gpu::probe_arm(cls) {
2703 crate::gpu::ProbeArm::Gpu => {
2704 if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2705 crate::gpu::probe_record(cls, true, t0.elapsed());
2706 return;
2707 }
2708 }
2709 crate::gpu::ProbeArm::CpuTimed => {
2710 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2711 crate::gpu::probe_record(cls, false, t0.elapsed());
2712 return;
2713 }
2714 crate::gpu::ProbeArm::Cpu => {}
2715 }
2716 }
2717 }
2718 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2719 }
2720
2721 fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2726 let (rows, cols) = (self.rows(), self.cols());
2727 if let Self::Mapped { model, idx, .. } = self {
2728 if b >= 2
2729 && b * rows * cols >= 33_554_432
2730 && !row_exact()
2731 && !crate::gpu::mm_killed()
2732 && crate::gpu::enabled_here()
2733 {
2734 let class = if b >= 128 {
2735 crate::gpu::OpClass::MatmatWide
2736 } else {
2737 crate::gpu::OpClass::Matmat
2738 };
2739 let t0 = std::time::Instant::now();
2740 match crate::gpu::probe_arm(class) {
2741 crate::gpu::ProbeArm::Gpu => {
2742 if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2743 crate::gpu::probe_record(class, true, t0.elapsed());
2744 return;
2745 }
2746 }
2747 crate::gpu::ProbeArm::CpuTimed => {
2748 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2749 crate::gpu::probe_record(class, false, t0.elapsed());
2750 return;
2751 }
2752 crate::gpu::ProbeArm::Cpu => {}
2753 }
2754 }
2755 }
2756 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2757 }
2758}
2759
2760impl QTensor {
2761 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2771 let (rows, cols) = (self.rows(), self.cols());
2772 let Self::Mapped {
2773 model,
2774 idx,
2775 dtype,
2776 row_scale,
2777 col_field,
2778 ..
2779 } = self
2780 else {
2781 return false;
2782 };
2783 if crate::prism::has_contract(model) {
2784 return false;
2785 }
2786 match *dtype {
2787 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2788 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2792 if *dtype == TensorDtype::Q8_2f
2795 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2796 && crate::gpu::q8_matmat_2f(
2797 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2798 )
2799 {
2800 return true;
2801 }
2802 let flat: Vec<f32> = (0..b)
2803 .flat_map(|bi| {
2804 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2805 })
2806 .collect();
2807 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2808 }
2809 _ => false,
2810 }
2811 }
2812
2813 pub fn matvec_many<const N: usize>(
2820 ts: [&QTensor; N],
2821 x: &[f32],
2822 mut outs: [&mut [f32]; N],
2823 pool: Option<&Pool>,
2824 ) {
2825 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2826 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2827 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2832 t.matvec(x, o, pool);
2833 }
2834 return;
2835 }
2836 let uniform_q8 = ts.iter().all(|t| {
2837 matches!(
2838 t,
2839 Self::Mapped {
2840 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2841 ..
2842 }
2843 )
2844 });
2845 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2846 if uniform_f32 && crate::f32_backend::active() {
2847 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2848 t.matvec(x, o, pool);
2849 }
2850 return;
2851 }
2852 let uniform_q4 = ts.iter().all(|t| {
2853 matches!(
2854 t,
2855 Self::Mapped {
2856 dtype: TensorDtype::Q4Block,
2857 ..
2858 }
2859 )
2860 });
2861 let uniform_vbit = ts.iter().all(|t| {
2862 matches!(
2863 t,
2864 Self::Mapped {
2865 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2866 ..
2867 }
2868 )
2869 });
2870 let uniform_q1 =
2873 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2874 let uniform_q1t = ts.iter().all(|t| {
2875 matches!(
2876 t,
2877 Self::Mapped {
2878 dtype: TensorDtype::Q1T,
2879 ..
2880 }
2881 )
2882 });
2883 let uniform_q4tp = ts.iter().all(|t| {
2888 matches!(
2889 t,
2890 Self::Mapped {
2891 dtype: TensorDtype::Q4TiledP,
2892 ..
2893 }
2894 )
2895 }) && ts
2896 .iter()
2897 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2898 let Some(pool) = pool else {
2899 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2900 t.matvec(x, o, None);
2901 }
2902 return;
2903 };
2904 if total_rows < 256
2905 || !(uniform_q8
2906 || uniform_f32
2907 || uniform_q4
2908 || uniform_vbit
2909 || uniform_q1
2910 || uniform_q1t
2911 || uniform_q4tp)
2912 {
2913 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2914 t.matvec(x, o, Some(pool));
2915 }
2916 return;
2917 }
2918
2919 if uniform_q4tp {
2920 let cols = ts[0].cols();
2926 let gpr = cols / GROUP_SIZE;
2927 let views: [Q4tpView; N] =
2928 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2929 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2930 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2931 let locate = |flat: usize| -> (usize, usize) {
2933 let mut acc = 0;
2934 for (i, &r) in rows_of.iter().enumerate() {
2935 if flat < acc + r {
2936 return (i, flat - acc);
2937 }
2938 acc += r;
2939 }
2940 (rows_of.len() - 1, 0)
2941 };
2942 let (views, outs_addr) = (&views, &outs_addr);
2943 if a8w8_enabled() {
2944 let act = split_act(x);
2945 let act = &act;
2946 let run = |start: usize, end: usize| {
2947 with_krow(gpr, |sc| {
2948 for flat in start..end {
2949 let (t, r) = locate(flat);
2950 let v = &views[t];
2951 v.scales_into(r, gpr, sc);
2952 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2953 for &(j, xv) in &act.outliers {
2954 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2955 acc += w * s * xv;
2956 }
2957 unsafe { *outs_addr[t].at(r) = acc };
2959 }
2960 });
2961 };
2962 pool.run_rows(total_rows, &run);
2963 } else {
2964 let run = |start: usize, end: usize| {
2965 with_krow(gpr, |sc| {
2966 for flat in start..end {
2967 let (t, r) = locate(flat);
2968 let v = &views[t];
2969 v.scales_into(r, gpr, sc);
2970 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2972 }
2973 });
2974 };
2975 pool.run_rows(total_rows, &run);
2976 }
2977 return;
2978 }
2979
2980 if uniform_q1 {
2981 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2984 if a8w8_enabled() {
2985 let act = split_act(x);
2986 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2987 let (act, gsum) = (&act, &gsum);
2988 let closures: [_; N] = std::array::from_fn(|i| {
2989 let (w, gpr, out) = (
2990 ts[i].bin_w().expect("uniform_q1 checked"),
2991 ts[i].cols() / GROUP_SIZE,
2992 outs_addr[i],
2993 );
2994 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2995 });
2996 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2997 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2998 pool.run_many(&parts);
2999 } else {
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_f32(w, gpr, x, 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 }
3012 return;
3013 }
3014
3015 if uniform_q1t {
3016 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3020 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3021 if a8w8_enabled() {
3022 let act = split_act(x);
3023 let act = &act;
3024 let x_ref = x;
3025 let closures: [_; N] = std::array::from_fn(|i| {
3026 let bytes = ts[i].quant_bytes();
3027 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3028 let gpr = cols / GROUP_SIZE;
3029 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3030 let out = outs_addr[i];
3031 move |s: usize, e: usize| {
3032 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3033 }
3034 });
3035 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037 pool.run_many(&parts);
3038 } else {
3039 let x_ref = x;
3040 let closures: [_; N] = std::array::from_fn(|i| {
3041 let bytes = ts[i].quant_bytes();
3042 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3043 let gpr = cols / GROUP_SIZE;
3044 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3045 let out = outs_addr[i];
3046 move |s: usize, e: usize| {
3047 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3048 }
3049 });
3050 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3051 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3052 pool.run_many(&parts);
3053 }
3054 return;
3055 }
3056
3057 if uniform_q4 || uniform_vbit {
3058 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3059 if a8w8_enabled() {
3061 let act = split_act(x);
3062 let act = &act;
3063 if uniform_q4 {
3064 let closures: [_; N] = std::array::from_fn(|i| {
3065 let (packed, scales) =
3066 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3067 let (gpr, cols, out) =
3068 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3069 move |s: usize, e: usize| {
3070 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3071 }
3072 });
3073 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3074 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3075 pool.run_many(&parts);
3076 } else {
3077 let closures: [_; N] = std::array::from_fn(|i| {
3078 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3079 unreachable!()
3080 };
3081 let (bytes, rows, cols, out) = (
3082 ts[i].quant_bytes(),
3083 ts[i].rows(),
3084 ts[i].cols(),
3085 outs_addr[i],
3086 );
3087 move |s: usize, e: usize| {
3088 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3089 }
3090 });
3091 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3092 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3093 pool.run_many(&parts);
3094 }
3095 return;
3096 }
3097 if uniform_q4 {
3098 let closures: [_; N] = std::array::from_fn(|i| {
3099 let (packed, scales) =
3100 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3101 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3102 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3103 });
3104 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3105 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3106 pool.run_many(&parts);
3107 } else {
3108 let closures: [_; N] = std::array::from_fn(|i| {
3109 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3110 unreachable!()
3111 };
3112 let (bytes, rows, cols, out) = (
3113 ts[i].quant_bytes(),
3114 ts[i].rows(),
3115 ts[i].cols(),
3116 outs_addr[i],
3117 );
3118 move |s: usize, e: usize| {
3119 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3120 }
3121 });
3122 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3123 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3124 pool.run_many(&parts);
3125 }
3126 return;
3127 }
3128
3129 if uniform_f32 {
3130 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3131 let closures: [_; N] = std::array::from_fn(|i| {
3132 let Self::F32 { data, cols, .. } = ts[i] else {
3133 unreachable!()
3134 };
3135 let out = outs_addr[i];
3136 move |start: usize, end: usize| {
3137 for o in start..end {
3138 let row = &data[o * cols..(o + 1) * cols];
3139 let mut sum = 0.0f32;
3140 for j in 0..*cols {
3141 sum += row[j] * x[j];
3142 }
3143 unsafe { *out.at(o) = sum };
3145 }
3146 }
3147 });
3148 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3149 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3150 pool.run_many(&parts);
3151 return;
3152 }
3153
3154 struct Ctx<'a> {
3157 bytes: &'a [u8],
3158 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3159 rep: &'a [u8],
3160 row_scale: &'a [f32],
3161 cols: usize,
3162 xs: std::borrow::Cow<'a, [f32]>,
3163 }
3164 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3165 let Self::Mapped {
3166 dtype,
3167 cols,
3168 row_scale,
3169 col_field,
3170 repack,
3171 ..
3172 } = ts[i]
3173 else {
3174 unreachable!()
3175 };
3176 Ctx {
3177 bytes: ts[i].quant_bytes(),
3178 rep: repack,
3179 row_scale,
3180 cols: *cols,
3181 xs: prescale(x, col_field, *dtype),
3182 }
3183 });
3184 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3185 #[cfg(target_arch = "aarch64")]
3186 if sdot_enabled() {
3187 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3188 let closures: [_; N] = std::array::from_fn(|i| {
3189 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3190 move |start: usize, end: usize| {
3191 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3192 }
3193 });
3194 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3195 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3196 pool.run_many(&parts);
3197 return;
3198 }
3199 #[cfg(target_arch = "x86_64")]
3200 if avx2_a8w8_enabled() {
3201 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3202 let closures: [_; N] = std::array::from_fn(|i| {
3203 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3204 move |start: usize, end: usize| {
3205 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3206 }
3207 });
3208 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3209 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3210 pool.run_many(&parts);
3211 return;
3212 }
3213 let closures: [_; N] = std::array::from_fn(|i| {
3214 let (c, out) = (&ctxs[i], outs_addr[i]);
3215 move |start: usize, end: usize| {
3216 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3217 }
3218 });
3219 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3220 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3221 pool.run_many(&parts);
3222 }
3223}
3224
3225impl QTensor {
3226 #[allow(clippy::needless_range_loop)]
3231 pub fn matvec2_many<const N: usize>(
3232 ts: [&QTensor; N],
3233 x1: &[f32],
3234 x2: &[f32],
3235 mut o1s: [&mut [f32]; N],
3236 mut o2s: [&mut [f32]; N],
3237 pool: Option<&Pool>,
3238 ) {
3239 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3240 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3241 for i in 0..N {
3242 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3243 }
3244 return;
3245 }
3246 let uniform_q8 = ts.iter().all(|t| {
3247 matches!(
3248 t,
3249 Self::Mapped {
3250 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3251 ..
3252 }
3253 )
3254 });
3255 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3256 let uniform_q4 = ts.iter().all(|t| {
3257 matches!(
3258 t,
3259 Self::Mapped {
3260 dtype: TensorDtype::Q4Block,
3261 ..
3262 }
3263 )
3264 });
3265 let uniform_vbit = ts.iter().all(|t| {
3266 matches!(
3267 t,
3268 Self::Mapped {
3269 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3270 ..
3271 }
3272 )
3273 });
3274 let fusable = pool.is_some()
3275 && total_rows >= 256
3276 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3277 if !fusable {
3278 for i in 0..N {
3279 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3280 }
3281 return;
3282 }
3283 let pool = pool.unwrap();
3284
3285 if uniform_q4 || uniform_vbit {
3286 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3287 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3288 if a8w8_enabled() {
3290 let a1 = split_act(x1);
3291 let a2 = split_act(x2);
3292 let (a1, a2) = (&a1, &a2);
3293 if uniform_q4 {
3294 let closures: [_; N] = std::array::from_fn(|i| {
3295 let (packed, scales) =
3296 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3297 let (gpr, cols, o1, o2) =
3298 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3299 move |s: usize, e: usize| {
3300 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3301 }
3302 });
3303 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3304 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3305 pool.run_many(&parts);
3306 } else {
3307 let closures: [_; N] = std::array::from_fn(|i| {
3308 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3309 unreachable!()
3310 };
3311 let (bytes, rows, cols, o1, o2) = (
3312 ts[i].quant_bytes(),
3313 ts[i].rows(),
3314 ts[i].cols(),
3315 p1[i],
3316 p2[i],
3317 );
3318 move |s: usize, e: usize| {
3319 vbit_range2_a8w8(
3320 bytes,
3321 vbit_offsets,
3322 x1,
3323 x2,
3324 a1,
3325 a2,
3326 rows,
3327 cols,
3328 o1,
3329 o2,
3330 s,
3331 e,
3332 )
3333 }
3334 });
3335 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3336 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3337 pool.run_many(&parts);
3338 }
3339 return;
3340 }
3341 if uniform_q4 {
3342 let closures: [_; N] = std::array::from_fn(|i| {
3343 let (packed, scales) =
3344 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3345 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3346 move |s: usize, e: usize| {
3347 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3348 }
3349 });
3350 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3351 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3352 pool.run_many(&parts);
3353 } else {
3354 let closures: [_; N] = std::array::from_fn(|i| {
3355 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3356 unreachable!()
3357 };
3358 let (bytes, rows, cols, o1, o2) = (
3359 ts[i].quant_bytes(),
3360 ts[i].rows(),
3361 ts[i].cols(),
3362 p1[i],
3363 p2[i],
3364 );
3365 move |s: usize, e: usize| {
3366 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3367 }
3368 });
3369 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3370 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3371 pool.run_many(&parts);
3372 }
3373 return;
3374 }
3375
3376 if uniform_f32 {
3377 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3378 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3379 let closures: [_; N] = std::array::from_fn(|i| {
3380 let Self::F32 { data, cols, .. } = ts[i] else {
3381 unreachable!()
3382 };
3383 let (o1, o2) = (p1[i], p2[i]);
3384 move |start: usize, end: usize| {
3385 for o in start..end {
3386 let row = &data[o * cols..(o + 1) * cols];
3387 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3388 for j in 0..*cols {
3389 s1 += row[j] * x1[j];
3390 s2 += row[j] * x2[j];
3391 }
3392 unsafe {
3394 *o1.at(o) = s1;
3395 *o2.at(o) = s2;
3396 }
3397 }
3398 }
3399 });
3400 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3401 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3402 pool.run_many(&parts);
3403 return;
3404 }
3405
3406 struct Ctx<'a> {
3407 bytes: &'a [u8],
3408 row_scale: &'a [f32],
3409 cols: usize,
3410 xs1: std::borrow::Cow<'a, [f32]>,
3411 xs2: std::borrow::Cow<'a, [f32]>,
3412 }
3413 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3414 let Self::Mapped {
3415 dtype,
3416 cols,
3417 row_scale,
3418 col_field,
3419 ..
3420 } = ts[i]
3421 else {
3422 unreachable!()
3423 };
3424 Ctx {
3425 bytes: ts[i].quant_bytes(),
3426 row_scale,
3427 cols: *cols,
3428 xs1: prescale(x1, col_field, *dtype),
3429 xs2: prescale(x2, col_field, *dtype),
3430 }
3431 });
3432 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3433 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3434 #[cfg(target_arch = "aarch64")]
3435 if sdot_enabled() {
3436 let acts: [(SplitAct, SplitAct); N] =
3437 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3438 let closures: [_; N] = std::array::from_fn(|i| {
3439 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3440 move |start: usize, end: usize| {
3441 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3442 }
3443 });
3444 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3445 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3446 pool.run_many(&parts);
3447 return;
3448 }
3449 #[cfg(target_arch = "x86_64")]
3450 if avx2_a8w8_enabled() {
3451 let acts: [(SplitAct, SplitAct); N] =
3452 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3453 let closures: [_; N] = std::array::from_fn(|i| {
3454 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3455 move |start: usize, end: usize| {
3456 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3457 }
3458 });
3459 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3460 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3461 pool.run_many(&parts);
3462 return;
3463 }
3464 let closures: [_; N] = std::array::from_fn(|i| {
3465 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3466 move |start: usize, end: usize| {
3467 q8_range2_f32(
3468 c.bytes,
3469 c.row_scale,
3470 &c.xs1,
3471 &c.xs2,
3472 c.cols,
3473 o1,
3474 o2,
3475 start,
3476 end,
3477 )
3478 }
3479 });
3480 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3481 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3482 pool.run_many(&parts);
3483 }
3484
3485 pub fn matvec_silu_mul(
3490 gate: &QTensor,
3491 up: &QTensor,
3492 x: &[f32],
3493 out: &mut [f32],
3494 pool: Option<&Pool>,
3495 ) -> bool {
3496 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3497 }
3498
3499 pub fn matvec_silu_mul_limited(
3505 gate: &QTensor,
3506 up: &QTensor,
3507 x: &[f32],
3508 out: &mut [f32],
3509 limit: f32,
3510 pool: Option<&Pool>,
3511 ) -> bool {
3512 if gate.has_prism_contract()
3513 || up.has_prism_contract()
3514 || gate.per_tensor_only()
3515 || up.per_tensor_only()
3516 {
3517 return false;
3521 }
3522 let inter = gate.rows();
3523 debug_assert_eq!(up.rows(), inter);
3524 debug_assert_eq!(out.len(), inter);
3525 debug_assert_eq!(gate.cols(), up.cols());
3526 if !a8w8_enabled() {
3527 return false;
3528 }
3529 let act = split_act(x);
3530 let act = &act;
3531 let x_ref = x;
3532 let out_addr = SendMut(out.as_mut_ptr());
3533
3534 match (gate, up) {
3535 (
3537 Self::Mapped {
3538 dtype: TensorDtype::Q4Block,
3539 ..
3540 },
3541 Self::Mapped {
3542 dtype: TensorDtype::Q4Block,
3543 ..
3544 },
3545 ) => {
3546 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3547 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3548 let gpr = gate.cols() / GROUP_SIZE;
3549 let cols = gate.cols();
3550 let run = move |start: usize, end: usize| {
3551 for r in start..end {
3552 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3553 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3554 for &(j, xv) in &act.outliers {
3555 let flat = r * cols + j;
3556 let gb = gp[flat / 2];
3557 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3558 let gsc = f16_to_f32(u16::from_le_bytes([
3559 gs[(flat / GROUP_SIZE) * 2],
3560 gs[(flat / GROUP_SIZE) * 2 + 1],
3561 ]));
3562 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3563 let ub = up_p[flat / 2];
3564 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3565 let usc = f16_to_f32(u16::from_le_bytes([
3566 up_s[(flat / GROUP_SIZE) * 2],
3567 up_s[(flat / GROUP_SIZE) * 2 + 1],
3568 ]));
3569 uv += ((un as i32 - 8) as f32) * usc * xv;
3570 }
3571 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3573 }
3574 };
3575 dispatch_rows(pool, inter, &run);
3576 true
3577 }
3578 (
3582 Self::Mapped {
3583 dtype: TensorDtype::Q4Tiled,
3584 ..
3585 },
3586 Self::Mapped {
3587 dtype: TensorDtype::Q4Tiled,
3588 ..
3589 },
3590 ) => {
3591 let g_bytes = gate.quant_bytes();
3592 let u_bytes = up.quant_bytes();
3593 let gpr = gate.cols() / GROUP_SIZE;
3594 let run = move |start: usize, end: usize| {
3595 for r in start..end {
3596 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3597 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3598 for &(j, xv) in &act.outliers {
3599 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3600 gv += w * s * xv;
3601 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3602 uv += w * s * xv;
3603 }
3604 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3606 }
3607 };
3608 dispatch_rows(pool, inter, &run);
3609 true
3610 }
3611 (
3614 Self::Mapped {
3615 dtype: TensorDtype::Q4TiledP,
3616 ..
3617 },
3618 Self::Mapped {
3619 dtype: TensorDtype::Q4TiledP,
3620 ..
3621 },
3622 ) => {
3623 let cols = gate.cols();
3624 let gpr = cols / GROUP_SIZE;
3625 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3626 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3627 let run = |start: usize, end: usize| {
3628 with_krows(gpr, |gsc, usc| {
3629 for r in start..end {
3630 gv_view.scales_into(r, gpr, gsc);
3631 uv_view.scales_into(r, gpr, usc);
3632 let mut gv =
3633 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3634 let mut uv =
3635 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3636 for &(j, xv) in &act.outliers {
3637 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3638 gv += w * s * xv;
3639 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3640 uv += w * s * xv;
3641 }
3642 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3644 }
3645 });
3646 };
3647 dispatch_rows(pool, inter, &run);
3648 true
3649 }
3650 (
3657 Self::Mapped {
3658 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3659 ..
3660 },
3661 Self::Mapped {
3662 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3663 ..
3664 },
3665 ) => {
3666 let gw = gate.bin_w().expect("binary dtype matched");
3667 let uw = up.bin_w().expect("binary dtype matched");
3668 let gpr = gate.cols() / GROUP_SIZE;
3669 let gsum = q1_group_sums(&act.xq, gpr);
3670 let gsum = &gsum;
3671 let run = move |start: usize, end: usize| {
3672 for r in start..end {
3673 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3674 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3675 for &(j, xv) in &act.outliers {
3676 let (w, s) = gw.outlier(r, gpr, j);
3677 gv += w * s * xv;
3678 let (w, s) = uw.outlier(r, gpr, j);
3679 uv += w * s * xv;
3680 }
3681 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3683 }
3684 };
3685 dispatch_rows(pool, inter, &run);
3686 true
3687 }
3688 (
3692 Self::Mapped {
3693 dtype: TensorDtype::Q2TiledP,
3694 ..
3695 },
3696 Self::Mapped {
3697 dtype: TensorDtype::Q2TiledP,
3698 ..
3699 },
3700 ) => {
3701 let cols = gate.cols();
3702 let gpr = cols / GROUP_SIZE;
3703 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3704 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3705 let gsum = q1_group_sums(&act.xq, gpr);
3706 let gsum = &gsum;
3707 let run = move |start: usize, end: usize| {
3708 with_krows(gpr, |gsc, usc| {
3709 for r in start..end {
3710 gv_view.scales_into(r, gpr, gsc);
3711 uv_view.scales_into(r, gpr, usc);
3712 let mut gv =
3713 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3714 let mut uv =
3715 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3716 for &(j, xv) in &act.outliers {
3717 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3718 gv += w * s * xv;
3719 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3720 uv += w * s * xv;
3721 }
3722 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3724 }
3725 });
3726 };
3727 dispatch_rows(pool, inter, &run);
3728 true
3729 }
3730 (
3735 Self::Mapped {
3736 dtype: TensorDtype::Q8Row,
3737 row_scale: g_rs,
3738 ..
3739 },
3740 Self::Mapped {
3741 dtype: TensorDtype::Q8Row,
3742 row_scale: u_rs,
3743 ..
3744 },
3745 ) => {
3746 let g_bytes = gate.quant_bytes();
3747 let u_bytes = up.quant_bytes();
3748 let cols = gate.cols();
3749 let run = move |start: usize, end: usize| {
3750 for r in start..end {
3751 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3752 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3753 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3755 }
3756 };
3757 dispatch_rows(pool, inter, &run);
3758 true
3759 }
3760 (
3762 Self::Mapped {
3763 dtype: TensorDtype::Q1T,
3764 ..
3765 },
3766 Self::Mapped {
3767 dtype: TensorDtype::Q1T,
3768 ..
3769 },
3770 ) => {
3771 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3772 let g_bytes = gate.quant_bytes();
3773 let u_bytes = up.quant_bytes();
3774 let gpr = gate.cols() / GROUP_SIZE;
3775 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3776 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3777 let run = move |start: usize, end: usize| {
3778 for r in start..end {
3779 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3780 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3781 for &(j, xv) in &act.outliers {
3782 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3783 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3784 }
3785 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3786 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3787 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3789 }
3790 };
3791 dispatch_rows(pool, inter, &run);
3792 true
3793 }
3794 _ => false,
3795 }
3796 }
3797
3798 pub fn moe_gate_up_many(
3813 pairs: &[(&QTensor, &QTensor)],
3814 x: &[f32],
3815 outs: &mut [Vec<f32>],
3816 pool: Option<&Pool>,
3817 ) -> bool {
3818 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3819 }
3820
3821 pub fn moe_gate_up_many_limited(
3825 pairs: &[(&QTensor, &QTensor)],
3826 x: &[f32],
3827 outs: &mut [Vec<f32>],
3828 limit: f32,
3829 pool: Option<&Pool>,
3830 ) -> bool {
3831 if pairs.is_empty() || pairs.len() != outs.len() {
3832 return false;
3833 }
3834 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3835 return false;
3836 }
3837 if !a8w8_enabled() {
3838 if limit > 0.0 {
3839 return false;
3842 }
3843 let groups = vec![vec![0]; pairs.len()];
3844 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3845 }
3846 let inter = pairs[0].0.rows();
3847 let cols = pairs[0].0.cols();
3848 if cols % GROUP_SIZE != 0 {
3849 return false;
3850 }
3851 let gpr = cols / GROUP_SIZE;
3852 let q2 = matches!(
3855 pairs[0].0,
3856 Self::Mapped {
3857 dtype: TensorDtype::Q2TiledP,
3858 ..
3859 }
3860 );
3861 let want = if q2 {
3862 TensorDtype::Q2TiledP
3863 } else {
3864 TensorDtype::Q4TiledP
3865 };
3866 let mut views = Vec::with_capacity(pairs.len() * 2);
3867 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3868 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3869 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3870 if !both
3871 || g.rows() != inter
3872 || u.rows() != inter
3873 || g.cols() != cols
3874 || u.cols() != cols
3875 || o.len() != inter
3876 {
3877 return false;
3878 }
3879 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3880 views.push(mk(g.quant_bytes(), inter, cols));
3881 views.push(mk(u.quant_bytes(), inter, cols));
3882 }
3883 let act = split_act(x);
3884 let gsum = if q2 {
3885 q1_group_sums(&act.xq, gpr)
3886 } else {
3887 Vec::new()
3888 };
3889 let (act, gsum) = (&act, &gsum);
3890 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3891 let (views, ptrs) = (&views, &ptrs);
3892 let run = |start: usize, end: usize| {
3893 with_krows(gpr, |gsc, usc| {
3894 for flat in start..end {
3895 let (e, r) = (flat / inter, flat % inter);
3896 let gv_view = &views[e * 2];
3897 let uv_view = &views[e * 2 + 1];
3898 gv_view.scales_into(r, gpr, gsc);
3899 uv_view.scales_into(r, gpr, usc);
3900 let (mut gv, mut uv) = if q2 {
3901 (
3902 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3903 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3904 )
3905 } else {
3906 (
3907 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3908 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3909 )
3910 };
3911 for &(j, xv) in &act.outliers {
3912 let (og, ou) = if q2 {
3913 (
3914 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3915 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3916 )
3917 } else {
3918 (
3919 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3920 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3921 )
3922 };
3923 gv += og.0 * og.1 * xv;
3924 uv += ou.0 * ou.1 * xv;
3925 }
3926 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3928 }
3929 });
3930 };
3931 dispatch_rows(pool, pairs.len() * inter, &run);
3932 true
3933 }
3934
3935 pub fn moe_down_many(
3944 downs: &[&QTensor],
3945 gs: &[Vec<f32>],
3946 weights: &[f32],
3947 out: &mut [f32],
3948 pool: Option<&Pool>,
3949 ) -> bool {
3950 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3951 return false;
3952 }
3953 if downs.iter().any(|d| d.per_tensor_only()) {
3954 return false;
3955 }
3956 if !a8w8_enabled() {
3957 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3958 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3959 return false;
3960 }
3961 out.fill(0.0);
3962 for (row, &w) in terms.iter().zip(weights) {
3963 for (o, &v) in out.iter_mut().zip(row) {
3964 *o += w * v;
3965 }
3966 }
3967 return true;
3968 }
3969 let rows = out.len();
3970 let cols = downs[0].cols();
3971 if cols % GROUP_SIZE != 0 {
3972 return false;
3973 }
3974 let gpr = cols / GROUP_SIZE;
3975 let mut views = Vec::with_capacity(downs.len());
3976 for (d, g) in downs.iter().zip(gs.iter()) {
3977 if !matches!(
3978 d,
3979 Self::Mapped {
3980 dtype: TensorDtype::Q4TiledP,
3981 ..
3982 }
3983 ) || d.rows() != rows
3984 || d.cols() != cols
3985 || g.len() != cols
3986 {
3987 return false;
3988 }
3989 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3990 }
3991 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3993 let out_addr = SendMut(out.as_mut_ptr());
4001 let (views, acts, weights) = (&views, &acts, &weights);
4002 let run = |start: usize, end: usize| {
4003 with_krow(gpr, |sc| {
4004 for r in start..end {
4005 let mut acc = 0f32;
4006 for (e, v) in views.iter().enumerate() {
4007 v.scales_into(r, gpr, sc);
4008 let a = &acts[e];
4009 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4010 for &(j, xv) in &a.outliers {
4011 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4012 d += w * s * xv;
4013 }
4014 acc += weights[e] * d;
4015 }
4016 unsafe { *out_addr.at(r) = acc };
4018 }
4019 });
4020 };
4021 dispatch_rows(pool, rows, &run);
4022 true
4023 }
4024
4025 pub fn moe_gate_up_rows(
4037 pairs: &[(&QTensor, &QTensor)],
4038 groups: &[Vec<usize>],
4039 xs: &[f32],
4040 outs: &mut [Vec<f32>],
4041 pool: Option<&Pool>,
4042 ) -> bool {
4043 if pairs.is_empty() || pairs.len() != groups.len() {
4044 return false;
4045 }
4046 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4047 return false;
4048 }
4049 let inter = pairs[0].0.rows();
4050 let cols = pairs[0].0.cols();
4051 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4052 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4053 return false;
4054 }
4055 let b = xs.len() / cols;
4056 let gpr = cols / GROUP_SIZE;
4057 let mut views = Vec::with_capacity(pairs.len() * 2);
4058 for (g, u) in pairs {
4059 let q4tp = |t: &QTensor| {
4060 matches!(
4061 t,
4062 Self::Mapped {
4063 dtype: TensorDtype::Q4TiledP,
4064 ..
4065 }
4066 )
4067 };
4068 if g.has_prism_contract()
4069 || u.has_prism_contract()
4070 || !q4tp(g)
4071 || !q4tp(u)
4072 || g.rows() != inter
4073 || u.rows() != inter
4074 || g.cols() != cols
4075 || u.cols() != cols
4076 {
4077 return false;
4078 }
4079 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4080 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4081 }
4082 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4083 return false;
4084 }
4085 let quantized = a8w8_enabled();
4086 let acts: Vec<SplitAct> = if quantized {
4087 (0..b)
4088 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4089 .collect()
4090 } else {
4091 Vec::new()
4092 };
4093 let mut offs = Vec::with_capacity(groups.len());
4094 let mut o = 0usize;
4095 for g in groups {
4096 offs.push(o);
4097 o += g.len();
4098 }
4099 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4100 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4101 let run = |start: usize, end: usize| {
4102 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4103 for flat in start..end {
4104 let (e, r) = (flat / inter, flat % inter);
4105 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4106 gv_view.scales_into(r, gpr, &mut gsc);
4107 uv_view.scales_into(r, gpr, &mut usc);
4108 for (k, &t) in groups[e].iter().enumerate() {
4109 if !quantized {
4110 let x = &xs[t * cols..(t + 1) * cols];
4111 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4112 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4113 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4114 continue;
4115 }
4116 let act = &acts[t];
4117 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4118 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4119 for &(j, xv) in &act.outliers {
4120 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4121 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4122 gv += og.0 * og.1 * xv;
4123 uv += ou.0 * ou.1 * xv;
4124 }
4125 let silu_g = gv / (1.0 + (-gv).exp());
4126 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4129 }
4130 }
4131 };
4132 dispatch_rows(pool, pairs.len() * inter, &run);
4133 true
4134 }
4135
4136 pub fn moe_down_rows(
4143 downs: &[&QTensor],
4144 group_lens: &[usize],
4145 gs: &[Vec<f32>],
4146 outs: &mut [Vec<f32>],
4147 pool: Option<&Pool>,
4148 ) -> bool {
4149 if downs.is_empty() || downs.len() != group_lens.len() {
4150 return false;
4151 }
4152 if downs.iter().any(|d| d.per_tensor_only()) {
4153 return false;
4154 }
4155 let rows = downs[0].rows();
4156 let cols = downs[0].cols();
4157 let n_pairs: usize = group_lens.iter().sum();
4158 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4159 return false;
4160 }
4161 let gpr = cols / GROUP_SIZE;
4162 let mut views = Vec::with_capacity(downs.len());
4163 for d in downs {
4164 if d.has_prism_contract()
4165 || !matches!(
4166 d,
4167 Self::Mapped {
4168 dtype: TensorDtype::Q4TiledP,
4169 ..
4170 }
4171 )
4172 || d.rows() != rows
4173 || d.cols() != cols
4174 {
4175 return false;
4176 }
4177 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4178 }
4179 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4180 return false;
4181 }
4182 let quantized = a8w8_enabled();
4183 let acts: Vec<SplitAct> = if quantized {
4184 gs.iter().map(|g| split_act(g)).collect()
4185 } else {
4186 Vec::new()
4187 };
4188 let mut offs = Vec::with_capacity(group_lens.len());
4189 let mut o = 0usize;
4190 for &l in group_lens {
4191 offs.push(o);
4192 o += l;
4193 }
4194 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4195 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4196 let run = |start: usize, end: usize| {
4197 let mut sc = vec![0f32; gpr];
4198 for flat in start..end {
4199 let (e, r) = (flat / rows, flat % rows);
4200 let v = &views[e];
4201 v.scales_into(r, gpr, &mut sc);
4202 for k in 0..group_lens[e] {
4203 if !quantized {
4204 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4205 unsafe { *ptrs[offs[e] + k].at(r) = d };
4206 continue;
4207 }
4208 let a = &acts[offs[e] + k];
4209 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4210 for &(j, xv) in &a.outliers {
4211 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4212 d += w * s * xv;
4213 }
4214 unsafe { *ptrs[offs[e] + k].at(r) = d };
4216 }
4217 }
4218 };
4219 dispatch_rows(pool, downs.len() * rows, &run);
4220 true
4221 }
4222}
4223
4224#[cfg(target_os = "macos")]
4229mod accel_blas {
4230 #[link(name = "Accelerate", kind = "framework")]
4231 unsafe extern "C" {
4232 pub fn cblas_sgemm(
4233 order: i32,
4234 trans_a: i32,
4235 trans_b: i32,
4236 m: i32,
4237 n: i32,
4238 k: i32,
4239 alpha: f32,
4240 a: *const f32,
4241 lda: i32,
4242 b: *const f32,
4243 ldb: i32,
4244 beta: f32,
4245 c: *mut f32,
4246 ldc: i32,
4247 );
4248 }
4249}
4250
4251#[cfg(target_os = "macos")]
4252pub(crate) fn accel_gemm_enabled() -> bool {
4253 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4254 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4255}
4256
4257#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4260pub(crate) fn accel_gemm_enabled() -> bool {
4261 true
4262}
4263
4264#[cfg(target_arch = "aarch64")]
4271#[allow(clippy::too_many_arguments)]
4272pub(crate) fn neon_gemm_rm(
4273 m: usize,
4274 n: usize,
4275 k: usize,
4276 alpha: f32,
4277 a: &[f32],
4278 lda: usize,
4279 b_mat: &[f32],
4280 ldb: usize,
4281 b_rows_are_n: bool,
4282 c: &mut [f32],
4283 ldc: usize,
4284) {
4285 debug_assert!(a.len() >= (m - 1) * lda + k);
4286 debug_assert!(c.len() >= (m - 1) * ldc + n);
4287 unsafe {
4289 use core::arch::aarch64::*;
4290 let mut i = 0usize;
4291 while i < m {
4292 let mi = (m - i).min(4);
4293 let mut j = 0usize;
4294 while j < n {
4295 let nj = (n - j).min(8);
4296 if mi == 4 && nj == 8 {
4297 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4298 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4299 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4300 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4301 for p in 0..k {
4302 let (b0, b1) = if b_rows_are_n {
4303 let base = b_mat.as_ptr().add(j * ldb + p);
4306 let g = |o: usize| *base.add(o * ldb);
4307 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4308 } else {
4309 let base = b_mat.as_ptr().add(p * ldb + j);
4310 (
4311 [*base, *base.add(1), *base.add(2), *base.add(3)],
4312 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4313 )
4314 };
4315 let bv0 = vld1q_f32(b0.as_ptr());
4316 let bv1 = vld1q_f32(b1.as_ptr());
4317 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4318 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4319 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4320 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4321 c0a = vfmaq_f32(c0a, a0, bv0);
4322 c0b = vfmaq_f32(c0b, a0, bv1);
4323 c1a = vfmaq_f32(c1a, a1, bv0);
4324 c1b = vfmaq_f32(c1b, a1, bv1);
4325 c2a = vfmaq_f32(c2a, a2, bv0);
4326 c2b = vfmaq_f32(c2b, a2, bv1);
4327 c3a = vfmaq_f32(c3a, a3, bv0);
4328 c3b = vfmaq_f32(c3b, a3, bv1);
4329 }
4330 let al = vdupq_n_f32(alpha);
4331 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4332 .iter()
4333 .enumerate()
4334 {
4335 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4336 vst1q_f32(dst, vmulq_f32(*ca, al));
4337 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4338 }
4339 } else {
4340 for r in 0..mi {
4341 for q in 0..nj {
4342 let mut acc = 0f32;
4343 for p in 0..k {
4344 let bv = if b_rows_are_n {
4345 b_mat[(j + q) * ldb + p]
4346 } else {
4347 b_mat[p * ldb + j + q]
4348 };
4349 acc += a[(i + r) * lda + p] * bv;
4350 }
4351 c[(i + r) * ldc + j + q] = acc * alpha;
4352 }
4353 }
4354 }
4355 j += nj;
4356 }
4357 i += mi;
4358 }
4359 }
4360}
4361
4362#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4364#[allow(clippy::too_many_arguments)]
4365pub(crate) fn sgemm_rm(
4366 m: usize,
4367 n: usize,
4368 k: usize,
4369 alpha: f32,
4370 a: &[f32],
4371 lda: usize,
4372 b_mat: &[f32],
4373 ldb: usize,
4374 b_rows_are_n: bool,
4375 c: &mut [f32],
4376 ldc: usize,
4377) {
4378 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4379}
4380
4381#[allow(clippy::too_many_arguments)]
4385pub fn sgemm_public(
4386 m: usize,
4387 n: usize,
4388 k: usize,
4389 alpha: f32,
4390 a: &[f32],
4391 lda: usize,
4392 b_mat: &[f32],
4393 ldb: usize,
4394 b_rows_are_n: bool,
4395 c: &mut [f32],
4396 ldc: usize,
4397) {
4398 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4399 {
4400 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4401 }
4402 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4407 {
4408 if m == 0 || n == 0 {
4409 return;
4410 }
4411 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4412 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4413 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4414 unsafe {
4416 matrixmultiply::sgemm(
4417 m,
4418 k,
4419 n,
4420 alpha,
4421 a.as_ptr(),
4422 lda as isize,
4423 1,
4424 b_mat.as_ptr(),
4425 rsb,
4426 csb,
4427 0.0,
4428 c.as_mut_ptr(),
4429 ldc as isize,
4430 1,
4431 );
4432 }
4433 }
4434}
4435
4436#[cfg(target_os = "macos")]
4439#[allow(clippy::too_many_arguments)]
4440pub(crate) fn sgemm_rm(
4441 m: usize,
4442 n: usize,
4443 k: usize,
4444 alpha: f32,
4445 a: &[f32],
4446 lda: usize,
4447 b_mat: &[f32],
4448 ldb: usize,
4449 b_rows_are_n: bool,
4450 c: &mut [f32],
4451 ldc: usize,
4452) {
4453 debug_assert!(a.len() >= (m - 1) * lda + k);
4454 debug_assert!(c.len() >= (m - 1) * ldc + n);
4455 #[cfg(target_arch = "aarch64")]
4460 if std::env::var("CMF_FORCE_NEON_GEMM")
4461 .map(|v| v == "1")
4462 .unwrap_or(false)
4463 {
4464 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4465 }
4466 unsafe {
4467 accel_blas::cblas_sgemm(
4468 101, 111, if b_rows_are_n { 112 } else { 111 },
4471 m as i32,
4472 n as i32,
4473 k as i32,
4474 alpha,
4475 a.as_ptr(),
4476 lda as i32,
4477 b_mat.as_ptr(),
4478 ldb as i32,
4479 0.0,
4480 c.as_mut_ptr(),
4481 ldc as i32,
4482 );
4483 }
4484}
4485
4486#[cfg(target_os = "macos")]
4493fn qmatmat_accel(
4494 q: &[u8],
4495 row_scale: &[f32],
4496 pre: &[std::borrow::Cow<'_, [f32]>],
4497 rows: usize,
4498 cols: usize,
4499 out: &mut [f32],
4500 pool: Option<&Pool>,
4501) {
4502 const TR: usize = 2048;
4507 let b = pre.len();
4508 thread_local! {
4509 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4510 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4511 }
4512 XPANEL.with(|xp| {
4513 WTILE.with(|wt| {
4514 let mut xpanel = xp.borrow_mut();
4515 xpanel.clear();
4516 for x in pre {
4517 xpanel.extend_from_slice(x);
4518 }
4519 let mut wtile = wt.borrow_mut();
4520 wtile.resize(TR * cols, 0.0);
4521 let mut r0 = 0usize;
4522 while r0 < rows {
4523 let tr = TR.min(rows - r0);
4524 let wt_addr = SendMut(wtile.as_mut_ptr());
4526 let run = |start: usize, end: usize| {
4527 for r in start..end {
4528 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4529 let s = row_scale[r0 + r];
4530 let dst =
4532 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4533 for (d, &v) in dst.iter_mut().zip(row) {
4534 *d = (v as i8) as f32 * s;
4535 }
4536 }
4537 };
4538 dispatch_rows(pool, tr, &run);
4539 unsafe {
4541 accel_blas::cblas_sgemm(
4542 101, 111, 112, b as i32,
4546 tr as i32,
4547 cols as i32,
4548 1.0,
4549 xpanel.as_ptr(),
4550 cols as i32,
4551 wtile.as_ptr(),
4552 cols as i32,
4553 0.0,
4554 out.as_mut_ptr().add(r0),
4555 rows as i32,
4556 );
4557 }
4558 r0 += tr;
4559 }
4560 })
4561 });
4562}
4563
4564fn qmatmat(
4565 q: &[u8],
4566 row_scale: &[f32],
4567 pre: &[std::borrow::Cow<'_, [f32]>],
4568 rows: usize,
4569 cols: usize,
4570 out: &mut [f32],
4571 pool: Option<&Pool>,
4572) {
4573 let b = pre.len();
4574 debug_assert_eq!(out.len(), b * rows);
4575 #[cfg(target_os = "macos")]
4580 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4581 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4582 return;
4583 }
4584 #[cfg(target_arch = "aarch64")]
4585 if sdot_enabled() {
4586 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4587 let out_addr = SendMut(out.as_mut_ptr());
4588 let blocked_ok = blocked_enabled();
4591 let use_i8mm = i8mm_enabled();
4592 if blocked_ok {
4593 let run = |start: usize, end: usize| {
4594 let mut o = start;
4595 while o < end {
4596 if o + 2 <= end {
4597 let r0 = &q[o * cols..(o + 1) * cols];
4598 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4599 let mut bi = 0usize;
4600 while bi + 4 <= acts.len() {
4601 let xs = [
4602 acts[bi].xq.as_slice(),
4603 acts[bi + 1].xq.as_slice(),
4604 acts[bi + 2].xq.as_slice(),
4605 acts[bi + 3].xq.as_slice(),
4606 ];
4607 let d = if use_i8mm {
4608 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4609 } else {
4610 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4611 };
4612 for (r, row) in [r0, r1].into_iter().enumerate() {
4613 for k in 0..4 {
4614 let act = &acts[bi + k];
4615 let mut v = d[r][k] as f32 * act.sx;
4616 for &(j, xv) in &act.outliers {
4617 v += (row[j] as i8) as f32 * xv;
4618 }
4619 unsafe {
4620 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4621 };
4622 }
4623 }
4624 bi += 4;
4625 }
4626 while bi < acts.len() {
4627 for (r, row) in [r0, r1].into_iter().enumerate() {
4628 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4629 unsafe { *out_addr.at(bi * rows + o + r) = v };
4630 }
4631 bi += 1;
4632 }
4633 o += 2;
4634 } else {
4635 let row = &q[o * cols..(o + 1) * cols];
4636 for (bi, act) in acts.iter().enumerate() {
4637 let v = row_dot_sdot(row, act) * row_scale[o];
4638 unsafe { *out_addr.at(bi * rows + o) = v };
4639 }
4640 o += 1;
4641 }
4642 }
4643 };
4644 dispatch_rows(pool, rows, &run);
4645 return;
4646 }
4647 let run = |start: usize, end: usize| {
4648 for o in start..end {
4649 let row = &q[o * cols..(o + 1) * cols];
4650 for (bi, act) in acts.iter().enumerate() {
4651 let v = row_dot_sdot(row, act) * row_scale[o];
4652 unsafe { *out_addr.at(bi * rows + o) = v };
4653 }
4654 }
4655 };
4656 dispatch_rows(pool, rows, &run);
4657 return;
4658 }
4659 #[cfg(target_arch = "x86_64")]
4664 if avx2_a8w8_enabled() {
4665 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4666 let out_addr = SendMut(out.as_mut_ptr());
4667 let blocked_ok = blocked_enabled();
4670 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4671 let run = |start: usize, end: usize| {
4672 let mut o = start;
4673 while o < end {
4674 if o + 2 <= end {
4675 let r0 = &q[o * cols..(o + 1) * cols];
4676 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4677 let mut bi = 0usize;
4678 while bi + 4 <= acts.len() {
4679 let xs = [
4680 acts[bi].xq.as_slice(),
4681 acts[bi + 1].xq.as_slice(),
4682 acts[bi + 2].xq.as_slice(),
4683 acts[bi + 3].xq.as_slice(),
4684 ];
4685 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4686 for (r, row) in [r0, r1].into_iter().enumerate() {
4687 for k in 0..4 {
4688 let act = &acts[bi + k];
4689 let mut v = d[r][k] as f32 * act.sx;
4690 for &(j, xv) in &act.outliers {
4691 v += (row[j] as i8) as f32 * xv;
4692 }
4693 unsafe {
4694 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4695 };
4696 }
4697 }
4698 bi += 4;
4699 }
4700 while bi < acts.len() {
4701 for (r, row) in [r0, r1].into_iter().enumerate() {
4702 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4703 unsafe { *out_addr.at(bi * rows + o + r) = v };
4704 }
4705 bi += 1;
4706 }
4707 o += 2;
4708 } else {
4709 let row = &q[o * cols..(o + 1) * cols];
4710 for (bi, act) in acts.iter().enumerate() {
4711 let v = row_dot_avx2(row, act) * row_scale[o];
4712 unsafe { *out_addr.at(bi * rows + o) = v };
4713 }
4714 o += 1;
4715 }
4716 }
4717 };
4718 dispatch_rows(pool, rows, &run);
4719 return;
4720 }
4721 let run = |start: usize, end: usize| {
4722 for o in start..end {
4723 let row = &q[o * cols..(o + 1) * cols];
4724 for (bi, act) in acts.iter().enumerate() {
4725 let v = row_dot_avx2(row, act) * row_scale[o];
4726 unsafe { *out_addr.at(bi * rows + o) = v };
4727 }
4728 }
4729 };
4730 dispatch_rows(pool, rows, &run);
4731 return;
4732 }
4733 let out_addr = SendMut(out.as_mut_ptr());
4734 let run = |start: usize, end: usize| {
4735 for o in start..end {
4736 let row = &q[o * cols..(o + 1) * cols];
4737 for (bi, x) in pre.iter().enumerate() {
4738 let mut acc = 0f32;
4739 for j in 0..cols {
4740 acc += (row[j] as i8) as f32 * x[j];
4741 }
4742 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4743 }
4744 }
4745 };
4746 dispatch_rows(pool, rows, &run);
4747}
4748
4749fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4752 match pool {
4753 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4754 _ => run(0, rows),
4755 }
4756}
4757
4758fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4760 let groups = rows * cols / GROUP_SIZE;
4761 bytes.split_at(groups * 16)
4762}
4763
4764#[inline]
4769fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4770 #[cfg(target_arch = "aarch64")]
4771 unsafe {
4772 return vbit_fill4_neon(data, buf);
4773 }
4774 #[cfg(target_arch = "x86_64")]
4775 if avx2_enabled() {
4776 return unsafe { vbit_fill4_avx2(data, buf) };
4777 }
4778 #[allow(unreachable_code)]
4779 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4780 let u = unpack8::<4>(&data[blk * 4..]);
4781 for k in 0..8 {
4782 chunk[k] = (u[k] - 7) as i8 as u8;
4783 }
4784 }
4785}
4786
4787#[cfg(target_arch = "aarch64")]
4788#[target_feature(enable = "neon")]
4789unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4790 unsafe {
4793 use core::arch::aarch64::*;
4794 let n = buf.len();
4795 let mask = vdupq_n_u8(0x0F);
4796 let seven = vdupq_n_s8(7);
4797 let mut g = 0usize;
4798 while g * 32 + 32 <= n {
4799 let b = vld1q_u8(data.as_ptr().add(g * 16));
4800 let hi = vshrq_n_u8::<4>(b);
4801 let lo = vandq_u8(b, mask);
4802 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4803 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4804 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4805 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4806 g += 1;
4807 }
4808 }
4809}
4810
4811#[cfg(target_arch = "x86_64")]
4812#[target_feature(enable = "avx2")]
4813unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4814 unsafe {
4816 use core::arch::x86_64::*;
4817 let n = buf.len();
4818 let mask = _mm_set1_epi8(0x0F);
4819 let seven = _mm256_set1_epi8(7);
4820 let mut g = 0usize;
4821 while g * 32 + 32 <= n {
4822 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4823 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4824 let lo = _mm_and_si128(b, mask);
4825 let z = _mm256_sub_epi8(
4826 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4827 seven,
4828 );
4829 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4830 g += 1;
4831 }
4832 }
4833}
4834
4835#[inline(always)]
4840fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4841 let mut acc = 0u64;
4842 for i in 0..B {
4843 acc = (acc << 8) | data[i] as u64;
4844 }
4845 let mask = (1u64 << B) - 1;
4846 let mut out = [0i32; 8];
4847 for (k, o) in out.iter_mut().enumerate() {
4848 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4849 }
4850 out
4851}
4852
4853#[allow(clippy::too_many_arguments)]
4859fn vbitmatvec(
4860 bytes: &[u8],
4861 offsets: &[usize],
4862 x: &[f32],
4863 rows: usize,
4864 cols: usize,
4865 out: &mut [f32],
4866 pool: Option<&Pool>,
4867) {
4868 debug_assert_eq!(out.len(), rows);
4869 debug_assert_eq!(offsets.len(), rows + 1);
4870
4871 if a8w8_enabled() {
4875 let act = split_act(x);
4876 let out_addr = SendMut(out.as_mut_ptr());
4877 let run = move |start: usize, end: usize| {
4878 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4879 };
4880 dispatch_rows(pool, rows, &run);
4881 return;
4882 }
4883
4884 let out_addr = SendMut(out.as_mut_ptr());
4885 let run = move |start: usize, end: usize| {
4886 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4887 };
4888 dispatch_rows(pool, rows, &run);
4889}
4890
4891#[allow(clippy::too_many_arguments)]
4895fn vbit_range_a8w8(
4896 bytes: &[u8],
4897 offsets: &[usize],
4898 x: &[f32],
4899 act: &SplitAct,
4900 rows: usize,
4901 cols: usize,
4902 out: SendMut,
4903 start: usize,
4904 end: usize,
4905) {
4906 let ng = cols / GROUP_SIZE;
4907 let bits = &bytes[..rows];
4908 let sc_off = rows;
4909 let row_dot = |r: usize| -> f32 {
4910 let b = bits[r] as usize;
4911 let l = (1i32 << (b - 1)) - 1;
4912 let mask = (1u64 << b) - 1;
4913 let data = &bytes[offsets[r]..offsets[r + 1]];
4914 if b == 8 {
4915 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4917 let mut dot = 0f32;
4918 for g in 0..ng {
4919 let so = (r * ng + g) * 2;
4920 let sgf = f16_to_f32(u16::from_le_bytes([
4921 bytes[sc_off + so],
4922 bytes[sc_off + so + 1],
4923 ]));
4924 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4925 let mut gd = 0f32;
4926 for &xv in xg.iter() {
4927 if nbits < 8 {
4928 acc = (acc << 8) | data[idx] as u64;
4929 idx += 1;
4930 nbits += 8;
4931 }
4932 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4933 nbits -= 8;
4934 gd += (u - l) as f32 * xv;
4935 }
4936 dot += gd * sgf;
4937 }
4938 return dot;
4939 }
4940 thread_local! {
4944 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4945 const { std::cell::RefCell::new(Vec::new()) };
4946 }
4947 #[inline(always)]
4948 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4949 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4950 let u = unpack8::<B>(&data[blk * B..]);
4951 for k in 0..8 {
4952 chunk[k] = (u[k] - l) as i8 as u8;
4953 }
4954 }
4955 }
4956 let _ = mask;
4957 VBIT_SCRATCH.with(|scratch| {
4958 let mut buf = scratch.borrow_mut();
4959 buf.resize(cols, 0);
4960 match b {
4961 3 => fill::<3>(data, l, &mut buf),
4962 4 => vbit_fill4(data, &mut buf),
4963 5 => fill::<5>(data, l, &mut buf),
4964 6 => fill::<6>(data, l, &mut buf),
4965 _ => unreachable!(),
4966 }
4967 let mut dot = 0f32;
4968 for g in 0..ng {
4969 let so = (r * ng + g) * 2;
4970 let s = f16_to_f32(u16::from_le_bytes([
4971 bytes[sc_off + so],
4972 bytes[sc_off + so + 1],
4973 ]));
4974 let d = dot_i8_i8(
4975 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4976 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4977 ) as f32
4978 * act.sx;
4979 dot += d * s;
4980 }
4981 for &(j, xv) in &act.outliers {
4982 let so = (r * ng + j / GROUP_SIZE) * 2;
4983 let s = f16_to_f32(u16::from_le_bytes([
4984 bytes[sc_off + so],
4985 bytes[sc_off + so + 1],
4986 ]));
4987 dot += (buf[j] as i8) as f32 * s * xv;
4989 }
4990 dot
4991 })
4992 };
4993 for r in start..end {
4994 unsafe { *out.at(r) = row_dot(r) };
4996 }
4997}
4998
4999#[allow(clippy::too_many_arguments)]
5001fn vbit_range_f32(
5002 bytes: &[u8],
5003 offsets: &[usize],
5004 x: &[f32],
5005 rows: usize,
5006 cols: usize,
5007 out: SendMut,
5008 start: usize,
5009 end: usize,
5010) {
5011 let ng = cols / GROUP_SIZE;
5012 let bits = &bytes[..rows];
5013 let sc_off = rows;
5014 #[inline(always)]
5018 fn dot_row<const B: usize>(
5019 data: &[u8],
5020 bytes: &[u8],
5021 sc_off: usize,
5022 r: usize,
5023 ng: usize,
5024 x: &[f32],
5025 ) -> f32 {
5026 let l = ((1i32 << (B - 1)) - 1) as f32;
5027 let gbytes = GROUP_SIZE * B / 8;
5028 let mut dot = 0f32;
5029 for g in 0..ng {
5030 let so = (r * ng + g) * 2;
5031 let s = f16_to_f32(u16::from_le_bytes([
5032 bytes[sc_off + so],
5033 bytes[sc_off + so + 1],
5034 ]));
5035 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5036 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5037 let mut gd = 0f32;
5038 for blk in 0..GROUP_SIZE / 8 {
5039 let u = unpack8::<B>(&gd0[blk * B..]);
5040 let xb = &xg[blk * 8..blk * 8 + 8];
5041 for k in 0..8 {
5042 gd += (u[k] as f32 - l) * xb[k];
5043 }
5044 }
5045 dot += gd * s;
5046 }
5047 dot
5048 }
5049 for r in start..end {
5050 let data = &bytes[offsets[r]..offsets[r + 1]];
5051 let v = match bits[r] {
5052 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5053 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5054 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5055 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5056 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5057 b => unreachable!("vbit bit-width {b} (validated at load)"),
5058 };
5059 unsafe { *out.at(r) = v };
5061 }
5062}
5063
5064#[allow(clippy::too_many_arguments)]
5069fn vbitmatvec2(
5070 bytes: &[u8],
5071 offsets: &[usize],
5072 x1: &[f32],
5073 x2: &[f32],
5074 rows: usize,
5075 cols: usize,
5076 o1: &mut [f32],
5077 o2: &mut [f32],
5078 pool: Option<&Pool>,
5079) {
5080 debug_assert_eq!(o1.len(), rows);
5081 debug_assert_eq!(o2.len(), rows);
5082
5083 if a8w8_enabled() {
5084 let a1 = split_act(x1);
5085 let a2 = split_act(x2);
5086 let p1 = SendMut(o1.as_mut_ptr());
5087 let p2 = SendMut(o2.as_mut_ptr());
5088 let run = move |start: usize, end: usize| {
5089 vbit_range2_a8w8(
5090 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5091 )
5092 };
5093 dispatch_rows(pool, rows, &run);
5094 return;
5095 }
5096
5097 let p1 = SendMut(o1.as_mut_ptr());
5098 let p2 = SendMut(o2.as_mut_ptr());
5099 let run = move |start: usize, end: usize| {
5100 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5101 };
5102 dispatch_rows(pool, rows, &run);
5103}
5104
5105#[allow(clippy::too_many_arguments)]
5109fn vbit_range2_a8w8(
5110 bytes: &[u8],
5111 offsets: &[usize],
5112 x1: &[f32],
5113 x2: &[f32],
5114 a1: &SplitAct,
5115 a2: &SplitAct,
5116 rows: usize,
5117 cols: usize,
5118 p1: SendMut,
5119 p2: SendMut,
5120 start: usize,
5121 end: usize,
5122) {
5123 let ng = cols / GROUP_SIZE;
5124 let bits = &bytes[..rows];
5125 let sc_off = rows;
5126 let row_dots = |r: usize| -> (f32, f32) {
5127 let b = bits[r] as usize;
5128 let l = (1i32 << (b - 1)) - 1;
5129 let data = &bytes[offsets[r]..offsets[r + 1]];
5130 if b == 8 {
5131 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5134 let (mut d1, mut d2) = (0f32, 0f32);
5135 for g in 0..ng {
5136 let so = (r * ng + g) * 2;
5137 let sgf = f16_to_f32(u16::from_le_bytes([
5138 bytes[sc_off + so],
5139 bytes[sc_off + so + 1],
5140 ]));
5141 let (mut g1, mut g2) = (0f32, 0f32);
5142 for k in 0..GROUP_SIZE {
5143 if nbits < 8 {
5144 acc = (acc << 8) | data[idx] as u64;
5145 idx += 1;
5146 nbits += 8;
5147 }
5148 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5149 nbits -= 8;
5150 let w = (u - l) as f32;
5151 g1 += w * x1[g * GROUP_SIZE + k];
5152 g2 += w * x2[g * GROUP_SIZE + k];
5153 }
5154 d1 += g1 * sgf;
5155 d2 += g2 * sgf;
5156 }
5157 return (d1, d2);
5158 }
5159 thread_local! {
5160 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5161 const { std::cell::RefCell::new(Vec::new()) };
5162 }
5163 #[inline(always)]
5164 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5165 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5166 let u = unpack8::<B>(&data[blk * B..]);
5167 for k in 0..8 {
5168 chunk[k] = (u[k] - l) as i8 as u8;
5169 }
5170 }
5171 }
5172 VBIT_SCRATCH2.with(|scratch| {
5173 let mut buf = scratch.borrow_mut();
5174 buf.resize(cols, 0);
5175 match b {
5176 3 => fill::<3>(data, l, &mut buf),
5177 4 => vbit_fill4(data, &mut buf),
5178 5 => fill::<5>(data, l, &mut buf),
5179 6 => fill::<6>(data, l, &mut buf),
5180 _ => unreachable!(),
5181 }
5182 let (mut d1, mut d2) = (0f32, 0f32);
5183 for g in 0..ng {
5184 let so = (r * ng + g) * 2;
5185 let s = f16_to_f32(u16::from_le_bytes([
5186 bytes[sc_off + so],
5187 bytes[sc_off + so + 1],
5188 ]));
5189 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5190 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5191 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5192 d1 += v1 * s;
5193 d2 += v2 * s;
5194 }
5195 for &(j, xv) in &a1.outliers {
5196 let so = (r * ng + j / GROUP_SIZE) * 2;
5197 let s = f16_to_f32(u16::from_le_bytes([
5198 bytes[sc_off + so],
5199 bytes[sc_off + so + 1],
5200 ]));
5201 d1 += (buf[j] as i8) as f32 * s * xv;
5202 }
5203 for &(j, xv) in &a2.outliers {
5204 let so = (r * ng + j / GROUP_SIZE) * 2;
5205 let s = f16_to_f32(u16::from_le_bytes([
5206 bytes[sc_off + so],
5207 bytes[sc_off + so + 1],
5208 ]));
5209 d2 += (buf[j] as i8) as f32 * s * xv;
5210 }
5211 (d1, d2)
5212 })
5213 };
5214 for r in start..end {
5215 let (v1, v2) = row_dots(r);
5216 unsafe {
5218 *p1.at(r) = v1;
5219 *p2.at(r) = v2;
5220 }
5221 }
5222}
5223
5224#[allow(clippy::too_many_arguments)]
5228fn vbit_range2_f32(
5229 bytes: &[u8],
5230 offsets: &[usize],
5231 x1: &[f32],
5232 x2: &[f32],
5233 rows: usize,
5234 cols: usize,
5235 p1: SendMut,
5236 p2: SendMut,
5237 start: usize,
5238 end: usize,
5239) {
5240 let ng = cols / GROUP_SIZE;
5241 let bits = &bytes[..rows];
5242 let sc_off = rows;
5243 #[inline(always)]
5244 #[allow(clippy::too_many_arguments)]
5245 fn dot_row2<const B: usize>(
5246 data: &[u8],
5247 bytes: &[u8],
5248 sc_off: usize,
5249 r: usize,
5250 ng: usize,
5251 x1: &[f32],
5252 x2: &[f32],
5253 ) -> (f32, f32) {
5254 let l = ((1i32 << (B - 1)) - 1) as f32;
5255 let gbytes = GROUP_SIZE * B / 8;
5256 let (mut d1, mut d2) = (0f32, 0f32);
5257 for g in 0..ng {
5258 let so = (r * ng + g) * 2;
5259 let s = f16_to_f32(u16::from_le_bytes([
5260 bytes[sc_off + so],
5261 bytes[sc_off + so + 1],
5262 ]));
5263 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5264 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5265 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5266 let (mut g1, mut g2) = (0f32, 0f32);
5267 for blk in 0..GROUP_SIZE / 8 {
5268 let u = unpack8::<B>(&gd0[blk * B..]);
5269 for k in 0..8 {
5270 let w = u[k] as f32 - l;
5271 g1 += w * x1g[blk * 8 + k];
5272 g2 += w * x2g[blk * 8 + k];
5273 }
5274 }
5275 d1 += g1 * s;
5276 d2 += g2 * s;
5277 }
5278 (d1, d2)
5279 }
5280 for r in start..end {
5281 let data = &bytes[offsets[r]..offsets[r + 1]];
5282 let (v1, v2) = match bits[r] {
5283 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5284 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5285 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5286 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5287 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5288 b => unreachable!("vbit bit-width {b} (validated at load)"),
5289 };
5290 unsafe {
5292 *p1.at(r) = v1;
5293 *p2.at(r) = v2;
5294 }
5295 }
5296}
5297
5298#[inline]
5305#[allow(unreachable_code)]
5306fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5307 #[cfg(target_arch = "aarch64")]
5308 unsafe {
5309 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5310 }
5311 #[cfg(target_arch = "x86_64")]
5312 unsafe {
5313 if vnni_tiles_enabled() {
5314 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5315 }
5316 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5317 }
5318 let mut acc = 0f32;
5319 for gi in 0..gpr {
5320 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5321 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5322 let mut d = 0i32;
5323 for (k, &b) in tile[2..].iter().enumerate() {
5324 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5325 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5326 }
5327 acc += d as f32 * s;
5328 }
5329 acc
5330}
5331
5332#[cfg(target_arch = "aarch64")]
5333#[target_feature(enable = "neon,dotprod")]
5334unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5335 unsafe {
5338 use core::arch::aarch64::*;
5339 use core::arch::asm;
5340 let lomask = vdupq_n_u8(0x0F);
5341 let eight = vdupq_n_s8(8);
5342 let mut acc = 0f32;
5343 for gi in 0..gpr {
5344 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5345 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5346 let b = vld1q_u8(t.add(2));
5347 let lo = vandq_u8(b, lomask);
5348 let hi = vshrq_n_u8::<4>(b);
5349 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5350 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5351 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5352 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5353 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5354 asm!(
5355 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5356 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5357 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5358 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5359 options(pure, nomem, nostack),
5360 );
5361 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5362 }
5363 acc
5364 }
5365}
5366
5367#[cfg(target_arch = "x86_64")]
5368#[target_feature(enable = "avx2")]
5369unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5370 unsafe {
5372 use core::arch::x86_64::*;
5373 let lomask = _mm_set1_epi8(0x0F);
5374 let eight = _mm256_set1_epi8(8);
5375 let ones = _mm256_set1_epi16(1);
5376 let mut acc = 0f32;
5377 for gi in 0..gpr {
5378 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5379 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5380 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5381 let lo = _mm_and_si128(b, lomask);
5382 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5383 let w = _mm256_sub_epi8(
5384 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5385 eight,
5386 );
5387 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5388 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5389 let d = _mm256_madd_epi16(p16, ones);
5390 let hi128 = _mm256_extracti128_si256::<1>(d);
5391 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5392 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5393 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5394 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5395 }
5396 acc
5397 }
5398}
5399
5400#[cfg(target_arch = "x86_64")]
5404#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5405unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5406 unsafe {
5408 use core::arch::x86_64::*;
5409 let lomask = _mm_set1_epi8(0x0F);
5410 let eight = _mm256_set1_epi8(8);
5411 let mut acc = 0f32;
5412 for gi in 0..gpr {
5413 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5414 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5415 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5416 let lo = _mm_and_si128(b, lomask);
5417 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5418 let w = _mm256_sub_epi8(
5419 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5420 eight,
5421 );
5422 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5423 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5424 acc += d as f32 * s;
5425 }
5426 acc
5427 }
5428}
5429
5430#[cfg(target_arch = "x86_64")]
5435#[target_feature(enable = "avx2,fma")]
5440unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5441 unsafe {
5443 use core::arch::x86_64::*;
5444 let lomask = _mm_set1_epi8(0x0F);
5445 let eight = _mm256_set1_epi8(8);
5446 let ones = _mm256_set1_epi16(1);
5447 let mut f0 = _mm256_setzero_ps();
5460 let mut f1 = _mm256_setzero_ps();
5461 let mut f2 = _mm256_setzero_ps();
5462 let mut f3 = _mm256_setzero_ps();
5463 for gi in 0..gpr {
5464 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5465 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5466 let sv = _mm256_set1_ps(s);
5467 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5468 let lo = _mm_and_si128(bb, lomask);
5469 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5470 let w = _mm256_sub_epi8(
5471 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5472 eight,
5473 );
5474 let aw = _mm256_abs_epi8(w);
5475 let off = gi * GROUP_SIZE;
5476 let dot = |xq: &[i8]| {
5477 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5478 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5479 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5480 };
5481 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5482 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5483 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5484 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5485 }
5486 [
5487 hsum256_ps(f0),
5488 hsum256_ps(f1),
5489 hsum256_ps(f2),
5490 hsum256_ps(f3),
5491 ]
5492 }
5493}
5494
5495#[cfg(target_arch = "x86_64")]
5498#[target_feature(enable = "avx2")]
5499#[inline]
5500unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5501 unsafe {
5503 use core::arch::x86_64::*;
5504 let hi = _mm256_extractf128_ps::<1>(v);
5505 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5506 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5507 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5508 _mm_cvtss_f32(s)
5509 }
5510}
5511
5512#[cfg(target_arch = "x86_64")]
5514#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5515unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5516 unsafe {
5518 use core::arch::x86_64::*;
5519 let lomask = _mm_set1_epi8(0x0F);
5520 let eight = _mm256_set1_epi8(8);
5521 let mut f0 = _mm256_setzero_ps();
5524 let mut f1 = _mm256_setzero_ps();
5525 let mut f2 = _mm256_setzero_ps();
5526 let mut f3 = _mm256_setzero_ps();
5527 for gi in 0..gpr {
5528 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5529 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5530 let sv = _mm256_set1_ps(s);
5531 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5532 let lo = _mm_and_si128(bb, lomask);
5533 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5534 let w = _mm256_sub_epi8(
5535 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5536 eight,
5537 );
5538 let aw = _mm256_abs_epi8(w);
5539 let off = gi * GROUP_SIZE;
5540 let dot = |xq: &[i8]| {
5541 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5542 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5543 _mm256_setzero_si256(),
5544 aw,
5545 _mm256_sign_epi8(x, w),
5546 ))
5547 };
5548 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5549 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5550 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5551 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5552 }
5553 let acc = [
5554 hsum256_ps(f0),
5555 hsum256_ps(f1),
5556 hsum256_ps(f2),
5557 hsum256_ps(f3),
5558 ];
5559 acc
5560 }
5561}
5562
5563#[cfg(target_arch = "aarch64")]
5568#[target_feature(enable = "neon,dotprod")]
5569unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5570 unsafe {
5572 use core::arch::aarch64::*;
5573 use core::arch::asm;
5574 let lomask = vdupq_n_u8(0x0F);
5575 let eight = vdupq_n_s8(8);
5576 let mut acc = [0f32; 4];
5577 for gi in 0..gpr {
5578 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5579 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5580 let b = vld1q_u8(t.add(2));
5581 let lo = vandq_u8(b, lomask);
5582 let hi = vshrq_n_u8::<4>(b);
5583 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5584 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5585 for (k, xq) in xs.iter().enumerate() {
5586 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5587 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5588 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5589 asm!(
5590 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5591 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5592 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5593 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5594 options(pure, nomem, nostack),
5595 );
5596 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5597 }
5598 }
5599 acc
5600 }
5601}
5602
5603#[inline]
5605fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5606 let gi = j / GROUP_SIZE;
5607 let k = j % GROUP_SIZE;
5608 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5609 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5610 let byte = tile[2 + k / 2];
5611 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5612 ((nib as i32 - 8) as f32, s)
5613}
5614
5615#[inline]
5618fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5619 let mut acc = 0f32;
5620 for gi in 0..gpr {
5621 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5622 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5623 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5624 let mut ga = 0f32;
5625 for (k, &b) in tile[2..].iter().enumerate() {
5626 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5627 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5628 }
5629 acc += ga * s;
5630 }
5631 acc
5632}
5633
5634struct Q4tpView<'a> {
5638 nib: &'a [u8],
5639 params: &'a [u8],
5640 codes: &'a [u8],
5641 stride: usize,
5642 zero_rung: bool,
5644}
5645
5646impl<'a> Q4tpView<'a> {
5647 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5648 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5649 Self {
5650 nib: &bytes[..params_off],
5651 params: &bytes[params_off..codes_off],
5652 codes: &bytes[codes_off..],
5653 stride,
5654 zero_rung: false,
5655 }
5656 }
5657
5658 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5660 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5661 Self {
5662 nib: &bytes[..params_off],
5663 params: &bytes[params_off..codes_off],
5664 codes: &bytes[codes_off..],
5665 stride,
5666 zero_rung: true,
5667 }
5668 }
5669
5670 #[inline]
5686 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5687 let tab = if self.zero_rung {
5688 q2tp_ladder(self.params, r)
5689 } else {
5690 q4tp_ladder(self.params, r)
5691 };
5692 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5693 let out = &mut out[..gpr];
5694 let mut chunks = out.chunks_exact_mut(8);
5695 let mut ci = 0usize;
5696 for c in &mut chunks {
5697 let w = u64::from(codes[ci])
5698 | u64::from(codes[ci + 1]) << 8
5699 | u64::from(codes[ci + 2]) << 16
5700 | u64::from(codes[ci + 3]) << 24
5701 | u64::from(codes[ci + 4]) << 32;
5702 for (k, o) in c.iter_mut().enumerate() {
5703 *o = tab[((w >> (5 * k)) & 31) as usize];
5704 }
5705 ci += 5;
5706 }
5707 let tail = &codes[ci..];
5710 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5711 *o = tab[q4tp_code(tail, k)];
5712 }
5713 }
5714}
5715
5716#[inline]
5717fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5718 #[cfg(target_arch = "aarch64")]
5719 unsafe {
5720 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5721 }
5722 #[cfg(target_arch = "x86_64")]
5723 unsafe {
5724 if vnni_tiles_enabled() {
5725 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5726 }
5727 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5728 }
5729 #[allow(unreachable_code)]
5730 {
5731 let mut acc = 0f32;
5732 for gi in 0..gpr {
5733 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5734 let s = scales[gi];
5735 let mut d = 0i32;
5736 for (k, &b) in tile.iter().enumerate() {
5737 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5738 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5739 }
5740 acc += d as f32 * s;
5741 }
5742 acc
5743 }
5744}
5745
5746#[cfg(target_arch = "aarch64")]
5749#[target_feature(enable = "neon,dotprod")]
5750unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5751 unsafe {
5754 use core::arch::aarch64::*;
5755 use core::arch::asm;
5756 let lomask = vdupq_n_u8(0x0F);
5757 let eight = vdupq_n_s8(8);
5758 let mut acc = 0f32;
5759 for gi in 0..gpr {
5760 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5761 let s = *scales.get_unchecked(gi);
5762 let b = vld1q_u8(t);
5763 let lo = vandq_u8(b, lomask);
5764 let hi = vshrq_n_u8::<4>(b);
5765 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5766 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5767 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5768 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5769 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5770 asm!(
5771 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5772 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5773 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5774 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5775 options(pure, nomem, nostack),
5776 );
5777 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5778 }
5779 acc
5780 }
5781}
5782
5783#[cfg(target_arch = "x86_64")]
5784#[target_feature(enable = "avx2")]
5785unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5786 unsafe {
5788 use core::arch::x86_64::*;
5789 let lomask = _mm_set1_epi8(0x0F);
5790 let eight = _mm256_set1_epi8(8);
5791 let ones = _mm256_set1_epi16(1);
5792 let mut acc = 0f32;
5793 for gi in 0..gpr {
5794 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5795 let s = *scales.get_unchecked(gi);
5796 let b = _mm_loadu_si128(t as *const __m128i);
5797 let lo = _mm_and_si128(b, lomask);
5798 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5799 let w = _mm256_sub_epi8(
5800 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5801 eight,
5802 );
5803 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5804 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5805 let d = _mm256_madd_epi16(p16, ones);
5806 let hi128 = _mm256_extracti128_si256::<1>(d);
5807 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5808 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5809 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5810 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5811 }
5812 acc
5813 }
5814}
5815
5816#[cfg(target_arch = "x86_64")]
5819#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5820unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5821 unsafe {
5823 use core::arch::x86_64::*;
5824 let lomask = _mm_set1_epi8(0x0F);
5825 let eight = _mm256_set1_epi8(8);
5826 let mut acc = 0f32;
5827 for gi in 0..gpr {
5828 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5829 let s = *scales.get_unchecked(gi);
5830 let b = _mm_loadu_si128(t as *const __m128i);
5831 let lo = _mm_and_si128(b, lomask);
5832 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5833 let w = _mm256_sub_epi8(
5834 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5835 eight,
5836 );
5837 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5838 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5839 }
5840 acc
5841 }
5842}
5843
5844#[inline]
5847fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5848 #[cfg(target_arch = "x86_64")]
5849 if avx2_enabled() {
5850 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5852 }
5853 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5854}
5855
5856#[inline]
5857fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5858 let mut acc = 0f32;
5859 for gi in 0..gpr {
5860 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5861 let s = scales[gi];
5862 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5863 let mut ga = 0f32;
5864 for (k, &b) in tile.iter().enumerate() {
5865 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5866 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5867 }
5868 acc += ga * s;
5869 }
5870 acc
5871}
5872
5873#[cfg(target_arch = "x86_64")]
5877#[target_feature(enable = "avx2")]
5878unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5879 unsafe {
5882 use core::arch::x86_64::*;
5883 let mask = _mm_set1_epi8(15);
5884 let eight = _mm_set1_epi8(8);
5885 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5886 let mut acc = 0.0f32;
5887 for gi in 0..gpr {
5888 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5889 let lo = _mm_and_si128(packed, mask);
5890 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5891 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5892 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5893 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5894 let a = _mm256_mul_ps(
5895 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5896 _mm256_loadu_ps(xp),
5897 );
5898 let b = _mm256_mul_ps(
5899 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5900 _mm256_loadu_ps(xp.add(8)),
5901 );
5902 let c = _mm256_mul_ps(
5903 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5904 _mm256_loadu_ps(xp.add(16)),
5905 );
5906 let d = _mm256_mul_ps(
5907 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5908 _mm256_loadu_ps(xp.add(24)),
5909 );
5910 let mut pairs = [0.0f32; 16];
5911 _mm256_storeu_ps(
5912 pairs.as_mut_ptr(),
5913 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5914 );
5915 _mm256_storeu_ps(
5916 pairs.as_mut_ptr().add(8),
5917 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5918 );
5919 let mut ga = 0.0f32;
5920 for v in pairs {
5921 ga += v;
5922 }
5923 acc += ga * scales[gi];
5924 }
5925 acc
5926 }
5927}
5928
5929#[inline]
5932fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5933 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5934 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5935 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5936 ((n as i32 - 8) as f32, scales[gi])
5937}
5938
5939fn q4tp_matvec(
5941 bytes: &[u8],
5942 x: &[f32],
5943 rows: usize,
5944 cols: usize,
5945 out: &mut [f32],
5946 pool: Option<&Pool>,
5947) {
5948 debug_assert_eq!(out.len(), rows);
5949 let gpr = cols / GROUP_SIZE;
5950 let v = Q4tpView::new(bytes, rows, cols);
5951 let out_addr = SendMut(out.as_mut_ptr());
5952 if a8w8_enabled() {
5953 let act = split_act(x);
5954 let run = |start: usize, end: usize| {
5955 with_krow(gpr, |sc| {
5957 for r in start..end {
5958 v.scales_into(r, gpr, sc);
5959 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5960 for &(j, xv) in &act.outliers {
5961 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5962 acc += w * s * xv;
5963 }
5964 unsafe { *out_addr.at(r) = acc };
5966 }
5967 })
5968 };
5969 dispatch_rows(pool, rows, &run);
5970 return;
5971 }
5972 let run = |start: usize, end: usize| {
5973 with_krow(gpr, |sc| {
5974 for r in start..end {
5975 v.scales_into(r, gpr, sc);
5976 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5978 }
5979 })
5980 };
5981 dispatch_rows(pool, rows, &run);
5982}
5983
5984#[allow(clippy::too_many_arguments)]
5987fn q4tp_matvec2(
5988 bytes: &[u8],
5989 x1: &[f32],
5990 x2: &[f32],
5991 rows: usize,
5992 cols: usize,
5993 o1: &mut [f32],
5994 o2: &mut [f32],
5995 pool: Option<&Pool>,
5996) {
5997 let gpr = cols / GROUP_SIZE;
5998 let v = Q4tpView::new(bytes, rows, cols);
5999 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6000 let run = |start: usize, end: usize| {
6001 let mut sc = vec![0f32; gpr];
6002 for r in start..end {
6003 v.scales_into(r, gpr, &mut sc);
6004 unsafe {
6006 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6007 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6008 }
6009 }
6010 };
6011 dispatch_rows(pool, rows, &run);
6012}
6013
6014#[inline]
6017fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6018 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6019 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6020 let c = (byte >> (2 * (k % 4))) & 3;
6021 (c as f32 - 1.5, scales[gi])
6022}
6023
6024#[cfg(target_arch = "x86_64")]
6025const Q2TP_DECODE_U32: [u32; 256] = {
6026 let mut tab = [0u32; 256];
6027 let mut b = 0usize;
6028 while b < 256 {
6029 tab[b] = ((b as u32) & 3)
6030 | ((((b as u32) >> 2) & 3) << 8)
6031 | ((((b as u32) >> 4) & 3) << 16)
6032 | ((((b as u32) >> 6) & 3) << 24);
6033 b += 1;
6034 }
6035 tab
6036};
6037
6038#[cfg(target_arch = "x86_64")]
6042#[target_feature(enable = "avx2")]
6043unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6044 use core::arch::x86_64::*;
6045 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6046 let codes = _mm256_setr_epi32(
6047 Q2TP_DECODE_U32[ch[0] as usize] as i32,
6048 Q2TP_DECODE_U32[ch[1] as usize] as i32,
6049 Q2TP_DECODE_U32[ch[2] as usize] as i32,
6050 Q2TP_DECODE_U32[ch[3] as usize] as i32,
6051 Q2TP_DECODE_U32[ch[4] as usize] as i32,
6052 Q2TP_DECODE_U32[ch[5] as usize] as i32,
6053 Q2TP_DECODE_U32[ch[6] as usize] as i32,
6054 Q2TP_DECODE_U32[ch[7] as usize] as i32,
6055 );
6056 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6057 let pair = _mm256_maddubs_epi16(codes, xv);
6058 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6059 let sum128 = _mm_add_epi32(
6060 _mm256_castsi256_si128(quad),
6061 _mm256_extracti128_si256(quad, 1),
6062 );
6063 let sum64 = _mm_hadd_epi32(sum128, sum128);
6064 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6065}
6066
6067#[cfg(target_arch = "x86_64")]
6068#[inline]
6069fn q2tp_avx2_enabled() -> bool {
6070 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6071 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6072}
6073
6074#[inline]
6081fn dot_q2tp_row_i8(
6082 chunks: &[u8],
6083 r: usize,
6084 gpr: usize,
6085 xq: &[i8],
6086 gsum: &[i32],
6087 scales: &[f32],
6088) -> f32 {
6089 let mut acc = 0f32;
6090 let base = r * gpr * Q2TP_CHUNK;
6091 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6092 let mut codes = [0i8; GROUP_SIZE];
6093 #[cfg(target_arch = "x86_64")]
6094 let avx2 = q2tp_avx2_enabled();
6099 for gi in 0..gpr {
6100 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
6101 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6102 #[cfg(target_arch = "aarch64")]
6103 let dot = unsafe {
6109 use core::arch::aarch64::*;
6110 let b = vld1_u8(ch.as_ptr());
6111 let three = vdup_n_u8(3);
6112 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
6113 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
6114 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
6115 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
6116 let x4 = vld4_s8(xg.as_ptr());
6117 let mut acc4 = vdupq_n_s32(0);
6118 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
6119 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
6120 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
6121 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
6122 vaddvq_s32(acc4)
6123 };
6124 #[cfg(target_arch = "x86_64")]
6125 let dot: i32 = if avx2 {
6126 unsafe { q2tp_code_dot_avx2(ch, xg) }
6129 } else {
6130 ch.iter()
6131 .enumerate()
6132 .map(|(k, &b)| {
6133 ((b & 3) as i32) * xg[k * 4] as i32
6134 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
6135 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
6136 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
6137 })
6138 .sum()
6139 };
6140 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6141 let dot: i32 = {
6142 for (k, &b) in ch.iter().enumerate() {
6143 codes[k * 4] = (b & 3) as i8;
6144 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6145 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6146 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6147 }
6148 codes
6149 .iter()
6150 .zip(xg)
6151 .map(|(&c, &x)| c as i32 * x as i32)
6152 .sum()
6153 };
6154 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6155 }
6156 acc
6157}
6158
6159fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6163 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6164}
6165
6166#[inline]
6170fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6171 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6172}
6173
6174#[inline]
6175fn q2tp_row_exact_center(
6176 chunks: &[u8],
6177 r: usize,
6178 gpr: usize,
6179 x: &[f32],
6180 scales: &[f32],
6181 center: f32,
6182) -> f32 {
6183 let mut acc = 0f32;
6184 for gi in 0..gpr {
6185 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6186 let s = scales[gi];
6187 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6188 let mut g = 0f32;
6189 for (k, &b) in ch.iter().enumerate() {
6190 g += ((b & 3) as f32 - center) * xb[k * 4]
6191 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6192 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6193 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6194 }
6195 acc += s * g;
6196 }
6197 acc
6198}
6199
6200fn q2tp_matvec(
6201 bytes: &[u8],
6202 x: &[f32],
6203 rows: usize,
6204 cols: usize,
6205 out: &mut [f32],
6206 pool: Option<&Pool>,
6207) {
6208 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6209}
6210
6211fn q2tp_affine_matvec(
6212 bytes: &[u8],
6213 x: &[f32],
6214 rows: usize,
6215 cols: usize,
6216 out: &mut [f32],
6217 pool: Option<&Pool>,
6218) {
6219 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6220}
6221
6222fn q2tp_matvec_mode(
6223 bytes: &[u8],
6224 x: &[f32],
6225 rows: usize,
6226 cols: usize,
6227 out: &mut [f32],
6228 pool: Option<&Pool>,
6229 affine: bool,
6230) {
6231 debug_assert_eq!(out.len(), rows);
6232 let gpr = cols / GROUP_SIZE;
6233 let v = Q4tpView::new_q2(bytes, rows, cols);
6234 let out_addr = SendMut(out.as_mut_ptr());
6235 if !affine && a8w8_enabled() {
6240 let act = split_act(x);
6241 let gsum = q1_group_sums(&act.xq, gpr);
6242 let (act, gsum) = (&act, &gsum);
6243 let run = move |start: usize, end: usize| {
6244 with_krow(gpr, |sc| {
6245 for r in start..end {
6246 v.scales_into(r, gpr, sc);
6247 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6248 for &(j, xv) in &act.outliers {
6249 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6250 acc += w * s * xv;
6251 }
6252 unsafe { *out_addr.at(r) = acc };
6254 }
6255 })
6256 };
6257 dispatch_rows(pool, rows, &run);
6258 return;
6259 }
6260 let run = |start: usize, end: usize| {
6261 with_krow(gpr, |sc| {
6262 for r in start..end {
6263 v.scales_into(r, gpr, sc);
6264 unsafe {
6266 *out_addr.at(r) = if affine {
6267 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6268 } else {
6269 q2tp_row_exact(v.nib, r, gpr, x, sc)
6270 }
6271 };
6272 }
6273 })
6274 };
6275 dispatch_rows(pool, rows, &run);
6276}
6277
6278#[allow(clippy::too_many_arguments)]
6280fn q2tp_matvec2(
6281 bytes: &[u8],
6282 x1: &[f32],
6283 x2: &[f32],
6284 rows: usize,
6285 cols: usize,
6286 o1: &mut [f32],
6287 o2: &mut [f32],
6288 pool: Option<&Pool>,
6289) {
6290 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6291}
6292
6293#[allow(clippy::too_many_arguments)]
6294fn q2tp_affine_matvec2(
6295 bytes: &[u8],
6296 x1: &[f32],
6297 x2: &[f32],
6298 rows: usize,
6299 cols: usize,
6300 o1: &mut [f32],
6301 o2: &mut [f32],
6302 pool: Option<&Pool>,
6303) {
6304 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6305}
6306
6307#[allow(clippy::too_many_arguments)]
6308fn q2tp_matvec2_mode(
6309 bytes: &[u8],
6310 x1: &[f32],
6311 x2: &[f32],
6312 rows: usize,
6313 cols: usize,
6314 o1: &mut [f32],
6315 o2: &mut [f32],
6316 pool: Option<&Pool>,
6317 affine: bool,
6318) {
6319 let gpr = cols / GROUP_SIZE;
6320 let v = Q4tpView::new_q2(bytes, rows, cols);
6321 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6322 let run = |start: usize, end: usize| {
6323 let mut sc = vec![0f32; gpr];
6324 for r in start..end {
6325 v.scales_into(r, gpr, &mut sc);
6326 unsafe {
6328 *p1.at(r) = if affine {
6329 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6330 } else {
6331 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6332 };
6333 *p2.at(r) = if affine {
6334 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6335 } else {
6336 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6337 };
6338 }
6339 }
6340 };
6341 dispatch_rows(pool, rows, &run);
6342}
6343
6344pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6351 let gpr = cols / GROUP_SIZE;
6356 let v = Q4tpView::new_q2(bytes, rows, cols);
6357 with_krow(gpr, |sc| {
6358 for r in 0..rows {
6359 v.scales_into(r, gpr, sc);
6360 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6361 }
6362 });
6363}
6364
6365pub fn q2tp_affine_matvec_for_test(
6368 bytes: &[u8],
6369 x: &[f32],
6370 rows: usize,
6371 cols: usize,
6372 out: &mut [f32],
6373) {
6374 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6375}
6376
6377pub fn q2tp_matmat_for_test(
6378 bytes: &[u8],
6379 xs_all: &[f32],
6380 b: usize,
6381 rows: usize,
6382 cols: usize,
6383 out: &mut [f32],
6384) {
6385 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6386}
6387
6388fn q2tp_matmat(
6389 bytes: &[u8],
6390 xs_all: &[f32],
6391 b: usize,
6392 rows: usize,
6393 cols: usize,
6394 out: &mut [f32],
6395 pool: Option<&Pool>,
6396) {
6397 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6398}
6399
6400fn q2tp_affine_matmat(
6401 bytes: &[u8],
6402 xs_all: &[f32],
6403 b: usize,
6404 rows: usize,
6405 cols: usize,
6406 out: &mut [f32],
6407 pool: Option<&Pool>,
6408) {
6409 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6410}
6411
6412fn q2tp_matmat_mode(
6413 bytes: &[u8],
6414 xs_all: &[f32],
6415 b: usize,
6416 rows: usize,
6417 cols: usize,
6418 out: &mut [f32],
6419 pool: Option<&Pool>,
6420 affine: bool,
6421) {
6422 debug_assert_eq!(out.len(), b * rows);
6423 let gpr = cols / GROUP_SIZE;
6424 let v = Q4tpView::new_q2(bytes, rows, cols);
6425 let out_addr = SendMut(out.as_mut_ptr());
6426 let run = |start: usize, end: usize| {
6427 let mut sc = vec![0f32; gpr];
6428 for r in start..end {
6429 v.scales_into(r, gpr, &mut sc);
6430 for bi in 0..b {
6431 let x = &xs_all[bi * cols..(bi + 1) * cols];
6432 unsafe {
6434 *out_addr.at(bi * rows + r) = if affine {
6435 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6436 } else {
6437 q2tp_row_exact(v.nib, r, gpr, x, &sc)
6438 }
6439 };
6440 }
6441 }
6442 };
6443 dispatch_rows(pool, rows, &run);
6444}
6445
6446#[cfg(target_arch = "aarch64")]
6454#[target_feature(enable = "neon,dotprod")]
6455unsafe fn dot_q4tp_row_1x4_sdot_v1(
6456 nib: &[u8],
6457 r: usize,
6458 gpr: usize,
6459 xs: [&[i8]; 4],
6460 scales: &[f32],
6461) -> [f32; 4] {
6462 unsafe {
6463 use core::arch::aarch64::*;
6464 use core::arch::asm;
6465 let lomask = vdupq_n_u8(0x0F);
6466 let eight = vdupq_n_s8(8);
6467 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6468 for gi in 0..gpr {
6469 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6470 let s = *scales.get_unchecked(gi);
6471 let bb = vld1q_u8(t);
6472 let lo = vandq_u8(bb, lomask);
6473 let hi = vshrq_n_u8::<4>(bb);
6474 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6475 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6476 let mut d = [0f32; 4];
6477 for (k, dk) in d.iter_mut().enumerate() {
6478 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6479 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6480 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6481 asm!(
6482 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6483 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6484 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6485 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6486 options(pure, nomem, nostack),
6487 );
6488 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6489 }
6490 f0 += d[0];
6491 f1 += d[1];
6492 f2 += d[2];
6493 f3 += d[3];
6494 }
6495 [f0, f1, f2, f3]
6496 }
6497}
6498
6499#[allow(dead_code)]
6506static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6507
6508#[cfg(test)]
6511static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6512
6513#[cfg(target_arch = "x86_64")]
6519fn q4tp_blocked_x86() -> bool {
6520 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6521 1 => false,
6522 2 => avx512vnni_enabled(),
6527 _ => blocked_enabled() && avx512vnni_enabled(),
6531 }
6532}
6533
6534#[cfg(target_arch = "aarch64")]
6536#[allow(dead_code)]
6537fn q4tp_v1() -> bool {
6538 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6539 1 => true,
6540 2 => false,
6541 _ => {
6542 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6543 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6544 }
6545 }
6546}
6547
6548#[cfg(target_arch = "x86_64")]
6560#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6561unsafe fn dot_q4tp_2x8_avx512(
6562 nib: &[u8],
6563 r0: usize,
6564 gpr: usize,
6565 xs: [&[i8]; 8],
6566 sc0: &[f32],
6567 sc1: &[f32],
6568) -> [[f32; 8]; 2] {
6569 unsafe {
6572 use core::arch::x86_64::*;
6573 let lomask = _mm256_set1_epi8(0x0F);
6574 let eight = _mm256_set1_epi8(8);
6575 let zero = _mm512_setzero_si512();
6576 let mut v0 = [_mm512_setzero_ps(); 8];
6577 let mut v1 = [_mm512_setzero_ps(); 8];
6578 let pairs = gpr / 2;
6579 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6580 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6581 let bb = _mm256_loadu_si256(t as *const __m256i);
6582 let lo = _mm256_and_si256(bb, lomask);
6583 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6584 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6585 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6586 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6587 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6588 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6589 };
6590 for gp in 0..pairs {
6591 let gi = gp * 2;
6592 let (wa0, neg0) = unpack(r0, gi);
6593 let (wa1, neg1) = unpack(r0 + 1, gi);
6594 let off = gi * GROUP_SIZE;
6595 let sv = |sc: &[f32]| {
6596 _mm512_insertf32x8::<1>(
6597 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6598 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6599 )
6600 };
6601 let s0 = sv(sc0);
6602 let s1 = sv(sc1);
6603 for k in 0..8 {
6604 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6605 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6606 zero,
6607 wa0,
6608 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6609 ));
6610 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6611 zero,
6612 wa1,
6613 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6614 ));
6615 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6616 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6617 }
6618 }
6619 let mut acc = [[0f32; 8]; 2];
6620 for k in 0..8 {
6621 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6622 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6623 }
6624 if gpr % 2 == 1 {
6625 let off = (gpr - 1) * GROUP_SIZE;
6626 for j in off..off + GROUP_SIZE {
6627 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6628 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6629 for k in 0..8 {
6630 let x = *xs[k].get_unchecked(j) as f32;
6631 acc[0][k] += w0 * sa * x;
6632 acc[1][k] += w1 * sb * x;
6633 }
6634 }
6635 }
6636 acc
6637 }
6638}
6639
6640#[cfg(target_arch = "x86_64")]
6645#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6646unsafe fn dot_q4tp_row_1x8_avx512(
6647 nib: &[u8],
6648 r: usize,
6649 gpr: usize,
6650 xs: [&[i8]; 8],
6651 scales: &[f32],
6652) -> [f32; 8] {
6653 unsafe {
6655 use core::arch::x86_64::*;
6656 let lomask = _mm256_set1_epi8(0x0F);
6657 let eight = _mm256_set1_epi8(8);
6658 let zero = _mm512_setzero_si512();
6659 let (mut v0, mut v1, mut v2, mut v3) = (
6660 _mm512_setzero_ps(),
6661 _mm512_setzero_ps(),
6662 _mm512_setzero_ps(),
6663 _mm512_setzero_ps(),
6664 );
6665 let (mut v4, mut v5, mut v6, mut v7) = (
6666 _mm512_setzero_ps(),
6667 _mm512_setzero_ps(),
6668 _mm512_setzero_ps(),
6669 _mm512_setzero_ps(),
6670 );
6671 let pairs = gpr / 2;
6672 for gp in 0..pairs {
6673 let gi = gp * 2;
6674 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6675 let bb = _mm256_loadu_si256(t as *const __m256i);
6676 let lo = _mm256_and_si256(bb, lomask);
6677 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6678 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6683 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6684 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6685 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6686 let wabs = _mm512_abs_epi8(w);
6687 let neg = _mm512_movepi8_mask(w);
6688 let off = gi * GROUP_SIZE;
6689 let sv = _mm512_insertf32x8::<1>(
6690 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6691 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6692 );
6693 let dot = |x: &[i8]| -> __m512 {
6694 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6695 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6696 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6697 };
6698 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6699 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6700 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6701 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6702 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6703 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6704 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6705 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6706 }
6707 let mut acc = [
6708 _mm512_reduce_add_ps(v0),
6709 _mm512_reduce_add_ps(v1),
6710 _mm512_reduce_add_ps(v2),
6711 _mm512_reduce_add_ps(v3),
6712 _mm512_reduce_add_ps(v4),
6713 _mm512_reduce_add_ps(v5),
6714 _mm512_reduce_add_ps(v6),
6715 _mm512_reduce_add_ps(v7),
6716 ];
6717 if gpr % 2 == 1 {
6720 let off = (gpr - 1) * GROUP_SIZE;
6721 for j in off..off + GROUP_SIZE {
6722 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6723 let ws = w * s;
6724 for k in 0..8 {
6725 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6726 }
6727 }
6728 }
6729 acc
6730 }
6731}
6732
6733#[cfg(target_arch = "x86_64")]
6746#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6747unsafe fn dot_q4tp_row_1x4_avx512(
6748 nib: &[u8],
6749 r: usize,
6750 gpr: usize,
6751 xs: [&[i8]; 4],
6752 scales: &[f32],
6753) -> [f32; 4] {
6754 unsafe {
6756 use core::arch::x86_64::*;
6757 let lomask = _mm256_set1_epi8(0x0F);
6758 let eight = _mm256_set1_epi8(8);
6759 let zero = _mm512_setzero_si512();
6760 let (mut v0, mut v1, mut v2, mut v3) = (
6761 _mm512_setzero_ps(),
6762 _mm512_setzero_ps(),
6763 _mm512_setzero_ps(),
6764 _mm512_setzero_ps(),
6765 );
6766 let pairs = gpr / 2;
6767 for gp in 0..pairs {
6768 let gi = gp * 2;
6769 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6770 let bb = _mm256_loadu_si256(t as *const __m256i);
6771 let lo = _mm256_and_si256(bb, lomask);
6772 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6773 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6778 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6779 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6780 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6781 let wabs = _mm512_abs_epi8(w);
6782 let neg = _mm512_movepi8_mask(w);
6783 let off = gi * GROUP_SIZE;
6784 let sv = _mm512_insertf32x8::<1>(
6785 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6786 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6787 );
6788 let dot = |x: &[i8]| -> __m512 {
6789 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6790 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6791 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6792 };
6793 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6794 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6795 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6796 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6797 }
6798 let mut acc = [
6799 _mm512_reduce_add_ps(v0),
6800 _mm512_reduce_add_ps(v1),
6801 _mm512_reduce_add_ps(v2),
6802 _mm512_reduce_add_ps(v3),
6803 ];
6804 if gpr % 2 == 1 {
6807 let off = (gpr - 1) * GROUP_SIZE;
6808 for j in off..off + GROUP_SIZE {
6809 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6810 let ws = w * s;
6811 for k in 0..4 {
6812 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6813 }
6814 }
6815 }
6816 acc
6817 }
6818}
6819
6820#[cfg(target_arch = "aarch64")]
6824#[target_feature(enable = "neon,dotprod")]
6825unsafe fn dot_q4tp_row_1x4_sdot(
6826 nib: &[u8],
6827 r: usize,
6828 gpr: usize,
6829 xs: [&[i8]; 4],
6830 scales: &[f32],
6831) -> [f32; 4] {
6832 unsafe {
6834 use core::arch::aarch64::*;
6835 use core::arch::asm;
6836 let lomask = vdupq_n_u8(0x0F);
6837 let eight = vdupq_n_s8(8);
6838 let (mut v0, mut v1, mut v2, mut v3) = (
6854 vdupq_n_f32(0.0),
6855 vdupq_n_f32(0.0),
6856 vdupq_n_f32(0.0),
6857 vdupq_n_f32(0.0),
6858 );
6859 for gi in 0..gpr {
6860 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6861 let s = *scales.get_unchecked(gi);
6862 let bb = vld1q_u8(t);
6863 let lo = vandq_u8(bb, lomask);
6864 let hi = vshrq_n_u8::<4>(bb);
6865 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6866 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6867 let off = gi * GROUP_SIZE;
6868 let dot4 = |x: &[i8]| -> int32x4_t {
6869 let x0 = vld1q_s8(x.as_ptr().add(off));
6870 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6871 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6872 asm!(
6873 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6874 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6875 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6876 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6877 options(pure, nomem, nostack),
6878 );
6879 vaddq_s32(a0, a1)
6880 };
6881 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6882 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6883 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6884 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6885 }
6886 [
6887 vaddvq_f32(v0),
6888 vaddvq_f32(v1),
6889 vaddvq_f32(v2),
6890 vaddvq_f32(v3),
6891 ]
6892 }
6893}
6894
6895fn q4tp_matmat(
6903 bytes: &[u8],
6904 xs_all: &[f32],
6905 b: usize,
6906 rows: usize,
6907 cols: usize,
6908 out: &mut [f32],
6909 pool: Option<&Pool>,
6910) {
6911 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6912}
6913
6914#[allow(clippy::too_many_arguments)]
6918fn q4tp_matmat_with(
6919 bytes: &[u8],
6920 xs_all: &[f32],
6921 b: usize,
6922 rows: usize,
6923 cols: usize,
6924 out: &mut [f32],
6925 pool: Option<&Pool>,
6926 exact: bool,
6927) {
6928 debug_assert_eq!(out.len(), b * rows);
6929 let gpr = cols / GROUP_SIZE;
6930 let v = Q4tpView::new(bytes, rows, cols);
6931
6932 #[cfg(target_os = "macos")]
6936 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6937 dequant_matmat_accel(
6938 &|r, dst| {
6939 let mut sc = [0f32; 32];
6940 let mut scv;
6941 let s: &[f32] = if gpr <= 32 {
6942 v.scales_into(r, gpr, &mut sc);
6943 &sc[..gpr]
6944 } else {
6945 scv = vec![0f32; gpr];
6946 v.scales_into(r, gpr, &mut scv);
6947 &scv
6948 };
6949 for gi in 0..gpr {
6950 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6951 for (k, &bb) in tile.iter().enumerate() {
6952 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6953 dst[gi * GROUP_SIZE + k * 2 + 1] =
6954 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6955 }
6956 }
6957 },
6958 xs_all,
6959 b,
6960 rows,
6961 cols,
6962 out,
6963 pool,
6964 );
6965 return;
6966 }
6967
6968 let out_addr = SendMut(out.as_mut_ptr());
6969 if a8w8_enabled() {
6970 let acts: Vec<SplitAct> = (0..b)
6971 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6972 .collect();
6973 let acts = &acts;
6974 #[cfg(target_arch = "aarch64")]
6981 let blocked_ok = sdot_enabled() && blocked_enabled();
6982 #[cfg(target_arch = "x86_64")]
6989 let blocked_ok = q4tp_blocked_x86() && !exact;
6990 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6991 let blocked_ok = {
6992 let _ = exact;
6993 false
6994 };
6995 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7005 .ok()
7006 .and_then(|v| v.parse().ok())
7007 .filter(|v| *v > 0)
7008 .unwrap_or(256);
7009 let run = |start: usize, end: usize| {
7010 for abase in (0..acts.len()).step_by(panel_cols) {
7011 let alen = (acts.len() - abase).min(panel_cols);
7012 let mut sc = vec![0f32; gpr];
7013 #[cfg(target_arch = "x86_64")]
7014 let mut r_lo = start;
7015 #[cfg(target_arch = "x86_64")]
7016 if blocked_ok && alen >= 8 {
7017 let mut sc1 = vec![0f32; gpr];
7018 while r_lo + 2 <= end {
7019 v.scales_into(r_lo, gpr, &mut sc);
7020 v.scales_into(r_lo + 1, gpr, &mut sc1);
7021 let mut bi = 0usize;
7022 while bi + 8 <= alen {
7023 let xs = [
7024 acts[abase + bi].xq.as_slice(),
7025 acts[abase + bi + 1].xq.as_slice(),
7026 acts[abase + bi + 2].xq.as_slice(),
7027 acts[abase + bi + 3].xq.as_slice(),
7028 acts[abase + bi + 4].xq.as_slice(),
7029 acts[abase + bi + 5].xq.as_slice(),
7030 acts[abase + bi + 6].xq.as_slice(),
7031 acts[abase + bi + 7].xq.as_slice(),
7032 ];
7033 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7034 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7035 for k in 0..8 {
7036 let act = &acts[abase + bi + k];
7037 let mut acc = dr[k] * act.sx;
7038 for &(j, xv) in &act.outliers {
7039 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7040 acc += w * s * xv;
7041 }
7042 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7044 }
7045 }
7046 bi += 8;
7047 }
7048 for row in [r_lo, r_lo + 1] {
7051 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7052 for b2 in bi..alen {
7053 let act = &acts[abase + b2];
7054 let xs4 = [
7055 act.xq.as_slice(),
7056 act.xq.as_slice(),
7057 act.xq.as_slice(),
7058 act.xq.as_slice(),
7059 ];
7060 let d =
7061 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7062 let mut acc = d[0] * act.sx;
7063 for &(j, xv) in &act.outliers {
7064 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7065 acc += w * s * xv;
7066 }
7067 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7069 }
7070 }
7071 r_lo += 2;
7072 }
7073 }
7074 #[cfg(target_arch = "x86_64")]
7075 let row_start = r_lo;
7076 #[cfg(not(target_arch = "x86_64"))]
7077 let row_start = start;
7078 for r in row_start..end {
7079 v.scales_into(r, gpr, &mut sc);
7080 let mut bi = 0usize;
7081 #[cfg(target_arch = "x86_64")]
7082 if blocked_ok {
7083 while bi + 8 <= alen {
7084 let xs = [
7085 acts[abase + bi].xq.as_slice(),
7086 acts[abase + bi + 1].xq.as_slice(),
7087 acts[abase + bi + 2].xq.as_slice(),
7088 acts[abase + bi + 3].xq.as_slice(),
7089 acts[abase + bi + 4].xq.as_slice(),
7090 acts[abase + bi + 5].xq.as_slice(),
7091 acts[abase + bi + 6].xq.as_slice(),
7092 acts[abase + bi + 7].xq.as_slice(),
7093 ];
7094 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
7095 for k in 0..8 {
7096 let act = &acts[abase + bi + k];
7097 let mut acc = d[k] * act.sx;
7098 for &(j, xv) in &act.outliers {
7099 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7100 acc += w * s * xv;
7101 }
7102 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7104 }
7105 bi += 8;
7106 }
7107 while bi + 4 <= alen {
7108 let xs = [
7109 acts[abase + bi].xq.as_slice(),
7110 acts[abase + bi + 1].xq.as_slice(),
7111 acts[abase + bi + 2].xq.as_slice(),
7112 acts[abase + bi + 3].xq.as_slice(),
7113 ];
7114 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
7115 for k in 0..4 {
7116 let act = &acts[abase + bi + k];
7117 let mut acc = d[k] * act.sx;
7118 for &(j, xv) in &act.outliers {
7119 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7120 acc += w * s * xv;
7121 }
7122 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7124 }
7125 bi += 4;
7126 }
7127 }
7128 #[cfg(target_arch = "aarch64")]
7129 if blocked_ok {
7130 while bi + 4 <= alen {
7131 let xs = [
7132 acts[abase + bi].xq.as_slice(),
7133 acts[abase + bi + 1].xq.as_slice(),
7134 acts[abase + bi + 2].xq.as_slice(),
7135 acts[abase + bi + 3].xq.as_slice(),
7136 ];
7137 let d = unsafe {
7138 if exact || q4tp_v1() {
7139 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
7140 } else {
7141 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7142 }
7143 };
7144 for k in 0..4 {
7145 let act = &acts[abase + bi + k];
7146 let mut acc = d[k] * act.sx;
7147 for &(j, xv) in &act.outliers {
7148 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7149 acc += w * s * xv;
7150 }
7151 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7153 }
7154 bi += 4;
7155 }
7156 }
7157 let _ = blocked_ok;
7158 while bi < alen {
7159 let act = &acts[abase + bi];
7160 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7161 for &(j, xv) in &act.outliers {
7162 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7163 acc += w * s * xv;
7164 }
7165 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7167 bi += 1;
7168 }
7169 }
7170 }
7171 };
7172 dispatch_rows(pool, rows, &run);
7173 return;
7174 }
7175
7176 let run = |start: usize, end: usize| {
7177 let mut sc = vec![0f32; gpr];
7178 for r in start..end {
7179 v.scales_into(r, gpr, &mut sc);
7180 for bi in 0..b {
7181 let x = &xs_all[bi * cols..(bi + 1) * cols];
7182 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7184 }
7185 }
7186 };
7187 dispatch_rows(pool, rows, &run);
7188}
7189
7190fn q4t_matvec(
7192 bytes: &[u8],
7193 x: &[f32],
7194 rows: usize,
7195 cols: usize,
7196 out: &mut [f32],
7197 pool: Option<&Pool>,
7198) {
7199 debug_assert_eq!(out.len(), rows);
7200 let gpr = cols / GROUP_SIZE;
7201 let out_addr = SendMut(out.as_mut_ptr());
7202 if a8w8_enabled() {
7203 let act = split_act(x);
7204 let run = move |start: usize, end: usize| {
7205 for r in start..end {
7206 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7207 for &(j, xv) in &act.outliers {
7208 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7209 acc += w * s * xv;
7210 }
7211 unsafe { *out_addr.at(r) = acc };
7213 }
7214 };
7215 dispatch_rows(pool, rows, &run);
7216 return;
7217 }
7218 let run = move |start: usize, end: usize| {
7219 for r in start..end {
7220 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7222 }
7223 };
7224 dispatch_rows(pool, rows, &run);
7225}
7226
7227#[allow(clippy::too_many_arguments)]
7229fn q4t_matvec2(
7230 bytes: &[u8],
7231 x1: &[f32],
7232 x2: &[f32],
7233 rows: usize,
7234 cols: usize,
7235 o1: &mut [f32],
7236 o2: &mut [f32],
7237 pool: Option<&Pool>,
7238) {
7239 let gpr = cols / GROUP_SIZE;
7240 let p1 = SendMut(o1.as_mut_ptr());
7241 let p2 = SendMut(o2.as_mut_ptr());
7242 if a8w8_enabled() {
7243 let a1 = split_act(x1);
7244 let a2 = split_act(x2);
7245 let run = move |start: usize, end: usize| {
7246 for r in start..end {
7247 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7248 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7249 for &(j, xv) in &a1.outliers {
7250 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7251 v1 += w * s * xv;
7252 }
7253 for &(j, xv) in &a2.outliers {
7254 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7255 v2 += w * s * xv;
7256 }
7257 unsafe {
7259 *p1.at(r) = v1;
7260 *p2.at(r) = v2;
7261 }
7262 }
7263 };
7264 dispatch_rows(pool, rows, &run);
7265 return;
7266 }
7267 let run = move |start: usize, end: usize| {
7268 for r in start..end {
7269 unsafe {
7271 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7272 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7273 }
7274 }
7275 };
7276 dispatch_rows(pool, rows, &run);
7277}
7278
7279#[allow(clippy::too_many_arguments)]
7281#[cfg(target_os = "macos")]
7287fn dequant_matmat_accel(
7288 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7289 xs_all: &[f32],
7290 b: usize,
7291 rows: usize,
7292 cols: usize,
7293 out: &mut [f32],
7294 pool: Option<&Pool>,
7295) {
7296 const TR: usize = 2048;
7297 thread_local! {
7298 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7299 }
7300 WTILE.with(|wt| {
7301 let mut wtile = wt.borrow_mut();
7302 wtile.resize(TR * cols, 0.0);
7303 let mut r0 = 0usize;
7304 while r0 < rows {
7305 let tr = TR.min(rows - r0);
7306 let wt_addr = SendMut(wtile.as_mut_ptr());
7307 let run = |start: usize, end: usize| {
7308 for r in start..end {
7309 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7311 dequant_row(r0 + r, dst);
7312 }
7313 };
7314 dispatch_rows(pool, tr, &run);
7315 unsafe {
7316 accel_blas::cblas_sgemm(
7317 101, 111, 112, b as i32,
7321 tr as i32,
7322 cols as i32,
7323 1.0,
7324 xs_all.as_ptr(),
7325 cols as i32,
7326 wtile.as_ptr(),
7327 cols as i32,
7328 0.0,
7329 out.as_mut_ptr().add(r0),
7330 rows as i32,
7331 );
7332 }
7333 r0 += tr;
7334 }
7335 });
7336}
7337
7338fn q4t_matmat(
7339 bytes: &[u8],
7340 xs_all: &[f32],
7341 b: usize,
7342 rows: usize,
7343 cols: usize,
7344 out: &mut [f32],
7345 pool: Option<&Pool>,
7346) {
7347 debug_assert_eq!(out.len(), b * rows);
7348 let gpr = cols / GROUP_SIZE;
7349 #[cfg(target_os = "macos")]
7353 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7354 dequant_matmat_accel(
7355 &|r, dst| {
7356 for gi in 0..gpr {
7357 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7358 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7359 for (k, &bb) in tile[2..].iter().enumerate() {
7360 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7361 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7362 }
7363 }
7364 },
7365 xs_all,
7366 b,
7367 rows,
7368 cols,
7369 out,
7370 pool,
7371 );
7372 return;
7373 }
7374 let out_addr = SendMut(out.as_mut_ptr());
7375 if a8w8_enabled() {
7376 let acts: Vec<SplitAct> = (0..b)
7377 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7378 .collect();
7379 let acts = &acts;
7380 #[cfg(target_arch = "x86_64")]
7381 let blocked_ok = avx2_enabled() && blocked_enabled();
7382 #[cfg(target_arch = "aarch64")]
7383 let blocked_ok = sdot_enabled() && blocked_enabled();
7384 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7385 let blocked_ok = false;
7386 let run = move |start: usize, end: usize| {
7387 for r in start..end {
7388 let mut bi = 0usize;
7389 #[cfg(target_arch = "aarch64")]
7390 if blocked_ok {
7391 while bi + 4 <= acts.len() {
7392 let xs = [
7393 acts[bi].xq.as_slice(),
7394 acts[bi + 1].xq.as_slice(),
7395 acts[bi + 2].xq.as_slice(),
7396 acts[bi + 3].xq.as_slice(),
7397 ];
7398 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7399 for k in 0..4 {
7400 let act = &acts[bi + k];
7401 let mut acc = d[k] * act.sx;
7402 for &(j, xv) in &act.outliers {
7403 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7404 acc += w * sc * xv;
7405 }
7406 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7408 }
7409 bi += 4;
7410 }
7411 }
7412 #[cfg(target_arch = "x86_64")]
7413 if blocked_ok {
7414 while bi + 4 <= acts.len() {
7415 let xs = [
7416 acts[bi].xq.as_slice(),
7417 acts[bi + 1].xq.as_slice(),
7418 acts[bi + 2].xq.as_slice(),
7419 acts[bi + 3].xq.as_slice(),
7420 ];
7421 let d = unsafe {
7422 if vnni_tiles_enabled() {
7423 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7424 } else {
7425 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7426 }
7427 };
7428 for k in 0..4 {
7429 let act = &acts[bi + k];
7430 let mut acc = d[k] * act.sx;
7431 for &(j, xv) in &act.outliers {
7432 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7433 acc += w * sc * xv;
7434 }
7435 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7437 }
7438 bi += 4;
7439 }
7440 }
7441 let _ = blocked_ok;
7442 while bi < acts.len() {
7443 let act = &acts[bi];
7444 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7445 for &(j, xv) in &act.outliers {
7446 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7447 acc += w * s * xv;
7448 }
7449 unsafe { *out_addr.at(bi * rows + r) = acc };
7451 bi += 1;
7452 }
7453 }
7454 };
7455 dispatch_rows(pool, rows, &run);
7456 return;
7457 }
7458 let run = move |start: usize, end: usize| {
7459 for r in start..end {
7460 for bi in 0..b {
7461 let x = &xs_all[bi * cols..(bi + 1) * cols];
7462 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7464 }
7465 }
7466 };
7467 dispatch_rows(pool, rows, &run);
7468}
7469
7470fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7479 (0..gpr)
7480 .map(|gi| {
7481 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7482 .iter()
7483 .map(|&v| v as i32)
7484 .sum()
7485 })
7486 .collect()
7487}
7488
7489#[inline]
7493#[allow(unreachable_code)]
7494#[cfg(target_arch = "x86_64")]
7499#[target_feature(enable = "avx2")]
7500unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7501 unsafe {
7503 use core::arch::x86_64::*;
7504 let expand = _mm256_setr_epi8(
7506 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,
7507 3, 3, 3,
7508 );
7509 let bitsel = _mm256_setr_epi8(
7510 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7511 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7512 );
7513 let ones8 = _mm256_set1_epi8(1);
7514 let ones16 = _mm256_set1_epi16(1);
7515 let mut acc = 0f32;
7516 for gi in 0..gpr {
7517 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7518 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7519 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7520 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7521 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7522 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7523 let sel = _mm256_and_si256(x, mask);
7524 let p16 = _mm256_maddubs_epi16(ones8, sel);
7526 let d32 = _mm256_madd_epi16(p16, ones16);
7527 let hi128 = _mm256_extracti128_si256::<1>(d32);
7528 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7529 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7530 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7531 let msum = _mm_cvtsi128_si32(s32);
7532 let d = 2 * msum - gsum[gi];
7535 acc += d as f32 * s;
7536 }
7537 acc
7538 }
7539}
7540
7541#[cfg(target_arch = "x86_64")]
7544#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7545unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7546 unsafe {
7548 use core::arch::x86_64::*;
7549 let expand = _mm256_setr_epi8(
7550 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,
7551 3, 3, 3,
7552 );
7553 let bitsel = _mm256_setr_epi8(
7554 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7555 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7556 );
7557 let ones8 = _mm256_set1_epi8(1);
7558 let mut acc = 0f32;
7559 for gi in 0..gpr {
7560 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7561 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7562 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7563 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7564 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7565 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7566 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7567 let d = 2 * msum - gsum[gi];
7568 acc += d as f32 * s;
7569 }
7570 acc
7571 }
7572}
7573
7574#[cfg(target_arch = "x86_64")]
7576#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7577unsafe fn dot_q1_row_1x4_vnni(
7578 bytes: &[u8],
7579 r: usize,
7580 gpr: usize,
7581 xs: [&[i8]; 4],
7582 gsums: [&[i32]; 4],
7583) -> [f32; 4] {
7584 unsafe {
7586 use core::arch::x86_64::*;
7587 let expand = _mm256_setr_epi8(
7588 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,
7589 3, 3, 3,
7590 );
7591 let bitsel = _mm256_setr_epi8(
7592 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7593 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7594 );
7595 let ones8 = _mm256_set1_epi8(1);
7596 let mut acc = [0f32; 4];
7597 for gi in 0..gpr {
7598 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7599 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7600 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7601 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7602 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7603 for (k, xq) in xs.iter().enumerate() {
7604 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7605 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7606 let d = 2 * msum - gsums[k][gi];
7607 acc[k] += d as f32 * s;
7608 }
7609 }
7610 acc
7611 }
7612}
7613
7614#[cfg(target_arch = "x86_64")]
7617#[target_feature(enable = "avx2")]
7618unsafe fn dot_q1_row_1x4_avx2(
7619 bytes: &[u8],
7620 r: usize,
7621 gpr: usize,
7622 xs: [&[i8]; 4],
7623 gsums: [&[i32]; 4],
7624) -> [f32; 4] {
7625 unsafe {
7627 use core::arch::x86_64::*;
7628 let expand = _mm256_setr_epi8(
7629 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,
7630 3, 3, 3,
7631 );
7632 let bitsel = _mm256_setr_epi8(
7633 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7634 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7635 );
7636 let ones8 = _mm256_set1_epi8(1);
7637 let ones16 = _mm256_set1_epi16(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 sel = _mm256_and_si256(x, mask);
7648 let p16 = _mm256_maddubs_epi16(ones8, sel);
7649 let d32 = _mm256_madd_epi16(p16, ones16);
7650 let hi128 = _mm256_extracti128_si256::<1>(d32);
7651 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7652 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7653 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7654 let msum = _mm_cvtsi128_si32(s32);
7655 let d = 2 * msum - gsums[k][gi];
7656 acc[k] += d as f32 * s;
7657 }
7658 }
7659 acc
7660 }
7661}
7662
7663#[allow(unreachable_code)]
7664fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7665 #[cfg(target_arch = "aarch64")]
7666 unsafe {
7667 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7668 }
7669 #[cfg(target_arch = "x86_64")]
7670 if avx2_enabled() {
7671 unsafe {
7672 if vnni_tiles_enabled() {
7673 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7674 }
7675 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7676 }
7677 }
7678 let _ = gsum;
7679 let mut acc = 0f32;
7680 for gi in 0..gpr {
7681 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7682 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7683 let mut d = 0i32;
7684 for (j, &b) in tile[2..].iter().enumerate() {
7685 for k in 0..8 {
7686 let w = ((b >> k) & 1) as i32 * 2 - 1;
7687 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7688 }
7689 }
7690 acc += d as f32 * s;
7691 }
7692 acc
7693}
7694
7695#[cfg(target_arch = "aarch64")]
7704#[target_feature(enable = "neon,dotprod")]
7705unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7706 unsafe {
7709 use core::arch::aarch64::*;
7710 use core::arch::asm;
7711 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7712 let m = vld1q_u8(MASKS.as_ptr());
7713 macro_rules! tile_dot {
7715 ($t:expr, $x:expr) => {{
7716 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7717 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7718 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7719 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7720 let x0 = vld1q_s8($x);
7721 let x1 = vld1q_s8($x.add(16));
7722 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7723 asm!(
7724 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7725 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7726 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7727 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7728 options(pure, nomem, nostack),
7729 );
7730 vaddq_s32(a0, a1)
7731 }};
7732 }
7733 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7742 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7743 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7744 const IW11: [u8; 16] = [
7745 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7746 ];
7747 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7748 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7749 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7750 let isc = vld1_u8(ISC.as_ptr());
7751 macro_rules! tile_dot_tbl {
7753 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7754 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7755 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7756 let x0 = vld1q_s8($x);
7757 let x1 = vld1q_s8($x.add(16));
7758 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7759 asm!(
7760 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7761 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7762 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7763 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7764 options(pure, nomem, nostack),
7765 );
7766 vaddq_s32(a0, a1)
7767 }};
7768 }
7769 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7770 let row_base = r * gpr * Q1_TILE;
7771 let abs_end = bytes.len();
7772 let xp = xq.as_ptr();
7773 let gp = gsum.as_ptr();
7774 let mut accv = vdupq_n_f32(0.0);
7775 let mut gi = 0;
7776 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7779 let t0 = base.add(gi * Q1_TILE);
7780 let ld_a = vld1q_u8(t0);
7781 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7782 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7783 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7784 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7785 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7786 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7788 let g = vld1q_s32(gp.add(gi));
7789 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7790 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7791 let scf: float32x4_t;
7792 asm!(
7793 "fcvtl {o:v}.4s, {i:v}.4h",
7794 o = out(vreg) scf, i = in(vreg) sc16,
7795 options(pure, nomem, nostack),
7796 );
7797 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7798 gi += 4;
7799 }
7800 let mut acc = vaddvq_f32(accv);
7801 while gi < gpr {
7802 let t = base.add(gi * Q1_TILE);
7803 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7804 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7805 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7806 gi += 1;
7807 }
7808 acc
7809 }
7810}
7811
7812#[cfg(target_arch = "aarch64")]
7817#[target_feature(enable = "neon,dotprod")]
7818unsafe fn dot_q1_row_1x4_sdot(
7819 bytes: &[u8],
7820 r: usize,
7821 gpr: usize,
7822 xs: [&[i8]; 4],
7823 gs: [&[i32]; 4],
7824) -> [f32; 4] {
7825 unsafe {
7827 use core::arch::aarch64::*;
7828 use core::arch::asm;
7829 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7830 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7831 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7832 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7833 const IW11: [u8; 16] = [
7834 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7835 ];
7836 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7837 let m = vld1q_u8(MASKS.as_ptr());
7838 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7839 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7840 let isc = vld1_u8(ISC.as_ptr());
7841 macro_rules! sdot2 {
7842 ($w0:expr, $w1:expr, $x:expr) => {{
7843 let x0 = vld1q_s8($x);
7844 let x1 = vld1q_s8($x.add(16));
7845 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7846 asm!(
7847 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7848 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7849 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7850 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7851 options(pure, nomem, nostack),
7852 );
7853 vaddq_s32(a0, a1)
7854 }};
7855 }
7856 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7857 let row_base = r * gpr * Q1_TILE;
7858 let abs_end = bytes.len();
7859 let mut accv = [vdupq_n_f32(0.0); 4];
7860 let mut gi = 0;
7861 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7862 let t0 = base.add(gi * Q1_TILE);
7863 let ld_a = vld1q_u8(t0);
7864 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7865 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7867 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7868 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7869 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7870 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7871 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7872 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7873 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7874 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7875 let scf: float32x4_t;
7876 asm!(
7877 "fcvtl {o:v}.4s, {i:v}.4h",
7878 o = out(vreg) scf, i = in(vreg) sc16,
7879 options(pure, nomem, nostack),
7880 );
7881 for k in 0..4 {
7882 let xp = xs[k].as_ptr();
7883 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7884 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7885 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7886 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7887 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7888 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7889 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7890 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7891 }
7892 gi += 4;
7893 }
7894 let mut acc = [
7895 vaddvq_f32(accv[0]),
7896 vaddvq_f32(accv[1]),
7897 vaddvq_f32(accv[2]),
7898 vaddvq_f32(accv[3]),
7899 ];
7900 while gi < gpr {
7901 let t = base.add(gi * Q1_TILE);
7902 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7903 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7904 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7905 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7906 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7907 for k in 0..4 {
7908 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7909 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7910 }
7911 gi += 1;
7912 }
7913 acc
7914 }
7915}
7916
7917#[inline]
7919fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7920 let gi = j / GROUP_SIZE;
7921 let k = j % GROUP_SIZE;
7922 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7923 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7924 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7925 ((bit as i32 * 2 - 1) as f32, s)
7926}
7927
7928#[inline]
7930fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7931 let mut acc = 0f32;
7932 for gi in 0..gpr {
7933 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7934 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7935 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7936 let mut ga = 0f32;
7937 for (j, &b) in tile[2..].iter().enumerate() {
7938 for k in 0..8 {
7939 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7940 }
7941 }
7942 acc += ga * s;
7943 }
7944 acc
7945}
7946
7947#[allow(clippy::too_many_arguments)]
7950fn bin_range_a8w8(
7951 w: BinW,
7952 gpr: usize,
7953 act: &SplitAct,
7954 gsum: &[i32],
7955 out: SendMut,
7956 start: usize,
7957 end: usize,
7958) {
7959 for r in start..end {
7960 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7961 for &(j, xv) in &act.outliers {
7962 let (wv, s) = w.outlier(r, gpr, j);
7963 acc += wv * s * xv;
7964 }
7965 unsafe { *out.at(r) = acc };
7967 }
7968}
7969
7970fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7972 for r in start..end {
7973 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7975 }
7976}
7977
7978#[derive(Clone, Copy)]
7986pub(crate) enum BinW<'a> {
7987 Q1(&'a [u8]),
7988 Q1g {
7989 bytes: &'a [u8],
7990 sc_off: usize,
7992 rows: usize,
7995 },
7996}
7997
7998impl<'a> BinW<'a> {
7999 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8000 Self::Q1g {
8001 bytes,
8002 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8003 rows,
8004 }
8005 }
8006
8007 #[inline]
8008 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8009 match self {
8010 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8011 Self::Q1g {
8012 bytes,
8013 sc_off,
8014 rows,
8015 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8016 }
8017 }
8018
8019 #[inline]
8020 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8021 match self {
8022 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8023 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8024 }
8025 }
8026
8027 #[inline]
8028 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8029 match self {
8030 Self::Q1(b) => q1_outlier(b, r, gpr, j),
8031 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8032 }
8033 }
8034
8035 #[cfg(target_arch = "aarch64")]
8040 #[inline]
8041 unsafe fn dot_1x4_sdot(
8042 self,
8043 r: usize,
8044 gpr: usize,
8045 xs: [&[i8]; 4],
8046 gs: [&[i32]; 4],
8047 ) -> [f32; 4] {
8048 unsafe {
8050 match self {
8051 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8052 Self::Q1g {
8053 bytes,
8054 sc_off,
8055 rows,
8056 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8057 }
8058 }
8059 }
8060
8061 #[cfg(target_arch = "x86_64")]
8066 #[inline]
8067 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
8068 unsafe {
8070 match (self, vnni_tiles_enabled()) {
8071 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
8072 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
8073 (Self::Q1g { bytes, sc_off, .. }, true) => {
8074 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
8075 }
8076 (Self::Q1g { bytes, sc_off, .. }, false) => {
8077 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
8078 }
8079 }
8080 }
8081 }
8082}
8083
8084#[inline]
8094fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
8095 u32::from_le_bytes([
8096 bytes[t * 4],
8097 bytes[t * 4 + 1],
8098 bytes[t * 4 + 2],
8099 bytes[t * 4 + 3],
8100 ])
8101}
8102
8103#[inline]
8105fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
8106 let o = sc_off + (t / 4) * 2;
8107 u16::from_le_bytes([bytes[o], bytes[o + 1]])
8108}
8109
8110#[allow(unreachable_code)]
8112#[allow(clippy::too_many_arguments)]
8113fn dot_q1g_row_i8(
8114 bytes: &[u8],
8115 sc_off: usize,
8116 rows: usize,
8117 r: usize,
8118 gpr: usize,
8119 xq: &[i8],
8120 gsum: &[i32],
8121) -> f32 {
8122 #[cfg(target_arch = "aarch64")]
8123 unsafe {
8124 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
8125 }
8126 #[cfg(target_arch = "x86_64")]
8127 if avx2_enabled() {
8128 unsafe {
8129 if vnni_tiles_enabled() {
8130 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
8131 }
8132 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
8133 }
8134 }
8135 let _ = (gsum, rows);
8136 let mut acc = 0f32;
8137 for gi in 0..gpr {
8138 let t = r * gpr + gi;
8139 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8140 let bits = q1g_bits(bytes, t);
8141 let mut d = 0i32;
8142 for k in 0..GROUP_SIZE {
8143 let w = ((bits >> k) & 1) as i32 * 2 - 1;
8144 d += w * xq[gi * GROUP_SIZE + k] as i32;
8145 }
8146 acc += d as f32 * s;
8147 }
8148 acc
8149}
8150
8151#[cfg(target_arch = "x86_64")]
8156#[target_feature(enable = "avx2")]
8157#[inline]
8158unsafe fn q1g_hsum4(
8159 d0: core::arch::x86_64::__m256i,
8160 d1: core::arch::x86_64::__m256i,
8161 d2: core::arch::x86_64::__m256i,
8162 d3: core::arch::x86_64::__m256i,
8163) -> core::arch::x86_64::__m128i {
8164 unsafe {
8166 use core::arch::x86_64::*;
8167 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8168 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8169 }
8170}
8171
8172#[cfg(target_arch = "x86_64")]
8176#[target_feature(enable = "avx2")]
8177#[inline]
8178unsafe fn q1g_group_acc(
8179 acc: f32,
8180 msum4: core::arch::x86_64::__m128i,
8181 gsum4: &[i32],
8182 s: f32,
8183) -> f32 {
8184 unsafe {
8186 use core::arch::x86_64::*;
8187 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8188 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8189 let mut f = [0f32; 4];
8190 _mm_storeu_ps(f.as_mut_ptr(), d);
8191 let mut acc = acc;
8192 for v in f {
8193 acc += v * s;
8194 }
8195 acc
8196 }
8197}
8198
8199#[cfg(target_arch = "x86_64")]
8204#[target_feature(enable = "avx2")]
8205unsafe fn dot_q1g_row_avx2(
8206 bytes: &[u8],
8207 sc_off: usize,
8208 r: usize,
8209 gpr: usize,
8210 xq: &[i8],
8211 gsum: &[i32],
8212) -> f32 {
8213 unsafe {
8216 use core::arch::x86_64::*;
8217 let expand = _mm256_setr_epi8(
8218 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,
8219 3, 3, 3,
8220 );
8221 let bitsel = _mm256_setr_epi8(
8222 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8223 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8224 );
8225 let ones8 = _mm256_set1_epi8(1);
8226 let ones16 = _mm256_set1_epi16(1);
8227 let bp = bytes.as_ptr().add(r * gpr * 4);
8228 let mut acc = 0f32;
8229 let mut gi = 0;
8230 while gi < gpr {
8231 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8232 let mut d = [_mm256_setzero_si256(); 4];
8233 for (k, dk) in d.iter_mut().enumerate() {
8234 let t = gi + k;
8235 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8236 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8237 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8238 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8239 let sel = _mm256_and_si256(x, mask);
8240 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8241 }
8242 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8243 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8244 gi += 4;
8245 }
8246 acc
8247 }
8248}
8249
8250#[cfg(target_arch = "x86_64")]
8253#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8254unsafe fn dot_q1g_row_vnni(
8255 bytes: &[u8],
8256 sc_off: usize,
8257 r: usize,
8258 gpr: usize,
8259 xq: &[i8],
8260 gsum: &[i32],
8261) -> f32 {
8262 unsafe {
8264 use core::arch::x86_64::*;
8265 let expand = _mm256_setr_epi8(
8266 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,
8267 3, 3, 3,
8268 );
8269 let bitsel = _mm256_setr_epi8(
8270 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8271 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8272 );
8273 let ones8 = _mm256_set1_epi8(1);
8274 let bp = bytes.as_ptr().add(r * gpr * 4);
8275 let mut acc = 0f32;
8276 let mut gi = 0;
8277 while gi < gpr {
8278 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8279 let mut d = [_mm256_setzero_si256(); 4];
8280 for (k, dk) in d.iter_mut().enumerate() {
8281 let t = gi + k;
8282 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8283 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8284 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8285 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8286 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8287 }
8288 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8289 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8290 gi += 4;
8291 }
8292 acc
8293 }
8294}
8295
8296#[cfg(target_arch = "x86_64")]
8300#[target_feature(enable = "avx2")]
8301unsafe fn dot_q1g_row_1x4_avx2(
8302 bytes: &[u8],
8303 sc_off: usize,
8304 r: usize,
8305 gpr: usize,
8306 xs: [&[i8]; 4],
8307 gsums: [&[i32]; 4],
8308) -> [f32; 4] {
8309 unsafe {
8311 use core::arch::x86_64::*;
8312 let expand = _mm256_setr_epi8(
8313 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,
8314 3, 3, 3,
8315 );
8316 let bitsel = _mm256_setr_epi8(
8317 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8318 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8319 );
8320 let ones8 = _mm256_set1_epi8(1);
8321 let ones16 = _mm256_set1_epi16(1);
8322 let bp = bytes.as_ptr().add(r * gpr * 4);
8323 let mut acc = [0f32; 4];
8324 let mut gi = 0;
8325 while gi < gpr {
8326 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8327 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8329 for t in 0..4 {
8330 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8331 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8332 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8333 for (k, xq) in xs.iter().enumerate() {
8334 let x = _mm256_loadu_si256(
8335 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8336 );
8337 let sel = _mm256_and_si256(x, mask);
8338 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8339 }
8340 }
8341 for k in 0..4 {
8342 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8343 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8344 }
8345 gi += 4;
8346 }
8347 acc
8348 }
8349}
8350
8351#[cfg(target_arch = "x86_64")]
8353#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8354unsafe fn dot_q1g_row_1x4_vnni(
8355 bytes: &[u8],
8356 sc_off: usize,
8357 r: usize,
8358 gpr: usize,
8359 xs: [&[i8]; 4],
8360 gsums: [&[i32]; 4],
8361) -> [f32; 4] {
8362 unsafe {
8364 use core::arch::x86_64::*;
8365 let expand = _mm256_setr_epi8(
8366 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,
8367 3, 3, 3,
8368 );
8369 let bitsel = _mm256_setr_epi8(
8370 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8371 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8372 );
8373 let ones8 = _mm256_set1_epi8(1);
8374 let bp = bytes.as_ptr().add(r * gpr * 4);
8375 let mut acc = [0f32; 4];
8376 let mut gi = 0;
8377 while gi < gpr {
8378 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8379 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8380 for t in 0..4 {
8381 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8382 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8383 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8384 for (k, xq) in xs.iter().enumerate() {
8385 let x = _mm256_loadu_si256(
8386 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8387 );
8388 d[k][t] = _mm256_dpbusd_epi32(
8389 _mm256_setzero_si256(),
8390 ones8,
8391 _mm256_and_si256(x, mask),
8392 );
8393 }
8394 }
8395 for k in 0..4 {
8396 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8397 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8398 }
8399 gi += 4;
8400 }
8401 acc
8402 }
8403}
8404
8405#[cfg(target_arch = "aarch64")]
8416#[target_feature(enable = "neon,dotprod")]
8417#[allow(clippy::too_many_arguments)]
8418unsafe fn dot_q1g_row_sdot(
8419 bytes: &[u8],
8420 sc_off: usize,
8421 rows: usize,
8422 r: usize,
8423 gpr: usize,
8424 xq: &[i8],
8425 gsum: &[i32],
8426) -> f32 {
8427 unsafe {
8430 use core::arch::aarch64::*;
8431 use core::arch::asm;
8432 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8433 const IW: [[u8; 16]; 8] = [
8435 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8436 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8437 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8438 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8439 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8440 [
8441 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8442 ],
8443 [
8444 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8445 ],
8446 [
8447 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8448 ],
8449 ];
8450 let m = vld1q_u8(MASKS.as_ptr());
8451 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8452 macro_rules! tile_dot {
8453 ($w0:expr, $w1:expr, $x:expr) => {{
8454 let x0 = vld1q_s8($x);
8455 let x1 = vld1q_s8($x.add(16));
8456 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8457 asm!(
8458 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8459 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8460 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8461 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8462 options(pure, nomem, nostack),
8463 );
8464 vaddq_s32(a0, a1)
8465 }};
8466 }
8467 macro_rules! tbl_mask {
8468 ($ld:expr, $i:expr) => {
8469 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8470 };
8471 }
8472 let bp = bytes.as_ptr().add(r * gpr * 4);
8473 let xp = xq.as_ptr();
8474 let gp = gsum.as_ptr();
8475 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8476 let mut accv = vdupq_n_f32(0.0);
8477 let mut gi = 0;
8478 while gi + 4 <= vec_end {
8479 let ld = vld1q_u8(bp.add(gi * 4));
8480 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8481 let d1 = tile_dot!(
8482 tbl_mask!(ld, 2),
8483 tbl_mask!(ld, 3),
8484 xp.add((gi + 1) * GROUP_SIZE)
8485 );
8486 let d2 = tile_dot!(
8487 tbl_mask!(ld, 4),
8488 tbl_mask!(ld, 5),
8489 xp.add((gi + 2) * GROUP_SIZE)
8490 );
8491 let d3 = tile_dot!(
8492 tbl_mask!(ld, 6),
8493 tbl_mask!(ld, 7),
8494 xp.add((gi + 3) * GROUP_SIZE)
8495 );
8496 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8497 let g = vld1q_s32(gp.add(gi));
8498 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8499 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8500 let scf: float32x4_t;
8501 asm!(
8502 "fcvtl {o:v}.4s, {i:v}.4h",
8503 o = out(vreg) scf, i = in(vreg) h,
8504 options(pure, nomem, nostack),
8505 );
8506 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8507 gi += 4;
8508 }
8509 let mut acc = vaddvq_f32(accv);
8510 while gi < gpr {
8511 let t = bp.add(gi * 4);
8512 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8513 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8514 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8515 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8516 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8517 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8518 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8519 gi += 1;
8520 }
8521 acc
8522 }
8523}
8524
8525#[cfg(target_arch = "aarch64")]
8529#[target_feature(enable = "neon,dotprod")]
8530#[allow(clippy::too_many_arguments)]
8531unsafe fn dot_q1g_row_1x4_sdot(
8532 bytes: &[u8],
8533 sc_off: usize,
8534 rows: usize,
8535 r: usize,
8536 gpr: usize,
8537 xs: [&[i8]; 4],
8538 gs: [&[i32]; 4],
8539) -> [f32; 4] {
8540 unsafe {
8542 use core::arch::aarch64::*;
8543 use core::arch::asm;
8544 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8545 const IW: [[u8; 16]; 8] = [
8546 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8547 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8548 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8549 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8550 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8551 [
8552 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8553 ],
8554 [
8555 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8556 ],
8557 [
8558 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8559 ],
8560 ];
8561 let m = vld1q_u8(MASKS.as_ptr());
8562 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8563 macro_rules! sdot2 {
8564 ($w0:expr, $w1:expr, $x:expr) => {{
8565 let x0 = vld1q_s8($x);
8566 let x1 = vld1q_s8($x.add(16));
8567 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8568 asm!(
8569 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8570 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8571 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8572 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8573 options(pure, nomem, nostack),
8574 );
8575 vaddq_s32(a0, a1)
8576 }};
8577 }
8578 let bp = bytes.as_ptr().add(r * gpr * 4);
8579 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8580 let mut accv = [vdupq_n_f32(0.0); 4];
8581 let mut gi = 0;
8582 while gi + 4 <= vec_end {
8583 let ld = vld1q_u8(bp.add(gi * 4));
8584 let w: [int8x16_t; 8] =
8586 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8587 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8588 let scf: float32x4_t;
8589 asm!(
8590 "fcvtl {o:v}.4s, {i:v}.4h",
8591 o = out(vreg) scf, i = in(vreg) h,
8592 options(pure, nomem, nostack),
8593 );
8594 for k in 0..4 {
8595 let xp = xs[k].as_ptr();
8596 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8597 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8598 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8599 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8600 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8601 let g = vld1q_s32(gs[k].as_ptr().add(gi));
8602 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8603 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8604 }
8605 gi += 4;
8606 }
8607 let mut acc = [
8608 vaddvq_f32(accv[0]),
8609 vaddvq_f32(accv[1]),
8610 vaddvq_f32(accv[2]),
8611 vaddvq_f32(accv[3]),
8612 ];
8613 while gi < gpr {
8614 let t = bp.add(gi * 4);
8615 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8616 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8617 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8618 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8619 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8620 for k in 0..4 {
8621 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8622 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8623 }
8624 gi += 1;
8625 }
8626 acc
8627 }
8628}
8629
8630#[inline]
8632fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8633 let t = r * gpr + j / GROUP_SIZE;
8634 let k = j % GROUP_SIZE;
8635 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8636 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8637 ((bit as i32 * 2 - 1) as f32, s)
8638}
8639
8640#[inline]
8642fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8643 let mut acc = 0f32;
8644 for gi in 0..gpr {
8645 let t = r * gpr + gi;
8646 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8647 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8648 let mut ga = 0f32;
8649 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8650 for k in 0..8 {
8651 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8652 }
8653 }
8654 acc += ga * s;
8655 }
8656 acc
8657}
8658
8659fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8664 let entries = base_len + (rows + 1) * 4;
8665 (base_len, entries, entries <= bytes.len())
8666}
8667
8668#[inline]
8670fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8671 let o = rp_off + r * 4;
8672 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8673}
8674
8675const SIGN5: [[f32; 5]; 256] = {
8679 let mut lut = [[0.0f32; 5]; 256];
8680 let pow3 = [1u16, 3, 9, 27, 81];
8681 let mut byte = 0usize;
8682 while byte < 256 {
8683 let mut i = 0usize;
8684 while i < 5 {
8685 let code = (byte as u16 / pow3[i]) % 3;
8686 lut[byte][i] = if code == 1 {
8687 1.0
8688 } else if code == 2 {
8689 -1.0
8690 } else {
8691 0.0
8692 };
8693 i += 1;
8694 }
8695 byte += 1;
8696 }
8697 lut
8698};
8699
8700const SIGN5_I8: [[i8; 5]; 256] = {
8702 let mut lut = [[0i8; 5]; 256];
8703 let pow3 = [1u16, 3, 9, 27, 81];
8704 let mut byte = 0usize;
8705 while byte < 256 {
8706 let mut i = 0usize;
8707 while i < 5 {
8708 let code = (byte as u16 / pow3[i]) % 3;
8709 lut[byte][i] = if code == 1 {
8710 1
8711 } else if code == 2 {
8712 -1
8713 } else {
8714 0
8715 };
8716 i += 1;
8717 }
8718 byte += 1;
8719 }
8720 lut
8721};
8722
8723const SIGN5_U64: [u64; 256] = {
8729 let mut lut = [0u64; 256];
8730 let pow3 = [1u16, 3, 9, 27, 81];
8731 let mut byte = 0usize;
8732 while byte < 256 {
8733 let mut v = 0u64;
8734 let mut i = 0usize;
8735 while i < 5 {
8736 let code = (byte as u16 / pow3[i]) % 3;
8737 let s: u8 = if code == 1 {
8738 1
8739 } else if code == 2 {
8740 0xFF
8741 } else {
8742 0
8743 };
8744 v |= (s as u64) << (i * 8);
8745 i += 1;
8746 }
8747 lut[byte] = v;
8748 byte += 1;
8749 }
8750 lut
8751};
8752
8753#[inline]
8758fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8759 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8760 let off = (r * gpr + j / GROUP_SIZE) * TILE;
8761 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8762 let within = j % GROUP_SIZE;
8763 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8764}
8765
8766#[cfg(target_arch = "aarch64")]
8769#[target_feature(enable = "neon,dotprod")]
8770#[inline]
8771unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8772 unsafe {
8774 use core::arch::aarch64::*;
8775 use core::arch::asm;
8776 let w0 = vld1q_s8(w);
8777 let w1 = vld1q_s8(w.add(16));
8778 let x0 = vld1q_s8(x);
8779 let x1 = vld1q_s8(x.add(16));
8780 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8781 asm!(
8782 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8783 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8784 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8785 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8786 options(pure, nomem, nostack),
8787 );
8788 vaddvq_s32(vaddq_s32(a0, a1))
8789 }
8790}
8791
8792#[cfg(target_arch = "x86_64")]
8795#[target_feature(enable = "avx2")]
8796#[inline]
8797unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8798 unsafe {
8800 use core::arch::x86_64::*;
8801 let wv = _mm256_loadu_si256(w as *const __m256i);
8802 let xv = _mm256_loadu_si256(x as *const __m256i);
8803 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8804 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8805 let hi128 = _mm256_extracti128_si256::<1>(d);
8806 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8807 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8808 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8809 _mm_cvtsi128_si32(s32)
8810 }
8811}
8812
8813#[inline]
8818fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8819 debug_assert!(dst.len() >= 40);
8820 unsafe {
8823 let p = dst.as_mut_ptr();
8824 for bi in 0..7 {
8825 core::ptr::write_unaligned(
8826 p.add(bi * 5) as *mut u64,
8827 SIGN5_U64[*codes.add(bi) as usize],
8828 );
8829 }
8830 }
8831}
8832
8833#[inline]
8838fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8839 #[cfg(target_arch = "aarch64")]
8840 unsafe {
8841 return sdot32_i8(w, x);
8842 }
8843 #[cfg(target_arch = "x86_64")]
8844 unsafe {
8845 return i8dot32_avx2(w, x);
8846 }
8847 #[allow(unreachable_code)]
8848 unsafe {
8849 let mut s = 0i32;
8850 for k in 0..GROUP_SIZE {
8851 s += *w.add(k) as i32 * *x.add(k) as i32;
8852 }
8853 s
8854 }
8855}
8856
8857#[inline]
8858unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8859 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8860 (
8861 SIGN5_U64[*codes as usize],
8862 SIGN5_U64[*codes.add(1) as usize],
8863 SIGN5_U64[*codes.add(2) as usize],
8864 SIGN5_U64[*codes.add(3) as usize],
8865 SIGN5_U64[*codes.add(4) as usize],
8866 SIGN5_U64[*codes.add(5) as usize],
8867 SIGN5_U64[*codes.add(6) as usize],
8868 )
8869 };
8870
8871 let u0 = s0 | (s1 << 40);
8872 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8873 let u2 = (s3 >> 8) | (s4 << 32);
8874 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8875
8876 (u0, u1, u2, u3)
8877}
8878
8879#[cfg(target_arch = "aarch64")]
8883#[target_feature(enable = "neon,dotprod")]
8884unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8885 use core::arch::aarch64::*;
8886 use core::arch::asm;
8887 unsafe {
8888 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8889 let mut acc = 0f32;
8890 let bytes_ptr = bytes.as_ptr();
8891 let xq_ptr = xq.as_ptr();
8892 let row_off = r * gpr * TILE;
8893
8894 let gpr2 = gpr & !1;
8895 let mut gi = 0;
8896 while gi < gpr2 {
8897 let off0 = row_off + gi * TILE;
8898 let off1 = off0 + TILE;
8899 let s0 = f16_to_f32(u16::from_le_bytes([
8900 *bytes_ptr.add(off0),
8901 *bytes_ptr.add(off0 + 1),
8902 ]));
8903 let s1 = f16_to_f32(u16::from_le_bytes([
8904 *bytes_ptr.add(off1),
8905 *bytes_ptr.add(off1 + 1),
8906 ]));
8907
8908 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8909 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8910
8911 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8912 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8913 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8914 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8915
8916 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8917 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8918 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8919 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8920
8921 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8922 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8923 asm!(
8924 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8925 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8926 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8927 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8928 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8929 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8930 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8931 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8932 options(pure, nomem, nostack),
8933 );
8934 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8935 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8936 acc += d0 as f32 * s0 + d1 as f32 * s1;
8937 gi += 2;
8938 }
8939
8940 if gi < gpr {
8941 let off = row_off + gi * TILE;
8942 let s = f16_to_f32(u16::from_le_bytes([
8943 *bytes_ptr.add(off),
8944 *bytes_ptr.add(off + 1),
8945 ]));
8946 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8947 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8948 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8949 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8950 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8951 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8952 asm!(
8953 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8954 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8955 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8956 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8957 options(pure, nomem, nostack),
8958 );
8959 let d = vaddvq_s32(vaddq_s32(a0, a1));
8960 acc += d as f32 * s;
8961 }
8962 acc
8963 }
8964}
8965
8966#[cfg(target_arch = "x86_64")]
8968#[target_feature(enable = "avx2")]
8969unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8970 use core::arch::x86_64::*;
8971 unsafe {
8972 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8973 let mut acc = 0f32;
8974 let bytes_ptr = bytes.as_ptr();
8975 let xq_ptr = xq.as_ptr();
8976 let row_off = r * gpr * TILE;
8977
8978 let ones = _mm256_set1_epi16(1);
8979 for gi in 0..gpr {
8980 let off = row_off + gi * TILE;
8981 let s = f16_to_f32(u16::from_le_bytes([
8982 *bytes_ptr.add(off),
8983 *bytes_ptr.add(off + 1),
8984 ]));
8985 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8986 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8987 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8988 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8989 let d256 = _mm256_madd_epi16(p16, ones);
8990 let d128 = _mm_add_epi32(
8991 _mm256_castsi256_si128(d256),
8992 _mm256_extracti128_si256(d256, 1),
8993 );
8994 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8995 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8996 acc += d32 as f32 * s;
8997 }
8998 acc
8999 }
9000}
9001
9002#[cfg(target_arch = "x86_64")]
9004#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9005unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9006 use core::arch::x86_64::*;
9007 unsafe {
9009 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9010 let mut acc = 0f32;
9011 let bytes_ptr = bytes.as_ptr();
9012 let xq_ptr = xq.as_ptr();
9013 let row_off = r * gpr * TILE;
9014 for gi in 0..gpr {
9015 let off = row_off + gi * TILE;
9016 let s = f16_to_f32(u16::from_le_bytes([
9017 *bytes_ptr.add(off),
9018 *bytes_ptr.add(off + 1),
9019 ]));
9020 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9021 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9022 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9023 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9024 acc += d as f32 * s;
9025 }
9026 acc
9027 }
9028}
9029
9030#[inline]
9034fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9035 #[cfg(target_arch = "aarch64")]
9036 unsafe {
9037 return q1t_dot_row_sdot(bytes, r, gpr, xq);
9038 }
9039 #[cfg(target_arch = "x86_64")]
9040 unsafe {
9041 if vnni_tiles_enabled() {
9042 return q1t_dot_row_vnni(bytes, r, gpr, xq);
9043 }
9044 return q1t_dot_row_avx2(bytes, r, gpr, xq);
9045 }
9046 #[allow(unreachable_code)]
9047 {
9048 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9049 let mut acc = 0f32;
9050 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
9052 let off = (r * gpr + gi) * TILE;
9053 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9054 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9055 let mut d = 0i32;
9056 for k in 0..GROUP_SIZE {
9057 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9058 }
9059 acc += d as f32 * s;
9060 }
9061 acc
9062 }
9063}
9064
9065fn q1t_row_outlier_correction(
9072 bytes: &[u8],
9073 r: usize,
9074 rp_off: usize,
9075 entries_off: usize,
9076 has_ov: bool,
9077 x: &[f32],
9078) -> f32 {
9079 if !has_ov {
9080 return 0.0;
9081 }
9082 let (c0, c1) = (
9083 q1t_rowptr(bytes, rp_off, r),
9084 q1t_rowptr(bytes, rp_off, r + 1),
9085 );
9086 let mut corr = 0f32;
9087 for p in c0..c1 {
9088 let e = entries_off + p * 4;
9089 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9090 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9091 corr += val * x[col];
9092 }
9093 corr
9094}
9095
9096fn q1t_dequant_row(
9100 bytes: &[u8],
9101 r: usize,
9102 gpr: usize,
9103 rp_off: usize,
9104 entries_off: usize,
9105 has_ov: bool,
9106 buf: &mut [f32],
9107) {
9108 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9109 for g in 0..gpr {
9110 let off = (r * gpr + g) * TILE;
9111 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9112 let codes = &bytes[off + 2..off + TILE];
9113 let bc = g * GROUP_SIZE;
9114 for bi in 0..6 {
9116 let lut = &SIGN5[codes[bi] as usize];
9117 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
9118 for i in 0..5 {
9119 d[i] = lut[i] * s;
9120 }
9121 }
9122 let lut = &SIGN5[codes[6] as usize];
9123 buf[bc + 30] = lut[0] * s;
9124 buf[bc + 31] = lut[1] * s;
9125 }
9126 if !has_ov {
9127 return;
9128 }
9129 let (c0, c1) = (
9130 q1t_rowptr(bytes, rp_off, r),
9131 q1t_rowptr(bytes, rp_off, r + 1),
9132 );
9133 for p in c0..c1 {
9134 let e = entries_off + p * 4;
9135 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9136 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9137 }
9138}
9139
9140fn q1t_add_overlay(
9144 bytes: &[u8],
9145 x: &[f32],
9146 rows: usize,
9147 cols: usize,
9148 out: &mut [f32],
9149 pool: Option<&Pool>,
9150) {
9151 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9152 let gpr = cols / GROUP_SIZE;
9153 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9154 if !has_ov {
9155 return;
9156 }
9157 let out_addr = SendMut(out.as_mut_ptr());
9158 let run = move |start: usize, end: usize| {
9159 for r in start..end {
9160 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9161 unsafe { *out_addr.at(r) += corr };
9163 }
9164 };
9165 dispatch_rows(pool, rows, &run);
9166}
9167
9168#[allow(clippy::too_many_arguments)]
9171fn q1t_range_a8w8(
9172 bytes: &[u8],
9173 gpr: usize,
9174 rp_off: usize,
9175 ent_off: usize,
9176 has_ov: bool,
9177 act: &SplitAct,
9178 x: &[f32],
9179 out: SendMut,
9180 start: usize,
9181 end: usize,
9182) {
9183 for r in start..end {
9184 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9185 for &(j, xv) in &act.outliers {
9186 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9187 }
9188 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9189 unsafe { *out.at(r) = acc };
9191 }
9192}
9193
9194#[allow(clippy::too_many_arguments)]
9197fn q1t_range_f32_batch(
9198 bytes: &[u8],
9199 gpr: usize,
9200 rp_off: usize,
9201 ent_off: usize,
9202 has_ov: bool,
9203 x: &[f32],
9204 out: SendMut,
9205 start: usize,
9206 end: usize,
9207) {
9208 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9209 let mut sg = [0f32; GROUP_SIZE];
9210 for r in start..end {
9211 let mut acc = 0f32;
9212 for g in 0..gpr {
9213 let off = (r * gpr + g) * TILE;
9214 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9215 let codes = &bytes[off + 2..off + TILE];
9216 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9217 for bi in 0..6 {
9218 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9219 }
9220 let lut = &SIGN5[codes[6] as usize];
9221 sg[30] = lut[0];
9222 sg[31] = lut[1];
9223 let mut gsum = 0f32;
9224 for k in 0..GROUP_SIZE {
9225 gsum += sg[k] * xg[k];
9226 }
9227 acc += s * gsum;
9228 }
9229 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9230 unsafe { *out.at(r) = acc };
9232 }
9233}
9234
9235fn q1t_matvec(
9239 bytes: &[u8],
9240 x: &[f32],
9241 rows: usize,
9242 cols: usize,
9243 out: &mut [f32],
9244 pool: Option<&Pool>,
9245) {
9246 debug_assert_eq!(out.len(), rows);
9247 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9248 let gpr = cols / GROUP_SIZE;
9249 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9250 let out_addr = SendMut(out.as_mut_ptr());
9251 if a8w8_enabled() {
9255 let act = split_act(x);
9256 let act = &act;
9257 let run = move |start: usize, end: usize| {
9258 for r in start..end {
9259 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9260 for &(j, xv) in &act.outliers {
9261 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9262 }
9263 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9264 unsafe { *out_addr.at(r) = acc };
9266 }
9267 };
9268 dispatch_rows(pool, rows, &run);
9269 return;
9270 }
9271 let run = move |start: usize, end: usize| {
9272 let mut sg = [0f32; GROUP_SIZE];
9276 for r in start..end {
9277 let mut acc = 0f32;
9278 for g in 0..gpr {
9279 let off = (r * gpr + g) * TILE;
9280 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9281 let codes = &bytes[off + 2..off + TILE];
9282 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9283 for bi in 0..6 {
9284 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9285 }
9286 let lut = &SIGN5[codes[6] as usize];
9287 sg[30] = lut[0];
9288 sg[31] = lut[1];
9289 let mut gsum = 0f32;
9290 for k in 0..GROUP_SIZE {
9291 gsum += sg[k] * xg[k];
9292 }
9293 acc += s * gsum;
9294 }
9295 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9296 unsafe { *out_addr.at(r) = acc };
9297 }
9298 };
9299 dispatch_rows(pool, rows, &run);
9300}
9301
9302#[cfg(target_arch = "aarch64")]
9308#[target_feature(enable = "neon,dotprod")]
9309unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9310 use core::arch::aarch64::*;
9311 use core::arch::asm;
9312 unsafe {
9314 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9315 let bytes_ptr = bytes.as_ptr();
9316 let row_off = r * gpr * TILE;
9317 let xp = [xa.as_ptr(), xb.as_ptr()];
9318 let mut acc = [0f32; 2];
9319 macro_rules! sdot2 {
9320 ($w0:expr, $w1:expr, $x:expr) => {{
9321 let x0 = vld1q_s8($x);
9322 let x1 = vld1q_s8($x.add(16));
9323 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9324 asm!(
9325 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9326 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9327 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9328 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9329 options(pure, nomem, nostack),
9330 );
9331 vaddvq_s32(vaddq_s32(a0, a1))
9332 }};
9333 }
9334 let gpr2 = gpr & !1;
9335 let mut gi = 0;
9336 while gi < gpr2 {
9337 let off0 = row_off + gi * TILE;
9338 let off1 = off0 + TILE;
9339 let s0 = f16_to_f32(u16::from_le_bytes([
9340 *bytes_ptr.add(off0),
9341 *bytes_ptr.add(off0 + 1),
9342 ]));
9343 let s1 = f16_to_f32(u16::from_le_bytes([
9344 *bytes_ptr.add(off1),
9345 *bytes_ptr.add(off1 + 1),
9346 ]));
9347 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9348 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9349 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9350 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9351 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9352 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9353 for k in 0..2 {
9354 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9355 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9356 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9357 }
9358 gi += 2;
9359 }
9360 if gi < gpr {
9361 let off = row_off + gi * TILE;
9362 let s = f16_to_f32(u16::from_le_bytes([
9363 *bytes_ptr.add(off),
9364 *bytes_ptr.add(off + 1),
9365 ]));
9366 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9367 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9368 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9369 for k in 0..2 {
9370 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9371 acc[k] += d as f32 * s;
9372 }
9373 }
9374 acc
9375 }
9376}
9377
9378fn q1t_matvec2(
9384 bytes: &[u8],
9385 x1: &[f32],
9386 x2: &[f32],
9387 rows: usize,
9388 cols: usize,
9389 o1: &mut [f32],
9390 o2: &mut [f32],
9391 pool: Option<&Pool>,
9392) {
9393 debug_assert_eq!(o1.len(), rows);
9394 debug_assert_eq!(o2.len(), rows);
9395 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9396 let gpr = cols / GROUP_SIZE;
9397 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9398 let out1 = SendMut(o1.as_mut_ptr());
9399 let out2 = SendMut(o2.as_mut_ptr());
9400 if a8w8_enabled() {
9401 let a1 = split_act(x1);
9402 let a2 = split_act(x2);
9403 let (a1, a2) = (&a1, &a2);
9404 let run = move |start: usize, end: usize| {
9405 for r in start..end {
9406 #[cfg(target_arch = "aarch64")]
9407 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9410 #[cfg(not(target_arch = "aarch64"))]
9411 let ds = [
9412 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9413 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9414 ];
9415 let mut acc1 = ds[0] * a1.sx;
9416 for &(j, xv) in &a1.outliers {
9417 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9418 }
9419 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9420 let mut acc2 = ds[1] * a2.sx;
9421 for &(j, xv) in &a2.outliers {
9422 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9423 }
9424 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9425 unsafe {
9427 *out1.at(r) = acc1;
9428 *out2.at(r) = acc2;
9429 }
9430 }
9431 };
9432 dispatch_rows(pool, rows, &run);
9433 return;
9434 }
9435 let run = move |start: usize, end: usize| {
9436 let mut sg = [0f32; GROUP_SIZE];
9439 for r in start..end {
9440 let mut acc1 = 0f32;
9441 let mut acc2 = 0f32;
9442 for g in 0..gpr {
9443 let off = (r * gpr + g) * TILE;
9444 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9445 let codes = &bytes[off + 2..off + TILE];
9446 for bi in 0..6 {
9447 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9448 }
9449 let lut = &SIGN5[codes[6] as usize];
9450 sg[30] = lut[0];
9451 sg[31] = lut[1];
9452 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9453 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9454 let mut gsum1 = 0f32;
9455 for k in 0..GROUP_SIZE {
9456 gsum1 += sg[k] * xg1[k];
9457 }
9458 acc1 += s * gsum1;
9459 let mut gsum2 = 0f32;
9460 for k in 0..GROUP_SIZE {
9461 gsum2 += sg[k] * xg2[k];
9462 }
9463 acc2 += s * gsum2;
9464 }
9465 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
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}
9476
9477fn q1t_matmat(
9480 bytes: &[u8],
9481 xs: &[f32],
9482 b: usize,
9483 rows: usize,
9484 cols: usize,
9485 out: &mut [f32],
9486 pool: Option<&Pool>,
9487) {
9488 debug_assert_eq!(out.len(), b * rows);
9489 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9490 let gpr = cols / GROUP_SIZE;
9491 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9492 let out_addr = SendMut(out.as_mut_ptr());
9493 if a8w8_enabled() {
9497 let acts: Vec<SplitAct> = (0..b)
9498 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9499 .collect();
9500 let acts = &acts;
9501 let run = move |start: usize, end: usize| {
9502 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
9506 for g in 0..gpr {
9507 let off = (r * gpr + g) * TILE;
9508 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9509 q1t_unpack_group_i8(
9510 bytes.as_ptr().wrapping_add(off + 2),
9511 &mut sg[g * GROUP_SIZE..],
9512 );
9513 }
9514 for bi in 0..b {
9515 let act = &acts[bi];
9516 let mut isum = 0f32;
9517 for g in 0..gpr {
9518 let d = q1t_i8dot32(
9519 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9520 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9521 );
9522 isum += d as f32 * sc[g];
9523 }
9524 let mut acc = isum * act.sx;
9525 for &(j, xv) in &act.outliers {
9526 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9527 }
9528 accs[bi] = acc;
9529 }
9530 if has_ov {
9534 let (c0, c1) = (
9535 q1t_rowptr(bytes, rp_off, r),
9536 q1t_rowptr(bytes, rp_off, r + 1),
9537 );
9538 for p in c0..c1 {
9539 let e = ent_off + p * 4;
9540 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9541 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9542 for bi in 0..b {
9543 accs[bi] += val * xs[bi * cols + col];
9544 }
9545 }
9546 }
9547 for bi in 0..b {
9548 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9549 }
9550 }
9551 };
9552 dispatch_rows(pool, rows, &run);
9553 return;
9554 }
9555 let run = move |start: usize, end: usize| {
9556 let mut buf = vec![0f32; cols];
9557 for r in start..end {
9558 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9559 for bi in 0..b {
9560 let xr = &xs[bi * cols..(bi + 1) * cols];
9561 let mut acc = 0f32;
9562 for j in 0..cols {
9563 acc += buf[j] * xr[j];
9564 }
9565 unsafe { *out_addr.at(bi * rows + r) = acc };
9566 }
9567 }
9568 };
9569 dispatch_rows(pool, rows, &run);
9570}
9571
9572#[cfg(test)]
9575fn q1_matvec(
9576 bytes: &[u8],
9577 x: &[f32],
9578 rows: usize,
9579 cols: usize,
9580 out: &mut [f32],
9581 pool: Option<&Pool>,
9582) {
9583 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9584}
9585
9586fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9589 debug_assert_eq!(out.len(), rows);
9590 let gpr = cols / GROUP_SIZE;
9591 let out_addr = SendMut(out.as_mut_ptr());
9592 if a8w8_enabled() {
9593 let act = split_act(x);
9594 let gsum = q1_group_sums(&act.xq, gpr);
9595 let (act, gsum) = (&act, &gsum);
9596 let run =
9597 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9598 dispatch_rows(pool, rows, &run);
9599 return;
9600 }
9601 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9602 dispatch_rows(pool, rows, &run);
9603}
9604
9605#[cfg(test)]
9607#[allow(clippy::too_many_arguments)]
9608fn q1_matvec2(
9609 bytes: &[u8],
9610 x1: &[f32],
9611 x2: &[f32],
9612 rows: usize,
9613 cols: usize,
9614 o1: &mut [f32],
9615 o2: &mut [f32],
9616 pool: Option<&Pool>,
9617) {
9618 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9619}
9620
9621#[allow(clippy::too_many_arguments)]
9623fn bin_matvec2(
9624 w: BinW,
9625 x1: &[f32],
9626 x2: &[f32],
9627 rows: usize,
9628 cols: usize,
9629 o1: &mut [f32],
9630 o2: &mut [f32],
9631 pool: Option<&Pool>,
9632) {
9633 let gpr = cols / GROUP_SIZE;
9634 let p1 = SendMut(o1.as_mut_ptr());
9635 let p2 = SendMut(o2.as_mut_ptr());
9636 if a8w8_enabled() {
9637 let a1 = split_act(x1);
9638 let a2 = split_act(x2);
9639 let g1 = q1_group_sums(&a1.xq, gpr);
9640 let g2 = q1_group_sums(&a2.xq, gpr);
9641 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9642 let run = move |start: usize, end: usize| {
9643 for r in start..end {
9644 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9645 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9646 for &(j, xv) in &a1.outliers {
9647 let (wv, s) = w.outlier(r, gpr, j);
9648 v1 += wv * s * xv;
9649 }
9650 for &(j, xv) in &a2.outliers {
9651 let (wv, s) = w.outlier(r, gpr, j);
9652 v2 += wv * s * xv;
9653 }
9654 unsafe {
9656 *p1.at(r) = v1;
9657 *p2.at(r) = v2;
9658 }
9659 }
9660 };
9661 dispatch_rows(pool, rows, &run);
9662 return;
9663 }
9664 let run = move |start: usize, end: usize| {
9665 for r in start..end {
9666 unsafe {
9668 *p1.at(r) = w.row_exact(r, gpr, x1);
9669 *p2.at(r) = w.row_exact(r, gpr, x2);
9670 }
9671 }
9672 };
9673 dispatch_rows(pool, rows, &run);
9674}
9675
9676#[cfg(test)]
9678#[allow(clippy::too_many_arguments)]
9679fn q1_matmat(
9680 bytes: &[u8],
9681 xs_all: &[f32],
9682 b: usize,
9683 rows: usize,
9684 cols: usize,
9685 out: &mut [f32],
9686 pool: Option<&Pool>,
9687) {
9688 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9689}
9690
9691#[allow(clippy::too_many_arguments)]
9693fn bin_matmat(
9694 w: BinW,
9695 xs_all: &[f32],
9696 b: usize,
9697 rows: usize,
9698 cols: usize,
9699 out: &mut [f32],
9700 pool: Option<&Pool>,
9701) {
9702 debug_assert_eq!(out.len(), b * rows);
9703 let gpr = cols / GROUP_SIZE;
9704 let out_addr = SendMut(out.as_mut_ptr());
9705 if a8w8_enabled() {
9706 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9707 .map(|bi| {
9708 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9709 let gsum = q1_group_sums(&act.xq, gpr);
9710 (act, gsum)
9711 })
9712 .collect();
9713 let acts = &acts;
9714 #[cfg(target_arch = "x86_64")]
9715 let blocked_ok = avx2_enabled() && blocked_enabled();
9716 #[cfg(target_arch = "aarch64")]
9717 let blocked_ok = sdot_enabled() && blocked_enabled();
9718 let run = move |start: usize, end: usize| {
9719 for r in start..end {
9720 let mut bi = 0usize;
9721 #[cfg(target_arch = "aarch64")]
9724 if blocked_ok {
9725 while bi + 4 <= acts.len() {
9726 let xs = [
9727 acts[bi].0.xq.as_slice(),
9728 acts[bi + 1].0.xq.as_slice(),
9729 acts[bi + 2].0.xq.as_slice(),
9730 acts[bi + 3].0.xq.as_slice(),
9731 ];
9732 let gs = [
9733 acts[bi].1.as_slice(),
9734 acts[bi + 1].1.as_slice(),
9735 acts[bi + 2].1.as_slice(),
9736 acts[bi + 3].1.as_slice(),
9737 ];
9738 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9740 for k in 0..4 {
9741 let (act, _) = &acts[bi + k];
9742 let mut acc = d[k] * act.sx;
9743 for &(j, xv) in &act.outliers {
9744 let (wv, sc) = w.outlier(r, gpr, j);
9745 acc += wv * sc * xv;
9746 }
9747 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9749 }
9750 bi += 4;
9751 }
9752 }
9753 #[cfg(target_arch = "x86_64")]
9754 if blocked_ok {
9755 while bi + 4 <= acts.len() {
9756 let xs = [
9757 acts[bi].0.xq.as_slice(),
9758 acts[bi + 1].0.xq.as_slice(),
9759 acts[bi + 2].0.xq.as_slice(),
9760 acts[bi + 3].0.xq.as_slice(),
9761 ];
9762 let gs = [
9763 acts[bi].1.as_slice(),
9764 acts[bi + 1].1.as_slice(),
9765 acts[bi + 2].1.as_slice(),
9766 acts[bi + 3].1.as_slice(),
9767 ];
9768 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9770 for k in 0..4 {
9771 let (act, _) = &acts[bi + k];
9772 let mut acc = d[k] * act.sx;
9773 for &(j, xv) in &act.outliers {
9774 let (wv, sc) = w.outlier(r, gpr, j);
9775 acc += wv * sc * xv;
9776 }
9777 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9779 }
9780 bi += 4;
9781 }
9782 }
9783 while bi < acts.len() {
9784 let (act, gsum) = &acts[bi];
9785 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9786 for &(j, xv) in &act.outliers {
9787 let (wv, s) = w.outlier(r, gpr, j);
9788 acc += wv * s * xv;
9789 }
9790 unsafe { *out_addr.at(bi * rows + r) = acc };
9792 bi += 1;
9793 }
9794 }
9795 };
9796 dispatch_rows(pool, rows, &run);
9797 return;
9798 }
9799 let run = move |start: usize, end: usize| {
9800 for r in start..end {
9801 for bi in 0..b {
9802 let x = &xs_all[bi * cols..(bi + 1) * cols];
9803 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9805 }
9806 }
9807 };
9808 dispatch_rows(pool, rows, &run);
9809}
9810
9811fn q4matvec(
9817 bytes: &[u8],
9818 x: &[f32],
9819 rows: usize,
9820 cols: usize,
9821 out: &mut [f32],
9822 pool: Option<&Pool>,
9823) {
9824 debug_assert_eq!(out.len(), rows);
9825 let (packed, scales) = q4_split(bytes, rows, cols);
9826 let gpr = cols / GROUP_SIZE;
9827 let out_addr = SendMut(out.as_mut_ptr());
9828
9829 if a8w8_enabled() {
9830 let act = split_act(x);
9831 let run = move |start: usize, end: usize| {
9832 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9833 };
9834 dispatch_rows(pool, rows, &run);
9835 return;
9836 }
9837
9838 let run =
9839 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9840 dispatch_rows(pool, rows, &run);
9841}
9842
9843#[inline]
9846#[allow(unreachable_code)]
9847#[cfg(target_arch = "x86_64")]
9852#[target_feature(enable = "avx2")]
9853unsafe fn dot_q4b_row_1x4_avx2(
9854 buf: &[u8],
9855 scales: &[u8],
9856 g0: usize,
9857 gpr: usize,
9858 xs: [&[i8]; 4],
9859) -> [f32; 4] {
9860 unsafe {
9862 use core::arch::x86_64::*;
9863 let ones = _mm256_set1_epi16(1);
9864 let mut acc = [0f32; 4];
9865 for gi in 0..gpr {
9866 let s = f16_to_f32(u16::from_le_bytes([
9867 scales[(g0 + gi) * 2],
9868 scales[(g0 + gi) * 2 + 1],
9869 ]));
9870 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9871 let aw = _mm256_abs_epi8(w);
9872 for (k, xq) in xs.iter().enumerate() {
9873 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9874 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9875 let d = _mm256_madd_epi16(p16, ones);
9876 let hi128 = _mm256_extracti128_si256::<1>(d);
9877 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9878 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9879 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9880 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9881 }
9882 }
9883 acc
9884 }
9885}
9886
9887#[cfg(target_arch = "x86_64")]
9889#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9890unsafe fn dot_q4b_row_1x4_vnni(
9891 buf: &[u8],
9892 scales: &[u8],
9893 g0: usize,
9894 gpr: usize,
9895 xs: [&[i8]; 4],
9896) -> [f32; 4] {
9897 unsafe {
9899 use core::arch::x86_64::*;
9900 let mut acc = [0f32; 4];
9901 for gi in 0..gpr {
9902 let s = f16_to_f32(u16::from_le_bytes([
9903 scales[(g0 + gi) * 2],
9904 scales[(g0 + gi) * 2 + 1],
9905 ]));
9906 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9907 let aw = _mm256_abs_epi8(w);
9908 for (k, xq) in xs.iter().enumerate() {
9909 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9910 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9911 acc[k] += d as f32 * s;
9912 }
9913 }
9914 acc
9915 }
9916}
9917
9918#[cfg(target_arch = "x86_64")]
9924#[target_feature(enable = "avx2")]
9925unsafe fn dot_q4b_row_1x4_sx_avx2(
9926 buf: &[u8],
9927 scales: &[u8],
9928 g0: usize,
9929 gpr: usize,
9930 xs: [&[i8]; 4],
9931 sxs: [f32; 4],
9932) -> [f32; 4] {
9933 unsafe {
9935 use core::arch::x86_64::*;
9936 let ones = _mm256_set1_epi16(1);
9937 let mut acc = [0f32; 4];
9938 for gi in 0..gpr {
9939 let s = f16_to_f32(u16::from_le_bytes([
9940 scales[(g0 + gi) * 2],
9941 scales[(g0 + gi) * 2 + 1],
9942 ]));
9943 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9944 let aw = _mm256_abs_epi8(w);
9945 for (k, xq) in xs.iter().enumerate() {
9946 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9947 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9948 let d = _mm256_madd_epi16(p16, ones);
9949 let hi128 = _mm256_extracti128_si256::<1>(d);
9950 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9951 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9952 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9953 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9954 }
9955 }
9956 acc
9957 }
9958}
9959
9960#[cfg(target_arch = "x86_64")]
9963#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9964unsafe fn dot_q4b_row_1x4_sx_vnni(
9965 buf: &[u8],
9966 scales: &[u8],
9967 g0: usize,
9968 gpr: usize,
9969 xs: [&[i8]; 4],
9970 sxs: [f32; 4],
9971) -> [f32; 4] {
9972 unsafe {
9974 use core::arch::x86_64::*;
9975 let mut acc = [0f32; 4];
9976 for gi in 0..gpr {
9977 let s = f16_to_f32(u16::from_le_bytes([
9978 scales[(g0 + gi) * 2],
9979 scales[(g0 + gi) * 2 + 1],
9980 ]));
9981 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9982 let aw = _mm256_abs_epi8(w);
9983 for (k, xq) in xs.iter().enumerate() {
9984 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9985 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9986 acc[k] += (d as f32 * sxs[k]) * s;
9987 }
9988 }
9989 acc
9990 }
9991}
9992
9993#[allow(unreachable_code)]
9994fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9995 #[cfg(target_arch = "aarch64")]
9996 unsafe {
9997 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9998 }
9999 #[cfg(target_arch = "x86_64")]
10000 unsafe {
10001 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10002 }
10003 let mut acc = 0f32;
10004 for gi in 0..gpr {
10005 let g = g0 + gi;
10006 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10007 let mut d = 0i32;
10008 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10009 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10010 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10011 }
10012 acc += d as f32 * s;
10013 }
10014 acc
10015}
10016
10017#[inline]
10019#[allow(unreachable_code)]
10020fn dot_q4_row_i8_2(
10021 packed: &[u8],
10022 scales: &[u8],
10023 g0: usize,
10024 gpr: usize,
10025 xq1: &[i8],
10026 xq2: &[i8],
10027) -> (f32, f32) {
10028 #[cfg(target_arch = "aarch64")]
10029 unsafe {
10030 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10031 }
10032 #[cfg(target_arch = "x86_64")]
10033 unsafe {
10034 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10035 }
10036 (
10037 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10038 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10039 )
10040}
10041
10042#[allow(clippy::too_many_arguments)]
10045fn q4_range_a8w8(
10046 packed: &[u8],
10047 scales: &[u8],
10048 gpr: usize,
10049 cols: usize,
10050 act: &SplitAct,
10051 out: SendMut,
10052 start: usize,
10053 end: usize,
10054) {
10055 for r in start..end {
10056 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10057 for &(j, xv) in &act.outliers {
10059 let flat = r * cols + j;
10060 let byte = packed[flat / 2];
10061 let nib = if flat & 1 == 0 {
10062 byte & 0x0F
10063 } else {
10064 byte >> 4
10065 };
10066 let s = f16_to_f32(u16::from_le_bytes([
10067 scales[(flat / GROUP_SIZE) * 2],
10068 scales[(flat / GROUP_SIZE) * 2 + 1],
10069 ]));
10070 acc += ((nib as i32 - 8) as f32) * s * xv;
10071 }
10072 unsafe { *out.at(r) = acc };
10074 }
10075}
10076
10077#[allow(clippy::too_many_arguments)]
10080fn q4_range2_a8w8(
10081 packed: &[u8],
10082 scales: &[u8],
10083 gpr: usize,
10084 cols: usize,
10085 a1: &SplitAct,
10086 a2: &SplitAct,
10087 p1: SendMut,
10088 p2: SendMut,
10089 start: usize,
10090 end: usize,
10091) {
10092 for r in start..end {
10093 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
10094 let mut acc1 = s1 * a1.sx;
10095 let mut acc2 = s2 * a2.sx;
10096 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
10098 for &(j, xv) in outliers {
10099 let flat = r * cols + j;
10100 let byte = packed[flat / 2];
10101 let nib = if flat & 1 == 0 {
10102 byte & 0x0F
10103 } else {
10104 byte >> 4
10105 };
10106 let s = f16_to_f32(u16::from_le_bytes([
10107 scales[(flat / GROUP_SIZE) * 2],
10108 scales[(flat / GROUP_SIZE) * 2 + 1],
10109 ]));
10110 *acc += ((nib as i32 - 8) as f32) * s * xv;
10111 }
10112 };
10113 fix(&a1.outliers, &mut acc1);
10114 fix(&a2.outliers, &mut acc2);
10115 unsafe {
10117 *p1.at(r) = acc1;
10118 *p2.at(r) = acc2;
10119 }
10120 }
10121}
10122
10123fn q4_range_f32(
10125 packed: &[u8],
10126 scales: &[u8],
10127 gpr: usize,
10128 x: &[f32],
10129 out: SendMut,
10130 start: usize,
10131 end: usize,
10132) {
10133 for r in start..end {
10134 let mut acc = 0f32;
10135 for gi in 0..gpr {
10136 let g = r * gpr + gi;
10137 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10138 let pk = &packed[g * 16..(g + 1) * 16];
10139 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10140 let mut ga = 0f32;
10141 for (k, &b) in pk.iter().enumerate() {
10142 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10143 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10144 }
10145 acc += ga * s;
10146 }
10147 unsafe { *out.at(r) = acc };
10149 }
10150}
10151
10152#[allow(clippy::too_many_arguments)]
10156fn q4matvec2(
10157 bytes: &[u8],
10158 x1: &[f32],
10159 x2: &[f32],
10160 rows: usize,
10161 cols: usize,
10162 o1: &mut [f32],
10163 o2: &mut [f32],
10164 pool: Option<&Pool>,
10165) {
10166 debug_assert_eq!(o1.len(), rows);
10167 debug_assert_eq!(o2.len(), rows);
10168 let (packed, scales) = q4_split(bytes, rows, cols);
10169 let gpr = cols / GROUP_SIZE;
10170
10171 if a8w8_enabled() {
10172 let a1 = split_act(x1);
10173 let a2 = split_act(x2);
10174 let p1 = SendMut(o1.as_mut_ptr());
10175 let p2 = SendMut(o2.as_mut_ptr());
10176 let run = move |start: usize, end: usize| {
10177 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10178 };
10179 dispatch_rows(pool, rows, &run);
10180 return;
10181 }
10182
10183 let p1 = SendMut(o1.as_mut_ptr());
10184 let p2 = SendMut(o2.as_mut_ptr());
10185 let run = move |start: usize, end: usize| {
10186 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10187 };
10188 dispatch_rows(pool, rows, &run);
10189}
10190
10191#[allow(clippy::too_many_arguments)]
10193fn q4_range2_f32(
10194 packed: &[u8],
10195 scales: &[u8],
10196 gpr: usize,
10197 x1: &[f32],
10198 x2: &[f32],
10199 p1: SendMut,
10200 p2: SendMut,
10201 start: usize,
10202 end: usize,
10203) {
10204 for r in start..end {
10205 let (mut acc1, mut acc2) = (0f32, 0f32);
10206 for gi in 0..gpr {
10207 let g = r * gpr + gi;
10208 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10209 let pk = &packed[g * 16..(g + 1) * 16];
10210 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10211 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10212 let (mut g1, mut g2) = (0f32, 0f32);
10213 for (k, &b) in pk.iter().enumerate() {
10214 let wl = (b & 0x0F) as f32 - 8.0;
10215 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10216 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10217 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10218 }
10219 acc1 += g1 * s;
10220 acc2 += g2 * s;
10221 }
10222 unsafe {
10224 *p1.at(r) = acc1;
10225 *p2.at(r) = acc2;
10226 }
10227 }
10228}
10229
10230thread_local! {
10231 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10234 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10235}
10236
10237#[allow(clippy::too_many_arguments)]
10243fn q4matmat(
10244 bytes: &[u8],
10245 xs_all: &[f32],
10246 b: usize,
10247 rows: usize,
10248 cols: usize,
10249 out: &mut [f32],
10250 pool: Option<&Pool>,
10251) {
10252 debug_assert_eq!(xs_all.len(), b * cols);
10253 debug_assert_eq!(out.len(), b * rows);
10254 let (packed, scales) = q4_split(bytes, rows, cols);
10255 let gpr = cols / GROUP_SIZE;
10256 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10257
10258 if a8w8_enabled() {
10259 let acts: Vec<SplitAct> = (0..b)
10260 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10261 .collect();
10262 let acts = &acts;
10263 let out_addr = SendMut(out.as_mut_ptr());
10264 let run = move |start: usize, end: usize| {
10265 ROW_I8.with(|rb| {
10266 let mut buf = rb.borrow_mut();
10267 buf.resize(cols, 0);
10268 for r in start..end {
10269 for gi in 0..gpr {
10273 let g = r * gpr + gi;
10274 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10275 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10276 buf[gi * GROUP_SIZE + k * 2 + 1] =
10277 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10278 }
10279 }
10280 let mut bi = 0usize;
10281 #[cfg(target_arch = "x86_64")]
10282 if avx2_enabled() && blocked_enabled() {
10283 while bi + 4 <= acts.len() {
10284 let xs = [
10285 acts[bi].xq.as_slice(),
10286 acts[bi + 1].xq.as_slice(),
10287 acts[bi + 2].xq.as_slice(),
10288 acts[bi + 3].xq.as_slice(),
10289 ];
10290 let d = unsafe {
10291 if vnni_tiles_enabled() {
10292 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10293 } else {
10294 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10295 }
10296 };
10297 for k in 0..4 {
10298 let act = &acts[bi + k];
10299 let mut acc = d[k] * act.sx;
10300 for &(j, xv) in &act.outliers {
10301 acc += (buf[j] as i8) as f32
10302 * gscale((r * cols + j) / GROUP_SIZE)
10303 * xv;
10304 }
10305 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10307 }
10308 bi += 4;
10309 }
10310 }
10311 while bi < acts.len() {
10312 let act = &acts[bi];
10313 let mut acc = 0f32;
10314 for gi in 0..gpr {
10315 let d = dot_i8_i8(
10316 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10317 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10318 );
10319 acc += d as f32 * gscale(r * gpr + gi);
10320 }
10321 acc *= act.sx;
10322 for &(j, xv) in &act.outliers {
10324 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10325 }
10326 unsafe { *out_addr.at(bi * rows + r) = acc };
10328 bi += 1;
10329 }
10330 }
10331 })
10332 };
10333 dispatch_rows(pool, rows, &run);
10334 return;
10335 }
10336
10337 let out_addr = SendMut(out.as_mut_ptr());
10338 let run = move |start: usize, end: usize| {
10339 ROW_F32.with(|rb| {
10340 let mut buf = rb.borrow_mut();
10341 buf.resize(cols, 0.0);
10342 for r in start..end {
10343 for gi in 0..gpr {
10346 let g = r * gpr + gi;
10347 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10348 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10349 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10350 }
10351 }
10352 for bi in 0..b {
10353 let x = &xs_all[bi * cols..(bi + 1) * cols];
10354 let mut acc = 0f32;
10355 for gi in 0..gpr {
10356 let mut ga = 0f32;
10357 for k in 0..GROUP_SIZE / 2 {
10362 let e = gi * GROUP_SIZE + k * 2;
10363 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10364 }
10365 acc += ga * gscale(r * gpr + gi);
10366 }
10367 unsafe { *out_addr.at(bi * rows + r) = acc };
10369 }
10370 }
10371 })
10372 };
10373 dispatch_rows(pool, rows, &run);
10374}
10375
10376#[allow(clippy::too_many_arguments)]
10381fn vbitmatmat(
10382 bytes: &[u8],
10383 offsets: &[usize],
10384 xs_all: &[f32],
10385 b: usize,
10386 rows: usize,
10387 cols: usize,
10388 out: &mut [f32],
10389 pool: Option<&Pool>,
10390) {
10391 debug_assert_eq!(xs_all.len(), b * cols);
10392 debug_assert_eq!(out.len(), b * rows);
10393 debug_assert_eq!(offsets.len(), rows + 1);
10394 let ng = cols / GROUP_SIZE;
10395 let bits = &bytes[..rows];
10396 let sc_off = rows;
10397 let gscale = |r: usize, g: usize| {
10398 let so = (r * ng + g) * 2;
10399 f16_to_f32(u16::from_le_bytes([
10400 bytes[sc_off + so],
10401 bytes[sc_off + so + 1],
10402 ]))
10403 };
10404
10405 let decode_f32 = |r: usize, dst: &mut [f32]| {
10407 let bw = bits[r] as usize;
10408 let l = ((1i32 << (bw - 1)) - 1) as f32;
10409 let data = &bytes[offsets[r]..offsets[r + 1]];
10410 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10411 for d in dst.iter_mut() {
10412 while nbits < bw {
10413 acc = (acc << 8) | data[idx] as u64;
10414 idx += 1;
10415 nbits += 8;
10416 }
10417 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10418 nbits -= bw;
10419 *d = u - l;
10420 }
10421 };
10422
10423 if a8w8_enabled() {
10424 let acts: Vec<SplitAct> = (0..b)
10425 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10426 .collect();
10427 let acts = &acts;
10428 let out_addr = SendMut(out.as_mut_ptr());
10429 let run = move |start: usize, end: usize| {
10430 for r in start..end {
10431 let bw = bits[r] as usize;
10432 if bw == 8 {
10433 ROW_F32.with(|rb| {
10436 let mut buf = rb.borrow_mut();
10437 buf.resize(cols, 0.0);
10438 decode_f32(r, &mut buf);
10439 for bi in 0..b {
10440 let x = &xs_all[bi * cols..(bi + 1) * cols];
10441 let mut dot = 0f32;
10442 for g in 0..ng {
10443 let mut gd = 0f32;
10444 for k in 0..GROUP_SIZE {
10445 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10446 }
10447 dot += gd * gscale(r, g);
10448 }
10449 unsafe { *out_addr.at(bi * rows + r) = dot };
10451 }
10452 });
10453 continue;
10454 }
10455 let l = (1i32 << (bw - 1)) - 1;
10456 let data = &bytes[offsets[r]..offsets[r + 1]];
10457 ROW_I8.with(|rb| {
10458 let mut buf = rb.borrow_mut();
10459 buf.resize(cols, 0);
10460 #[inline(always)]
10461 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10462 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10463 let u = unpack8::<B>(&data[blk * B..]);
10464 for k in 0..8 {
10465 chunk[k] = (u[k] - l) as i8 as u8;
10466 }
10467 }
10468 }
10469 match bw {
10470 3 => fill::<3>(data, l, &mut buf),
10471 4 => vbit_fill4(data, &mut buf),
10472 5 => fill::<5>(data, l, &mut buf),
10473 6 => fill::<6>(data, l, &mut buf),
10474 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10475 }
10476 let mut bi = 0usize;
10477 #[cfg(target_arch = "x86_64")]
10481 if avx2_enabled() && blocked_enabled() {
10482 while bi + 4 <= acts.len() {
10483 let xs = [
10484 acts[bi].xq.as_slice(),
10485 acts[bi + 1].xq.as_slice(),
10486 acts[bi + 2].xq.as_slice(),
10487 acts[bi + 3].xq.as_slice(),
10488 ];
10489 let sxs = [
10490 acts[bi].sx,
10491 acts[bi + 1].sx,
10492 acts[bi + 2].sx,
10493 acts[bi + 3].sx,
10494 ];
10495 let d = unsafe {
10496 if vnni_tiles_enabled() {
10497 dot_q4b_row_1x4_sx_vnni(
10498 &buf,
10499 &bytes[sc_off..],
10500 r * ng,
10501 ng,
10502 xs,
10503 sxs,
10504 )
10505 } else {
10506 dot_q4b_row_1x4_sx_avx2(
10507 &buf,
10508 &bytes[sc_off..],
10509 r * ng,
10510 ng,
10511 xs,
10512 sxs,
10513 )
10514 }
10515 };
10516 for k in 0..4 {
10517 let act = &acts[bi + k];
10518 let mut dot = d[k];
10519 for &(j, xv) in &act.outliers {
10520 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10521 }
10522 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10524 }
10525 bi += 4;
10526 }
10527 }
10528 while bi < acts.len() {
10529 let act = &acts[bi];
10530 let mut dot = 0f32;
10531 for g in 0..ng {
10532 let d = dot_i8_i8(
10533 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10534 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10535 ) as f32
10536 * act.sx;
10537 dot += d * gscale(r, g);
10538 }
10539 for &(j, xv) in &act.outliers {
10540 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10541 }
10542 unsafe { *out_addr.at(bi * rows + r) = dot };
10544 bi += 1;
10545 }
10546 });
10547 }
10548 };
10549 dispatch_rows(pool, rows, &run);
10550 return;
10551 }
10552
10553 let out_addr = SendMut(out.as_mut_ptr());
10554 let run = move |start: usize, end: usize| {
10555 ROW_F32.with(|rb| {
10556 let mut buf = rb.borrow_mut();
10557 buf.resize(cols, 0.0);
10558 for r in start..end {
10559 decode_f32(r, &mut buf);
10560 for bi in 0..b {
10561 let x = &xs_all[bi * cols..(bi + 1) * cols];
10562 let mut dot = 0f32;
10563 for g in 0..ng {
10564 let mut gd = 0f32;
10565 for k in 0..GROUP_SIZE {
10566 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10567 }
10568 dot += gd * gscale(r, g);
10569 }
10570 unsafe { *out_addr.at(bi * rows + r) = dot };
10572 }
10573 }
10574 })
10575 };
10576 dispatch_rows(pool, rows, &run);
10577}
10578
10579pub(crate) fn gpu_batch_job<'a>(
10583 t: &'a QTensor,
10584 x: &[f32],
10585) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10586 match t {
10587 QTensor::Mapped {
10588 model,
10589 idx,
10590 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10591 rows,
10592 cols,
10593 row_scale,
10594 col_field,
10595 ..
10596 } => Some((
10597 model.clone(),
10598 crate::gpu::BatchJob {
10599 idx: *idx,
10600 rows: *rows,
10601 cols: *cols,
10602 row_scale,
10603 xs: prescale(x, col_field, *dt).into_owned(),
10604 layout: crate::gpu::BatchLayout::Q8,
10605 },
10606 )),
10607 QTensor::Mapped {
10609 model,
10610 idx,
10611 dtype: TensorDtype::Q1,
10612 rows,
10613 cols,
10614 ..
10615 } => Some((
10616 model.clone(),
10617 crate::gpu::BatchJob {
10618 idx: *idx,
10619 rows: *rows,
10620 cols: *cols,
10621 row_scale: &[],
10622 xs: x.to_vec(),
10623 layout: crate::gpu::BatchLayout::Q1,
10624 },
10625 )),
10626 QTensor::Mapped {
10628 model,
10629 idx,
10630 dtype: TensorDtype::Q1G,
10631 rows,
10632 cols,
10633 ..
10634 } => Some((
10635 model.clone(),
10636 crate::gpu::BatchJob {
10637 idx: *idx,
10638 rows: *rows,
10639 cols: *cols,
10640 row_scale: &[],
10641 xs: x.to_vec(),
10642 layout: crate::gpu::BatchLayout::Q1g,
10643 },
10644 )),
10645 QTensor::Mapped {
10650 model,
10651 idx,
10652 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10653 rows,
10654 cols,
10655 ..
10656 } => Some((
10657 model.clone(),
10658 crate::gpu::BatchJob {
10659 idx: *idx,
10660 rows: *rows,
10661 cols: *cols,
10662 row_scale: &[],
10663 xs: x.to_vec(),
10664 layout: if *dt == TensorDtype::Q4Tiled {
10665 crate::gpu::BatchLayout::Q4t
10666 } else {
10667 crate::gpu::BatchLayout::Q4tp
10668 },
10669 },
10670 )),
10671 _ => None,
10672 }
10673}
10674
10675thread_local! {
10676 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10677 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10678}
10679
10680pub(crate) fn prescale<'a>(
10681 x: &'a [f32],
10682 col_field: &[f32],
10683 dtype: TensorDtype,
10684) -> std::borrow::Cow<'a, [f32]> {
10685 if dtype == TensorDtype::Q8_2f {
10686 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10687 } else {
10688 std::borrow::Cow::Borrowed(x)
10689 }
10690}
10691
10692pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10695 x: &[f32],
10696 col_field: &[f32],
10697 dtype: TensorDtype,
10698 buf_id: u8,
10699 f: F,
10700) -> R {
10701 if dtype == TensorDtype::Q8_2f {
10702 if buf_id == 1 {
10703 PRESCALE_BUF1.with(|b| {
10704 let mut buf = b.borrow_mut();
10705 buf.clear();
10706 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10707 f(&buf)
10708 })
10709 } else {
10710 PRESCALE_BUF2.with(|b| {
10711 let mut buf = b.borrow_mut();
10712 buf.clear();
10713 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10714 f(&buf)
10715 })
10716 }
10717 } else {
10718 f(x)
10719 }
10720}
10721
10722#[cfg(target_arch = "x86_64")]
10727pub(crate) fn avx2_enabled() -> bool {
10728 use std::sync::OnceLock;
10729 static ON: OnceLock<bool> = OnceLock::new();
10730 *ON.get_or_init(|| {
10731 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10732 && std::arch::is_x86_feature_detected!("avx2")
10733 && std::arch::is_x86_feature_detected!("fma")
10734 })
10735}
10736
10737#[cfg(target_arch = "x86_64")]
10742fn avx2_a8w8_enabled() -> bool {
10743 if FLOAT_ACTIVATIONS.get() {
10744 return false;
10745 }
10746 use std::sync::OnceLock;
10747 static ON: OnceLock<bool> = OnceLock::new();
10748 *ON.get_or_init(|| {
10749 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10750 })
10751}
10752
10753thread_local! {
10754 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10755}
10756
10757pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10760 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10761 impl Drop for Restore {
10762 fn drop(&mut self) {
10763 FULL_GPU_Q8.set(self.0);
10764 }
10765 }
10766 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10767}
10768
10769thread_local! {
10773 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10774}
10775
10776pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
10781 float_activations_scope(f)
10782}
10783
10784#[doc(hidden)]
10789pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10790 struct Restore(bool);
10791 impl Drop for Restore {
10792 fn drop(&mut self) {
10793 FLOAT_ACTIVATIONS.set(self.0);
10794 }
10795 }
10796 let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10797 f()
10798}
10799
10800static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10812
10813pub(crate) fn row_exact() -> bool {
10814 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10815}
10816
10817fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10818 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10819 impl Drop for Restore<'_> {
10820 fn drop(&mut self) {
10821 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10822 }
10823 }
10824 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10825 let _restore = Restore(active);
10826 f()
10827}
10828
10829pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10831 counted_row_exact_scope(&ROW_EXACT, f)
10832}
10833
10834#[inline]
10838pub(crate) fn a8w8_enabled() -> bool {
10839 #[cfg(target_arch = "aarch64")]
10840 {
10841 sdot_enabled()
10842 }
10843 #[cfg(target_arch = "x86_64")]
10844 {
10845 avx2_a8w8_enabled()
10846 }
10847 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10848 {
10849 false
10850 }
10851}
10852
10853#[inline]
10856#[allow(unreachable_code)]
10857fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10858 #[cfg(target_arch = "aarch64")]
10859 unsafe {
10860 return dot_i8_sdot(w, xq);
10861 }
10862 #[cfg(target_arch = "x86_64")]
10863 unsafe {
10864 if avx512vnni_enabled() {
10865 return dot_i8_i8_vnni(w, xq);
10866 }
10867 return dot_i8_i8_avx2(w, xq);
10868 }
10869 w.iter()
10870 .zip(xq)
10871 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10872 .sum()
10873}
10874
10875#[cfg(target_arch = "x86_64")]
10879fn avx512vnni_enabled() -> bool {
10880 use std::sync::OnceLock;
10881 static ON: OnceLock<bool> = OnceLock::new();
10882 *ON.get_or_init(|| {
10883 std::env::var("CMF_AVX512")
10884 .map(|v| v != "0")
10885 .unwrap_or(true)
10886 && std::arch::is_x86_feature_detected!("avx512f")
10887 && std::arch::is_x86_feature_detected!("avx512bw")
10888 && std::arch::is_x86_feature_detected!("avx512vl")
10889 && std::arch::is_x86_feature_detected!("avx512vnni")
10890 })
10891}
10892
10893#[cfg(target_arch = "x86_64")]
10901fn vnni_tiles_enabled() -> bool {
10902 use std::sync::OnceLock;
10903 static ON: OnceLock<bool> = OnceLock::new();
10904 *ON.get_or_init(|| {
10905 std::env::var("CMF_VNNI_TILES")
10906 .map(|v| v != "0")
10907 .unwrap_or(true)
10908 && avx512vnni_enabled()
10909 })
10910}
10911
10912#[cfg(target_arch = "x86_64")]
10917#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10918#[inline]
10919unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10920 unsafe {
10922 use core::arch::x86_64::*;
10923 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10924 let hi128 = _mm256_extracti128_si256::<1>(d);
10925 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10926 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10927 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10928 _mm_cvtsi128_si32(s32)
10929 }
10930}
10931
10932#[cfg(target_arch = "x86_64")]
10937#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10938unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10939 unsafe {
10941 use core::arch::x86_64::*;
10942 let n = w.len();
10943 let mut j = 0usize;
10944 let mut total: i32;
10945 {
10950 #[inline(always)]
10951 unsafe fn step(
10952 w: *const u8,
10953 x: *const i8,
10954 acc: core::arch::x86_64::__m512i,
10955 ) -> core::arch::x86_64::__m512i {
10956 unsafe {
10957 use core::arch::x86_64::*;
10958 let wv = _mm512_loadu_si512(w as *const _);
10959 let xv = _mm512_loadu_si512(x as *const _);
10960 let aw = _mm512_abs_epi8(wv);
10961 let neg = _mm512_movepi8_mask(wv);
10962 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10963 _mm512_dpbusd_epi32(acc, aw, sx)
10964 }
10965 }
10966 let (mut a0, mut a1, mut a2, mut a3) = (
10967 _mm512_setzero_si512(),
10968 _mm512_setzero_si512(),
10969 _mm512_setzero_si512(),
10970 _mm512_setzero_si512(),
10971 );
10972 while j + 256 <= n {
10973 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10974 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10975 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10976 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10977 j += 256;
10978 }
10979 while j + 64 <= n {
10980 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10981 j += 64;
10982 }
10983 let s01 = _mm512_add_epi32(a0, a1);
10984 let s23 = _mm512_add_epi32(a2, a3);
10985 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10986 }
10987 if j + 32 <= n {
10989 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10990 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10991 let d = _mm256_dpbusd_epi32(
10992 _mm256_setzero_si256(),
10993 _mm256_abs_epi8(wv),
10994 _mm256_sign_epi8(xv, wv),
10995 );
10996 let hi128 = _mm256_extracti128_si256::<1>(d);
10997 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10998 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10999 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11000 total += _mm_cvtsi128_si32(s32);
11001 j += 32;
11002 }
11003 while j < n {
11004 total += (w[j] as i8) as i32 * xq[j] as i32;
11005 j += 1;
11006 }
11007 total
11008 }
11009}
11010
11011#[cfg(target_arch = "x86_64")]
11013#[target_feature(enable = "avx2,fma")]
11014unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
11015 unsafe {
11017 use core::arch::x86_64::*;
11018 let n = x.len();
11019 let wp = w.as_ptr();
11020 let xp = x.as_ptr();
11021 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
11022 let mut j = 0usize;
11023 while j + 16 <= n {
11024 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
11025 let lo = _mm256_cvtepi8_epi32(wb);
11026 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
11027 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
11028 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
11029 j += 16;
11030 }
11031 let acc = _mm256_add_ps(a0, a1);
11032 let hi128 = _mm256_extractf128_ps::<1>(acc);
11033 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
11034 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
11035 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
11036 let mut sum = _mm_cvtss_f32(s32);
11037 while j < n {
11038 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
11039 j += 1;
11040 }
11041 sum
11042 }
11043}
11044
11045#[cfg(target_arch = "x86_64")]
11050#[target_feature(enable = "avx2")]
11051unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
11052 unsafe {
11054 use core::arch::x86_64::*;
11055 let n = w.len();
11056 let ones = _mm256_set1_epi16(1);
11057 let mut acc = _mm256_setzero_si256();
11058 let mut j = 0usize;
11059 while j + 32 <= n {
11060 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11061 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11062 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
11063 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
11064 j += 32;
11065 }
11066 let hi128 = _mm256_extracti128_si256::<1>(acc);
11067 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
11068 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11069 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11070 let mut s = _mm_cvtsi128_si32(s32);
11071 while j < n {
11072 s += (w[j] as i8) as i32 * xq[j] as i32;
11073 j += 1;
11074 }
11075 s
11076 }
11077}
11078
11079#[cfg(target_arch = "aarch64")]
11083#[target_feature(enable = "neon,i8mm")]
11084unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11085 unsafe {
11087 use core::arch::aarch64::*;
11088 use core::arch::asm;
11089 let n = w0.len();
11090 let w0p = w0.as_ptr() as *const i8;
11091 let w1p = w1.as_ptr() as *const i8;
11092 let mut acc01 = vdupq_n_s32(0);
11095 let mut acc23 = vdupq_n_s32(0);
11096 let mut i = 0usize;
11097 while i + 8 <= n {
11098 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
11099 let xb01 = vcombine_s8(
11100 vld1_s8(xs[0].as_ptr().add(i)),
11101 vld1_s8(xs[1].as_ptr().add(i)),
11102 );
11103 let xb23 = vcombine_s8(
11104 vld1_s8(xs[2].as_ptr().add(i)),
11105 vld1_s8(xs[3].as_ptr().add(i)),
11106 );
11107 asm!(
11108 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
11109 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
11110 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
11111 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
11112 options(pure, nomem, nostack),
11113 );
11114 i += 8;
11115 }
11116 let mut out = [[0i32; 4]; 2];
11117 let a01: [i32; 4] = core::mem::transmute(acc01);
11118 let a23: [i32; 4] = core::mem::transmute(acc23);
11119 out[0][0] = a01[0];
11120 out[0][1] = a01[1];
11121 out[1][0] = a01[2];
11122 out[1][1] = a01[3];
11123 out[0][2] = a23[0];
11124 out[0][3] = a23[1];
11125 out[1][2] = a23[2];
11126 out[1][3] = a23[3];
11127 if i < n {
11128 for (k, x) in xs.iter().enumerate() {
11129 for j in i..n {
11130 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11131 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11132 }
11133 }
11134 }
11135 out
11136 }
11137}
11138
11139#[cfg(target_arch = "aarch64")]
11143#[target_feature(enable = "neon,dotprod")]
11144unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11145 unsafe {
11147 use core::arch::aarch64::*;
11148 use core::arch::asm;
11149 let n = w0.len();
11150 let w0p = w0.as_ptr() as *const i8;
11151 let w1p = w1.as_ptr() as *const i8;
11152 let mut acc = [[vdupq_n_s32(0); 4]; 2];
11153 let mut i = 0usize;
11154 while i + 16 <= n {
11155 let wv0 = vld1q_s8(w0p.add(i));
11156 let wv1 = vld1q_s8(w1p.add(i));
11157 for (k, x) in xs.iter().enumerate() {
11158 let xv = vld1q_s8(x.as_ptr().add(i));
11159 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11160 asm!(
11161 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11162 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11163 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11164 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11165 options(pure, nomem, nostack),
11166 );
11167 acc[0][k] = a0;
11168 acc[1][k] = a1;
11169 }
11170 i += 16;
11171 }
11172 let mut out = [[0i32; 4]; 2];
11173 for r in 0..2 {
11174 for k in 0..4 {
11175 out[r][k] = vaddvq_s32(acc[r][k]);
11176 }
11177 }
11178 if i < n {
11179 for (k, x) in xs.iter().enumerate() {
11180 for j in i..n {
11181 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11182 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11183 }
11184 }
11185 }
11186 out
11187 }
11188}
11189
11190#[cfg(target_arch = "x86_64")]
11196#[target_feature(enable = "avx2")]
11197unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11198 unsafe {
11200 use core::arch::x86_64::*;
11201 let n = w0.len();
11202 let ones = _mm256_set1_epi16(1);
11203 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11204 let mut j = 0usize;
11205 while j + 32 <= n {
11206 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11207 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11208 let aw0 = _mm256_abs_epi8(wv0);
11209 let aw1 = _mm256_abs_epi8(wv1);
11210 for (k, x) in xs.iter().enumerate() {
11211 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11212 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11213 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11214 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11215 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11216 }
11217 j += 32;
11218 }
11219 let mut out = [[0i32; 4]; 2];
11220 for r in 0..2 {
11221 for k in 0..4 {
11222 let a = acc[r][k];
11223 let hi128 = _mm256_extracti128_si256::<1>(a);
11224 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11225 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11226 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11227 out[r][k] = _mm_cvtsi128_si32(s32);
11228 }
11229 }
11230 if j < n {
11231 for (k, x) in xs.iter().enumerate() {
11232 for i in j..n {
11233 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11234 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11235 }
11236 }
11237 }
11238 out
11239 }
11240}
11241
11242#[cfg(target_arch = "x86_64")]
11247#[inline]
11248fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11249 let dot = if avx512vnni_enabled() && row.len() >= 64 {
11250 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11251 } else {
11252 unsafe { dot_i8_i8_avx2(row, &act.xq) }
11253 };
11254 let mut acc = dot as f32 * act.sx;
11255 for &(j, xv) in &act.outliers {
11256 acc += (row[j] as i8) as f32 * xv;
11257 }
11258 acc
11259}
11260
11261#[cfg(target_arch = "x86_64")]
11265#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11266unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11267 unsafe {
11269 use core::arch::x86_64::*;
11270 let n = w.len();
11271 let flip = _mm512_set1_epi8(-128); #[inline(always)]
11273 unsafe fn step(
11274 w: *const u8,
11275 x: *const i8,
11276 flip: core::arch::x86_64::__m512i,
11277 acc: core::arch::x86_64::__m512i,
11278 ) -> core::arch::x86_64::__m512i {
11279 unsafe {
11280 use core::arch::x86_64::*;
11281 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11282 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11283 }
11284 }
11285 let (mut a0, mut a1, mut a2, mut a3) = (
11286 _mm512_setzero_si512(),
11287 _mm512_setzero_si512(),
11288 _mm512_setzero_si512(),
11289 _mm512_setzero_si512(),
11290 );
11291 let mut j = 0usize;
11292 while j + 256 <= n {
11293 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11294 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11295 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11296 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11297 j += 256;
11298 }
11299 while j + 64 <= n {
11300 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11301 j += 64;
11302 }
11303 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11304 _mm512_add_epi32(a0, a1),
11305 _mm512_add_epi32(a2, a3),
11306 ));
11307 while j < n {
11309 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11310 j += 1;
11311 }
11312 total
11313 }
11314}
11315
11316#[cfg(target_arch = "x86_64")]
11322#[target_feature(enable = "avx2")]
11323unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11324 unsafe {
11327 use core::arch::x86_64::*;
11328 let lomask = _mm_set1_epi8(0x0F);
11329 let eight = _mm256_set1_epi8(8);
11330 let ones = _mm256_set1_epi16(1);
11331 let mut acc = 0f32;
11332 for gi in 0..gpr {
11333 let g = g0 + gi;
11334 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11335 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11336 let lo = _mm_and_si128(b, lomask);
11337 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11338 let w = _mm256_sub_epi8(
11339 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11340 eight,
11341 );
11342 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11343 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11344 let d = _mm256_madd_epi16(p16, ones);
11345 let hi128 = _mm256_extracti128_si256::<1>(d);
11346 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11347 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11348 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11349 acc += _mm_cvtsi128_si32(s32) as f32 * s;
11350 }
11351 acc
11352 }
11353}
11354
11355#[cfg(target_arch = "x86_64")]
11358#[target_feature(enable = "avx2")]
11359unsafe fn dot_q4_row_avx2_2(
11360 packed: &[u8],
11361 scales: &[u8],
11362 g0: usize,
11363 gpr: usize,
11364 xq1: &[i8],
11365 xq2: &[i8],
11366) -> (f32, f32) {
11367 unsafe {
11369 use core::arch::x86_64::*;
11370 let lomask = _mm_set1_epi8(0x0F);
11371 let eight = _mm256_set1_epi8(8);
11372 let ones = _mm256_set1_epi16(1);
11373 let (mut acc1, mut acc2) = (0f32, 0f32);
11374 #[inline(always)]
11375 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11376 unsafe {
11377 use core::arch::x86_64::*;
11378 let hi128 = _mm256_extracti128_si256::<1>(d);
11379 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11380 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11381 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11382 _mm_cvtsi128_si32(s32)
11383 }
11384 }
11385 for gi in 0..gpr {
11386 let g = g0 + gi;
11387 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11388 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11389 let lo = _mm_and_si128(b, lomask);
11390 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11391 let w = _mm256_sub_epi8(
11392 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11393 eight,
11394 );
11395 let aw = _mm256_abs_epi8(w);
11396 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11397 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11398 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11399 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11400 acc1 += hsum(d1) as f32 * s;
11401 acc2 += hsum(d2) as f32 * s;
11402 }
11403 (acc1, acc2)
11404 }
11405}
11406
11407#[cfg(target_arch = "x86_64")]
11409fn q8_range_avx2(
11410 q: &[u8],
11411 row_scale: &[f32],
11412 act: &SplitAct,
11413 cols: usize,
11414 out_addr: SendMut,
11415 start: usize,
11416 end: usize,
11417) {
11418 for o in start..end {
11419 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11420 unsafe { *out_addr.at(o) = v };
11422 }
11423}
11424
11425#[cfg(target_arch = "x86_64")]
11427#[allow(clippy::too_many_arguments)]
11428fn q8_range2_avx2(
11429 q: &[u8],
11430 row_scale: &[f32],
11431 a1: &SplitAct,
11432 a2: &SplitAct,
11433 cols: usize,
11434 p1: SendMut,
11435 p2: SendMut,
11436 start: usize,
11437 end: usize,
11438) {
11439 for o in start..end {
11440 let row = &q[o * cols..(o + 1) * cols];
11441 unsafe {
11443 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11444 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11445 }
11446 }
11447}
11448
11449#[cfg(target_arch = "aarch64")]
11460fn i8mm_enabled() -> bool {
11461 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11462 *ON.get_or_init(|| {
11463 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11464 && std::arch::is_aarch64_feature_detected!("i8mm")
11465 })
11466}
11467
11468#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11472fn sdot_enabled() -> bool {
11473 if FLOAT_ACTIVATIONS.get() {
11474 return false;
11475 }
11476 use std::sync::OnceLock;
11477 static ON: OnceLock<bool> = OnceLock::new();
11478 *ON.get_or_init(|| {
11479 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11480 if !want {
11481 return false;
11482 }
11483
11484 #[cfg(target_arch = "aarch64")]
11485 {
11486 if std::arch::is_aarch64_feature_detected!("dotprod") {
11487 return true;
11488 }
11489 #[cfg(target_os = "android")]
11490 {
11491 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11492 if cpuinfo.lines().any(|l| {
11493 (l.starts_with("Features") || l.starts_with("features"))
11494 && l.contains("asimddp")
11495 }) {
11496 return true;
11497 }
11498 }
11499 }
11500 false
11501 }
11502 #[cfg(not(target_arch = "aarch64"))]
11503 {
11504 false
11505 }
11506 })
11507}
11508
11509struct SplitAct {
11514 xq: Vec<i8>,
11515 sx: f32,
11516 outliers: Vec<(usize, f32)>,
11517 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11520 xsum: i32,
11521}
11522
11523thread_local! {
11524 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11527 const { std::cell::RefCell::new(Vec::new()) };
11528}
11529
11530impl Drop for SplitAct {
11531 fn drop(&mut self) {
11532 let buf = std::mem::take(&mut self.xq);
11533 if buf.capacity() > 0 {
11534 XQ_FREE.with(|f| {
11535 let mut f = f.borrow_mut();
11536 if f.len() < 16 {
11537 f.push(buf);
11538 }
11539 });
11540 }
11541 }
11542}
11543
11544thread_local! {
11545 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11552 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11557 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11558}
11559
11560#[inline]
11563fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11564 KROW.with(|s| {
11565 let b = unsafe { &mut *s.get() };
11568 if b.len() < n {
11569 b.resize(n, 0.0);
11570 }
11571 f(&mut b[..n])
11572 })
11573}
11574
11575#[inline(always)]
11586fn q8_round(t: f32) -> i8 {
11587 let t = t.clamp(-127.0, 127.0);
11588 let i = t as i32;
11589 let f = t - i as f32;
11590 let r = if f >= 0.5 {
11591 i + 1
11592 } else if f <= -0.5 {
11593 i - 1
11594 } else {
11595 i
11596 };
11597 r as i8
11598}
11599
11600#[inline]
11601fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11602 KROWS.with(|s| {
11603 let b = unsafe { &mut *s.get() };
11606 for row in b.iter_mut() {
11607 if row.len() < n {
11608 row.resize(n, 0.0);
11609 }
11610 }
11611 let (a, b) = b.split_at_mut(1);
11612 f(&mut a[0][..n], &mut b[0][..n])
11613 })
11614}
11615
11616#[inline]
11617fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11618 if limit > 0.0 {
11619 up = up.clamp(-limit, limit);
11620 gate = gate.min(limit);
11621 }
11622 gate / (1.0 + (-gate).exp()) * up
11623}
11624
11625fn split_act(x: &[f32]) -> SplitAct {
11626 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11627 let n = x.len();
11628 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11629 let thr = 8.0 * rms;
11630 let mut outliers: Vec<(usize, f32)> = Vec::new();
11634 let mut amax = 0f32;
11635 for (j, &v) in x.iter().enumerate() {
11636 let a = v.abs();
11637 if a > thr {
11638 outliers.push((j, v));
11639 } else if a > amax {
11640 amax = a;
11641 }
11642 }
11643 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11644 let inv = 1.0 / sx;
11645 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11646 xq.clear();
11647 xq.reserve(n);
11648 if outliers.is_empty() {
11649 xq.extend(
11650 x.iter()
11651 .map(|&v| q8_round(v * inv)),
11652 );
11653 } else {
11654 xq.extend(x.iter().map(|&v| {
11656 if v.abs() > thr {
11657 0
11658 } else {
11659 q8_round(v * inv)
11660 }
11661 }));
11662 }
11663 let xsum = xq.iter().map(|&v| v as i32).sum();
11664 SplitAct {
11665 xq,
11666 sx,
11667 outliers,
11668 xsum,
11669 }
11670}
11671
11672fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11673 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11674 let n = x.len();
11675 let rms = (x
11676 .iter()
11677 .zip(col)
11678 .map(|(&a, &c)| {
11679 let v = a * c;
11680 (v * v) as f64
11681 })
11682 .sum::<f64>()
11683 / n.max(1) as f64)
11684 .sqrt() as f32;
11685 let thr = 8.0 * rms;
11686
11687 let mut outliers = Vec::new();
11688 let mut amax = 0f32;
11689 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11690 let v = a * c;
11691 let s = v.abs();
11692 if s > thr {
11693 outliers.push((j, v));
11694 } else if s > amax {
11695 amax = s;
11696 }
11697 }
11698
11699 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11700 let inv = 1.0 / sx;
11701 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11702 xq.clear();
11703 xq.reserve(n);
11704 if outliers.is_empty() {
11705 xq.extend(
11706 x.iter()
11707 .zip(col)
11708 .map(|(&a, &c)| q8_round((a * c) * inv)),
11709 );
11710 } else {
11711 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11712 let v = a * c;
11713 if v.abs() > thr {
11714 0
11715 } else {
11716 q8_round(v * inv)
11717 }
11718 }));
11719 }
11720 let xsum = xq.iter().map(|&v| v as i32).sum();
11721 SplitAct {
11722 xq,
11723 sx,
11724 outliers,
11725 xsum,
11726 }
11727}
11728
11729#[cfg(target_arch = "aarch64")]
11732#[target_feature(enable = "neon,dotprod")]
11733unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11734 unsafe {
11736 use core::arch::aarch64::*;
11737 use core::arch::asm;
11738 let wp = w.as_ptr() as *const i8;
11739 let n = w.len();
11740 let (mut a0, mut a1, mut a2, mut a3) = (
11741 vdupq_n_s32(0),
11742 vdupq_n_s32(0),
11743 vdupq_n_s32(0),
11744 vdupq_n_s32(0),
11745 );
11746 let mut i = 0;
11747 while i + 64 <= n {
11748 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11749 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11750 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11751 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11752 asm!(
11753 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11754 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11755 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11756 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11757 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11758 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11759 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11760 options(pure, nomem, nostack),
11761 );
11762 i += 64;
11763 }
11764 while i + 16 <= n {
11765 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11766 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11767 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11768 i += 16;
11769 }
11770 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11771 while i < n {
11772 s += (*wp.add(i)) as i32 * xq[i] as i32;
11773 i += 1;
11774 }
11775 s
11776 }
11777}
11778
11779#[cfg(target_arch = "aarch64")]
11783#[target_feature(enable = "neon,dotprod")]
11784unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11785 unsafe {
11787 use core::arch::aarch64::*;
11788 use core::arch::asm;
11789 let n = xq.len();
11790 let px = xq.as_ptr();
11791 let (p0, p1, p2, p3) = (
11792 w0.as_ptr() as *const i8,
11793 w1.as_ptr() as *const i8,
11794 w2.as_ptr() as *const i8,
11795 w3.as_ptr() as *const i8,
11796 );
11797 let (mut a0, mut a1, mut a2, mut a3) = (
11798 vdupq_n_s32(0),
11799 vdupq_n_s32(0),
11800 vdupq_n_s32(0),
11801 vdupq_n_s32(0),
11802 );
11803 let mut i = 0;
11804 while i + 16 <= n {
11805 let x = vld1q_s8(px.add(i));
11806 let v0 = vld1q_s8(p0.add(i));
11807 let v1 = vld1q_s8(p1.add(i));
11808 let v2 = vld1q_s8(p2.add(i));
11809 let v3 = vld1q_s8(p3.add(i));
11810 asm!(
11811 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11812 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11813 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11814 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11815 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11816 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11817 options(pure, nomem, nostack),
11818 );
11819 i += 16;
11820 }
11821 let mut r = [
11822 vaddvq_s32(a0),
11823 vaddvq_s32(a1),
11824 vaddvq_s32(a2),
11825 vaddvq_s32(a3),
11826 ];
11827 while i < n {
11828 let xi = *px.add(i) as i32;
11829 r[0] += (*p0.add(i)) as i32 * xi;
11830 r[1] += (*p1.add(i)) as i32 * xi;
11831 r[2] += (*p2.add(i)) as i32 * xi;
11832 r[3] += (*p3.add(i)) as i32 * xi;
11833 i += 1;
11834 }
11835 r
11836 }
11837}
11838
11839#[cfg(target_arch = "aarch64")]
11846#[target_feature(enable = "neon,dotprod")]
11847unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11848 unsafe {
11851 use core::arch::aarch64::*;
11852 use core::arch::asm;
11853 let n = xq.len();
11854 let px = xq.as_ptr();
11855 let pg = g.as_ptr() as *const i8;
11856 let (mut a0, mut a1, mut a2, mut a3) = (
11857 vdupq_n_s32(0),
11858 vdupq_n_s32(0),
11859 vdupq_n_s32(0),
11860 vdupq_n_s32(0),
11861 );
11862 let mut i = 0;
11863 while i + 16 <= n {
11864 let x = vld1q_s8(px.add(i));
11865 let base = pg.add(4 * i);
11866 let v0 = vld1q_s8(base);
11867 let v1 = vld1q_s8(base.add(16));
11868 let v2 = vld1q_s8(base.add(32));
11869 let v3 = vld1q_s8(base.add(48));
11870 asm!(
11871 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11872 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11873 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11874 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11875 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11876 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11877 options(pure, nomem, nostack),
11878 );
11879 i += 16;
11880 }
11881 [
11882 vaddvq_s32(a0),
11883 vaddvq_s32(a1),
11884 vaddvq_s32(a2),
11885 vaddvq_s32(a3),
11886 ]
11887 }
11888}
11889
11890#[cfg(target_arch = "aarch64")]
11896fn q8_range_sdot(
11897 q: &[u8],
11898 rep: &[u8],
11899 row_scale: &[f32],
11900 act: &SplitAct,
11901 cols: usize,
11902 out_addr: SendMut,
11903 start: usize,
11904 end: usize,
11905) {
11906 let mut o = start;
11907 if !rep.is_empty() {
11910 while o < end && o % 4 != 0 {
11911 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11912 unsafe { *out_addr.at(o) = v };
11913 o += 1;
11914 }
11915 }
11916 while o + 4 <= end {
11917 let r = if rep.is_empty() {
11918 unsafe {
11919 dot_i8_sdot_4rows(
11920 &q[o * cols..(o + 1) * cols],
11921 &q[(o + 1) * cols..(o + 2) * cols],
11922 &q[(o + 2) * cols..(o + 3) * cols],
11923 &q[(o + 3) * cols..(o + 4) * cols],
11924 &act.xq,
11925 )
11926 }
11927 } else {
11928 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11929 };
11930 for k in 0..4 {
11931 let mut acc = r[k] as f32 * act.sx;
11932 for &(j, xv) in &act.outliers {
11933 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11934 }
11935 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11937 }
11938 o += 4;
11939 }
11940 while o < end {
11941 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11942 unsafe { *out_addr.at(o) = v };
11943 o += 1;
11944 }
11945}
11946
11947#[cfg(target_arch = "aarch64")]
11950#[allow(clippy::too_many_arguments)]
11951fn q8_range2_sdot(
11952 q: &[u8],
11953 row_scale: &[f32],
11954 a1: &SplitAct,
11955 a2: &SplitAct,
11956 cols: usize,
11957 p1: SendMut,
11958 p2: SendMut,
11959 start: usize,
11960 end: usize,
11961) {
11962 for o in start..end {
11963 let row = &q[o * cols..(o + 1) * cols];
11964 unsafe {
11966 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11967 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11968 }
11969 }
11970}
11971
11972#[allow(clippy::too_many_arguments)]
11974fn q8_range2_f32(
11975 q: &[u8],
11976 row_scale: &[f32],
11977 x1: &[f32],
11978 x2: &[f32],
11979 cols: usize,
11980 p1: SendMut,
11981 p2: SendMut,
11982 start: usize,
11983 end: usize,
11984) {
11985 for o in start..end {
11986 let row = &q[o * cols..(o + 1) * cols];
11987 unsafe {
11989 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11990 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11991 }
11992 }
11993}
11994
11995fn q8_range_f32(
11997 q: &[u8],
11998 row_scale: &[f32],
11999 xs: &[f32],
12000 cols: usize,
12001 out_addr: SendMut,
12002 start: usize,
12003 end: usize,
12004) {
12005 for o in start..end {
12006 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12007 unsafe { *out_addr.at(o) = v };
12009 }
12010}
12011
12012#[inline]
12016fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
12017 #[cfg(target_arch = "aarch64")]
12018 return row_dot_sdot(row, act);
12019 #[cfg(target_arch = "x86_64")]
12020 return row_dot_avx2(row, act);
12021 #[allow(unreachable_code)]
12022 q8_row_dot_scalar(row, act)
12023}
12024
12025#[allow(dead_code)]
12026fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
12027 let mut acc = 0i32;
12028 for (k, &b) in row.iter().enumerate() {
12029 acc += (b as i8) as i32 * act.xq[k] as i32;
12030 }
12031 let mut acc = acc as f32 * act.sx;
12032 for &(j, xv) in &act.outliers {
12033 acc += (row[j] as i8) as f32 * xv;
12034 }
12035 acc
12036}
12037
12038#[cfg(target_arch = "aarch64")]
12041#[inline]
12042fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
12043 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
12044 for &(j, xv) in &act.outliers {
12045 acc += (row[j] as i8) as f32 * xv;
12046 }
12047 acc
12048}
12049
12050#[cfg(target_arch = "aarch64")]
12058#[target_feature(enable = "neon,dotprod")]
12059unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12060 unsafe {
12063 use core::arch::aarch64::*;
12064 use core::arch::asm;
12065 let lomask = vdupq_n_u8(0x0F);
12066 let eight = vdupq_n_s8(8);
12067 let mut acc = 0f32;
12068 for gi in 0..gpr {
12069 let g = g0 + gi;
12070 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12071 let b = vld1q_u8(packed.as_ptr().add(g * 16));
12072 let lo = vandq_u8(b, lomask);
12073 let hi = vshrq_n_u8::<4>(b);
12074 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12075 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12076 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
12077 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
12078 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
12079 asm!(
12080 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
12081 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
12082 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12083 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
12084 options(pure, nomem, nostack),
12085 );
12086 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12087 }
12088 acc
12089 }
12090}
12091
12092#[cfg(target_arch = "aarch64")]
12097#[target_feature(enable = "neon,dotprod")]
12098unsafe fn dot_q4_row_sdot2(
12099 packed: &[u8],
12100 scales: &[u8],
12101 g0: usize,
12102 gpr: usize,
12103 xq1: &[i8],
12104 xq2: &[i8],
12105) -> (f32, f32) {
12106 unsafe {
12109 use core::arch::aarch64::*;
12110 use core::arch::asm;
12111 let lomask = vdupq_n_u8(0x0F);
12112 let eight = vdupq_n_s8(8);
12113 let (mut acc1, mut acc2) = (0f32, 0f32);
12114 for gi in 0..gpr {
12115 let g = g0 + gi;
12116 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12117 let b = vld1q_u8(packed.as_ptr().add(g * 16));
12118 let lo = vandq_u8(b, lomask);
12119 let hi = vshrq_n_u8::<4>(b);
12120 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12121 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12122 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
12123 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
12124 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
12125 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
12126 let (mut a0, mut a1, mut b0, mut b1) = (
12127 vdupq_n_s32(0),
12128 vdupq_n_s32(0),
12129 vdupq_n_s32(0),
12130 vdupq_n_s32(0),
12131 );
12132 asm!(
12133 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
12134 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
12135 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
12136 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
12137 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12138 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
12139 e0 = in(vreg) e0, e1 = in(vreg) e1,
12140 x10 = in(vreg) x10, x11 = in(vreg) x11,
12141 x20 = in(vreg) x20, x21 = in(vreg) x21,
12142 options(pure, nomem, nostack),
12143 );
12144 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12145 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
12146 }
12147 (acc1, acc2)
12148 }
12149}
12150
12151#[inline]
12156pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12157 #[cfg(target_arch = "aarch64")]
12158 unsafe {
12159 return axpy_i8_f32_neon(acc, row, w);
12160 }
12161 #[cfg(target_arch = "x86_64")]
12162 if avx2_enabled() {
12163 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12164 }
12165 #[allow(unreachable_code)]
12166 {
12167 for (a, &b) in acc.iter_mut().zip(row) {
12168 *a += w * b as f32;
12169 }
12170 }
12171}
12172
12173#[cfg(target_arch = "x86_64")]
12175#[target_feature(enable = "avx2,fma")]
12176unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12177 unsafe {
12179 use core::arch::x86_64::*;
12180 let n = acc.len().min(row.len());
12181 let ap = acc.as_mut_ptr();
12182 let rp = row.as_ptr();
12183 let wv = _mm256_set1_ps(w);
12184 let mut j = 0usize;
12185 while j + 16 <= n {
12186 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12187 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12188 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12189 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12190 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12191 _mm256_storeu_ps(ap.add(j), v0);
12192 _mm256_storeu_ps(ap.add(j + 8), v1);
12193 j += 16;
12194 }
12195 while j < n {
12196 *ap.add(j) += w * (*rp.add(j)) as f32;
12197 j += 1;
12198 }
12199 }
12200}
12201
12202#[cfg(target_arch = "aarch64")]
12203#[target_feature(enable = "neon")]
12204unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12205 unsafe {
12207 use core::arch::aarch64::*;
12208 let n = acc.len().min(row.len());
12209 let ap = acc.as_mut_ptr();
12210 let rp = row.as_ptr();
12211 let wv = vdupq_n_f32(w);
12212 let mut j = 0usize;
12213 while j + 16 <= n {
12214 let rb = vld1q_s8(rp.add(j));
12215 let lo = vmovl_s8(vget_low_s8(rb));
12216 let hi = vmovl_s8(vget_high_s8(rb));
12217 for (off, half) in [(0, lo), (8, hi)] {
12218 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12219 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12220 let o = j + off;
12221 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12222 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12223 }
12224 j += 16;
12225 }
12226 while j < n {
12227 *ap.add(j) += w * (*rp.add(j)) as f32;
12228 j += 1;
12229 }
12230 }
12231}
12232
12233#[inline]
12236pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12237 #[cfg(target_arch = "aarch64")]
12238 unsafe {
12239 return dot_i8_f32_neon(w, x);
12240 }
12241 #[cfg(target_arch = "x86_64")]
12242 if avx2_enabled() {
12243 return unsafe { dot_i8_f32_avx2(w, x) };
12244 }
12245 #[allow(unreachable_code)]
12246 {
12247 let mut sum = 0.0f32;
12248 for (j, &b) in w.iter().enumerate() {
12249 sum += (b as i8) as f32 * x[j];
12250 }
12251 sum
12252 }
12253}
12254
12255#[inline]
12259fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12260 #[cfg(target_arch = "aarch64")]
12261 unsafe {
12262 return dot_i8_col_f32_neon(w, x, col);
12263 }
12264 #[allow(unreachable_code)]
12265 {
12266 let mut sum = 0.0f32;
12267 for (j, &b) in w.iter().enumerate() {
12268 sum += (b as i8) as f32 * x[j] * col[j];
12269 }
12270 sum
12271 }
12272}
12273
12274#[cfg(target_arch = "aarch64")]
12275#[target_feature(enable = "neon")]
12276unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12277 unsafe {
12279 use core::arch::aarch64::*;
12280 let n = x.len();
12281 let wp = w.as_ptr() as *const i8;
12282 let xp = x.as_ptr();
12283 let cp = col.as_ptr();
12284 let (mut a0, mut a1, mut a2, mut a3) = (
12285 vdupq_n_f32(0.0),
12286 vdupq_n_f32(0.0),
12287 vdupq_n_f32(0.0),
12288 vdupq_n_f32(0.0),
12289 );
12290 let mut j = 0usize;
12291 while j + 16 <= n {
12292 let wb = vld1q_s8(wp.add(j));
12293 let lo = vmovl_s8(vget_low_s8(wb));
12294 let hi = vmovl_s8(vget_high_s8(wb));
12295 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12296 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12297 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12298 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12299 a0 = vfmaq_f32(
12300 a0,
12301 w0,
12302 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12303 );
12304 a1 = vfmaq_f32(
12305 a1,
12306 w1,
12307 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12308 );
12309 a2 = vfmaq_f32(
12310 a2,
12311 w2,
12312 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12313 );
12314 a3 = vfmaq_f32(
12315 a3,
12316 w3,
12317 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12318 );
12319 j += 16;
12320 }
12321 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12322 while j < n {
12323 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12324 j += 1;
12325 }
12326 sum
12327 }
12328}
12329
12330#[cfg(target_arch = "aarch64")]
12331#[target_feature(enable = "neon")]
12332unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12333 unsafe {
12335 use core::arch::aarch64::*;
12336 let n = x.len();
12337 let wp = w.as_ptr() as *const i8;
12338 let xp = x.as_ptr();
12339 let (mut a0, mut a1, mut a2, mut a3) = (
12340 vdupq_n_f32(0.0),
12341 vdupq_n_f32(0.0),
12342 vdupq_n_f32(0.0),
12343 vdupq_n_f32(0.0),
12344 );
12345 let mut j = 0usize;
12346 while j + 16 <= n {
12347 let wb = vld1q_s8(wp.add(j));
12348 let lo = vmovl_s8(vget_low_s8(wb));
12349 let hi = vmovl_s8(vget_high_s8(wb));
12350 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12351 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12352 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12353 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12354 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12355 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12356 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12357 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12358 j += 16;
12359 }
12360 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12361 while j < n {
12362 sum += (*wp.add(j)) as f32 * *xp.add(j);
12363 j += 1;
12364 }
12365 sum
12366 }
12367}
12368
12369#[allow(clippy::too_many_arguments)]
12370fn qmatvec(
12371 q: &[u8],
12372 rep: &[u8],
12373 row_scale: &[f32],
12374 x: &[f32],
12375 col_field: &[f32],
12376 dtype: TensorDtype,
12377 rows: usize,
12378 cols: usize,
12379 out: &mut [f32],
12380 pool: Option<&Pool>,
12381) {
12382 debug_assert_eq!(out.len(), rows);
12383 #[cfg(not(target_arch = "aarch64"))]
12384 let _ = rep;
12385
12386 #[cfg(target_arch = "aarch64")]
12387 if sdot_enabled() {
12388 let act = if dtype == TensorDtype::Q8_2f {
12389 split_act_q8_2f(x, col_field)
12390 } else {
12391 split_act(x)
12392 };
12393 let out_addr = SendMut(out.as_mut_ptr());
12394 let run_range = |start: usize, end: usize| {
12395 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12396 };
12397 match pool {
12398 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12399 _ => run_range(0, rows),
12400 }
12401 return;
12402 }
12403 #[cfg(target_arch = "x86_64")]
12406 if avx2_a8w8_enabled() {
12407 let act = if dtype == TensorDtype::Q8_2f {
12408 split_act_q8_2f(x, col_field)
12409 } else {
12410 split_act(x)
12411 };
12412 let out_addr = SendMut(out.as_mut_ptr());
12413 let run_range = |start: usize, end: usize| {
12414 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12415 };
12416 match pool {
12417 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12418 _ => run_range(0, rows),
12419 }
12420 return;
12421 }
12422
12423 prescale_with(x, col_field, dtype, 1, |xs| {
12424 let out_addr = SendMut(out.as_mut_ptr());
12425 let run_range = move |start: usize, end: usize| {
12426 for o in start..end {
12427 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12428 unsafe { *out_addr.at(o) = v };
12430 }
12431 };
12432 match pool {
12433 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12434 _ => run_range(0, rows),
12435 }
12436 });
12437}
12438
12439#[allow(clippy::too_many_arguments)]
12440fn qmatvec2(
12441 q: &[u8],
12442 row_scale: &[f32],
12443 x1: &[f32],
12444 x2: &[f32],
12445 col_field: &[f32],
12446 dtype: TensorDtype,
12447 rows: usize,
12448 cols: usize,
12449 o1: &mut [f32],
12450 o2: &mut [f32],
12451 pool: Option<&Pool>,
12452) {
12453 #[cfg(target_arch = "aarch64")]
12454 if sdot_enabled() {
12455 let a1s = if dtype == TensorDtype::Q8_2f {
12456 split_act_q8_2f(x1, col_field)
12457 } else {
12458 split_act(x1)
12459 };
12460 let a2s = if dtype == TensorDtype::Q8_2f {
12461 split_act_q8_2f(x2, col_field)
12462 } else {
12463 split_act(x2)
12464 };
12465 let p1 = SendMut(o1.as_mut_ptr());
12466 let p2 = SendMut(o2.as_mut_ptr());
12467 let run_range = |start: usize, end: usize| {
12468 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12469 };
12470 match pool {
12471 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12472 _ => run_range(0, rows),
12473 }
12474 return;
12475 }
12476 #[cfg(target_arch = "x86_64")]
12477 if avx2_a8w8_enabled() {
12478 let a1s = if dtype == TensorDtype::Q8_2f {
12479 split_act_q8_2f(x1, col_field)
12480 } else {
12481 split_act(x1)
12482 };
12483 let a2s = if dtype == TensorDtype::Q8_2f {
12484 split_act_q8_2f(x2, col_field)
12485 } else {
12486 split_act(x2)
12487 };
12488 let p1 = SendMut(o1.as_mut_ptr());
12489 let p2 = SendMut(o2.as_mut_ptr());
12490 let run_range = |start: usize, end: usize| {
12491 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12492 };
12493 match pool {
12494 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12495 _ => run_range(0, rows),
12496 }
12497 return;
12498 }
12499
12500 prescale_with(x1, col_field, dtype, 1, |x1s| {
12501 prescale_with(x2, col_field, dtype, 2, |x2s| {
12502 let p1 = SendMut(o1.as_mut_ptr());
12503 let p2 = SendMut(o2.as_mut_ptr());
12504 let run_range = move |start: usize, end: usize| {
12505 for o in start..end {
12506 let row = &q[o * cols..(o + 1) * cols];
12507 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12508 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12509 unsafe {
12511 *p1.at(o) = s1;
12512 *p2.at(o) = s2;
12513 }
12514 }
12515 };
12516 match pool {
12517 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12518 _ => run_range(0, rows),
12519 }
12520 });
12521 });
12522}
12523
12524#[derive(Clone, Copy)]
12525struct SendMut(*mut f32);
12526unsafe impl Send for SendMut {}
12527unsafe impl Sync for SendMut {}
12528
12529impl SendMut {
12530 #[inline]
12531 fn at(self, i: usize) -> *mut f32 {
12532 unsafe { self.0.add(i) }
12533 }
12534}
12535
12536#[cfg(test)]
12537mod tests {
12538 #[test]
12542 fn q8_round_is_round_clamp() {
12543 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12544 let mut probe = vec![
12545 0.0f32,
12546 -0.0,
12547 f32::NAN,
12548 f32::INFINITY,
12549 f32::NEG_INFINITY,
12550 f32::MAX,
12551 f32::MIN,
12552 1e30,
12553 -1e30,
12554 f32::MIN_POSITIVE,
12555 -f32::MIN_POSITIVE,
12556 ];
12557 for k in -300i32..=300 {
12558 let h = k as f32 * 0.5;
12559 let up = f32::from_bits(h.to_bits() + 1);
12560 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12561 for t in [h, up, down] {
12562 probe.push(t);
12563 probe.push(-t);
12564 }
12565 }
12566 let mut t = -140.0f32;
12567 while t < 140.0 {
12568 probe.push(t);
12569 t += 0.000_731;
12570 }
12571 for t in probe {
12572 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12573 }
12574 }
12575
12576 use super::*;
12577
12578 #[test]
12582 fn f32_matmat_equals_matvec_bitwise() {
12583 let (rows, cols, b) = (5usize, 37usize, 7usize);
12584 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12585 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12586 let t = QTensor::from_f32(w, rows, cols);
12587 let mut out = vec![0f32; b * rows];
12588 t.matmat(&xs, b, &mut out, None);
12589 for bi in 0..b {
12590 let mut one = vec![0f32; rows];
12591 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12592 for o in 0..rows {
12593 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12594 }
12595 }
12596 }
12597
12598 #[test]
12599 fn q2tp_i8_dot_matches_exact_on_grid() {
12600 let (rows, cols) = (5, 64);
12604 let gpr = cols / GROUP_SIZE;
12605 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12609 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12610 .collect();
12611 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12612 let x: Vec<f32> = (0..cols)
12613 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12614 .collect();
12615 let act = split_act(&x);
12616 assert!(
12617 act.outliers.is_empty(),
12618 "on-grid input must have no outliers"
12619 );
12620 let gsum = q1_group_sums(&act.xq, gpr);
12621 for r in 0..rows {
12622 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12623 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12624 assert!(
12625 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12626 "row {r}: exact {exact} vs i8 {fast}"
12627 );
12628 }
12629 }
12630
12631 #[cfg(target_arch = "x86_64")]
12632 #[test]
12633 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12634 if !std::arch::is_x86_feature_detected!("avx2") {
12639 return;
12640 }
12641 let mut seed = 0x9e3779b9u32;
12642 let mut next = || {
12643 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12644 seed
12645 };
12646 for _ in 0..20_000 {
12647 let mut ch = [0u8; Q2TP_CHUNK];
12648 let mut x = [0i8; GROUP_SIZE];
12649 for b in &mut ch {
12650 *b = next() as u8;
12651 }
12652 for v in &mut x {
12653 *v = (next() >> 24) as i8;
12654 }
12655 let mut reference = 0i32;
12656 for (k, &b) in ch.iter().enumerate() {
12657 reference += (b & 3) as i32 * x[k * 4] as i32;
12658 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12659 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12660 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12661 }
12662 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12665 assert_eq!(got, reference, "packed q2 lane mismatch");
12666 }
12667 }
12668
12669 #[test]
12670 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12671 let (rows, cols) = (1usize, GROUP_SIZE);
12672 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12673 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12677 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12678 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
12680 let mut raw = vec![0.0f32; rows];
12681 let mut affine = vec![0.0f32; rows];
12682 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12683 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12684 assert_eq!(raw, vec![-24.0]);
12685 assert_eq!(affine, vec![-8.0]);
12686 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12687 }
12688
12689 #[test]
12690 fn q8_row_dot_fast_matches_scalar() {
12691 let cols = 96;
12694 let row: Vec<u8> = (0..cols)
12695 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12696 .collect();
12697 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12698 let act = split_act(&x);
12699 let fast = q8_row_dot(&row, &act);
12700 let scalar = q8_row_dot_scalar(&row, &act);
12701 assert!(
12702 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12703 "fast {fast} vs scalar {scalar}"
12704 );
12705 }
12706
12707 #[test]
12708 fn f32_matvec_matches_matvec_rows_bitexact() {
12709 let (rows, cols) = (300, 40);
12710 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12711 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12712 let qt = QTensor::from_f32(w.clone(), rows, cols);
12713
12714 let mut a = vec![0.0f32; rows];
12715 matvec_rows(None, &w, &x, &mut a);
12716 let mut b = vec![0.0f32; rows];
12717 qt.matvec(&x, &mut b, None);
12718 assert_eq!(a, b);
12719 }
12720
12721 #[test]
12722 fn sdot_kernel_exact_on_grid() {
12723 eprintln!("sdot_enabled = {}", sdot_enabled());
12728 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
12730 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12731 .collect();
12732 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12733 let x: Vec<f32> = (0..cols)
12734 .map(|i| match i % 3 {
12735 0 => 1.0,
12736 1 => -1.0,
12737 _ => 0.0,
12738 })
12739 .collect();
12740 let mut a = vec![0.0f32; rows];
12741 qmatvec(
12742 &w,
12743 &[],
12744 &scales,
12745 &x,
12746 &[],
12747 TensorDtype::Q8Row,
12748 rows,
12749 cols,
12750 &mut a,
12751 None,
12752 );
12753 for o in 0..rows {
12754 let mut acc = 0.0f32;
12755 for j in 0..cols {
12756 acc += (w[o * cols + j] as i8) as f32 * x[j];
12757 }
12758 let expect = acc * scales[o];
12759 assert!(
12760 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12761 "row {o}: {} vs {expect}",
12762 a[o]
12763 );
12764 }
12765 }
12766
12767 #[test]
12768 fn q1_tbl_fast_path_matches_reference() {
12769 let (rows, cols) = (5, 256);
12774 let gpr = cols / GROUP_SIZE;
12775 let mut bytes = Vec::new();
12776 for t in 0..rows * gpr {
12777 let s = 0.007 + (t % 11) as f32 * 0.004;
12778 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12779 for j in 0..4 {
12780 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12781 }
12782 }
12783 let x: Vec<f32> = (0..cols)
12784 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12785 .collect();
12786 let mut w = vec![0.0f32; rows * cols];
12787 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12788 let mut got = vec![0.0f32; rows];
12789 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12790 for o in 0..rows {
12791 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12792 assert!(
12793 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12794 "row {o}: {} vs {expect}",
12795 got[o]
12796 );
12797 }
12798 let b = 5usize;
12801 let mut xs_all = Vec::new();
12802 for bi in 0..b {
12803 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12804 }
12805 let mut mm = vec![0.0f32; b * rows];
12806 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12807 for bi in 0..b {
12808 let mut single = vec![0.0f32; rows];
12809 q1_matvec(
12810 &bytes,
12811 &xs_all[bi * cols..(bi + 1) * cols],
12812 rows,
12813 cols,
12814 &mut single,
12815 None,
12816 );
12817 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12818 }
12819 }
12820
12821 #[test]
12822 fn q1_kernels_match_exact_reference() {
12823 let (rows, cols) = (7, 96);
12825 let gpr = cols / GROUP_SIZE;
12826 let mut bytes = Vec::new();
12827 for t in 0..rows * gpr {
12828 let s = 0.01 + (t % 13) as f32 * 0.003;
12829 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12830 for j in 0..4 {
12831 bytes.push(((t * 31 + j * 97) % 251) as u8);
12832 }
12833 }
12834 let x: Vec<f32> = (0..cols)
12836 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12837 .collect();
12838 let mut w = vec![0.0f32; rows * cols];
12840 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12841 let mut expect = vec![0.0f32; rows];
12842 for o in 0..rows {
12843 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12844 }
12845 let mut got = vec![0.0f32; rows];
12846 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12847 for o in 0..rows {
12848 assert!(
12849 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12850 "row {o}: {} vs {}",
12851 got[o],
12852 expect[o]
12853 );
12854 }
12855 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12857 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12858 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12859 assert_eq!(a1, got);
12860 let mut xs = x.clone();
12861 xs.extend_from_slice(&x2);
12862 let mut mm = vec![0.0f32; 2 * rows];
12863 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12864 assert_eq!(&mm[..rows], got.as_slice());
12865 assert_eq!(&mm[rows..], a2.as_slice());
12866 }
12867
12868 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12871 let tiles = rows * cols / GROUP_SIZE;
12872 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12873 for t in 0..tiles {
12874 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12875 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12876 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12877 q1.extend_from_slice(&bits.to_le_bytes());
12878 }
12879 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12880 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12881 (q1, rp.data)
12882 }
12883
12884 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12887 let mut x: Vec<f32> = (0..cols)
12888 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12889 .collect();
12890 x[3 % cols] = 41.0;
12891 x[(cols / 2 + 5) % cols] = -37.5;
12892 x
12893 }
12894
12895 #[test]
12902 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12903 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12904 let gpr = cols / GROUP_SIZE;
12905 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12906 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12907 let x = q1g_test_x(cols, rows);
12908 let act = split_act(&x);
12909 assert!(!act.outliers.is_empty(), "test wants the outlier path");
12910 let gsum = q1_group_sums(&act.xq, gpr);
12911 for r in 0..rows {
12912 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12913 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12914 let (a, b) = (
12915 wa.dot_i8(r, gpr, &act.xq, &gsum),
12916 wb.dot_i8(r, gpr, &act.xq, &gsum),
12917 );
12918 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12919 for j in [0, 31, 32, 127, cols - 1] {
12920 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12921 assert_eq!(
12922 (a.0.to_bits(), a.1.to_bits()),
12923 (b.0.to_bits(), b.1.to_bits())
12924 );
12925 }
12926 #[cfg(target_arch = "x86_64")]
12927 if avx2_enabled() {
12928 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12929 unsafe {
12931 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12932 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12933 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12934 }
12935 if std::is_x86_feature_detected!("avx512vnni")
12936 && std::is_x86_feature_detected!("avx512bw")
12937 && std::is_x86_feature_detected!("avx512vl")
12938 {
12939 unsafe {
12941 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12942 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12943 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12944 }
12945 }
12946 }
12947 }
12948 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12950 bin_matvec(wa, &x, rows, cols, &mut ga, None);
12951 bin_matvec(wb, &x, rows, cols, &mut gb, None);
12952 assert_eq!(ga, gb, "{rows}x{cols} matvec");
12953 let x2 = q1g_test_x(cols, rows + 1);
12954 let (mut a1, mut a2, mut b1, mut b2) = (
12955 vec![0f32; rows],
12956 vec![0f32; rows],
12957 vec![0f32; rows],
12958 vec![0f32; rows],
12959 );
12960 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12961 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12962 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12963 let b = 5usize;
12964 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12965 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12966 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12967 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12968 assert_eq!(ma, mb, "{rows}x{cols} matmat");
12969 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12972 .map(|k| {
12973 let a = split_act(&xs[k * cols..(k + 1) * cols]);
12974 let g = q1_group_sums(&a.xq, gpr);
12975 (a, g)
12976 })
12977 .collect();
12978 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12979 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12980 for r in 0..rows {
12981 #[cfg(target_arch = "x86_64")]
12982 if avx2_enabled() {
12983 let (a, b) = unsafe {
12985 (
12986 wa.dot_1x4_x86(r, gpr, xq4, gs4),
12987 wb.dot_1x4_x86(r, gpr, xq4, gs4),
12988 )
12989 };
12990 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12991 }
12992 #[cfg(target_arch = "aarch64")]
12993 if sdot_enabled() {
12994 let (a, b) = unsafe {
12996 (
12997 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12998 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12999 )
13000 };
13001 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13002 }
13003 let _ = (r, &xq4, &gs4);
13004 }
13005 }
13006 }
13007
13008 #[test]
13014 #[ignore]
13015 fn q1g_cpu_matvec_speed() {
13016 use crate::pool::Pool;
13017 let threads = std::env::var("CMF_THREADS")
13018 .ok()
13019 .and_then(|v| v.parse().ok())
13020 .unwrap_or(8usize);
13021 let pool = Pool::new(threads);
13022 for &(name, rows, cols, reps) in &[
13023 ("gate/up_proj", 6144usize, 2048usize, 41usize),
13024 ("down_proj", 2048, 6144, 41),
13025 ("lm_head", 151_669, 2048, 9),
13026 ] {
13027 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
13028 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13029 let x = q1g_test_x(cols, 1);
13030 let mut out = vec![0f32; rows];
13031 let mut time = |w: BinW| {
13032 let t0 = std::time::Instant::now();
13033 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
13034 t0.elapsed().as_secs_f64()
13035 };
13036 let _ = (time(wa), time(wb));
13037 let (mut ta, mut tb) = (Vec::new(), Vec::new());
13038 for r in 0..reps {
13039 if r % 2 == 0 {
13040 ta.push(time(wa));
13041 tb.push(time(wb));
13042 } else {
13043 tb.push(time(wb));
13044 ta.push(time(wa));
13045 }
13046 }
13047 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
13048 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
13049 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
13050 eprintln!(
13051 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
13052 a8w8_enabled(),
13053 ma * 1e6,
13054 mb * 1e6,
13055 mb / ma
13056 );
13057 }
13058 }
13059
13060 #[test]
13065 fn q1g_mapped_tensor_paths_equal_q1() {
13066 use crate::pool::Pool;
13067 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13068 let (rows, cols) = (300usize, 256usize);
13069 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
13070 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
13071 let header: CmfHeader = serde_json::from_value(serde_json::json!({
13072 "format": "cmf", "version": cortiq_core::CMF_VERSION,
13073 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
13074 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13075 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
13076 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13077 "max_position_embeddings": 8},
13078 "quant_type": "VBIT"
13079 }))
13080 .unwrap();
13081 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
13082 name: name.into(),
13083 dtype,
13084 shape: vec![rows, cols],
13085 data: data.clone(),
13086 };
13087 let specs = [
13088 spec("g1", TensorDtype::Q1, &gq1),
13089 spec("u1", TensorDtype::Q1, &uq1),
13090 spec("gg", TensorDtype::Q1G, &gq1g),
13091 spec("ug", TensorDtype::Q1G, &uq1g),
13092 ];
13093 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
13094 std::fs::create_dir_all(&dir).unwrap();
13095 let path = dir.join("m.cmf");
13096 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13097 let model = Arc::new(CmfModel::open(&path).unwrap());
13098 assert!(model.verify().is_empty(), "{:?}", model.verify());
13099 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13100 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
13101 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
13102 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
13103 let pool = Pool::new(3);
13104 let x = q1g_test_x(cols, 4);
13105 let run = |a: &QTensor| {
13106 let mut o = vec![0f32; rows];
13107 a.matvec(&x, &mut o, Some(&pool));
13108 o
13109 };
13110 assert_eq!(run(&g1), run(&gg), "matvec");
13111 let many = |a: &QTensor, b: &QTensor| {
13112 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
13113 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
13114 (oa, ob)
13115 };
13116 let want = many(&g1, &u1);
13117 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
13118 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
13119 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
13120 let fused = |g: &QTensor, u: &QTensor| {
13121 let mut o = vec![0f32; rows];
13122 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
13123 (ok, o)
13124 };
13125 let (ok_a, fa) = fused(&g1, &u1);
13126 let (ok_b, fb) = fused(&gg, &ug);
13127 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
13128 assert_eq!(fa, fb, "fused gate/up");
13129 let b = 6usize;
13130 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
13131 let mm = |a: &QTensor| {
13132 let mut o = vec![0f32; b * rows];
13133 a.matmat(&xs, b, &mut o, Some(&pool));
13134 o
13135 };
13136 assert_eq!(mm(&g1), mm(&gg), "matmat");
13137 let x2 = q1g_test_x(cols, 77);
13138 let pair = |a: &QTensor| {
13139 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13140 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13141 (o1, o2)
13142 };
13143 assert_eq!(pair(&g1), pair(&gg), "matvec2");
13144 for r in [0, 1, rows / 2, rows - 1] {
13145 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
13146 g1.row_f32(r, &mut ra);
13147 gg.row_f32(r, &mut rb);
13148 assert_eq!(ra, rb, "embedding row {r}");
13149 }
13150 drop((g1, u1, gg, ug, model));
13151 let _ = std::fs::remove_dir_all(&dir);
13152 }
13153
13154 #[test]
13161 fn blr_mapped_tensor_paths_equal_dense() {
13162 use crate::pool::Pool;
13163 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13164 let (rows, cols, r) = (96usize, 128usize, 16usize);
13165 let mut rng = crate::blr::tests::Rng(41);
13166 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13167 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13168 let mut q8: Vec<u8> = (0..rows * cols)
13170 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13171 .collect();
13172 for o in 0..rows {
13173 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13174 }
13175 let header: CmfHeader = serde_json::from_value(serde_json::json!({
13176 "format": "cmf", "version": cortiq_core::CMF_VERSION,
13177 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13178 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13179 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13180 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13181 "max_position_embeddings": 8},
13182 "quant_type": "Q8_2F"
13183 }))
13184 .unwrap();
13185 let specs = [
13186 TensorSpec {
13187 name: "g".into(),
13188 dtype: TensorDtype::BinaryLowRank,
13189 shape: vec![rows, cols],
13190 data: gate.clone(),
13191 },
13192 TensorSpec {
13193 name: "u".into(),
13194 dtype: TensorDtype::BinaryLowRank,
13195 shape: vec![rows, cols],
13196 data: up.clone(),
13197 },
13198 TensorSpec {
13199 name: "q".into(),
13200 dtype: TensorDtype::Q8Row,
13201 shape: vec![rows, cols],
13202 data: q8,
13203 },
13204 ];
13205 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13206 std::fs::create_dir_all(&dir).unwrap();
13207 let path = dir.join("m.cmf");
13208 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13209 let model = Arc::new(CmfModel::open(&path).unwrap());
13210 assert!(model.verify().is_empty(), "{:?}", model.verify());
13211 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13212 let (g, u, q) = (t("g"), t("u"), t("q"));
13213 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13214 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13215 assert!(g.as_f32().is_none() && g.model_idx().is_none());
13216 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13217 assert!(!g.sparse_col_ok());
13218 assert_eq!((g.rows(), g.cols()), (rows, cols));
13219 let deq = |bytes: &[u8]| {
13220 let mut w = vec![0f32; rows * cols];
13221 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13222 w
13223 };
13224 let (wg, wu) = (deq(&gate), deq(&up));
13225 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13226 (0..rows)
13227 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13228 .collect()
13229 };
13230 let close = |a: &[f32], b: &[f32], what: &str| {
13231 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13232 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13233 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13234 };
13235 let pool = Pool::new(3);
13236 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13237 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13238 let mut o = vec![0f32; rows];
13239 g.matvec(&x, &mut o, Some(&pool));
13240 close(&o, &mv(&wg, &x), "matvec");
13241 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13242 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13243 close(&o1, &mv(&wg, &x), "matvec2 a");
13244 close(&o2, &mv(&wg, &x2), "matvec2 b");
13245 for b in [2usize, 5] {
13246 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13247 let mut om = vec![0f32; b * rows];
13248 g.matmat(&xs, b, &mut om, Some(&pool));
13249 for t in 0..b {
13250 close(
13251 &om[t * rows..(t + 1) * rows],
13252 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13253 "matmat",
13254 );
13255 }
13256 }
13257 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13260 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13261 close(&a, &mv(&wg, &x), "matvec_many blr gate");
13262 close(&c, &mv(&wu, &x), "matvec_many blr up");
13263 let mut qo = vec![0f32; rows];
13264 q.matvec(&x, &mut qo, Some(&pool));
13265 assert_eq!(bq, qo, "matvec_many q8_row sibling");
13266 let mut f = vec![0f32; rows];
13267 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13268 let mut row = vec![0f32; cols];
13269 g.row_f32(rows / 3, &mut row);
13270 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13271 let mut scratch = vec![0f32; cols];
13272 let d = g.row_dot(7, &x, &mut scratch);
13273 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13274 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13275 let mut col = vec![0f32; rows];
13276 g.add_col_scaled(cols / 2, 1.5, &mut col);
13277 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13278 close(&col, &want, "add_col_scaled");
13279 drop((g, u, q, model));
13280 let _ = std::fs::remove_dir_all(&dir);
13281 }
13282
13283 #[test]
13284 fn repack_is_bit_identical() {
13285 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
13291 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13292 .collect();
13293 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13294 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13295 let rep = q8_repack_layout(&w, rows, cols);
13296 for g in 0..rows / 4 {
13298 for c in 0..cols / 16 {
13299 for lane in 0..4 {
13300 assert_eq!(
13301 &rep[g * 4 * cols + c * 64 + lane * 16
13302 ..g * 4 * cols + c * 64 + lane * 16 + 16],
13303 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13304 );
13305 }
13306 }
13307 }
13308 let mut a = vec![0.0f32; rows];
13309 qmatvec(
13310 &w,
13311 &[],
13312 &scales,
13313 &x,
13314 &[],
13315 TensorDtype::Q8Row,
13316 rows,
13317 cols,
13318 &mut a,
13319 None,
13320 );
13321 let mut b = vec![0.0f32; rows];
13322 qmatvec(
13323 &w,
13324 &rep,
13325 &scales,
13326 &x,
13327 &[],
13328 TensorDtype::Q8Row,
13329 rows,
13330 cols,
13331 &mut b,
13332 None,
13333 );
13334 assert_eq!(a, b, "full-range repack output diverged");
13335
13336 #[cfg(target_arch = "aarch64")]
13337 if sdot_enabled() {
13338 let act = split_act(&x);
13340 let mut c1 = vec![0.0f32; rows];
13341 let mut c2 = vec![0.0f32; rows];
13342 q8_range_sdot(
13343 &w,
13344 &[],
13345 &scales,
13346 &act,
13347 cols,
13348 SendMut(c1.as_mut_ptr()),
13349 3,
13350 rows - 2,
13351 );
13352 q8_range_sdot(
13353 &w,
13354 &rep,
13355 &scales,
13356 &act,
13357 cols,
13358 SendMut(c2.as_mut_ptr()),
13359 3,
13360 rows - 2,
13361 );
13362 assert_eq!(c1, c2, "unaligned-range repack output diverged");
13363 }
13364 }
13365
13366 #[test]
13367 fn sdot_a8w8_noise_is_bounded() {
13368 let (rows, cols) = (16, 512);
13372 let w: Vec<u8> = (0..rows * cols)
13373 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13374 .collect();
13375 let scales = vec![0.01f32; rows];
13376 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13377 let mut a = vec![0.0f32; rows];
13378 qmatvec(
13379 &w,
13380 &[],
13381 &scales,
13382 &x,
13383 &[],
13384 TensorDtype::Q8Row,
13385 rows,
13386 cols,
13387 &mut a,
13388 None,
13389 );
13390 let (mut num, mut den) = (0f64, 0f64);
13391 for o in 0..rows {
13392 let mut acc = 0.0f32;
13393 for j in 0..cols {
13394 acc += (w[o * cols + j] as i8) as f32 * x[j];
13395 }
13396 let expect = acc * scales[o];
13397 num += ((a[o] - expect) as f64).powi(2);
13398 den += (expect as f64).powi(2);
13399 }
13400 let rel = (num / den.max(1e-12)).sqrt();
13401 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13402 }
13403
13404 #[test]
13405 fn i8_dot_neon_matches_scalar() {
13406 let n = 100;
13407 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13408 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13409 let mut scalar = 0.0f32;
13410 for j in 0..n {
13411 scalar += (w[j] as i8) as f32 * x[j];
13412 }
13413 let fast = dot_i8_f32(&w, &x);
13414 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13415 }
13416
13417 #[test]
13419 fn vbitmatvec_matches_full_dequant() {
13420 let (rows, cols) = (6, 64);
13421 let ng = cols / GROUP_SIZE;
13422 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13424 let mut bytes = bits.clone();
13425 for g in 0..rows * ng {
13426 let s = 0.02 + 0.001 * g as f32;
13427 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13428 }
13429 for r in 0..rows {
13430 let b = bits[r] as usize;
13431 let (mut acc, mut nb) = (0u64, 0usize);
13432 let mut rowbytes = Vec::new();
13433 for i in 0..cols {
13434 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13435 acc = (acc << b) | v;
13436 nb += b;
13437 while nb >= 8 {
13438 nb -= 8;
13439 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13440 }
13441 }
13442 if nb > 0 {
13443 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13444 }
13445 bytes.extend_from_slice(&rowbytes);
13446 }
13447 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13448
13449 let mut reference = vec![0f32; rows * cols];
13450 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13451 let mut expect = vec![0f32; rows];
13452 for r in 0..rows {
13453 expect[r] = reference[r * cols..(r + 1) * cols]
13454 .iter()
13455 .zip(&x)
13456 .map(|(w, xv)| w * xv)
13457 .sum();
13458 }
13459 let mut got = vec![0f32; rows];
13460 let offsets = vbit_row_offsets(&bytes, rows, cols);
13461 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13462 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13466 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13467 for r in 0..rows {
13468 assert!(
13469 (got[r] - expect[r]).abs() < tol * scale,
13470 "row {r}: {} vs {}",
13471 got[r],
13472 expect[r]
13473 );
13474 }
13475 }
13476
13477 #[test]
13482 #[cfg(target_arch = "x86_64")]
13483 fn vbit_matmat_blocked_matches_per_row() {
13484 let (rows, cols, b) = (64usize, 128usize, 9usize);
13485 let ng = cols / GROUP_SIZE;
13486 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13487 let mut bytes = bits.clone();
13488 for g in 0..rows * ng {
13489 let sc = 0.02 + 0.0005 * g as f32;
13490 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13491 }
13492 for r in 0..rows {
13493 let bw = bits[r] as usize;
13494 let (mut acc, mut nb) = (0u64, 0usize);
13495 let mut rowbytes = Vec::new();
13496 for i in 0..cols {
13497 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13498 acc = (acc << bw) | v;
13499 nb += bw;
13500 while nb >= 8 {
13501 nb -= 8;
13502 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13503 }
13504 }
13505 if nb > 0 {
13506 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13507 }
13508 bytes.extend_from_slice(&rowbytes);
13509 }
13510 let x: Vec<f32> = (0..b * cols)
13511 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13512 .collect();
13513 let offsets = vbit_row_offsets(&bytes, rows, cols);
13514 let mut y_a = vec![0f32; b * rows];
13515 let mut y_b = vec![0f32; b * rows];
13516 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13517 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13518 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13519 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13520 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13521 let max_d = y_a
13522 .iter()
13523 .zip(&y_b)
13524 .map(|(p, q)| (p - q).abs())
13525 .fold(0.0f32, f32::max);
13526 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13527 }
13528
13529 #[test]
13537 fn q4t_matmat_blocked_matches_per_row() {
13538 let (rows, cols, b) = (16usize, 64usize, 9usize);
13539 let gpr = cols / GROUP_SIZE;
13540 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13541 for r in 0..rows {
13542 for g in 0..gpr {
13543 let t = (r * gpr + g) * Q4_TILE;
13544 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13545 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13546 for k in 0..16 {
13547 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13548 }
13549 }
13550 }
13551 let x: Vec<f32> = (0..b * cols)
13552 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13553 .collect();
13554 let mut y_blk = vec![0f32; b * rows];
13555 let mut y_row = vec![0f32; b * rows];
13556 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13557 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13558 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13559 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13560 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13561 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13562 }
13563
13564 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13571 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13572 let gpr = cols / GROUP_SIZE;
13573 let stride = q4tp_code_stride(gpr);
13574 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13575 let mut b = vec![0u8; codes_off + rows * stride];
13576 for r in 0..rows {
13577 for g in 0..gpr {
13578 let t = (r * gpr + g) * Q4TP_NIB;
13579 for k in 0..16 {
13580 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13581 }
13582 }
13583 let lo = -6.0 - 0.03 * (r % 17) as f32;
13584 let step = 0.01 + 0.004 * (r % 11) as f32;
13585 let p = params_off + r * 4;
13586 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13587 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13588 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13589 for g in 0..gpr {
13590 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13591 }
13592 }
13593 b
13594 }
13595
13596 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13600 let gpr = cols / GROUP_SIZE;
13601 let v = Q4tpView::new(bytes, rows, cols);
13602 let mut out = vec![0u8; rows * gpr * Q4_TILE];
13603 let mut sc = vec![0f32; gpr];
13604 for r in 0..rows {
13605 v.scales_into(r, gpr, &mut sc);
13606 for g in 0..gpr {
13607 let t = (r * gpr + g) * Q4_TILE;
13608 let s = sc[g];
13609 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13610 let src = (r * gpr + g) * Q4TP_NIB;
13611 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13612 }
13613 }
13614 out
13615 }
13616
13617 #[test]
13623 fn q4tp_exact_path_matches_dequant_reference() {
13624 let (rows, cols) = (256usize, 512usize);
13625 let gpr = cols / GROUP_SIZE;
13626 let bytes = synth_q4tp(rows, cols);
13627 let mut w = vec![0f32; rows * cols];
13628 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13629
13630 let x: Vec<f32> = (0..cols)
13631 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13632 .collect();
13633 let v = Q4tpView::new(&bytes, rows, cols);
13634 let mut sc = vec![0f32; gpr];
13635 for r in 0..rows {
13636 v.scales_into(r, gpr, &mut sc);
13637 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13638 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13639 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13643 assert!(
13644 (got - want).abs() <= 1e-5 * mag,
13645 "row {r}: kernel {got} vs dequant {want}"
13646 );
13647 }
13648 }
13649
13650 #[test]
13656 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13657 let (rows, cols) = (256usize, 512usize);
13658 let bytes = synth_q4tp(rows, cols);
13659 let twin = q4tp_as_q4t(&bytes, rows, cols);
13660 let x: Vec<f32> = (0..cols)
13661 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13662 .collect();
13663
13664 let mut got = vec![0f32; rows];
13665 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13666 let mut want = vec![0f32; rows];
13667 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13668
13669 let mut w = vec![0f32; rows * cols];
13672 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13673 for r in 0..rows {
13674 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13675 assert!(
13676 (got[r] - want[r]).abs() <= 1e-3 * mag,
13677 "row {r}: q4tp {} vs q4t {}",
13678 got[r],
13679 want[r]
13680 );
13681 }
13682 }
13683
13684 #[test]
13689 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13690 let (rows, cols, b) = (256usize, 512usize, 5usize);
13691 let bytes = synth_q4tp(rows, cols);
13692 let twin = q4tp_as_q4t(&bytes, rows, cols);
13693 let xs: Vec<f32> = (0..b * cols)
13694 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13695 .collect();
13696
13697 let mut got = vec![0f32; b * rows];
13698 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13699 let mut want = vec![0f32; b * rows];
13700 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13701
13702 let mut w = vec![0f32; rows * cols];
13703 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13704 for t in 0..b {
13705 for r in 0..rows {
13706 let mag: f32 = (0..cols)
13707 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13708 .sum();
13709 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13710 assert!(
13711 (g - wa).abs() <= 1e-3 * mag,
13712 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13713 );
13714 }
13715 }
13716 }
13717
13718 #[test]
13719 fn q4tp_matvec2_matches_the_single_stream_kernel() {
13720 let (rows, cols) = (128usize, 256usize);
13721 let gpr = cols / GROUP_SIZE;
13722 let bytes = synth_q4tp(rows, cols);
13723 let xs: Vec<f32> = (0..2 * cols)
13724 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13725 .collect();
13726
13727 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13728 q4tp_matvec2(
13729 &bytes,
13730 &xs[..cols],
13731 &xs[cols..],
13732 rows,
13733 cols,
13734 &mut o1,
13735 &mut o2,
13736 None,
13737 );
13738
13739 let v = Q4tpView::new(&bytes, rows, cols);
13742 let mut sc = vec![0f32; gpr];
13743 for r in 0..rows {
13744 v.scales_into(r, gpr, &mut sc);
13745 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13746 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13747 }
13748 }
13749
13750 #[test]
13757 fn q4tp_matvec_keeps_pace_with_q4t() {
13758 let (rows, cols) = (4096usize, 3072usize);
13759 let bytes = synth_q4tp(rows, cols);
13760 let twin = q4tp_as_q4t(&bytes, rows, cols);
13761 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13762 let mut o = vec![0f32; rows];
13763 let n = 12;
13764 let mut best = (f64::MAX, f64::MAX);
13765 for _ in 0..3 {
13768 let t0 = std::time::Instant::now();
13769 for _ in 0..n {
13770 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13771 }
13772 best.0 = best.0.min(t0.elapsed().as_secs_f64());
13773 let t0 = std::time::Instant::now();
13774 for _ in 0..n {
13775 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13776 }
13777 best.1 = best.1.min(t0.elapsed().as_secs_f64());
13778 }
13779 let ratio = best.1 / best.0;
13780 println!(
13781 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13782 best.0 * 1e3 / n as f64,
13783 best.1 * 1e3 / n as f64
13784 );
13785 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13786 }
13787
13788 #[cfg(target_os = "macos")]
13789 #[test]
13790 fn q4t_matmat_accel_matches_dequant_reference() {
13791 if !accel_gemm_enabled() {
13792 return; }
13794 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
13796 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13797 for r in 0..rows {
13798 for g in 0..gpr {
13799 let t = (r * gpr + g) * Q4_TILE;
13800 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13801 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13802 for k in 0..16 {
13803 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13804 }
13805 }
13806 }
13807 let x: Vec<f32> = (0..b * cols)
13808 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13809 .collect();
13810 let mut got = vec![0f32; b * rows];
13811 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13812 let mut w = vec![0f32; rows * cols];
13814 for r in 0..rows {
13815 for g in 0..gpr {
13816 let t = (r * gpr + g) * Q4_TILE;
13817 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13818 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13819 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13820 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13821 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13822 }
13823 }
13824 }
13825 for bi in 0..b {
13826 for r in 0..rows {
13827 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13828 let d = (got[bi * rows + r] - want).abs();
13829 assert!(
13830 d <= want.abs().max(1.0) * 1e-4,
13831 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13832 got[bi * rows + r]
13833 );
13834 }
13835 }
13836 }
13837
13838 #[test]
13839 fn q4matvec_matches_full_dequant() {
13840 let (rows, cols) = (8, 64);
13841 let groups = rows * cols / GROUP_SIZE;
13842 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13844 for i in 0..groups * 16 {
13845 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13846 }
13847 for g in 0..groups {
13848 let s = 0.01 + 0.003 * g as f32;
13849 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13850 }
13851 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13852
13853 let mut reference = vec![0.0f32; rows * cols];
13854 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13855 let mut expect = vec![0.0f32; rows];
13856 for r in 0..rows {
13857 expect[r] = reference[r * cols..(r + 1) * cols]
13858 .iter()
13859 .zip(&x)
13860 .map(|(w, xv)| w * xv)
13861 .sum();
13862 }
13863
13864 let mut got = vec![0.0f32; rows];
13865 q4matvec(&bytes, &x, rows, cols, &mut got, None);
13866 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13870 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13871 for r in 0..rows {
13872 assert!(
13873 (got[r] - expect[r]).abs() < tol * scale,
13874 "row {r}: {} vs {}",
13875 got[r],
13876 expect[r]
13877 );
13878 }
13879 }
13880
13881 #[test]
13884 fn vbitmatvec2_equals_two_singles() {
13885 let (rows, cols) = (6, 64);
13886 let ng = cols / GROUP_SIZE;
13887 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13888 let mut bytes = bits.clone();
13889 for g in 0..rows * ng {
13890 let s = 0.02 + 0.001 * g as f32;
13891 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13892 }
13893 for r in 0..rows {
13894 let b = bits[r] as usize;
13895 let (mut acc, mut nb) = (0u64, 0usize);
13896 let mut rowbytes = Vec::new();
13897 for i in 0..cols {
13898 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13899 acc = (acc << b) | v;
13900 nb += b;
13901 while nb >= 8 {
13902 nb -= 8;
13903 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13904 }
13905 }
13906 if nb > 0 {
13907 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13908 }
13909 bytes.extend_from_slice(&rowbytes);
13910 }
13911 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13912 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13913 let offsets = vbit_row_offsets(&bytes, rows, cols);
13914
13915 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13916 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13917 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13918 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13919 vbitmatvec2(
13920 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13921 );
13922 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13923 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13924 }
13925
13926 #[test]
13928 fn q4matvec2_equals_two_singles() {
13929 let (rows, cols) = (8, 128);
13930 let groups = rows * cols / GROUP_SIZE;
13931 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13932 for i in 0..groups * 16 {
13933 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13934 }
13935 for g in 0..groups {
13936 let s = 0.01 + 0.003 * g as f32;
13937 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13938 }
13939 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13942 x1[9] = 250.0;
13943 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13944
13945 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13946 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13947 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13948 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13949 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13950 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13951 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13952 }
13953
13954 #[test]
13957 fn matvec_many_equals_separate_matvecs() {
13958 use crate::pool::Pool;
13959 let (r1, r2, cols) = (300, 200, 64);
13960 let mk = |salt: usize, rows: usize| {
13961 QTensor::from_f32(
13962 (0..rows * cols)
13963 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13964 .collect(),
13965 rows,
13966 cols,
13967 )
13968 };
13969 let (a, b) = (mk(1, r1), mk(5, r2));
13970 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13971 let pool = Pool::new(3);
13972
13973 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13974 a.matvec(&x, &mut ea, Some(&pool));
13975 b.matvec(&x, &mut eb, Some(&pool));
13976 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13977 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13978 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13979 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13980 }
13981
13982 #[test]
13987 fn q4tp_matvec_many_equals_separate_matvecs() {
13988 use crate::pool::Pool;
13989 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13990
13991 let (r1, r2, cols) = (300usize, 200usize, 64usize);
13992 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13993 "arch_name": "tiny-q4tp",
13994 "hidden_size": cols,
13995 "intermediate_size": cols * 2,
13996 "num_layers": 1,
13997 "num_attention_heads": 2,
13998 "num_kv_heads": 1,
13999 "head_dim": 32,
14000 "vocab_size": r1,
14001 "layer_types": ["FullAttention"],
14002 "rms_norm_eps": 1e-6,
14003 "max_position_embeddings": 8,
14004 "linear_conv_kernel_dim": 0,
14005 "linear_num_key_heads": 0,
14006 "linear_num_value_heads": 0
14007 }))
14008 .unwrap();
14009 let header = CmfHeader {
14010 format: "cmf".into(),
14011 version: CMF_VERSION,
14012 arch,
14013 quant_type: QuantType::Q4Block,
14014 provenance: None,
14015 tokenizer_config: None,
14016 section_hashes: None,
14017 skills: Vec::new(),
14018 shard: None,
14019 calibration: None,
14020 routing: None,
14021 genome: None,
14022 lineage: Vec::new(),
14023 router: None,
14024 segments: Vec::new(),
14025 };
14026 let specs = [
14027 TensorSpec {
14028 name: "q".into(),
14029 dtype: TensorDtype::Q4TiledP,
14030 shape: vec![r1, cols],
14031 data: synth_q4tp(r1, cols),
14032 },
14033 TensorSpec {
14034 name: "kv".into(),
14035 dtype: TensorDtype::Q4TiledP,
14036 shape: vec![r2, cols],
14037 data: synth_q4tp(r2, cols),
14038 },
14039 ];
14040 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
14041 std::fs::create_dir_all(&dir).unwrap();
14042 let path = dir.join("m.cmf");
14043 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14044 let model = Arc::new(CmfModel::open(&path).unwrap());
14045 let (a, b) = (
14046 QTensor::from_model(&model, "q").unwrap(),
14047 QTensor::from_model(&model, "kv").unwrap(),
14048 );
14049 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
14050 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
14051 let x: Vec<f32> = (0..cols)
14052 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
14053 .collect();
14054 let pool = Pool::new(3);
14055 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14056 a.matvec(&x, &mut ea, Some(&pool));
14057 b.matvec(&x, &mut eb, Some(&pool));
14058 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14059 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14060 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
14061 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
14062 let _ = std::fs::remove_dir_all(&dir);
14063 }
14064
14065 #[test]
14071 fn multi_token_moe_rows_equal_single_token_decode() {
14072 use crate::pool::Pool;
14073 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14074
14075 let (h, inter, ne) = (64usize, 128usize, 3usize);
14076 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14077 "arch_name": "tiny-q4tp-moe",
14078 "hidden_size": h,
14079 "intermediate_size": inter,
14080 "num_layers": 1,
14081 "num_attention_heads": 2,
14082 "num_kv_heads": 1,
14083 "head_dim": 32,
14084 "vocab_size": 8,
14085 "layer_types": ["FullAttention"],
14086 "rms_norm_eps": 1e-6,
14087 "max_position_embeddings": 8,
14088 "linear_conv_kernel_dim": 0,
14089 "linear_num_key_heads": 0,
14090 "linear_num_value_heads": 0
14091 }))
14092 .unwrap();
14093 let header = CmfHeader {
14094 format: "cmf".into(),
14095 version: CMF_VERSION,
14096 arch,
14097 quant_type: QuantType::Q4Block,
14098 provenance: None,
14099 tokenizer_config: None,
14100 section_hashes: None,
14101 skills: Vec::new(),
14102 shard: None,
14103 calibration: None,
14104 routing: None,
14105 genome: None,
14106 lineage: Vec::new(),
14107 router: None,
14108 segments: Vec::new(),
14109 };
14110 let mut specs = Vec::new();
14111 for e in 0..ne {
14112 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
14113 .into_iter()
14114 .enumerate()
14115 {
14116 let mut data = synth_q4tp(r, c);
14119 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
14120 .iter_mut()
14121 .enumerate()
14122 {
14123 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
14124 }
14125 specs.push(TensorSpec {
14126 name: format!("{n}{e}"),
14127 dtype: TensorDtype::Q4TiledP,
14128 shape: vec![r, c],
14129 data,
14130 });
14131 }
14132 }
14133 let dir = std::env::temp_dir().join(format!(
14134 "cmf-moe-rows-{}-{}",
14135 std::process::id(),
14136 FLOAT_ACTIVATIONS.get()
14137 ));
14138 std::fs::create_dir_all(&dir).unwrap();
14139 let path = dir.join("m.cmf");
14140 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14141 let model = Arc::new(CmfModel::open(&path).unwrap());
14142 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
14143 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
14144 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
14145 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
14146 let b = 4usize;
14147 let mut xs: Vec<f32> = (0..b * h)
14148 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
14149 .collect();
14150 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
14153 (vec![2, 0], vec![0.6, 0.4]),
14154 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14155 (vec![1], vec![1.0]),
14156 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14157 ];
14158 let pool = Pool::new(3);
14159 let mut want = vec![0f32; b * h];
14161 for (tk, (idx, w)) in routes.iter().enumerate() {
14162 let x = &xs[tk * h..(tk + 1) * h];
14163 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14164 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14165 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14166 if FLOAT_ACTIVATIONS.get() {
14167 for (slot, &e) in idx.iter().enumerate() {
14168 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14169 g[e].matvec(x, &mut gate, Some(&pool));
14170 u[e].matvec(x, &mut up, Some(&pool));
14171 for (v, u) in gate.iter_mut().zip(up) {
14172 *v = (*v / (1.0 + (-*v).exp())) * u;
14173 }
14174 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14175 }
14176 }
14177 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14178 assert!(QTensor::moe_down_many(
14179 &downs,
14180 &gs,
14181 w,
14182 &mut want[tk * h..(tk + 1) * h],
14183 Some(&pool)
14184 ));
14185 }
14186 if FLOAT_ACTIVATIONS.get() {
14187 for (tk, (idx, w)) in routes.iter().enumerate() {
14188 let mut scalar = vec![0.0; h];
14189 for (&e, &weight) in idx.iter().zip(w) {
14190 let (mut gate, mut up, mut down) =
14191 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14192 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14193 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14194 for (v, u) in gate.iter_mut().zip(up) {
14195 *v = (*v / (1.0 + (-*v).exp())) * u;
14196 }
14197 d[e].matvec(&gate, &mut down, Some(&pool));
14198 for (v, d) in scalar.iter_mut().zip(down) {
14199 *v += weight * d;
14200 }
14201 }
14202 assert_eq!(
14203 &want[tk * h..(tk + 1) * h],
14204 scalar,
14205 "float many equals scalar experts"
14206 );
14207 }
14208 }
14209 let mut experts: Vec<usize> = Vec::new();
14211 let mut groups: Vec<Vec<usize>> = Vec::new();
14212 for (tk, (idx, _)) in routes.iter().enumerate() {
14213 for &e in idx {
14214 match experts.iter().position(|&x| x == e) {
14215 Some(k) => groups[k].push(tk),
14216 None => {
14217 experts.push(e);
14218 groups.push(vec![tk]);
14219 }
14220 }
14221 }
14222 }
14223 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14224 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14225 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14226 assert!(QTensor::moe_gate_up_rows(
14227 &pairs,
14228 &groups,
14229 &xs,
14230 &mut gs,
14231 Some(&pool)
14232 ));
14233 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14234 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14235 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14236 assert!(QTensor::moe_down_rows(
14237 &downs,
14238 &lens,
14239 &gs,
14240 &mut ds,
14241 Some(&pool)
14242 ));
14243 let slot = |tk: usize, e: usize| {
14244 let k = experts.iter().position(|&x| x == e).unwrap();
14245 groups[..k].iter().map(|g| g.len()).sum::<usize>()
14246 + groups[k].iter().position(|&x| x == tk).unwrap()
14247 };
14248 let mut got = vec![0f32; b * h];
14249 for (tk, (idx, w)) in routes.iter().enumerate() {
14250 for i in 0..h {
14251 let mut acc = 0f32;
14252 for (&e, &we) in idx.iter().zip(w) {
14253 acc += we * ds[slot(tk, e)][i];
14254 }
14255 got[tk * h + i] = acc;
14256 }
14257 }
14258 assert!(want.iter().any(|v| *v != 0.0));
14259 assert_eq!(
14260 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14261 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14262 "multi-token MoE must equal decode bit for bit"
14263 );
14264
14265 let b5 = 5usize;
14268 let x5: Vec<f32> = (0..b5 * h)
14269 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14270 .collect();
14271 let mut mm = vec![0f32; b5 * inter];
14272 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14273 for tk in 0..b5 {
14274 let mut mv = vec![0f32; inter];
14275 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14276 assert_eq!(
14277 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14278 mm[tk * inter..(tk + 1) * inter]
14279 .iter()
14280 .map(|v| v.to_bits())
14281 .collect::<Vec<_>>(),
14282 "row-exact matmat token {tk}"
14283 );
14284 }
14285 let _ = std::fs::remove_dir_all(&dir);
14288 }
14289
14290 #[test]
14298 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14299 use crate::pool::Pool;
14300 use std::sync::atomic::Ordering::Relaxed;
14301 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14302 Q4TP_ALT.store(2, Relaxed);
14305 let pool = Pool::new(3);
14306 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14310 let bytes = synth_q4tp(rows, cols);
14311 let mut xs: Vec<f32> = (0..b * cols)
14312 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14313 .collect();
14314 xs[3] = 7.5; let run = |exact: bool| {
14316 let mut out = vec![0f32; b * rows];
14317 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14318 out
14319 };
14320 let (fast, exact) = (run(false), run(true));
14321 #[cfg(not(target_arch = "aarch64"))]
14322 let _ = fast;
14323 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14324 let mut matvecs = vec![0f32; b * rows];
14325 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14326 q4tp_matvec(
14327 &bytes,
14328 &xs[bi * cols..(bi + 1) * cols],
14329 rows,
14330 cols,
14331 o,
14332 Some(&pool),
14333 );
14334 }
14335 assert!(matvecs.iter().any(|v| *v != 0.0));
14336 assert_eq!(
14337 bits(&exact),
14338 bits(&matvecs),
14339 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14340 );
14341 #[cfg(target_arch = "aarch64")]
14342 {
14343 let gpr = cols / GROUP_SIZE;
14345 let v = Q4tpView::new(&bytes, rows, cols);
14346 let mut old = vec![0f32; b * rows];
14347 let mut sc = vec![0f32; gpr];
14348 let a8w8 = a8w8_enabled();
14349 let blocked = sdot_enabled() && blocked_enabled();
14350 let acts: Vec<SplitAct> = (0..b)
14351 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14352 .collect();
14353 for r in 0..rows {
14354 v.scales_into(r, gpr, &mut sc);
14355 if !a8w8 {
14356 for bi in 0..b {
14357 let x = &xs[bi * cols..(bi + 1) * cols];
14358 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14359 }
14360 continue;
14361 }
14362 let finish = |d: f32, act: &SplitAct| {
14363 let mut acc = d * act.sx;
14364 for &(j, xv) in &act.outliers {
14365 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14366 acc += w * s * xv;
14367 }
14368 acc
14369 };
14370 let mut bi = 0usize;
14371 while blocked && bi + 4 <= b {
14372 let xs4 = [
14373 acts[bi].xq.as_slice(),
14374 acts[bi + 1].xq.as_slice(),
14375 acts[bi + 2].xq.as_slice(),
14376 acts[bi + 3].xq.as_slice(),
14377 ];
14378 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14379 for k in 0..4 {
14380 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14381 }
14382 bi += 4;
14383 }
14384 for (bi, act) in acts.iter().enumerate().skip(bi) {
14385 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14386 old[bi * rows + r] = finish(d, act);
14387 }
14388 }
14389 assert_eq!(
14390 bits(&fast),
14391 bits(&old),
14392 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14393 );
14394 if blocked {
14397 assert_ne!(
14398 bits(&fast),
14399 bits(&matvecs),
14400 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14401 );
14402 }
14403 }
14404 }
14405 Q4TP_ALT.store(0, Relaxed);
14406 }
14407
14408 #[test]
14409 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14410 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14411 let active = AtomicUsize::new(0);
14414 counted_row_exact_scope(&active, || {
14415 assert_eq!(active.load(Ordering::Acquire), 1);
14416 counted_row_exact_scope(&active, || {
14417 assert_eq!(active.load(Ordering::Acquire), 2);
14418 });
14419 assert_eq!(active.load(Ordering::Acquire), 1);
14420 });
14421 assert_eq!(active.load(Ordering::Acquire), 0);
14422
14423 let both_entered = Barrier::new(2);
14424 let release_last = Barrier::new(2);
14425 std::thread::scope(|s| {
14426 let first = s.spawn(|| counted_row_exact_scope(&active, || {
14427 both_entered.wait();
14428 }));
14429 let last = s.spawn(|| counted_row_exact_scope(&active, || {
14430 both_entered.wait();
14431 release_last.wait();
14432 }));
14433 first.join().unwrap();
14434 let after_first = active.load(Ordering::Acquire);
14435 release_last.wait();
14436 last.join().unwrap();
14437 assert_eq!(after_first, 1, "second request must remain exact");
14438 });
14439 assert_eq!(active.load(Ordering::Acquire), 0);
14440 let panic = std::panic::catch_unwind(|| {
14441 counted_row_exact_scope(&active, || panic!("scope unwind"));
14442 });
14443 assert!(panic.is_err());
14444 assert_eq!(active.load(Ordering::Acquire), 0);
14445 }
14446
14447 #[test]
14448 #[cfg(target_arch = "x86_64")]
14449 fn q4tp_float_avx2_is_bitwise_scalar() {
14450 if !avx2_enabled() {
14451 return;
14452 }
14453 for cols in [32, 64, 96, 2048, 4096] {
14454 let rows = 9;
14455 let bytes = synth_q4tp(rows, cols);
14456 let v = Q4tpView::new(&bytes, rows, cols);
14457 let gpr = cols / GROUP_SIZE;
14458 let mut sc = vec![0.0; gpr];
14459 for seed in 1..=5 {
14460 let xs: Vec<f32> = (0..cols)
14461 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14462 .collect();
14463 for r in 0..rows {
14464 v.scales_into(r, gpr, &mut sc);
14465 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14466 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14467 assert_eq!(
14468 scalar.to_bits(),
14469 vector.to_bits(),
14470 "cols={cols} row={r} seed={seed}"
14471 );
14472 }
14473 }
14474 }
14475 }
14476
14477 #[test]
14478 fn multi_token_moe_rows_float_equal_single_token_decode() {
14479 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14480 }
14481
14482 #[test]
14483 fn full_gpu_q8_scope_is_nested_and_thread_local() {
14484 assert!(!FULL_GPU_Q8.get());
14485 let before = gpu_split_frac();
14486 {
14487 let _guard = enter_full_gpu_q8_scope();
14488 assert_eq!(gpu_split_frac(), 1.0);
14489 {
14490 let _nested = enter_full_gpu_q8_scope();
14491 }
14492 assert_eq!(gpu_split_frac(), 1.0);
14493 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14494 }
14495 assert!(!FULL_GPU_Q8.get());
14496 assert_eq!(gpu_split_frac(), before);
14497 }
14498
14499 #[test]
14500 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14501 assert!(!FLOAT_ACTIVATIONS.get());
14502 let before = a8w8_enabled();
14503 float_activations_scope(|| {
14504 assert!(!a8w8_enabled());
14505 float_activations_scope(|| assert!(!a8w8_enabled()));
14506 assert!(FLOAT_ACTIVATIONS.get());
14507 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14508 .join()
14509 .unwrap();
14510 });
14511 assert!(!FLOAT_ACTIVATIONS.get());
14512 assert_eq!(a8w8_enabled(), before);
14513 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14514 assert!(!FLOAT_ACTIVATIONS.get());
14515 }
14516
14517 #[test]
14520 fn batched_matmat_equals_per_position_matvec() {
14521 let (rows, cols, b) = (8, 64, 5);
14522 let groups = rows * cols / GROUP_SIZE;
14524 let mut q4 = Vec::new();
14525 for i in 0..groups * 16 {
14526 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14527 }
14528 for g in 0..groups {
14529 q4.extend_from_slice(
14530 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14531 );
14532 }
14533 let ng = cols / GROUP_SIZE;
14535 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14536 let mut vb = bits.clone();
14537 for g in 0..rows * ng {
14538 vb.extend_from_slice(
14539 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14540 );
14541 }
14542 for r in 0..rows {
14543 let bw = bits[r] as usize;
14544 let (mut acc, mut nb) = (0u64, 0usize);
14545 let mut rowbytes = Vec::new();
14546 for i in 0..cols {
14547 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14548 acc = (acc << bw) | v;
14549 nb += bw;
14550 while nb >= 8 {
14551 nb -= 8;
14552 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14553 }
14554 }
14555 if nb > 0 {
14556 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14557 }
14558 vb.extend_from_slice(&rowbytes);
14559 }
14560 let offsets = vbit_row_offsets(&vb, rows, cols);
14561
14562 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14563
14564 let mut got = vec![0f32; b * rows];
14566 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14567 for bi in 0..b {
14568 let mut expect = vec![0f32; rows];
14569 q4matvec(
14570 &q4,
14571 &xs[bi * cols..(bi + 1) * cols],
14572 rows,
14573 cols,
14574 &mut expect,
14575 None,
14576 );
14577 assert_eq!(
14578 &got[bi * rows..(bi + 1) * rows],
14579 &expect[..],
14580 "q4 batch pos {bi}"
14581 );
14582 }
14583
14584 let mut got = vec![0f32; b * rows];
14586 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14587 for bi in 0..b {
14588 let mut expect = vec![0f32; rows];
14589 vbitmatvec(
14590 &vb,
14591 &offsets,
14592 &xs[bi * cols..(bi + 1) * cols],
14593 rows,
14594 cols,
14595 &mut expect,
14596 None,
14597 );
14598 assert_eq!(
14599 &got[bi * rows..(bi + 1) * rows],
14600 &expect[..],
14601 "vbit batch pos {bi}"
14602 );
14603 }
14604 }
14605
14606 #[test]
14610 fn q4_tiled_matches_q4_block_bitexact() {
14611 let (rows, cols, b) = (8usize, 128usize, 3usize);
14612 let groups = rows * cols / GROUP_SIZE;
14613 let mut split = Vec::with_capacity(groups * 18);
14614 for i in 0..groups * 16 {
14615 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14616 }
14617 for g in 0..groups {
14618 split.extend_from_slice(
14619 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14620 );
14621 }
14622 let (packed, scales) = split.split_at(groups * 16);
14624 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14625 for g in 0..groups {
14626 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14627 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14628 }
14629
14630 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14631 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14633
14634 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14635 q4matvec(&split, &x1, rows, cols, &mut a, None);
14636 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14637 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14638
14639 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14640 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14641 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14642 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14643 assert_eq!(a1, t1);
14644 assert_eq!(a2, t2);
14645
14646 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14647 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14648 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14649 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14650 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14651 }
14652
14653 #[test]
14660 fn q4matvec_sdot_outlier_exact() {
14661 let (rows, cols) = (4, 128);
14662 let groups = rows * cols / GROUP_SIZE;
14663 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14664 for i in 0..groups * 16 {
14665 bytes.push(((i * 11 + 5) % 256) as u8);
14666 }
14667 for g in 0..groups {
14668 let s = 0.02 + 0.002 * g as f32;
14669 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14670 }
14671 let mut x: Vec<f32> = (0..cols)
14672 .map(|i| match i % 3 {
14673 0 => 1.0,
14674 1 => -1.0,
14675 _ => 0.0,
14676 })
14677 .collect();
14678 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
14681 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14682 let mut expect = vec![0.0f32; rows];
14683 for r in 0..rows {
14684 expect[r] = reference[r * cols..(r + 1) * cols]
14685 .iter()
14686 .zip(&x)
14687 .map(|(w, xv)| w * xv)
14688 .sum();
14689 }
14690 let mut got = vec![0.0f32; rows];
14691 q4matvec(&bytes, &x, rows, cols, &mut got, None);
14692 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14693 for r in 0..rows {
14694 assert!(
14695 (got[r] - expect[r]).abs() < 2e-3 * scale,
14696 "row {r}: {} vs {} (outlier term must be exact)",
14697 got[r],
14698 expect[r]
14699 );
14700 }
14701 }
14702
14703 #[test]
14707 fn q1t_matvec_matches_reference() {
14708 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14709 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
14711 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14712 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14714 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14715 let mut bytes = Vec::new();
14716 for r in 0..rows {
14717 for g in 0..gpr {
14718 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14719 let mut c = [0u8; 7];
14720 for k in 0..GROUP_SIZE {
14721 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14723 0
14724 } else {
14725 ((k + r * 3 + g) % 3) as u8 };
14727 cortiq_core::quant::q1t_pack(&mut c, k, code);
14728 }
14729 bytes.extend_from_slice(&c);
14730 }
14731 }
14732 let mut row_ptr = vec![0u32; rows + 1];
14735 for &(idx, _) in &outliers {
14736 row_ptr[idx as usize / cols + 1] += 1;
14737 }
14738 for r in 0..rows {
14739 row_ptr[r + 1] += row_ptr[r];
14740 }
14741 for &p in &row_ptr {
14742 bytes.extend_from_slice(&p.to_le_bytes());
14743 }
14744 for &(idx, v) in &outliers {
14745 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14746 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14747 }
14748
14749 let mut refw = vec![0f32; rows * cols];
14750 dequant_q1t(&bytes, rows, cols, &mut refw);
14751 let x: Vec<f32> = (0..cols)
14754 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14755 .collect();
14756 let mut expect = vec![0f32; rows];
14757 for r in 0..rows {
14758 let mut a = 0.0f32;
14759 for j in 0..cols {
14760 a += refw[r * cols + j] * x[j];
14761 }
14762 expect[r] = a;
14763 }
14764 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14765 let mut got = vec![0f32; rows];
14766 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14767 for r in 0..rows {
14768 assert!(
14769 (got[r] - expect[r]).abs() < tol(expect[r]),
14770 "row {r}: {} vs {}",
14771 got[r],
14772 expect[r]
14773 );
14774 }
14775 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14777 let mut gm = vec![0f32; 2 * rows];
14778 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14779 for r in 0..rows {
14780 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14781 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14782 }
14783 let xb: Vec<f32> = (0..cols)
14787 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14788 .collect();
14789 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14790 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14791 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14792 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14793 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14794 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14795 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14796 }
14797
14798 #[test]
14801 fn q1t_matvec2_odd_gpr_matches_singles() {
14802 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14803 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
14805 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14806 for r in 0..rows {
14807 for g in 0..gpr {
14808 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14809 let mut c = [0u8; 7];
14810 for k in 0..GROUP_SIZE {
14811 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14812 }
14813 bytes.extend_from_slice(&c);
14814 }
14815 }
14816 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14817 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14818 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14819 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14820 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14821 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14822 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14823 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14824 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14825 }
14826
14827 #[test]
14831 #[ignore]
14832 fn q1t_matvec2_speed() {
14833 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14834 use std::time::Instant;
14835 let (rows, cols) = (8192usize, 4096usize);
14836 let gpr = cols / GROUP_SIZE;
14837 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14838 for r in 0..rows {
14839 for g in 0..gpr {
14840 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14841 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14842 let mut c = [0u8; 7];
14843 for k in 0..GROUP_SIZE {
14844 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14845 }
14846 bytes.extend_from_slice(&c);
14847 }
14848 }
14849 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14850 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14851 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14852 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14853 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14855 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14856 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14857 for _ in 0..8 {
14858 let t0 = Instant::now();
14859 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14860 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14861 let t1 = Instant::now();
14862 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14863 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14864 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14865 }
14866 assert_eq!(p1, s1);
14867 assert_eq!(p2, s2);
14868 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14869 }
14870
14871 #[test]
14875 #[ignore]
14876 fn q1t_matvec_speed() {
14877 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14878 use std::time::Instant;
14879 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
14881 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14882 for r in 0..rows {
14883 for g in 0..gpr {
14884 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14885 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14886 let mut c = [0u8; 7];
14887 for k in 0..GROUP_SIZE {
14888 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14889 }
14890 bytes.extend_from_slice(&c);
14891 }
14892 }
14893 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
14895 let mut idx = 0usize;
14896 while idx < n {
14897 row_ptr[idx / cols + 1] += 1;
14898 idx += stride;
14899 }
14900 for r in 0..rows {
14901 row_ptr[r + 1] += row_ptr[r];
14902 }
14903 for &p in &row_ptr {
14904 bytes.extend_from_slice(&p.to_le_bytes());
14905 }
14906 let mut idx = 0usize;
14907 while idx < n {
14908 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14909 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14910 idx += stride;
14911 }
14912 let x: Vec<f32> = (0..cols)
14915 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14916 .collect();
14917 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14918
14919 let slow = |out: &mut [f32]| {
14921 let mut buf = vec![0f32; cols];
14922 for r in 0..rows {
14923 for g in 0..gpr {
14924 let off = (r * gpr + g) * Q1T_TILE;
14925 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14926 let codes = &bytes[off + 2..off + Q1T_TILE];
14927 for k in 0..GROUP_SIZE {
14928 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14929 1 => s,
14930 2 => -s,
14931 _ => 0.0,
14932 };
14933 }
14934 }
14935 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14936 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14937 }
14938 };
14939 let iters = 5;
14940 let mut a = vec![0f32; rows];
14941 slow(&mut a); let t = Instant::now();
14943 for _ in 0..iters {
14944 slow(&mut a);
14945 }
14946 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14947
14948 let mut b = vec![0f32; rows];
14949 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
14951 for _ in 0..iters {
14952 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14953 }
14954 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14955
14956 for r in 0..rows {
14957 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14958 }
14959 println!(
14960 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
14961 slow_ms / fast_ms
14962 );
14963 }
14964}
14965
14966#[cfg(test)]
14967mod gemm_bench {
14968 #[test]
14979 #[ignore]
14980 fn q4tp_matmat_throughput() {
14981 let _alt = super::Q4TP_ALT_TEST_LOCK
14982 .lock()
14983 .unwrap_or_else(|e| e.into_inner());
14984 let b: usize = std::env::var("CMF_BENCH_B")
14988 .ok()
14989 .and_then(|v| v.parse().ok())
14990 .unwrap_or(296);
14991 let (rows, cols) = (9216usize, 2304usize);
14992 let (_, _, _) = (rows, cols, b);
14993 let total =
14994 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14995 .unwrap();
14996 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15001 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15002 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15003 let step = cortiq_core::quant::f32_to_f16(0.1);
15004 for r in 0..rows {
15005 let o = params_off + r * 4;
15006 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15007 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15008 }
15009 let _ = codes_off;
15010 let xs: Vec<f32> = (0..b * cols)
15011 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15012 .collect();
15013 let mut out = vec![0f32; b * rows];
15014 let pool = crate::pool::Pool::from_env();
15015 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15021 let reps: usize = std::env::var("CMF_BENCH_REPS")
15022 .ok()
15023 .and_then(|v| v.parse().ok())
15024 .unwrap_or(10);
15025 let mut best = [f64::MAX; 2];
15026 let mut sums = [0f32; 2];
15027 for _ in 0..reps {
15028 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
15029 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
15030 let t = std::time::Instant::now();
15031 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15032 best[k] = best[k].min(t.elapsed().as_secs_f64());
15033 sums[k] = out.iter().take(64).sum::<f32>();
15034 }
15035 }
15036 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15037 for (k, name) in ["previous", "tuned "].iter().enumerate() {
15038 println!(
15039 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15040 best[k] * 1e3,
15041 flops / best[k] / 1e9,
15042 sums[k]
15043 );
15044 }
15045 assert!(
15046 (sums[0] - sums[1]).abs() < 1e-2,
15047 "the tuned kernel changed the result: {} vs {}",
15048 sums[0],
15049 sums[1]
15050 );
15051 }
15052
15053 #[test]
15059 fn q4tp_matmat_blocked_matches_scalar() {
15060 use std::sync::atomic::Ordering::Relaxed;
15061 let _alt = super::Q4TP_ALT_TEST_LOCK
15062 .lock()
15063 .unwrap_or_else(|e| e.into_inner());
15064 for &(rows, cols, b) in &[
15072 (64usize, 128usize, 7usize),
15073 (33, 96, 4),
15074 (16, 256, 9),
15075 (192, 2304, 37),
15076 ] {
15077 let total = cortiq_core::quant::expected_nbytes(
15078 cortiq_core::TensorDtype::Q4TiledP,
15079 &[rows, cols],
15080 )
15081 .unwrap();
15082 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15083 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
15084 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15085 let step = cortiq_core::quant::f32_to_f16(0.1);
15086 for r in 0..rows {
15087 let o = params_off + r * 4;
15088 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15089 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15090 }
15091 let xs: Vec<f32> = (0..b * cols)
15092 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
15093 .collect();
15094 let mut got = vec![0f32; b * rows];
15095 let mut want = vec![0f32; b * rows];
15096 let gpr = cols / 32;
15097 let view = super::Q4tpView::new(&bytes, rows, cols);
15098 let pool = crate::pool::Pool::from_env();
15099 super::Q4TP_ALT.store(2, Relaxed);
15100 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
15101 super::Q4TP_ALT.store(1, Relaxed);
15102 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
15103 super::Q4TP_ALT.store(0, Relaxed);
15104 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15111 let (mut worst, mut at) = (0f32, 0usize);
15112 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
15113 if (g - w).abs() > worst {
15114 worst = (g - w).abs();
15115 at = i;
15116 }
15117 }
15118 assert!(
15119 worst <= 1e-4 * scale,
15120 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
15121 (scale {scale:.3e}) at cell {at}: {} vs {}",
15122 got[at],
15123 want[at]
15124 );
15125
15126 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
15133 for bi in 0..b {
15134 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
15135 for r in 0..rows {
15136 let mut sc = vec![0f32; gpr];
15137 view.scales_into(r, gpr, &mut sc);
15138 let mut exact = 0f64;
15139 for j in 0..cols {
15140 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15141 exact += w as f64 * sq as f64 * act.xq[j] as f64;
15142 }
15143 exact *= act.sx as f64;
15144 for &(j, xv) in &act.outliers {
15145 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15146 exact += w as f64 * sq as f64 * xv as f64;
15147 }
15148 let i = bi * rows + r;
15149 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
15150 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
15151 }
15152 }
15153 println!(
15154 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15155 per-column {e_scalar:.3e}"
15156 );
15157 assert!(
15162 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15163 "{rows}x{cols} b={b}: error against f64 too large — blocked \
15164 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15165 );
15166 }
15167 }
15168
15169 #[test]
15175 #[ignore]
15176 fn q4tp_matmat_row_exact_cost() {
15177 let _alt = super::Q4TP_ALT_TEST_LOCK
15178 .lock()
15179 .unwrap_or_else(|e| e.into_inner());
15180 let pool = crate::pool::Pool::from_env();
15181 let reps: usize = std::env::var("CMF_BENCH_REPS")
15182 .ok()
15183 .and_then(|v| v.parse().ok())
15184 .unwrap_or(30);
15185 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15186 let total = cortiq_core::quant::expected_nbytes(
15187 cortiq_core::TensorDtype::Q4TiledP,
15188 &[rows, cols],
15189 )
15190 .unwrap();
15191 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15192 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15193 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15194 let step = cortiq_core::quant::f32_to_f16(0.1);
15195 for r in 0..rows {
15196 let o = params_off + r * 4;
15197 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15198 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15199 }
15200 for &b in &[2usize, 4, 5, 8] {
15201 let xs: Vec<f32> = (0..b * cols)
15202 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15203 .collect();
15204 let mut out = vec![0f32; b * rows];
15205 let mut best = [f64::MAX; 3];
15206 for _ in 0..reps {
15207 for (k, best_k) in best.iter_mut().enumerate() {
15208 let t = std::time::Instant::now();
15209 match k {
15210 0 | 1 => super::q4tp_matmat_with(
15211 &bytes,
15212 &xs,
15213 b,
15214 rows,
15215 cols,
15216 &mut out,
15217 pool.as_deref(),
15218 k == 1,
15219 ),
15220 _ => {
15221 for (bi, o) in out.chunks_mut(rows).enumerate() {
15222 super::q4tp_matvec(
15223 &bytes,
15224 &xs[bi * cols..(bi + 1) * cols],
15225 rows,
15226 cols,
15227 o,
15228 pool.as_deref(),
15229 );
15230 }
15231 }
15232 }
15233 *best_k = best_k.min(t.elapsed().as_secs_f64());
15234 }
15235 }
15236 println!(
15237 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15238 {b} matvecs {:.3} ms",
15239 best[0] * 1e3,
15240 best[1] * 1e3,
15241 best[2] * 1e3
15242 );
15243 }
15244 }
15245 }
15246
15247 #[test]
15251 #[ignore]
15252 fn q4t_matmat_throughput() {
15253 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15254 let total =
15255 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15256 .unwrap();
15257 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15260 let sc = cortiq_core::quant::f32_to_f16(0.02);
15261 for t in bytes.chunks_mut(super::Q4_TILE) {
15262 t[..2].copy_from_slice(&sc.to_le_bytes());
15263 }
15264 let xs: Vec<f32> = (0..b * cols)
15265 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15266 .collect();
15267 let mut out = vec![0f32; b * rows];
15268 let pool = crate::pool::Pool::from_env();
15269 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15270 let reps: usize = std::env::var("CMF_BENCH_REPS")
15271 .ok()
15272 .and_then(|v| v.parse().ok())
15273 .unwrap_or(10);
15274 let mut best = f64::MAX;
15275 for _ in 0..reps {
15276 let t = std::time::Instant::now();
15277 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15278 best = best.min(t.elapsed().as_secs_f64());
15279 }
15280 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15281 println!(
15282 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15283 best * 1e3,
15284 flops / best / 1e9,
15285 out.iter().take(64).sum::<f32>()
15286 );
15287 }
15288}