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
23pub enum QTensor {
24 F32 {
25 data: Vec<f32>,
26 rows: usize,
27 cols: usize,
28 },
29 Mapped {
30 model: Arc<CmfModel>,
31 idx: usize,
33 dtype: TensorDtype,
34 rows: usize,
35 cols: usize,
36 row_scale: Vec<f32>,
38 col_field: Vec<f32>,
40 vbit_offsets: Vec<usize>,
44 repack: Vec<u8>,
52 },
53 Blr(Box<crate::blr::BlrTensor>),
59}
60
61fn repack_enabled() -> bool {
68 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
69 *ON.get_or_init(|| {
70 std::env::var("CMF_REPACK")
71 .map(|v| v == "1")
72 .unwrap_or(cfg!(target_os = "android"))
73 })
74}
75
76fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
80 #[cfg(target_arch = "aarch64")]
81 let arch_ok = sdot_enabled();
82 #[cfg(not(target_arch = "aarch64"))]
83 let arch_ok = false;
84 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
85 return Vec::new();
86 }
87 q8_repack_layout(bytes, rows, cols)
88}
89
90fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
93 let groups = rows / 4;
94 let mut rep = vec![0u8; groups * 4 * cols];
95 for g in 0..groups {
96 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
97 for c in 0..cols / 16 {
98 for lane in 0..4 {
99 let src = (g * 4 + lane) * cols + c * 16;
100 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
101 .copy_from_slice(&bytes[src..src + 16]);
102 }
103 }
104 }
105 rep
106}
107
108fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
111 let ng = cols / GROUP_SIZE;
112 let bits = &bytes[..rows];
113 let mut offsets = Vec::with_capacity(rows + 1);
114 let mut off = rows + rows * ng * 2;
115 for r in 0..rows {
116 offsets.push(off);
117 off += (cols * bits[r] as usize).div_ceil(8);
118 }
119 offsets.push(off);
120 offsets
121}
122
123fn blocked_enabled() -> bool {
129 use std::sync::atomic::Ordering::Relaxed;
130 match BLOCKED_OVERRIDE.load(Relaxed) {
131 1 => false,
132 2 => true,
133 _ => {
134 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
135 *ON.get_or_init(|| {
136 std::env::var("CMF_X86_BLOCKED")
137 .map(|v| v != "0")
138 .unwrap_or(true)
139 })
140 }
141 }
142}
143
144static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
145
146pub fn set_blocked_override(on: Option<bool>) {
153 let v = match on {
154 None => 0,
155 Some(false) => 1,
156 Some(true) => 2,
157 };
158 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
159}
160
161fn gpu_lmhead_enabled() -> bool {
162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *ON.get_or_init(|| {
164 std::env::var("CMF_GPU_LMHEAD")
165 .map(|v| v != "0")
166 .unwrap_or(true)
167 })
168}
169
170fn gpu_split_frac() -> f32 {
171 if FULL_GPU_Q8.get() {
174 return 1.0;
175 }
176 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
177 *F.get_or_init(|| {
178 std::env::var("CMF_GPU_SPLIT")
179 .ok()
180 .and_then(|v| v.parse::<f32>().ok())
181 .unwrap_or(0.5)
182 .clamp(0.0, 1.0)
183 })
184}
185
186impl QTensor {
187 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
188 debug_assert_eq!(data.len(), rows * cols);
189 Self::F32 { data, rows, cols }
190 }
191
192 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
195 let idx = model
198 .tensor_index(name)
199 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
200 let entry = &model.tensors[idx];
201 if entry.shape.len() != 2 {
202 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
203 }
204 let (rows, cols) = (entry.shape[0], entry.shape[1]);
205 let bytes = model.entry_bytes(entry);
206
207 match entry.dtype {
208 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
209 let n = rows * cols;
210 let scales_off = n;
211 let row_scale: Vec<f32> = (0..rows)
212 .map(|o| {
213 f16_to_f32(u16::from_le_bytes([
214 bytes[scales_off + o * 2],
215 bytes[scales_off + o * 2 + 1],
216 ]))
217 })
218 .collect();
219 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
220 let col_off = n + rows * 2;
221 (0..cols)
222 .map(|i| {
223 f16_to_f32(u16::from_le_bytes([
224 bytes[col_off + i * 2],
225 bytes[col_off + i * 2 + 1],
226 ]))
227 })
228 .collect()
229 } else {
230 Vec::new()
231 };
232 Ok(Self::Mapped {
233 model: model.clone(),
234 idx,
235 dtype: entry.dtype,
236 rows,
237 cols,
238 row_scale,
239 col_field,
240 vbit_offsets: Vec::new(),
241 repack: q8_repack(bytes, rows, cols),
242 })
243 }
244 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
246 model: model.clone(),
247 idx,
248 dtype: entry.dtype,
249 rows,
250 cols,
251 row_scale: Vec::new(),
252 col_field: Vec::new(),
253 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
254 repack: Vec::new(),
255 }),
256 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
260 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
261 let offsets: Vec<usize> = (0..=rows)
262 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
263 .collect();
264 Ok(Self::Mapped {
265 model: model.clone(),
266 idx,
267 dtype: entry.dtype,
268 rows,
269 cols,
270 row_scale: Vec::new(),
271 col_field: Vec::new(),
272 vbit_offsets: offsets,
273 repack: Vec::new(),
274 })
275 }
276 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
282 model: model.clone(),
283 idx,
284 dtype: entry.dtype,
285 rows,
286 cols,
287 row_scale: Vec::new(),
288 col_field: Vec::new(),
289 vbit_offsets: Vec::new(),
290 repack: Vec::new(),
291 }),
292 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295 model: model.clone(),
296 idx,
297 dtype: entry.dtype,
298 rows,
299 cols,
300 row_scale: Vec::new(),
301 col_field: Vec::new(),
302 vbit_offsets: Vec::new(),
303 repack: Vec::new(),
304 }),
305 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
307 model: model.clone(),
308 idx,
309 dtype: entry.dtype,
310 rows,
311 cols,
312 row_scale: Vec::new(),
313 col_field: Vec::new(),
314 vbit_offsets: Vec::new(),
315 repack: Vec::new(),
316 }),
317 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318 model: model.clone(),
319 idx,
320 dtype: entry.dtype,
321 rows,
322 cols,
323 row_scale: Vec::new(),
324 col_field: Vec::new(),
325 vbit_offsets: Vec::new(),
326 repack: Vec::new(),
327 }),
328 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
331 model: model.clone(),
332 idx,
333 dtype: entry.dtype,
334 rows,
335 cols,
336 row_scale: Vec::new(),
337 col_field: Vec::new(),
338 vbit_offsets: Vec::new(),
339 repack: Vec::new(),
340 }),
341 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
342 model: model.clone(),
343 idx,
344 dtype: entry.dtype,
345 rows,
346 cols,
347 row_scale: Vec::new(),
348 col_field: Vec::new(),
349 vbit_offsets: Vec::new(),
350 repack: Vec::new(),
351 }),
352 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
356 model: model.clone(),
357 idx,
358 dtype: entry.dtype,
359 rows,
360 cols,
361 row_scale: Vec::new(),
362 col_field: Vec::new(),
363 vbit_offsets: Vec::new(),
364 repack: Vec::new(),
365 }),
366 TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
369 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
370 model, idx,
371 )?)))
372 }
373 _ => {
375 let mut data = vec![0.0f32; rows * cols];
376 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
377 Ok(Self::from_f32(data, rows, cols))
378 }
379 }
380 }
381
382 pub fn is_blr(&self) -> bool {
385 matches!(self, Self::Blr(_))
386 }
387
388 pub(crate) fn is_q1(&self) -> bool {
392 matches!(
393 self,
394 Self::Mapped {
395 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
396 ..
397 }
398 )
399 }
400
401 fn bin_w(&self) -> Option<BinW<'_>> {
404 match self {
405 Self::Mapped {
406 dtype: TensorDtype::Q1,
407 ..
408 } => Some(BinW::Q1(self.quant_bytes())),
409 Self::Mapped {
410 dtype: TensorDtype::Q1G,
411 rows,
412 cols,
413 ..
414 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
415 _ => None,
416 }
417 }
418
419 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
422 match self {
423 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
424 _ => None,
425 }
426 }
427
428 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
435 if self.has_prism_contract() {
436 return None;
437 }
438 match self {
439 #[cfg(target_os = "macos")]
440 Self::Mapped {
441 dtype: TensorDtype::Q1T,
442 ..
443 } if !crate::gpu::metal_q1t_enabled() => None,
444 Self::Mapped {
445 idx,
446 dtype:
447 TensorDtype::Q1
448 | TensorDtype::Q1G
449 | TensorDtype::Q1T
450 | TensorDtype::Q4Block
451 | TensorDtype::Q4Tiled
452 | TensorDtype::Q4TiledP
456 | TensorDtype::Q8Row
457 | TensorDtype::Q8_2f,
458 rows,
459 cols,
460 ..
461 } => Some((*idx, *rows, *cols)),
462 _ => None,
463 }
464 }
465
466 #[cfg(target_os = "macos")]
473 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
474 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
475 let name = &model.tensors[idx].name;
476 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
477 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
478 if forward && affine {
479 let e = model.tensors.get(idx)?;
480 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
481 }
482 }
483 self.q1_parts()
484 }
485
486 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
492 if self.has_prism_contract() {
493 return None;
494 }
495 match self {
496 Self::Mapped {
497 idx,
498 dtype: TensorDtype::Q4Tiled,
499 rows,
500 cols,
501 ..
502 } => Some((*idx, *rows, *cols)),
503 _ => None,
504 }
505 }
506
507 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
511 if self.has_prism_contract() {
512 return None;
513 }
514 match self {
515 Self::Mapped {
516 idx,
517 dtype: TensorDtype::Q4TiledP,
518 rows,
519 cols,
520 ..
521 } => Some((*idx, *rows, *cols)),
522 _ => None,
523 }
524 }
525
526 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
531 if self.has_prism_contract() {
532 return None;
533 }
534 match self {
535 Self::Mapped {
536 idx,
537 dtype: TensorDtype::Q8Row,
538 rows,
539 cols,
540 row_scale,
541 col_field,
542 ..
543 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
544 _ => None,
545 }
546 }
547
548 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
551 if self.has_prism_contract() {
552 return None;
553 }
554 match self {
555 Self::Mapped {
556 idx,
557 dtype: TensorDtype::Q8_2f,
558 rows,
559 cols,
560 row_scale,
561 col_field,
562 ..
563 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
564 _ => None,
565 }
566 }
567
568 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
572 match self {
573 Self::Mapped { dtype, .. } => Some(*dtype),
574 _ => None,
575 }
576 }
577
578 pub fn model_idx(&self) -> Option<usize> {
583 match self {
584 Self::Mapped { idx, .. } => Some(*idx),
585 _ => None,
586 }
587 }
588
589 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
594 match self {
595 Self::Mapped { model, .. } => Some(model.clone()),
596 _ => None,
597 }
598 }
599
600 pub(crate) fn has_prism_contract(&self) -> bool {
605 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
606 }
607
608 pub fn rows(&self) -> usize {
609 match self {
610 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
611 Self::Blr(t) => t.rows,
612 }
613 }
614
615 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
618 if self.has_prism_contract() {
619 return None;
620 }
621 match self {
622 Self::Mapped {
623 model,
624 idx,
625 dtype: TensorDtype::Q4Tiled,
626 ..
627 } => Some((model, *idx)),
628 _ => None,
629 }
630 }
631
632 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
635 if self.has_prism_contract() {
636 return None;
637 }
638 match self {
639 Self::Mapped {
640 model,
641 idx,
642 dtype: TensorDtype::Q4TiledP,
643 ..
644 } => Some((model, *idx)),
645 _ => None,
646 }
647 }
648
649 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
657 if self.has_prism_contract() {
658 return None;
659 }
660 match self {
661 Self::Mapped {
662 model,
663 idx,
664 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
665 ..
666 } => Some((model, *idx)),
667 _ => None,
668 }
669 }
670
671 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
674 if self.has_prism_contract() {
675 return None;
676 }
677 match self {
678 Self::Mapped {
679 model,
680 idx,
681 dtype: TensorDtype::Q8_2f,
682 ..
683 } => Some((model, *idx)),
684 _ => None,
685 }
686 }
687
688 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
691 if self.has_prism_contract() {
692 return None;
693 }
694 match self {
695 Self::Mapped {
696 model,
697 idx,
698 dtype: TensorDtype::Q2TiledP,
699 ..
700 } => Some((model, *idx)),
701 _ => None,
702 }
703 }
704
705 pub fn cols(&self) -> usize {
706 match self {
707 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
708 Self::Blr(t) => t.cols,
709 }
710 }
711
712 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
715 if self.has_prism_contract() {
716 return None;
717 }
718 match self {
719 Self::Mapped {
720 model,
721 idx,
722 dtype: TensorDtype::Q1,
723 ..
724 } => Some((model, *idx)),
725 _ => None,
726 }
727 }
728
729 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
740 if self.has_prism_contract() {
741 return None;
742 }
743 self.graph_weight_descriptor()
744 }
745
746 pub(crate) fn graph_weight_descriptor(
751 &self,
752 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
753 match self {
754 Self::Mapped {
755 model,
756 idx,
757 dtype: TensorDtype::Q8Row,
758 row_scale,
759 ..
760 } => Some((model, *idx, 0, row_scale.as_slice())),
761 Self::Mapped {
762 model,
763 idx,
764 dtype: TensorDtype::Q1,
765 ..
766 } => Some((model, *idx, 1, &[])),
767 Self::Mapped {
770 model,
771 idx,
772 dtype: TensorDtype::Q1G,
773 ..
774 } => Some((model, *idx, 10, &[])),
775 Self::Mapped {
780 model,
781 idx,
782 dtype: TensorDtype::Q4Tiled,
783 ..
784 } => Some((model, *idx, 5, &[])),
785 Self::Mapped {
789 model,
790 idx,
791 dtype: TensorDtype::Q4TiledP,
792 ..
793 } => Some((model, *idx, 6, &[])),
794 Self::Mapped {
795 model,
796 idx,
797 dtype: TensorDtype::Q4Block,
798 ..
799 } => Some((model, *idx, 2, &[])),
800 Self::Mapped {
805 model,
806 idx,
807 dtype: TensorDtype::Q8_2f,
808 ..
809 } => Some((model, *idx, 7, &[])),
810 Self::Mapped {
811 model,
812 idx,
813 dtype: TensorDtype::Q1T,
814 ..
815 } => Some((model, *idx, 3, &[])),
816 Self::Mapped {
820 model,
821 idx,
822 dtype: TensorDtype::Q2TiledP,
823 ..
824 } => Some((model, *idx, 9, &[])),
825 _ => None,
826 }
827 }
828
829 pub fn as_f32(&self) -> Option<&[f32]> {
832 match self {
833 Self::F32 { data, .. } => Some(data),
834 Self::Mapped { .. } | Self::Blr(_) => None,
835 }
836 }
837
838 fn quant_bytes(&self) -> &[u8] {
839 match self {
840 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
841 Self::Blr(t) => t.bytes(),
842 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
843 }
844 }
845
846 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
848 let cols = self.cols();
849 debug_assert_eq!(dst.len(), cols);
850 match self {
851 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
852 Self::Blr(t) => t.row_f32(r, dst),
853 Self::Mapped {
854 model,
855 idx,
856 dtype,
857 row_scale,
858 col_field,
859 vbit_offsets,
860 ..
861 } => {
862 if *dtype == TensorDtype::Q4Tiled {
863 let bytes = self.quant_bytes();
864 let gpr = cols / GROUP_SIZE;
865 for gi in 0..gpr {
866 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
867 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
868 for (k, &b) in tile[2..].iter().enumerate() {
869 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
870 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
871 }
872 }
873 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
874 crate::prism::inverse_embedding(model, dst);
875 }
876 return;
877 }
878 if *dtype == TensorDtype::Q4TiledP {
879 let bytes = self.quant_bytes();
880 let gpr = cols / GROUP_SIZE;
881 let v = Q4tpView::new(bytes, self.rows(), cols);
882 let mut sc = vec![0f32; gpr];
883 v.scales_into(r, gpr, &mut sc);
884 for gi in 0..gpr {
885 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
886 let s = sc[gi];
887 for (k, &b) in tile.iter().enumerate() {
888 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
889 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
890 }
891 }
892 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
893 crate::prism::inverse_embedding(model, dst);
894 }
895 return;
896 }
897 if *dtype == TensorDtype::Q2TiledP {
898 let bytes = self.quant_bytes();
899 let gpr = cols / GROUP_SIZE;
900 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
901 let mut sc = vec![0f32; gpr];
902 v.scales_into(r, gpr, &mut sc);
903 for gi in 0..gpr {
904 let ch =
905 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
906 let s = sc[gi];
907 for (k, &b) in ch.iter().enumerate() {
908 for j in 0..4 {
909 let center = if crate::prism::is_affine_target(
910 model,
911 &model.tensors[*idx].name,
912 ) {
913 1.0
914 } else {
915 1.5
916 };
917 dst[gi * GROUP_SIZE + k * 4 + j] =
918 (((b >> (2 * j)) & 3) as f32 - center) * s;
919 }
920 }
921 }
922 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
923 crate::prism::inverse_embedding(model, dst);
924 }
925 return;
926 }
927 if *dtype == TensorDtype::Q4Block {
928 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
929 let gpr = cols / GROUP_SIZE;
930 for gi in 0..gpr {
931 let g = r * gpr + gi;
932 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
933 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
934 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
935 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
936 }
937 }
938 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
939 crate::prism::inverse_embedding(model, dst);
940 }
941 return;
942 }
943 if *dtype == TensorDtype::Q1 {
944 let bytes = self.quant_bytes();
945 let gpr = cols / GROUP_SIZE;
946 for gi in 0..gpr {
947 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
948 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
949 for (j, &b) in tile[2..].iter().enumerate() {
950 for k in 0..8 {
951 dst[gi * GROUP_SIZE + j * 8 + k] =
952 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
953 }
954 }
955 }
956 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
957 crate::prism::inverse_embedding(model, dst);
958 }
959 return;
960 }
961 if *dtype == TensorDtype::Q1G {
962 let bytes = self.quant_bytes();
965 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
966 let gpr = cols / GROUP_SIZE;
967 for gi in 0..gpr {
968 let t = r * gpr + gi;
969 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
970 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
971 for k in 0..8 {
972 dst[gi * GROUP_SIZE + j * 8 + k] =
973 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
974 }
975 }
976 }
977 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
978 crate::prism::inverse_embedding(model, dst);
979 }
980 return;
981 }
982 if *dtype == TensorDtype::Q1T {
983 let bytes = self.quant_bytes();
984 let gpr = cols / GROUP_SIZE;
985 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
986 for gi in 0..gpr {
987 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
988 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
989 bytes[off],
990 bytes[off + 1],
991 ]));
992 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
993 for k in 0..GROUP_SIZE {
994 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
995 {
996 1 => s,
997 2 => -s,
998 _ => 0.0,
999 };
1000 }
1001 }
1002 let rows = self.rows();
1004 let entries = base_len + (rows + 1) * 4;
1005 if entries <= bytes.len() {
1006 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1007 let r0 = u32::from_le_bytes([
1008 ptrs[r * 4],
1009 ptrs[r * 4 + 1],
1010 ptrs[r * 4 + 2],
1011 ptrs[r * 4 + 3],
1012 ]) as usize;
1013 let r1 = u32::from_le_bytes([
1014 ptrs[(r + 1) * 4],
1015 ptrs[(r + 1) * 4 + 1],
1016 ptrs[(r + 1) * 4 + 2],
1017 ptrs[(r + 1) * 4 + 3],
1018 ]) as usize;
1019 let off = entries + r0 * 4;
1020 for i in 0..r1 - r0 {
1021 let item = &bytes[off + i * 4..off + i * 4 + 4];
1022 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1023 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1024 item[2], item[3],
1025 ]));
1026 if c < cols {
1027 dst[c] = v;
1028 }
1029 }
1030 }
1031 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1032 crate::prism::inverse_embedding(model, dst);
1033 }
1034 return;
1035 }
1036 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1037 let bytes = self.quant_bytes();
1038 let rows = self.rows();
1039 let ng = cols / GROUP_SIZE;
1040 let bits = &bytes[..rows];
1041 let sc_off = rows;
1042 let off = vbit_offsets[r];
1045 let b = bits[r] as usize;
1046 let l = ((1usize << (b - 1)) - 1) as f32;
1047 let data = &bytes[off..];
1048 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1049 for (i, d) in dst.iter_mut().enumerate() {
1050 while nbits < b {
1051 acc = (acc << 8) | data[byte_idx] as u64;
1052 byte_idx += 1;
1053 nbits += 8;
1054 }
1055 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1056 nbits -= b;
1057 let so = (r * ng + i / GROUP_SIZE) * 2;
1058 let sv = f16_to_f32(u16::from_le_bytes([
1059 bytes[sc_off + so],
1060 bytes[sc_off + so + 1],
1061 ]));
1062 *d = (u - l) * sv;
1063 }
1064 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1065 crate::prism::inverse_embedding(model, dst);
1066 }
1067 return;
1068 }
1069 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1070 let s = row_scale[r];
1071 match dtype {
1072 TensorDtype::Q8Row => {
1073 for (d, &b) in dst.iter_mut().zip(q) {
1074 *d = (b as i8) as f32 * s;
1075 }
1076 }
1077 TensorDtype::Q8_2f => {
1078 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1079 *d = (b as i8) as f32 * s * col_field[i];
1080 }
1081 }
1082 _ => unreachable!(),
1083 }
1084 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1085 crate::prism::inverse_embedding(model, dst);
1086 }
1087 }
1088 }
1089 }
1090
1091 pub fn sparse_col_ok(&self) -> bool {
1096 match self {
1097 Self::F32 { .. } => true,
1098 Self::Blr(_) => false,
1099 Self::Mapped { dtype, .. } => {
1100 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1101 }
1102 }
1103 }
1104
1105 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1109 let inter = self.cols();
1110 let hidden = self.rows();
1111 debug_assert_eq!(out.len(), hidden);
1112 match self {
1113 Self::F32 { data, .. } => {
1114 for (k, o) in out.iter_mut().enumerate() {
1115 *o += w * data[k * inter + c];
1116 }
1117 }
1118 Self::Blr(t) => t.add_col_scaled(c, w, out),
1119 Self::Mapped {
1120 dtype,
1121 row_scale,
1122 col_field,
1123 ..
1124 } => {
1125 let q = self.quant_bytes();
1126 let colf = if *dtype == TensorDtype::Q8_2f {
1127 col_field[c]
1128 } else {
1129 1.0
1130 };
1131 let wc = w * colf;
1132 for (k, o) in out.iter_mut().enumerate() {
1133 let b = q[k * inter + c] as i8 as f32;
1134 *o += wc * b * row_scale[k];
1135 }
1136 }
1137 }
1138 }
1139
1140 #[inline]
1149 pub fn prefetch_row(&self, r: usize) {
1150 let Self::Mapped { dtype, .. } = self else {
1151 return;
1152 };
1153 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1154 return;
1155 }
1156 let cols = self.cols();
1157 let q = self.quant_bytes();
1158 let (a, b) = (r * cols, (r + 1) * cols);
1159 if b > q.len() {
1160 return;
1161 }
1162 let mut j = a;
1163 while j < b {
1164 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1165 j += 512;
1166 }
1167 }
1168
1169 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1178 let cols = self.cols();
1179 debug_assert_eq!(out.len(), cols);
1180 match self {
1181 Self::F32 { data, .. } => {
1182 let row = &data[r * cols..(r + 1) * cols];
1183 for (o, v) in out.iter_mut().zip(row) {
1184 *o += w * v;
1185 }
1186 }
1187 Self::Blr(t) => {
1188 let tmp = &mut scratch[..cols];
1189 t.row_f32(r, tmp);
1190 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1191 *o += w * v;
1192 }
1193 }
1194 Self::Mapped {
1195 dtype,
1196 row_scale,
1197 col_field,
1198 ..
1199 } => match dtype {
1200 TensorDtype::Q8Row => {
1201 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1202 let ws = w * row_scale[r];
1203 let row: &[i8] =
1204 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1205 axpy_i8_f32(out, row, ws);
1206 }
1207 TensorDtype::Q8_2f => {
1208 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1209 let ws = w * row_scale[r];
1210 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1211 *o += ws * c * (*b as i8 as f32);
1212 }
1213 }
1214 _ => {
1215 self.row_f32(r, scratch);
1216 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1217 *o += w * v;
1218 }
1219 }
1220 },
1221 }
1222 }
1223
1224 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1228 let cols = self.cols();
1229 match self {
1230 Self::F32 { data, .. } => {
1231 let row = &data[r * cols..(r + 1) * cols];
1232 row.iter().zip(x).map(|(w, v)| w * v).sum()
1233 }
1234 Self::Blr(t) => {
1235 let tmp = &mut scratch[..cols];
1236 t.row_f32(r, tmp);
1237 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1238 }
1239 Self::Mapped {
1240 model,
1241 idx,
1242 dtype,
1243 row_scale,
1244 col_field,
1245 ..
1246 } => {
1247 let prism_forward =
1248 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1249 if prism_forward {
1250 let transformed = crate::prism::forward(model, &x[..cols]);
1251 let gpr = cols / GROUP_SIZE;
1252 match dtype {
1253 TensorDtype::Q2TiledP => {
1254 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1255 let mut sc = vec![0f32; gpr];
1256 v.scales_into(r, gpr, &mut sc);
1257 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1258 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1259 }
1260 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1261 }
1262 _ => {
1263 self.row_f32(r, scratch);
1264 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1265 }
1266 }
1267 }
1268 match dtype {
1269 TensorDtype::Q8Row => {
1270 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1271 dot_i8_f32(q, x) * row_scale[r]
1272 }
1273 TensorDtype::Q8_2f => {
1274 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1275 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1276 }
1277 _ => {
1278 self.row_f32(r, scratch);
1279 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1280 }
1281 }
1282 }
1283 }
1284 }
1285
1286 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1289 match self {
1290 Self::F32 { data, .. } => {
1294 if !crate::f32_backend::matvec(data, x, out) {
1295 matvec_rows(pool, data, x, out);
1296 }
1297 }
1298 Self::Blr(t) => {
1299 if crate::gpu::enabled_here()
1300 && crate::blr::gpu_wanted()
1301 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1302 {
1303 return;
1304 }
1305 t.matvec(x, out, pool)
1306 }
1307 Self::Mapped {
1308 model,
1309 idx,
1310 dtype,
1311 rows,
1312 cols,
1313 row_scale,
1314 col_field,
1315 vbit_offsets,
1316 repack,
1317 } => {
1318 let _ = (model, idx);
1319 assert!(
1329 out.len() >= *rows && x.len() >= *cols,
1330 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1331 out.len(),
1332 x.len(),
1333 );
1334 let prism_forward =
1335 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1336 if *dtype == TensorDtype::Q2TiledP
1337 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1338 {
1339 use std::sync::atomic::{AtomicUsize, Ordering};
1340 static N: AtomicUsize = AtomicUsize::new(0);
1341 let n = N.fetch_add(1, Ordering::Relaxed);
1342 if n < 128 {
1343 eprintln!(
1344 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1345 model.tensors[*idx].name,
1346 prism_forward,
1347 rows,
1348 cols,
1349 crate::gpu::enabled_here(),
1350 crate::gpu::q2tp_gpu_opt_in(),
1351 crate::gpu::cur_layer(),
1352 );
1353 }
1354 }
1355 if prism_forward {
1360 let transformed = crate::prism::forward(model, &x[..*cols]);
1361 match dtype {
1362 TensorDtype::Q4Block => {
1363 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1364 }
1365 TensorDtype::Q4Tiled => {
1366 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1367 }
1368 TensorDtype::Q4TiledP => {
1369 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1370 }
1371 TensorDtype::Q2TiledP => {
1372 let affine =
1373 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1374 if *rows * *cols >= 8_388_608
1375 && crate::gpu::enabled_here()
1376 && crate::gpu::q2tp_gpu_opt_in()
1377 {
1378 let gpu_ok = if affine {
1379 crate::gpu::q2tp_affine_matvec(
1380 model,
1381 *idx,
1382 &transformed,
1383 *rows,
1384 *cols,
1385 out,
1386 )
1387 } else {
1388 crate::gpu::q2tp_matvec(
1389 model,
1390 *idx,
1391 &transformed,
1392 *rows,
1393 *cols,
1394 out,
1395 )
1396 };
1397 if gpu_ok {
1398 return;
1399 }
1400 }
1401 if affine {
1402 q2tp_affine_matvec(
1403 self.quant_bytes(),
1404 &transformed,
1405 *rows,
1406 *cols,
1407 out,
1408 pool,
1409 )
1410 } else {
1411 q2tp_matvec(
1412 self.quant_bytes(),
1413 &transformed,
1414 *rows,
1415 *cols,
1416 out,
1417 pool,
1418 )
1419 }
1420 }
1421 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1422 self.bin_w().expect("binary dtype"),
1423 &transformed,
1424 *rows,
1425 *cols,
1426 out,
1427 pool,
1428 ),
1429 TensorDtype::Q1T => {
1430 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1431 }
1432 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1433 self.quant_bytes(),
1434 vbit_offsets,
1435 &transformed,
1436 *rows,
1437 *cols,
1438 out,
1439 pool,
1440 ),
1441 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1442 self.quant_bytes(),
1443 repack,
1444 row_scale,
1445 &transformed,
1446 col_field,
1447 *dtype,
1448 *rows,
1449 *cols,
1450 out,
1451 pool,
1452 ),
1453 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1454 }
1455 return;
1456 }
1457 if *dtype == TensorDtype::Q4Block {
1458 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1462 let t0 = std::time::Instant::now();
1463 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1464 crate::gpu::ProbeArm::Gpu => {
1465 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1466 crate::gpu::probe_record(
1467 crate::gpu::OpClass::Matvec,
1468 true,
1469 t0.elapsed(),
1470 );
1471 return;
1472 }
1473 }
1474 crate::gpu::ProbeArm::CpuTimed => {
1475 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1476 crate::gpu::probe_record(
1477 crate::gpu::OpClass::Matvec,
1478 false,
1479 t0.elapsed(),
1480 );
1481 return;
1482 }
1483 crate::gpu::ProbeArm::Cpu => {}
1484 }
1485 }
1486 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1487 return;
1488 }
1489 if *dtype == TensorDtype::Q4Tiled {
1490 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1495 let t0 = std::time::Instant::now();
1496 let cls = crate::gpu::matvec_class(*rows, *cols);
1497 match crate::gpu::probe_arm(cls) {
1498 crate::gpu::ProbeArm::Gpu => {
1499 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1500 crate::gpu::probe_record(cls, true, t0.elapsed());
1501 return;
1502 }
1503 }
1504 crate::gpu::ProbeArm::CpuTimed => {
1505 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1506 crate::gpu::probe_record(cls, false, t0.elapsed());
1507 return;
1508 }
1509 crate::gpu::ProbeArm::Cpu => {}
1510 }
1511 }
1512 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1513 return;
1514 }
1515 if *dtype == TensorDtype::Q4TiledP {
1516 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1522 let t0 = std::time::Instant::now();
1523 let cls = crate::gpu::matvec_class(*rows, *cols);
1524 match crate::gpu::probe_arm(cls) {
1525 crate::gpu::ProbeArm::Gpu => {
1526 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1527 crate::gpu::probe_record(cls, true, t0.elapsed());
1528 return;
1529 }
1530 }
1531 crate::gpu::ProbeArm::CpuTimed => {
1532 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1533 crate::gpu::probe_record(cls, false, t0.elapsed());
1534 return;
1535 }
1536 crate::gpu::ProbeArm::Cpu => {}
1537 }
1538 }
1539 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1540 return;
1541 }
1542 if *dtype == TensorDtype::Q2TiledP {
1543 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1544 return;
1545 }
1546 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1547 let w = self.bin_w().expect("binary dtype");
1552 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1553 let t0 = std::time::Instant::now();
1554 let arm = if crate::gpu::q1_force() {
1555 crate::gpu::ProbeArm::Gpu
1556 } else {
1557 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1558 };
1559 match arm {
1560 crate::gpu::ProbeArm::Gpu => {
1561 let ok = if *dtype == TensorDtype::Q1G {
1562 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1563 } else {
1564 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1565 };
1566 if ok {
1567 crate::gpu::probe_record(
1568 crate::gpu::OpClass::Matvec,
1569 true,
1570 t0.elapsed(),
1571 );
1572 return;
1573 }
1574 }
1575 crate::gpu::ProbeArm::CpuTimed => {
1576 bin_matvec(w, x, *rows, *cols, out, pool);
1577 crate::gpu::probe_record(
1578 crate::gpu::OpClass::Matvec,
1579 false,
1580 t0.elapsed(),
1581 );
1582 return;
1583 }
1584 crate::gpu::ProbeArm::Cpu => {}
1585 }
1586 }
1587 bin_matvec(w, x, *rows, *cols, out, pool);
1588 return;
1589 }
1590 if *dtype == TensorDtype::Q1T {
1591 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1595 let t0 = std::time::Instant::now();
1596 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1597 crate::gpu::ProbeArm::Gpu => {
1598 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1599 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1600 crate::gpu::probe_record(
1601 crate::gpu::OpClass::Matvec,
1602 true,
1603 t0.elapsed(),
1604 );
1605 return;
1606 }
1607 }
1608 crate::gpu::ProbeArm::CpuTimed => {
1609 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1610 crate::gpu::probe_record(
1611 crate::gpu::OpClass::Matvec,
1612 false,
1613 t0.elapsed(),
1614 );
1615 return;
1616 }
1617 crate::gpu::ProbeArm::Cpu => {}
1618 }
1619 }
1620 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1621 return;
1622 }
1623 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1624 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1625 return;
1626 }
1627 let xs = prescale(x, col_field, *dtype);
1628 if *rows >= crate::gpu::min_rows()
1633 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1634 && gpu_lmhead_enabled()
1635 && crate::gpu::enabled_here()
1636 {
1637 let t0 = std::time::Instant::now();
1640 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1641 crate::gpu::ProbeArm::Gpu => {}
1642 crate::gpu::ProbeArm::CpuTimed => {
1643 qmatvec(
1644 self.quant_bytes(),
1645 repack,
1646 row_scale,
1647 x,
1648 col_field,
1649 *dtype,
1650 *rows,
1651 *cols,
1652 out,
1653 pool,
1654 );
1655 crate::gpu::probe_record(
1656 crate::gpu::OpClass::Matvec,
1657 false,
1658 t0.elapsed(),
1659 );
1660 return;
1661 }
1662 crate::gpu::ProbeArm::Cpu => {
1663 qmatvec(
1664 self.quant_bytes(),
1665 repack,
1666 row_scale,
1667 x,
1668 col_field,
1669 *dtype,
1670 *rows,
1671 *cols,
1672 out,
1673 pool,
1674 );
1675 return;
1676 }
1677 }
1678 let frac = gpu_split_frac();
1679 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1680 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1681 let bytes = self.quant_bytes();
1682 let ok = std::thread::scope(|sc| {
1683 let g = sc.spawn(|| {
1684 crate::gpu::q8_matvec_range(
1685 model,
1686 *idx,
1687 cpu_rows,
1688 &row_scale[cpu_rows..],
1689 &xs,
1690 *rows - cpu_rows,
1691 *cols,
1692 out_gpu,
1693 )
1694 });
1695 if cpu_rows > 0 {
1696 let rep_cpu = if repack.is_empty() {
1699 &[][..]
1700 } else {
1701 &repack[..(cpu_rows / 4) * 4 * *cols]
1702 };
1703 qmatvec(
1704 &bytes[..cpu_rows * *cols],
1705 rep_cpu,
1706 &row_scale[..cpu_rows],
1707 x,
1708 col_field,
1709 *dtype,
1710 cpu_rows,
1711 *cols,
1712 out_cpu,
1713 pool,
1714 );
1715 }
1716 g.join().unwrap_or(false)
1717 });
1718 if ok {
1719 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1720 return;
1721 }
1722 qmatvec(
1725 &bytes[cpu_rows * *cols..(*rows) * *cols],
1726 &[],
1727 &row_scale[cpu_rows..],
1728 x,
1729 col_field,
1730 *dtype,
1731 *rows - cpu_rows,
1732 *cols,
1733 out_gpu,
1734 pool,
1735 );
1736 return;
1737 }
1738 qmatvec(
1739 self.quant_bytes(),
1740 repack,
1741 row_scale,
1742 x,
1743 col_field,
1744 *dtype,
1745 *rows,
1746 *cols,
1747 out,
1748 pool,
1749 );
1750 }
1751 }
1752 }
1753
1754 pub fn matvec2(
1756 &self,
1757 x1: &[f32],
1758 x2: &[f32],
1759 o1: &mut [f32],
1760 o2: &mut [f32],
1761 pool: Option<&Pool>,
1762 ) {
1763 match self {
1764 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1765 Self::Blr(_) => {
1766 self.matvec(x1, o1, pool);
1767 self.matvec(x2, o2, pool);
1768 }
1769 Self::Mapped {
1770 model,
1771 idx,
1772 dtype,
1773 rows,
1774 cols,
1775 row_scale,
1776 col_field,
1777 vbit_offsets,
1778 ..
1779 } => {
1780 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1781 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1782 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1783 match dtype {
1784 TensorDtype::Q4Block => {
1785 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1786 }
1787 TensorDtype::Q4Tiled => {
1788 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1789 }
1790 TensorDtype::Q4TiledP => {
1791 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1792 }
1793 TensorDtype::Q2TiledP => {
1794 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1795 q2tp_affine_matvec2(
1796 self.quant_bytes(),
1797 &tx1,
1798 &tx2,
1799 *rows,
1800 *cols,
1801 o1,
1802 o2,
1803 pool,
1804 )
1805 } else {
1806 q2tp_matvec2(
1807 self.quant_bytes(),
1808 &tx1,
1809 &tx2,
1810 *rows,
1811 *cols,
1812 o1,
1813 o2,
1814 pool,
1815 )
1816 }
1817 }
1818 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1819 self.bin_w().expect("binary dtype"),
1820 &tx1,
1821 &tx2,
1822 *rows,
1823 *cols,
1824 o1,
1825 o2,
1826 pool,
1827 ),
1828 TensorDtype::Q1T => {
1829 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1830 }
1831 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1832 self.quant_bytes(),
1833 vbit_offsets,
1834 &tx1,
1835 &tx2,
1836 *rows,
1837 *cols,
1838 o1,
1839 o2,
1840 pool,
1841 ),
1842 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1843 self.quant_bytes(),
1844 row_scale,
1845 &tx1,
1846 &tx2,
1847 col_field,
1848 *dtype,
1849 *rows,
1850 *cols,
1851 o1,
1852 o2,
1853 pool,
1854 ),
1855 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1856 }
1857 return;
1858 }
1859 if *dtype == TensorDtype::Q4Block {
1860 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1861 return;
1862 }
1863 if *dtype == TensorDtype::Q4Tiled {
1864 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1865 return;
1866 }
1867 if *dtype == TensorDtype::Q4TiledP {
1868 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1869 return;
1870 }
1871 if *dtype == TensorDtype::Q2TiledP {
1872 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1873 return;
1874 }
1875 if let Some(w) = self.bin_w() {
1876 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1877 return;
1878 }
1879 if *dtype == TensorDtype::Q1T {
1880 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1886 return;
1887 }
1888 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1889 vbitmatvec2(
1890 self.quant_bytes(),
1891 vbit_offsets,
1892 x1,
1893 x2,
1894 *rows,
1895 *cols,
1896 o1,
1897 o2,
1898 pool,
1899 );
1900 return;
1901 }
1902 qmatvec2(
1903 self.quant_bytes(),
1904 row_scale,
1905 x1,
1906 x2,
1907 col_field,
1908 *dtype,
1909 *rows,
1910 *cols,
1911 o1,
1912 o2,
1913 pool,
1914 );
1915 }
1916 }
1917 }
1918}
1919
1920impl QTensor {
1921 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1929 if self.has_prism_contract() {
1930 return None;
1931 }
1932 match self {
1933 Self::Mapped {
1934 model, idx, dtype, ..
1935 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1936 _ => None,
1937 }
1938 }
1939
1940 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1941 let cols = self.cols();
1942 let rows = self.rows();
1943 debug_assert_eq!(xs_all.len(), b * cols);
1944 debug_assert_eq!(out.len(), b * rows);
1945 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1946 if crate::gptq_capture::capturing() {
1950 if let Self::Mapped { model, idx, .. } = self {
1951 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1952 }
1953 if let Self::Blr(t) = self {
1954 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
1955 }
1956 }
1957 match self {
1958 Self::Blr(t) => {
1959 if crate::gpu::enabled_here()
1960 && crate::blr::gpu_wanted()
1961 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
1962 {
1963 return;
1964 }
1965 t.matmat(xs_all, b, out, pool)
1966 }
1967 Self::F32 { data, .. } => {
1968 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1969 return;
1970 }
1971 let out_addr = SendMut(out.as_mut_ptr());
1972 let run = |start: usize, end: usize| {
1979 let mut bi = start;
1980 while bi < end {
1981 let n = (end - bi).min(4);
1982 for o in 0..rows {
1983 let row = &data[o * cols..(o + 1) * cols];
1984 if n == 4 {
1985 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1986 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
1987 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
1988 for j in 0..cols {
1989 let w = row[j];
1990 a0 += w * x0[j];
1991 a1 += w * x1[j];
1992 a2 += w * x2[j];
1993 a3 += w * x3[j];
1994 }
1995 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
1996 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1997 }
1998 } else {
1999 for k in 0..n {
2000 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2001 let mut acc = 0f32;
2002 for j in 0..cols {
2003 acc += row[j] * x[j];
2004 }
2005 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2006 }
2007 }
2008 }
2009 bi += n;
2010 }
2011 };
2012 dispatch_rows(pool, b, &run);
2013 }
2014 Self::Mapped {
2015 model,
2016 idx,
2017 dtype,
2018 row_scale,
2019 col_field,
2020 vbit_offsets,
2021 ..
2022 } => {
2023 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2024 let mut transformed = Vec::with_capacity(xs_all.len());
2025 for bi in 0..b {
2026 transformed.extend_from_slice(&crate::prism::forward(
2027 model,
2028 &xs_all[bi * cols..(bi + 1) * cols],
2029 ));
2030 }
2031 match dtype {
2032 TensorDtype::Q4Block => {
2033 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2034 }
2035 TensorDtype::Q4Tiled => {
2036 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2037 }
2038 TensorDtype::Q4TiledP => {
2039 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2040 }
2041 TensorDtype::Q2TiledP => {
2042 let affine =
2043 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2044 let gpu_batch_ok = if affine {
2050 b >= 2
2051 } else {
2052 b >= 32 && b * rows * cols >= 128_000_000
2053 };
2054 if gpu_batch_ok
2055 && cols % 32 == 0
2056 && crate::gpu::enabled_here()
2057 && crate::gpu::q2tp_gpu_opt_in()
2058 {
2059 let gpu_ok = if affine {
2060 crate::gpu::q2tp_affine_matmat(
2061 model,
2062 *idx,
2063 &transformed,
2064 b,
2065 rows,
2066 cols,
2067 out,
2068 )
2069 } else {
2070 crate::gpu::q2tp_matmat(
2071 model,
2072 *idx,
2073 &transformed,
2074 b,
2075 rows,
2076 cols,
2077 out,
2078 )
2079 };
2080 if gpu_ok {
2081 return;
2082 }
2083 }
2084 if affine
2088 && b == 1
2089 && cols % 32 == 0
2090 && crate::gpu::enabled_here()
2091 && crate::gpu::q2tp_gpu_opt_in()
2092 && crate::gpu::q2tp_affine_matvec(
2093 model,
2094 *idx,
2095 &transformed[..cols],
2096 rows,
2097 cols,
2098 &mut out[..rows],
2099 )
2100 {
2101 return;
2102 }
2103 if affine {
2104 q2tp_affine_matmat(
2105 self.quant_bytes(),
2106 &transformed,
2107 b,
2108 rows,
2109 cols,
2110 out,
2111 pool,
2112 )
2113 } else {
2114 q2tp_matmat(
2115 self.quant_bytes(),
2116 &transformed,
2117 b,
2118 rows,
2119 cols,
2120 out,
2121 pool,
2122 )
2123 }
2124 }
2125 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2126 self.bin_w().expect("binary dtype"),
2127 &transformed,
2128 b,
2129 rows,
2130 cols,
2131 out,
2132 pool,
2133 ),
2134 TensorDtype::Q1T => {
2135 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2136 }
2137 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2138 self.quant_bytes(),
2139 vbit_offsets,
2140 &transformed,
2141 b,
2142 rows,
2143 cols,
2144 out,
2145 pool,
2146 ),
2147 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2148 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2149 .map(|bi| {
2150 prescale(
2151 &transformed[bi * cols..(bi + 1) * cols],
2152 col_field,
2153 *dtype,
2154 )
2155 })
2156 .collect();
2157 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2158 }
2159 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2160 }
2161 return;
2162 }
2163 if *dtype == TensorDtype::Q4Block {
2164 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2165 return;
2166 }
2167 if *dtype == TensorDtype::Q4TiledP {
2168 if b >= 32
2182 && b * rows * cols >= 128_000_000
2183 && cols % 32 == 0
2184 && !row_exact()
2185 && !crate::gpu::mm_killed()
2186 && crate::gpu::enabled_here()
2187 {
2188 let class = if b >= 128 {
2189 crate::gpu::OpClass::MatmatWide
2190 } else {
2191 crate::gpu::OpClass::Matmat
2192 };
2193 if let Self::Mapped { model, idx, .. } = self {
2194 if crate::mm_ab::on() {
2206 let mut g = vec![0f32; b * rows];
2207 let t = std::time::Instant::now();
2208 let took = crate::gpu::q4tp_matmat(
2209 model, *idx, xs_all, b, rows, cols, &mut g,
2210 );
2211 let dg = t.elapsed();
2212 let t = std::time::Instant::now();
2213 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2214 let dc = t.elapsed();
2215 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2216 return;
2217 }
2218 let t0 = std::time::Instant::now();
2219 let resident = crate::gpu::weight_is_resident(model, *idx);
2223 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2224 crate::gpu::ProbeArm::Gpu => {
2225 if crate::gpu::q4tp_matmat(
2226 model, *idx, xs_all, b, rows, cols, out,
2227 ) {
2228 let el = t0.elapsed();
2229 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2239 let budget = std::time::Duration::from_secs_f64(
2240 flops / 1.5e12 * 8.0 + 0.020,
2241 );
2242 crate::gpu::mm_budget_check(
2243 "q4tp matmat",
2244 el,
2245 budget,
2246 crate::gpu::probe_was_cold() || !resident,
2247 );
2248 crate::gpu::probe_record(class, true, el);
2249 return;
2250 }
2251 }
2252 crate::gpu::ProbeArm::CpuTimed => {
2253 q4tp_matmat(
2254 self.quant_bytes(),
2255 xs_all,
2256 b,
2257 rows,
2258 cols,
2259 out,
2260 pool,
2261 );
2262 crate::gpu::probe_record(class, false, t0.elapsed());
2263 return;
2264 }
2265 crate::gpu::ProbeArm::Cpu => {}
2266 }
2267 }
2268 }
2269 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2270 return;
2271 }
2272 if *dtype == TensorDtype::Q2TiledP {
2273 if b >= 32
2279 && b * rows * cols >= 128_000_000
2280 && cols % 32 == 0
2281 && !crate::gpu::mm_killed()
2282 && crate::gpu::enabled_here()
2283 {
2284 let class = if b >= 128 {
2285 crate::gpu::OpClass::MatmatWide
2286 } else {
2287 crate::gpu::OpClass::Matmat
2288 };
2289 if let Self::Mapped { model, idx, .. } = self {
2290 let t0 = std::time::Instant::now();
2291 match crate::gpu::probe_arm(class) {
2292 crate::gpu::ProbeArm::Gpu => {
2293 if crate::gpu::q2tp_matmat(
2294 model, *idx, xs_all, b, rows, cols, out,
2295 ) {
2296 crate::gpu::probe_record(class, true, t0.elapsed());
2297 return;
2298 }
2299 }
2300 crate::gpu::ProbeArm::CpuTimed => {
2301 q2tp_matmat(
2302 self.quant_bytes(),
2303 xs_all,
2304 b,
2305 rows,
2306 cols,
2307 out,
2308 pool,
2309 );
2310 crate::gpu::probe_record(class, false, t0.elapsed());
2311 return;
2312 }
2313 crate::gpu::ProbeArm::Cpu => {}
2314 }
2315 }
2316 }
2317 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2322 return;
2323 }
2324 if *dtype == TensorDtype::Q4Tiled {
2325 if b >= 32
2337 && b * rows * cols >= 128_000_000
2338 && cols % 32 == 0
2339 && !crate::gpu::mm_killed()
2340 && crate::gpu::enabled_here()
2341 {
2342 let class = if b >= 128 {
2343 crate::gpu::OpClass::MatmatWide
2344 } else {
2345 crate::gpu::OpClass::Matmat
2346 };
2347 if let Self::Mapped { model, idx, .. } = self {
2348 let t0 = std::time::Instant::now();
2349 match crate::gpu::probe_arm(class) {
2350 crate::gpu::ProbeArm::Gpu => {
2351 if crate::gpu::q4t_matmat(
2352 model, *idx, xs_all, b, rows, cols, out,
2353 ) {
2354 let el = t0.elapsed();
2355 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2365 let budget = std::time::Duration::from_secs_f64(
2366 flops / 1.5e12 * 8.0 + 0.020,
2367 );
2368 crate::gpu::mm_budget_check(
2369 "q4t matmat",
2370 el,
2371 budget,
2372 crate::gpu::probe_was_cold(),
2373 );
2374 crate::gpu::probe_record(class, true, el);
2375 return;
2376 }
2377 }
2378 crate::gpu::ProbeArm::CpuTimed => {
2379 q4t_matmat(
2380 self.quant_bytes(),
2381 xs_all,
2382 b,
2383 rows,
2384 cols,
2385 out,
2386 pool,
2387 );
2388 crate::gpu::probe_record(class, false, t0.elapsed());
2389 return;
2390 }
2391 crate::gpu::ProbeArm::Cpu => {}
2392 }
2393 }
2394 }
2395 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2396 return;
2397 }
2398 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2399 let w = self.bin_w().expect("binary dtype");
2403 if b >= 32
2404 && b * rows * cols >= 128_000_000
2405 && cols % 64 == 0
2406 && crate::gpu::enabled_here()
2407 {
2408 if let Self::Mapped { model, idx, .. } = self {
2409 let t0 = std::time::Instant::now();
2410 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2411 crate::gpu::ProbeArm::Gpu => {
2412 let ok = if *dtype == TensorDtype::Q1G {
2413 crate::gpu::q1g_matmat(
2414 model, *idx, xs_all, b, rows, cols, out,
2415 )
2416 } else {
2417 crate::gpu::q1_matmat(
2418 model, *idx, xs_all, b, rows, cols, out,
2419 )
2420 };
2421 if ok {
2422 crate::gpu::probe_record(
2423 crate::gpu::OpClass::Matmat,
2424 true,
2425 t0.elapsed(),
2426 );
2427 return;
2428 }
2429 }
2430 crate::gpu::ProbeArm::CpuTimed => {
2431 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2432 crate::gpu::probe_record(
2433 crate::gpu::OpClass::Matmat,
2434 false,
2435 t0.elapsed(),
2436 );
2437 return;
2438 }
2439 crate::gpu::ProbeArm::Cpu => {}
2440 }
2441 }
2442 }
2443 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2444 return;
2445 }
2446 if *dtype == TensorDtype::Q1T {
2447 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2450 if let Self::Mapped { model, idx, .. } = self {
2451 let t0 = std::time::Instant::now();
2452 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2453 crate::gpu::ProbeArm::Gpu => {
2454 if crate::gpu::q1t_matmat(
2455 model, *idx, xs_all, b, rows, cols, out,
2456 ) {
2457 crate::gpu::probe_record(
2458 crate::gpu::OpClass::Matmat,
2459 true,
2460 t0.elapsed(),
2461 );
2462 return;
2463 }
2464 }
2465 crate::gpu::ProbeArm::CpuTimed => {
2466 q1t_matmat(
2467 self.quant_bytes(),
2468 xs_all,
2469 b,
2470 rows,
2471 cols,
2472 out,
2473 pool,
2474 );
2475 crate::gpu::probe_record(
2476 crate::gpu::OpClass::Matmat,
2477 false,
2478 t0.elapsed(),
2479 );
2480 return;
2481 }
2482 crate::gpu::ProbeArm::Cpu => {}
2483 }
2484 }
2485 }
2486 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2487 return;
2488 }
2489 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2490 vbitmatmat(
2491 self.quant_bytes(),
2492 vbit_offsets,
2493 xs_all,
2494 b,
2495 rows,
2496 cols,
2497 out,
2498 pool,
2499 );
2500 return;
2501 }
2502 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2503 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2504 .collect();
2505 if row_exact()
2511 && (1..=4).contains(&b)
2512 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2513 && crate::gpu::enabled_here()
2514 && crate::gpu::wgpu_active()
2515 {
2516 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2517 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2518 return;
2519 }
2520 }
2521 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2527 if let Self::Mapped { model, idx, .. } = self {
2528 let t0 = std::time::Instant::now();
2529 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2530 crate::gpu::ProbeArm::Gpu
2531 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2532 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2533 {
2534 let q = self.quant_bytes();
2538 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2539 return;
2540 }
2541 crate::gpu::ProbeArm::Gpu => {
2542 if *dtype == TensorDtype::Q8_2f
2550 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2551 && crate::gpu::q8_matmat_2f(
2552 model, *idx, row_scale, col_field, xs_all, b, rows,
2553 cols, out,
2554 )
2555 {
2556 crate::gpu::probe_record(
2557 crate::gpu::OpClass::Matmat,
2558 true,
2559 t0.elapsed(),
2560 );
2561 return;
2562 }
2563 let flat: Vec<f32> =
2564 pre.iter().flat_map(|v| v.iter().copied()).collect();
2565 if crate::gpu::q8_matmat(
2566 model, *idx, row_scale, &flat, b, rows, cols, out,
2567 ) {
2568 crate::gpu::probe_record(
2569 crate::gpu::OpClass::Matmat,
2570 true,
2571 t0.elapsed(),
2572 );
2573 return;
2574 }
2575 }
2576 crate::gpu::ProbeArm::CpuTimed => {
2577 let q = self.quant_bytes();
2578 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2579 crate::gpu::probe_record(
2580 crate::gpu::OpClass::Matmat,
2581 false,
2582 t0.elapsed(),
2583 );
2584 return;
2585 }
2586 crate::gpu::ProbeArm::Cpu => {}
2587 }
2588 }
2589 }
2590 let q = self.quant_bytes();
2591 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2592 }
2593 }
2594 }
2595}
2596
2597impl QTensor {
2598 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2608 let (rows, cols) = (self.rows(), self.cols());
2609 let Self::Mapped {
2610 model,
2611 idx,
2612 dtype,
2613 row_scale,
2614 col_field,
2615 ..
2616 } = self
2617 else {
2618 return false;
2619 };
2620 if crate::prism::has_contract(model) {
2621 return false;
2622 }
2623 match *dtype {
2624 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2625 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2629 if *dtype == TensorDtype::Q8_2f
2632 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2633 && crate::gpu::q8_matmat_2f(
2634 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2635 )
2636 {
2637 return true;
2638 }
2639 let flat: Vec<f32> = (0..b)
2640 .flat_map(|bi| {
2641 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2642 })
2643 .collect();
2644 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2645 }
2646 _ => false,
2647 }
2648 }
2649
2650 pub fn matvec_many<const N: usize>(
2657 ts: [&QTensor; N],
2658 x: &[f32],
2659 mut outs: [&mut [f32]; N],
2660 pool: Option<&Pool>,
2661 ) {
2662 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2663 if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
2664 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2669 t.matvec(x, o, pool);
2670 }
2671 return;
2672 }
2673 let uniform_q8 = ts.iter().all(|t| {
2674 matches!(
2675 t,
2676 Self::Mapped {
2677 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2678 ..
2679 }
2680 )
2681 });
2682 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2683 if uniform_f32 && crate::f32_backend::active() {
2684 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2685 t.matvec(x, o, pool);
2686 }
2687 return;
2688 }
2689 let uniform_q4 = ts.iter().all(|t| {
2690 matches!(
2691 t,
2692 Self::Mapped {
2693 dtype: TensorDtype::Q4Block,
2694 ..
2695 }
2696 )
2697 });
2698 let uniform_vbit = ts.iter().all(|t| {
2699 matches!(
2700 t,
2701 Self::Mapped {
2702 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2703 ..
2704 }
2705 )
2706 });
2707 let uniform_q1 =
2710 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2711 let uniform_q1t = ts.iter().all(|t| {
2712 matches!(
2713 t,
2714 Self::Mapped {
2715 dtype: TensorDtype::Q1T,
2716 ..
2717 }
2718 )
2719 });
2720 let uniform_q4tp = ts.iter().all(|t| {
2725 matches!(
2726 t,
2727 Self::Mapped {
2728 dtype: TensorDtype::Q4TiledP,
2729 ..
2730 }
2731 )
2732 }) && ts
2733 .iter()
2734 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2735 let Some(pool) = pool else {
2736 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2737 t.matvec(x, o, None);
2738 }
2739 return;
2740 };
2741 if total_rows < 256
2742 || !(uniform_q8
2743 || uniform_f32
2744 || uniform_q4
2745 || uniform_vbit
2746 || uniform_q1
2747 || uniform_q1t
2748 || uniform_q4tp)
2749 {
2750 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2751 t.matvec(x, o, Some(pool));
2752 }
2753 return;
2754 }
2755
2756 if uniform_q4tp {
2757 let cols = ts[0].cols();
2763 let gpr = cols / GROUP_SIZE;
2764 let views: [Q4tpView; N] =
2765 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2766 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2767 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2768 let locate = |flat: usize| -> (usize, usize) {
2770 let mut acc = 0;
2771 for (i, &r) in rows_of.iter().enumerate() {
2772 if flat < acc + r {
2773 return (i, flat - acc);
2774 }
2775 acc += r;
2776 }
2777 (rows_of.len() - 1, 0)
2778 };
2779 let (views, outs_addr) = (&views, &outs_addr);
2780 if a8w8_enabled() {
2781 let act = split_act(x);
2782 let act = &act;
2783 let run = |start: usize, end: usize| {
2784 with_krow(gpr, |sc| {
2785 for flat in start..end {
2786 let (t, r) = locate(flat);
2787 let v = &views[t];
2788 v.scales_into(r, gpr, sc);
2789 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2790 for &(j, xv) in &act.outliers {
2791 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2792 acc += w * s * xv;
2793 }
2794 unsafe { *outs_addr[t].at(r) = acc };
2796 }
2797 });
2798 };
2799 pool.run_rows(total_rows, &run);
2800 } else {
2801 let run = |start: usize, end: usize| {
2802 with_krow(gpr, |sc| {
2803 for flat in start..end {
2804 let (t, r) = locate(flat);
2805 let v = &views[t];
2806 v.scales_into(r, gpr, sc);
2807 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2809 }
2810 });
2811 };
2812 pool.run_rows(total_rows, &run);
2813 }
2814 return;
2815 }
2816
2817 if uniform_q1 {
2818 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2821 if a8w8_enabled() {
2822 let act = split_act(x);
2823 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2824 let (act, gsum) = (&act, &gsum);
2825 let closures: [_; N] = std::array::from_fn(|i| {
2826 let (w, gpr, out) = (
2827 ts[i].bin_w().expect("uniform_q1 checked"),
2828 ts[i].cols() / GROUP_SIZE,
2829 outs_addr[i],
2830 );
2831 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2832 });
2833 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2834 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2835 pool.run_many(&parts);
2836 } else {
2837 let closures: [_; N] = std::array::from_fn(|i| {
2838 let (w, gpr, out) = (
2839 ts[i].bin_w().expect("uniform_q1 checked"),
2840 ts[i].cols() / GROUP_SIZE,
2841 outs_addr[i],
2842 );
2843 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
2844 });
2845 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2846 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2847 pool.run_many(&parts);
2848 }
2849 return;
2850 }
2851
2852 if uniform_q1t {
2853 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2857 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2858 if a8w8_enabled() {
2859 let act = split_act(x);
2860 let act = &act;
2861 let x_ref = x;
2862 let closures: [_; N] = std::array::from_fn(|i| {
2863 let bytes = ts[i].quant_bytes();
2864 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2865 let gpr = cols / GROUP_SIZE;
2866 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2867 let out = outs_addr[i];
2868 move |s: usize, e: usize| {
2869 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2870 }
2871 });
2872 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2873 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2874 pool.run_many(&parts);
2875 } else {
2876 let x_ref = x;
2877 let closures: [_; N] = std::array::from_fn(|i| {
2878 let bytes = ts[i].quant_bytes();
2879 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2880 let gpr = cols / GROUP_SIZE;
2881 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2882 let out = outs_addr[i];
2883 move |s: usize, e: usize| {
2884 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2885 }
2886 });
2887 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2888 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2889 pool.run_many(&parts);
2890 }
2891 return;
2892 }
2893
2894 if uniform_q4 || uniform_vbit {
2895 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2896 if a8w8_enabled() {
2898 let act = split_act(x);
2899 let act = &act;
2900 if uniform_q4 {
2901 let closures: [_; N] = std::array::from_fn(|i| {
2902 let (packed, scales) =
2903 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2904 let (gpr, cols, out) =
2905 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2906 move |s: usize, e: usize| {
2907 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2908 }
2909 });
2910 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2911 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2912 pool.run_many(&parts);
2913 } else {
2914 let closures: [_; N] = std::array::from_fn(|i| {
2915 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2916 unreachable!()
2917 };
2918 let (bytes, rows, cols, out) = (
2919 ts[i].quant_bytes(),
2920 ts[i].rows(),
2921 ts[i].cols(),
2922 outs_addr[i],
2923 );
2924 move |s: usize, e: usize| {
2925 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2926 }
2927 });
2928 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2929 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2930 pool.run_many(&parts);
2931 }
2932 return;
2933 }
2934 if uniform_q4 {
2935 let closures: [_; N] = std::array::from_fn(|i| {
2936 let (packed, scales) =
2937 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2938 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2939 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2940 });
2941 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2942 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2943 pool.run_many(&parts);
2944 } else {
2945 let closures: [_; N] = std::array::from_fn(|i| {
2946 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2947 unreachable!()
2948 };
2949 let (bytes, rows, cols, out) = (
2950 ts[i].quant_bytes(),
2951 ts[i].rows(),
2952 ts[i].cols(),
2953 outs_addr[i],
2954 );
2955 move |s: usize, e: usize| {
2956 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2957 }
2958 });
2959 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2960 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2961 pool.run_many(&parts);
2962 }
2963 return;
2964 }
2965
2966 if uniform_f32 {
2967 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2968 let closures: [_; N] = std::array::from_fn(|i| {
2969 let Self::F32 { data, cols, .. } = ts[i] else {
2970 unreachable!()
2971 };
2972 let out = outs_addr[i];
2973 move |start: usize, end: usize| {
2974 for o in start..end {
2975 let row = &data[o * cols..(o + 1) * cols];
2976 let mut sum = 0.0f32;
2977 for j in 0..*cols {
2978 sum += row[j] * x[j];
2979 }
2980 unsafe { *out.at(o) = sum };
2982 }
2983 }
2984 });
2985 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2986 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2987 pool.run_many(&parts);
2988 return;
2989 }
2990
2991 struct Ctx<'a> {
2994 bytes: &'a [u8],
2995 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2996 rep: &'a [u8],
2997 row_scale: &'a [f32],
2998 cols: usize,
2999 xs: std::borrow::Cow<'a, [f32]>,
3000 }
3001 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3002 let Self::Mapped {
3003 dtype,
3004 cols,
3005 row_scale,
3006 col_field,
3007 repack,
3008 ..
3009 } = ts[i]
3010 else {
3011 unreachable!()
3012 };
3013 Ctx {
3014 bytes: ts[i].quant_bytes(),
3015 rep: repack,
3016 row_scale,
3017 cols: *cols,
3018 xs: prescale(x, col_field, *dtype),
3019 }
3020 });
3021 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3022 #[cfg(target_arch = "aarch64")]
3023 if sdot_enabled() {
3024 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3025 let closures: [_; N] = std::array::from_fn(|i| {
3026 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3027 move |start: usize, end: usize| {
3028 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3029 }
3030 });
3031 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3032 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3033 pool.run_many(&parts);
3034 return;
3035 }
3036 #[cfg(target_arch = "x86_64")]
3037 if avx2_a8w8_enabled() {
3038 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3039 let closures: [_; N] = std::array::from_fn(|i| {
3040 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3041 move |start: usize, end: usize| {
3042 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3043 }
3044 });
3045 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3046 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3047 pool.run_many(&parts);
3048 return;
3049 }
3050 let closures: [_; N] = std::array::from_fn(|i| {
3051 let (c, out) = (&ctxs[i], outs_addr[i]);
3052 move |start: usize, end: usize| {
3053 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3054 }
3055 });
3056 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3057 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3058 pool.run_many(&parts);
3059 }
3060}
3061
3062impl QTensor {
3063 #[allow(clippy::needless_range_loop)]
3068 pub fn matvec2_many<const N: usize>(
3069 ts: [&QTensor; N],
3070 x1: &[f32],
3071 x2: &[f32],
3072 mut o1s: [&mut [f32]; N],
3073 mut o2s: [&mut [f32]; N],
3074 pool: Option<&Pool>,
3075 ) {
3076 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3077 if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
3078 for i in 0..N {
3079 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3080 }
3081 return;
3082 }
3083 let uniform_q8 = ts.iter().all(|t| {
3084 matches!(
3085 t,
3086 Self::Mapped {
3087 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3088 ..
3089 }
3090 )
3091 });
3092 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3093 let uniform_q4 = ts.iter().all(|t| {
3094 matches!(
3095 t,
3096 Self::Mapped {
3097 dtype: TensorDtype::Q4Block,
3098 ..
3099 }
3100 )
3101 });
3102 let uniform_vbit = ts.iter().all(|t| {
3103 matches!(
3104 t,
3105 Self::Mapped {
3106 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3107 ..
3108 }
3109 )
3110 });
3111 let fusable = pool.is_some()
3112 && total_rows >= 256
3113 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3114 if !fusable {
3115 for i in 0..N {
3116 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3117 }
3118 return;
3119 }
3120 let pool = pool.unwrap();
3121
3122 if uniform_q4 || uniform_vbit {
3123 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3124 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3125 if a8w8_enabled() {
3127 let a1 = split_act(x1);
3128 let a2 = split_act(x2);
3129 let (a1, a2) = (&a1, &a2);
3130 if uniform_q4 {
3131 let closures: [_; N] = std::array::from_fn(|i| {
3132 let (packed, scales) =
3133 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3134 let (gpr, cols, o1, o2) =
3135 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3136 move |s: usize, e: usize| {
3137 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3138 }
3139 });
3140 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3141 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3142 pool.run_many(&parts);
3143 } else {
3144 let closures: [_; N] = std::array::from_fn(|i| {
3145 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3146 unreachable!()
3147 };
3148 let (bytes, rows, cols, o1, o2) = (
3149 ts[i].quant_bytes(),
3150 ts[i].rows(),
3151 ts[i].cols(),
3152 p1[i],
3153 p2[i],
3154 );
3155 move |s: usize, e: usize| {
3156 vbit_range2_a8w8(
3157 bytes,
3158 vbit_offsets,
3159 x1,
3160 x2,
3161 a1,
3162 a2,
3163 rows,
3164 cols,
3165 o1,
3166 o2,
3167 s,
3168 e,
3169 )
3170 }
3171 });
3172 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3173 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3174 pool.run_many(&parts);
3175 }
3176 return;
3177 }
3178 if uniform_q4 {
3179 let closures: [_; N] = std::array::from_fn(|i| {
3180 let (packed, scales) =
3181 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3182 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3183 move |s: usize, e: usize| {
3184 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3185 }
3186 });
3187 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3188 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3189 pool.run_many(&parts);
3190 } else {
3191 let closures: [_; N] = std::array::from_fn(|i| {
3192 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3193 unreachable!()
3194 };
3195 let (bytes, rows, cols, o1, o2) = (
3196 ts[i].quant_bytes(),
3197 ts[i].rows(),
3198 ts[i].cols(),
3199 p1[i],
3200 p2[i],
3201 );
3202 move |s: usize, e: usize| {
3203 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3204 }
3205 });
3206 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3207 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3208 pool.run_many(&parts);
3209 }
3210 return;
3211 }
3212
3213 if uniform_f32 {
3214 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3215 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3216 let closures: [_; N] = std::array::from_fn(|i| {
3217 let Self::F32 { data, cols, .. } = ts[i] else {
3218 unreachable!()
3219 };
3220 let (o1, o2) = (p1[i], p2[i]);
3221 move |start: usize, end: usize| {
3222 for o in start..end {
3223 let row = &data[o * cols..(o + 1) * cols];
3224 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3225 for j in 0..*cols {
3226 s1 += row[j] * x1[j];
3227 s2 += row[j] * x2[j];
3228 }
3229 unsafe {
3231 *o1.at(o) = s1;
3232 *o2.at(o) = s2;
3233 }
3234 }
3235 }
3236 });
3237 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3238 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3239 pool.run_many(&parts);
3240 return;
3241 }
3242
3243 struct Ctx<'a> {
3244 bytes: &'a [u8],
3245 row_scale: &'a [f32],
3246 cols: usize,
3247 xs1: std::borrow::Cow<'a, [f32]>,
3248 xs2: std::borrow::Cow<'a, [f32]>,
3249 }
3250 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3251 let Self::Mapped {
3252 dtype,
3253 cols,
3254 row_scale,
3255 col_field,
3256 ..
3257 } = ts[i]
3258 else {
3259 unreachable!()
3260 };
3261 Ctx {
3262 bytes: ts[i].quant_bytes(),
3263 row_scale,
3264 cols: *cols,
3265 xs1: prescale(x1, col_field, *dtype),
3266 xs2: prescale(x2, col_field, *dtype),
3267 }
3268 });
3269 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3270 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3271 #[cfg(target_arch = "aarch64")]
3272 if sdot_enabled() {
3273 let acts: [(SplitAct, SplitAct); N] =
3274 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3275 let closures: [_; N] = std::array::from_fn(|i| {
3276 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3277 move |start: usize, end: usize| {
3278 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3279 }
3280 });
3281 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3282 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3283 pool.run_many(&parts);
3284 return;
3285 }
3286 #[cfg(target_arch = "x86_64")]
3287 if avx2_a8w8_enabled() {
3288 let acts: [(SplitAct, SplitAct); N] =
3289 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3290 let closures: [_; N] = std::array::from_fn(|i| {
3291 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3292 move |start: usize, end: usize| {
3293 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3294 }
3295 });
3296 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3297 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3298 pool.run_many(&parts);
3299 return;
3300 }
3301 let closures: [_; N] = std::array::from_fn(|i| {
3302 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3303 move |start: usize, end: usize| {
3304 q8_range2_f32(
3305 c.bytes,
3306 c.row_scale,
3307 &c.xs1,
3308 &c.xs2,
3309 c.cols,
3310 o1,
3311 o2,
3312 start,
3313 end,
3314 )
3315 }
3316 });
3317 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3318 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3319 pool.run_many(&parts);
3320 }
3321
3322 pub fn matvec_silu_mul(
3327 gate: &QTensor,
3328 up: &QTensor,
3329 x: &[f32],
3330 out: &mut [f32],
3331 pool: Option<&Pool>,
3332 ) -> bool {
3333 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3334 }
3335
3336 pub fn matvec_silu_mul_limited(
3342 gate: &QTensor,
3343 up: &QTensor,
3344 x: &[f32],
3345 out: &mut [f32],
3346 limit: f32,
3347 pool: Option<&Pool>,
3348 ) -> bool {
3349 if gate.has_prism_contract() || up.has_prism_contract() || gate.is_blr() || up.is_blr() {
3350 return false;
3354 }
3355 let inter = gate.rows();
3356 debug_assert_eq!(up.rows(), inter);
3357 debug_assert_eq!(out.len(), inter);
3358 debug_assert_eq!(gate.cols(), up.cols());
3359 if !a8w8_enabled() {
3360 return false;
3361 }
3362 let act = split_act(x);
3363 let act = &act;
3364 let x_ref = x;
3365 let out_addr = SendMut(out.as_mut_ptr());
3366
3367 match (gate, up) {
3368 (
3370 Self::Mapped {
3371 dtype: TensorDtype::Q4Block,
3372 ..
3373 },
3374 Self::Mapped {
3375 dtype: TensorDtype::Q4Block,
3376 ..
3377 },
3378 ) => {
3379 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3380 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3381 let gpr = gate.cols() / GROUP_SIZE;
3382 let cols = gate.cols();
3383 let run = move |start: usize, end: usize| {
3384 for r in start..end {
3385 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3386 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3387 for &(j, xv) in &act.outliers {
3388 let flat = r * cols + j;
3389 let gb = gp[flat / 2];
3390 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3391 let gsc = f16_to_f32(u16::from_le_bytes([
3392 gs[(flat / GROUP_SIZE) * 2],
3393 gs[(flat / GROUP_SIZE) * 2 + 1],
3394 ]));
3395 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3396 let ub = up_p[flat / 2];
3397 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3398 let usc = f16_to_f32(u16::from_le_bytes([
3399 up_s[(flat / GROUP_SIZE) * 2],
3400 up_s[(flat / GROUP_SIZE) * 2 + 1],
3401 ]));
3402 uv += ((un as i32 - 8) as f32) * usc * xv;
3403 }
3404 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3406 }
3407 };
3408 dispatch_rows(pool, inter, &run);
3409 true
3410 }
3411 (
3415 Self::Mapped {
3416 dtype: TensorDtype::Q4Tiled,
3417 ..
3418 },
3419 Self::Mapped {
3420 dtype: TensorDtype::Q4Tiled,
3421 ..
3422 },
3423 ) => {
3424 let g_bytes = gate.quant_bytes();
3425 let u_bytes = up.quant_bytes();
3426 let gpr = gate.cols() / GROUP_SIZE;
3427 let run = move |start: usize, end: usize| {
3428 for r in start..end {
3429 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3430 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3431 for &(j, xv) in &act.outliers {
3432 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3433 gv += w * s * xv;
3434 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3435 uv += w * s * xv;
3436 }
3437 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3439 }
3440 };
3441 dispatch_rows(pool, inter, &run);
3442 true
3443 }
3444 (
3447 Self::Mapped {
3448 dtype: TensorDtype::Q4TiledP,
3449 ..
3450 },
3451 Self::Mapped {
3452 dtype: TensorDtype::Q4TiledP,
3453 ..
3454 },
3455 ) => {
3456 let cols = gate.cols();
3457 let gpr = cols / GROUP_SIZE;
3458 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3459 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3460 let run = |start: usize, end: usize| {
3461 with_krows(gpr, |gsc, usc| {
3462 for r in start..end {
3463 gv_view.scales_into(r, gpr, gsc);
3464 uv_view.scales_into(r, gpr, usc);
3465 let mut gv =
3466 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3467 let mut uv =
3468 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3469 for &(j, xv) in &act.outliers {
3470 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3471 gv += w * s * xv;
3472 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3473 uv += w * s * xv;
3474 }
3475 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3477 }
3478 });
3479 };
3480 dispatch_rows(pool, inter, &run);
3481 true
3482 }
3483 (
3490 Self::Mapped {
3491 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3492 ..
3493 },
3494 Self::Mapped {
3495 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3496 ..
3497 },
3498 ) => {
3499 let gw = gate.bin_w().expect("binary dtype matched");
3500 let uw = up.bin_w().expect("binary dtype matched");
3501 let gpr = gate.cols() / GROUP_SIZE;
3502 let gsum = q1_group_sums(&act.xq, gpr);
3503 let gsum = &gsum;
3504 let run = move |start: usize, end: usize| {
3505 for r in start..end {
3506 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3507 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3508 for &(j, xv) in &act.outliers {
3509 let (w, s) = gw.outlier(r, gpr, j);
3510 gv += w * s * xv;
3511 let (w, s) = uw.outlier(r, gpr, j);
3512 uv += w * s * xv;
3513 }
3514 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3516 }
3517 };
3518 dispatch_rows(pool, inter, &run);
3519 true
3520 }
3521 (
3525 Self::Mapped {
3526 dtype: TensorDtype::Q2TiledP,
3527 ..
3528 },
3529 Self::Mapped {
3530 dtype: TensorDtype::Q2TiledP,
3531 ..
3532 },
3533 ) => {
3534 let cols = gate.cols();
3535 let gpr = cols / GROUP_SIZE;
3536 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3537 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3538 let gsum = q1_group_sums(&act.xq, gpr);
3539 let gsum = &gsum;
3540 let run = move |start: usize, end: usize| {
3541 with_krows(gpr, |gsc, usc| {
3542 for r in start..end {
3543 gv_view.scales_into(r, gpr, gsc);
3544 uv_view.scales_into(r, gpr, usc);
3545 let mut gv =
3546 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3547 let mut uv =
3548 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3549 for &(j, xv) in &act.outliers {
3550 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3551 gv += w * s * xv;
3552 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3553 uv += w * s * xv;
3554 }
3555 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3557 }
3558 });
3559 };
3560 dispatch_rows(pool, inter, &run);
3561 true
3562 }
3563 (
3568 Self::Mapped {
3569 dtype: TensorDtype::Q8Row,
3570 row_scale: g_rs,
3571 ..
3572 },
3573 Self::Mapped {
3574 dtype: TensorDtype::Q8Row,
3575 row_scale: u_rs,
3576 ..
3577 },
3578 ) => {
3579 let g_bytes = gate.quant_bytes();
3580 let u_bytes = up.quant_bytes();
3581 let cols = gate.cols();
3582 let run = move |start: usize, end: usize| {
3583 for r in start..end {
3584 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3585 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3586 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3588 }
3589 };
3590 dispatch_rows(pool, inter, &run);
3591 true
3592 }
3593 (
3595 Self::Mapped {
3596 dtype: TensorDtype::Q1T,
3597 ..
3598 },
3599 Self::Mapped {
3600 dtype: TensorDtype::Q1T,
3601 ..
3602 },
3603 ) => {
3604 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3605 let g_bytes = gate.quant_bytes();
3606 let u_bytes = up.quant_bytes();
3607 let gpr = gate.cols() / GROUP_SIZE;
3608 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3609 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3610 let run = move |start: usize, end: usize| {
3611 for r in start..end {
3612 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3613 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3614 for &(j, xv) in &act.outliers {
3615 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3616 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3617 }
3618 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3619 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3620 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3622 }
3623 };
3624 dispatch_rows(pool, inter, &run);
3625 true
3626 }
3627 _ => false,
3628 }
3629 }
3630
3631 pub fn moe_gate_up_many(
3646 pairs: &[(&QTensor, &QTensor)],
3647 x: &[f32],
3648 outs: &mut [Vec<f32>],
3649 pool: Option<&Pool>,
3650 ) -> bool {
3651 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3652 }
3653
3654 pub fn moe_gate_up_many_limited(
3658 pairs: &[(&QTensor, &QTensor)],
3659 x: &[f32],
3660 outs: &mut [Vec<f32>],
3661 limit: f32,
3662 pool: Option<&Pool>,
3663 ) -> bool {
3664 if pairs.is_empty() || pairs.len() != outs.len() {
3665 return false;
3666 }
3667 if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3668 return false;
3669 }
3670 if !a8w8_enabled() {
3671 if limit > 0.0 {
3672 return false;
3675 }
3676 let groups = vec![vec![0]; pairs.len()];
3677 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3678 }
3679 let inter = pairs[0].0.rows();
3680 let cols = pairs[0].0.cols();
3681 if cols % GROUP_SIZE != 0 {
3682 return false;
3683 }
3684 let gpr = cols / GROUP_SIZE;
3685 let q2 = matches!(
3688 pairs[0].0,
3689 Self::Mapped {
3690 dtype: TensorDtype::Q2TiledP,
3691 ..
3692 }
3693 );
3694 let want = if q2 {
3695 TensorDtype::Q2TiledP
3696 } else {
3697 TensorDtype::Q4TiledP
3698 };
3699 let mut views = Vec::with_capacity(pairs.len() * 2);
3700 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3701 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3702 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3703 if !both
3704 || g.rows() != inter
3705 || u.rows() != inter
3706 || g.cols() != cols
3707 || u.cols() != cols
3708 || o.len() != inter
3709 {
3710 return false;
3711 }
3712 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3713 views.push(mk(g.quant_bytes(), inter, cols));
3714 views.push(mk(u.quant_bytes(), inter, cols));
3715 }
3716 let act = split_act(x);
3717 let gsum = if q2 {
3718 q1_group_sums(&act.xq, gpr)
3719 } else {
3720 Vec::new()
3721 };
3722 let (act, gsum) = (&act, &gsum);
3723 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3724 let (views, ptrs) = (&views, &ptrs);
3725 let run = |start: usize, end: usize| {
3726 with_krows(gpr, |gsc, usc| {
3727 for flat in start..end {
3728 let (e, r) = (flat / inter, flat % inter);
3729 let gv_view = &views[e * 2];
3730 let uv_view = &views[e * 2 + 1];
3731 gv_view.scales_into(r, gpr, gsc);
3732 uv_view.scales_into(r, gpr, usc);
3733 let (mut gv, mut uv) = if q2 {
3734 (
3735 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3736 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3737 )
3738 } else {
3739 (
3740 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3741 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3742 )
3743 };
3744 for &(j, xv) in &act.outliers {
3745 let (og, ou) = if q2 {
3746 (
3747 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3748 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3749 )
3750 } else {
3751 (
3752 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3753 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3754 )
3755 };
3756 gv += og.0 * og.1 * xv;
3757 uv += ou.0 * ou.1 * xv;
3758 }
3759 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3761 }
3762 });
3763 };
3764 dispatch_rows(pool, pairs.len() * inter, &run);
3765 true
3766 }
3767
3768 pub fn moe_down_many(
3777 downs: &[&QTensor],
3778 gs: &[Vec<f32>],
3779 weights: &[f32],
3780 out: &mut [f32],
3781 pool: Option<&Pool>,
3782 ) -> bool {
3783 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3784 return false;
3785 }
3786 if downs.iter().any(|d| d.is_blr()) {
3787 return false;
3788 }
3789 if !a8w8_enabled() {
3790 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3791 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3792 return false;
3793 }
3794 out.fill(0.0);
3795 for (row, &w) in terms.iter().zip(weights) {
3796 for (o, &v) in out.iter_mut().zip(row) {
3797 *o += w * v;
3798 }
3799 }
3800 return true;
3801 }
3802 let rows = out.len();
3803 let cols = downs[0].cols();
3804 if cols % GROUP_SIZE != 0 {
3805 return false;
3806 }
3807 let gpr = cols / GROUP_SIZE;
3808 let mut views = Vec::with_capacity(downs.len());
3809 for (d, g) in downs.iter().zip(gs.iter()) {
3810 if !matches!(
3811 d,
3812 Self::Mapped {
3813 dtype: TensorDtype::Q4TiledP,
3814 ..
3815 }
3816 ) || d.rows() != rows
3817 || d.cols() != cols
3818 || g.len() != cols
3819 {
3820 return false;
3821 }
3822 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3823 }
3824 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3826 let out_addr = SendMut(out.as_mut_ptr());
3834 let (views, acts, weights) = (&views, &acts, &weights);
3835 let run = |start: usize, end: usize| {
3836 with_krow(gpr, |sc| {
3837 for r in start..end {
3838 let mut acc = 0f32;
3839 for (e, v) in views.iter().enumerate() {
3840 v.scales_into(r, gpr, sc);
3841 let a = &acts[e];
3842 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3843 for &(j, xv) in &a.outliers {
3844 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3845 d += w * s * xv;
3846 }
3847 acc += weights[e] * d;
3848 }
3849 unsafe { *out_addr.at(r) = acc };
3851 }
3852 });
3853 };
3854 dispatch_rows(pool, rows, &run);
3855 true
3856 }
3857
3858 pub fn moe_gate_up_rows(
3870 pairs: &[(&QTensor, &QTensor)],
3871 groups: &[Vec<usize>],
3872 xs: &[f32],
3873 outs: &mut [Vec<f32>],
3874 pool: Option<&Pool>,
3875 ) -> bool {
3876 if pairs.is_empty() || pairs.len() != groups.len() {
3877 return false;
3878 }
3879 if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3880 return false;
3881 }
3882 let inter = pairs[0].0.rows();
3883 let cols = pairs[0].0.cols();
3884 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3885 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3886 return false;
3887 }
3888 let b = xs.len() / cols;
3889 let gpr = cols / GROUP_SIZE;
3890 let mut views = Vec::with_capacity(pairs.len() * 2);
3891 for (g, u) in pairs {
3892 let q4tp = |t: &QTensor| {
3893 matches!(
3894 t,
3895 Self::Mapped {
3896 dtype: TensorDtype::Q4TiledP,
3897 ..
3898 }
3899 )
3900 };
3901 if g.has_prism_contract()
3902 || u.has_prism_contract()
3903 || !q4tp(g)
3904 || !q4tp(u)
3905 || g.rows() != inter
3906 || u.rows() != inter
3907 || g.cols() != cols
3908 || u.cols() != cols
3909 {
3910 return false;
3911 }
3912 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3913 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3914 }
3915 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3916 return false;
3917 }
3918 let quantized = a8w8_enabled();
3919 let acts: Vec<SplitAct> = if quantized {
3920 (0..b)
3921 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3922 .collect()
3923 } else {
3924 Vec::new()
3925 };
3926 let mut offs = Vec::with_capacity(groups.len());
3927 let mut o = 0usize;
3928 for g in groups {
3929 offs.push(o);
3930 o += g.len();
3931 }
3932 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3933 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3934 let run = |start: usize, end: usize| {
3935 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3936 for flat in start..end {
3937 let (e, r) = (flat / inter, flat % inter);
3938 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3939 gv_view.scales_into(r, gpr, &mut gsc);
3940 uv_view.scales_into(r, gpr, &mut usc);
3941 for (k, &t) in groups[e].iter().enumerate() {
3942 if !quantized {
3943 let x = &xs[t * cols..(t + 1) * cols];
3944 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3945 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3946 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3947 continue;
3948 }
3949 let act = &acts[t];
3950 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3951 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3952 for &(j, xv) in &act.outliers {
3953 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3954 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3955 gv += og.0 * og.1 * xv;
3956 uv += ou.0 * ou.1 * xv;
3957 }
3958 let silu_g = gv / (1.0 + (-gv).exp());
3959 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3962 }
3963 }
3964 };
3965 dispatch_rows(pool, pairs.len() * inter, &run);
3966 true
3967 }
3968
3969 pub fn moe_down_rows(
3976 downs: &[&QTensor],
3977 group_lens: &[usize],
3978 gs: &[Vec<f32>],
3979 outs: &mut [Vec<f32>],
3980 pool: Option<&Pool>,
3981 ) -> bool {
3982 if downs.is_empty() || downs.len() != group_lens.len() {
3983 return false;
3984 }
3985 if downs.iter().any(|d| d.is_blr()) {
3986 return false;
3987 }
3988 let rows = downs[0].rows();
3989 let cols = downs[0].cols();
3990 let n_pairs: usize = group_lens.iter().sum();
3991 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3992 return false;
3993 }
3994 let gpr = cols / GROUP_SIZE;
3995 let mut views = Vec::with_capacity(downs.len());
3996 for d in downs {
3997 if d.has_prism_contract()
3998 || !matches!(
3999 d,
4000 Self::Mapped {
4001 dtype: TensorDtype::Q4TiledP,
4002 ..
4003 }
4004 )
4005 || d.rows() != rows
4006 || d.cols() != cols
4007 {
4008 return false;
4009 }
4010 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4011 }
4012 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4013 return false;
4014 }
4015 let quantized = a8w8_enabled();
4016 let acts: Vec<SplitAct> = if quantized {
4017 gs.iter().map(|g| split_act(g)).collect()
4018 } else {
4019 Vec::new()
4020 };
4021 let mut offs = Vec::with_capacity(group_lens.len());
4022 let mut o = 0usize;
4023 for &l in group_lens {
4024 offs.push(o);
4025 o += l;
4026 }
4027 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4028 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4029 let run = |start: usize, end: usize| {
4030 let mut sc = vec![0f32; gpr];
4031 for flat in start..end {
4032 let (e, r) = (flat / rows, flat % rows);
4033 let v = &views[e];
4034 v.scales_into(r, gpr, &mut sc);
4035 for k in 0..group_lens[e] {
4036 if !quantized {
4037 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4038 unsafe { *ptrs[offs[e] + k].at(r) = d };
4039 continue;
4040 }
4041 let a = &acts[offs[e] + k];
4042 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4043 for &(j, xv) in &a.outliers {
4044 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4045 d += w * s * xv;
4046 }
4047 unsafe { *ptrs[offs[e] + k].at(r) = d };
4049 }
4050 }
4051 };
4052 dispatch_rows(pool, downs.len() * rows, &run);
4053 true
4054 }
4055}
4056
4057#[cfg(target_os = "macos")]
4062mod accel_blas {
4063 #[link(name = "Accelerate", kind = "framework")]
4064 unsafe extern "C" {
4065 pub fn cblas_sgemm(
4066 order: i32,
4067 trans_a: i32,
4068 trans_b: i32,
4069 m: i32,
4070 n: i32,
4071 k: i32,
4072 alpha: f32,
4073 a: *const f32,
4074 lda: i32,
4075 b: *const f32,
4076 ldb: i32,
4077 beta: f32,
4078 c: *mut f32,
4079 ldc: i32,
4080 );
4081 }
4082}
4083
4084#[cfg(target_os = "macos")]
4085pub(crate) fn accel_gemm_enabled() -> bool {
4086 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4087 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4088}
4089
4090#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4093pub(crate) fn accel_gemm_enabled() -> bool {
4094 true
4095}
4096
4097#[cfg(target_arch = "aarch64")]
4104#[allow(clippy::too_many_arguments)]
4105pub(crate) fn neon_gemm_rm(
4106 m: usize,
4107 n: usize,
4108 k: usize,
4109 alpha: f32,
4110 a: &[f32],
4111 lda: usize,
4112 b_mat: &[f32],
4113 ldb: usize,
4114 b_rows_are_n: bool,
4115 c: &mut [f32],
4116 ldc: usize,
4117) {
4118 debug_assert!(a.len() >= (m - 1) * lda + k);
4119 debug_assert!(c.len() >= (m - 1) * ldc + n);
4120 unsafe {
4122 use core::arch::aarch64::*;
4123 let mut i = 0usize;
4124 while i < m {
4125 let mi = (m - i).min(4);
4126 let mut j = 0usize;
4127 while j < n {
4128 let nj = (n - j).min(8);
4129 if mi == 4 && nj == 8 {
4130 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4131 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4132 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4133 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4134 for p in 0..k {
4135 let (b0, b1) = if b_rows_are_n {
4136 let base = b_mat.as_ptr().add(j * ldb + p);
4139 let g = |o: usize| *base.add(o * ldb);
4140 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4141 } else {
4142 let base = b_mat.as_ptr().add(p * ldb + j);
4143 (
4144 [*base, *base.add(1), *base.add(2), *base.add(3)],
4145 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4146 )
4147 };
4148 let bv0 = vld1q_f32(b0.as_ptr());
4149 let bv1 = vld1q_f32(b1.as_ptr());
4150 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4151 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4152 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4153 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4154 c0a = vfmaq_f32(c0a, a0, bv0);
4155 c0b = vfmaq_f32(c0b, a0, bv1);
4156 c1a = vfmaq_f32(c1a, a1, bv0);
4157 c1b = vfmaq_f32(c1b, a1, bv1);
4158 c2a = vfmaq_f32(c2a, a2, bv0);
4159 c2b = vfmaq_f32(c2b, a2, bv1);
4160 c3a = vfmaq_f32(c3a, a3, bv0);
4161 c3b = vfmaq_f32(c3b, a3, bv1);
4162 }
4163 let al = vdupq_n_f32(alpha);
4164 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4165 .iter()
4166 .enumerate()
4167 {
4168 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4169 vst1q_f32(dst, vmulq_f32(*ca, al));
4170 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4171 }
4172 } else {
4173 for r in 0..mi {
4174 for q in 0..nj {
4175 let mut acc = 0f32;
4176 for p in 0..k {
4177 let bv = if b_rows_are_n {
4178 b_mat[(j + q) * ldb + p]
4179 } else {
4180 b_mat[p * ldb + j + q]
4181 };
4182 acc += a[(i + r) * lda + p] * bv;
4183 }
4184 c[(i + r) * ldc + j + q] = acc * alpha;
4185 }
4186 }
4187 }
4188 j += nj;
4189 }
4190 i += mi;
4191 }
4192 }
4193}
4194
4195#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4197#[allow(clippy::too_many_arguments)]
4198pub(crate) fn sgemm_rm(
4199 m: usize,
4200 n: usize,
4201 k: usize,
4202 alpha: f32,
4203 a: &[f32],
4204 lda: usize,
4205 b_mat: &[f32],
4206 ldb: usize,
4207 b_rows_are_n: bool,
4208 c: &mut [f32],
4209 ldc: usize,
4210) {
4211 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4212}
4213
4214#[allow(clippy::too_many_arguments)]
4218pub fn sgemm_public(
4219 m: usize,
4220 n: usize,
4221 k: usize,
4222 alpha: f32,
4223 a: &[f32],
4224 lda: usize,
4225 b_mat: &[f32],
4226 ldb: usize,
4227 b_rows_are_n: bool,
4228 c: &mut [f32],
4229 ldc: usize,
4230) {
4231 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4232 {
4233 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4234 }
4235 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4240 {
4241 if m == 0 || n == 0 {
4242 return;
4243 }
4244 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4245 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4246 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4247 unsafe {
4249 matrixmultiply::sgemm(
4250 m,
4251 k,
4252 n,
4253 alpha,
4254 a.as_ptr(),
4255 lda as isize,
4256 1,
4257 b_mat.as_ptr(),
4258 rsb,
4259 csb,
4260 0.0,
4261 c.as_mut_ptr(),
4262 ldc as isize,
4263 1,
4264 );
4265 }
4266 }
4267}
4268
4269#[cfg(target_os = "macos")]
4272#[allow(clippy::too_many_arguments)]
4273pub(crate) fn sgemm_rm(
4274 m: usize,
4275 n: usize,
4276 k: usize,
4277 alpha: f32,
4278 a: &[f32],
4279 lda: usize,
4280 b_mat: &[f32],
4281 ldb: usize,
4282 b_rows_are_n: bool,
4283 c: &mut [f32],
4284 ldc: usize,
4285) {
4286 debug_assert!(a.len() >= (m - 1) * lda + k);
4287 debug_assert!(c.len() >= (m - 1) * ldc + n);
4288 #[cfg(target_arch = "aarch64")]
4293 if std::env::var("CMF_FORCE_NEON_GEMM")
4294 .map(|v| v == "1")
4295 .unwrap_or(false)
4296 {
4297 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4298 }
4299 unsafe {
4300 accel_blas::cblas_sgemm(
4301 101, 111, if b_rows_are_n { 112 } else { 111 },
4304 m as i32,
4305 n as i32,
4306 k as i32,
4307 alpha,
4308 a.as_ptr(),
4309 lda as i32,
4310 b_mat.as_ptr(),
4311 ldb as i32,
4312 0.0,
4313 c.as_mut_ptr(),
4314 ldc as i32,
4315 );
4316 }
4317}
4318
4319#[cfg(target_os = "macos")]
4326fn qmatmat_accel(
4327 q: &[u8],
4328 row_scale: &[f32],
4329 pre: &[std::borrow::Cow<'_, [f32]>],
4330 rows: usize,
4331 cols: usize,
4332 out: &mut [f32],
4333 pool: Option<&Pool>,
4334) {
4335 const TR: usize = 2048;
4340 let b = pre.len();
4341 thread_local! {
4342 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4343 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4344 }
4345 XPANEL.with(|xp| {
4346 WTILE.with(|wt| {
4347 let mut xpanel = xp.borrow_mut();
4348 xpanel.clear();
4349 for x in pre {
4350 xpanel.extend_from_slice(x);
4351 }
4352 let mut wtile = wt.borrow_mut();
4353 wtile.resize(TR * cols, 0.0);
4354 let mut r0 = 0usize;
4355 while r0 < rows {
4356 let tr = TR.min(rows - r0);
4357 let wt_addr = SendMut(wtile.as_mut_ptr());
4359 let run = |start: usize, end: usize| {
4360 for r in start..end {
4361 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4362 let s = row_scale[r0 + r];
4363 let dst =
4365 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4366 for (d, &v) in dst.iter_mut().zip(row) {
4367 *d = (v as i8) as f32 * s;
4368 }
4369 }
4370 };
4371 dispatch_rows(pool, tr, &run);
4372 unsafe {
4374 accel_blas::cblas_sgemm(
4375 101, 111, 112, b as i32,
4379 tr as i32,
4380 cols as i32,
4381 1.0,
4382 xpanel.as_ptr(),
4383 cols as i32,
4384 wtile.as_ptr(),
4385 cols as i32,
4386 0.0,
4387 out.as_mut_ptr().add(r0),
4388 rows as i32,
4389 );
4390 }
4391 r0 += tr;
4392 }
4393 })
4394 });
4395}
4396
4397fn qmatmat(
4398 q: &[u8],
4399 row_scale: &[f32],
4400 pre: &[std::borrow::Cow<'_, [f32]>],
4401 rows: usize,
4402 cols: usize,
4403 out: &mut [f32],
4404 pool: Option<&Pool>,
4405) {
4406 let b = pre.len();
4407 debug_assert_eq!(out.len(), b * rows);
4408 #[cfg(target_os = "macos")]
4413 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4414 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4415 return;
4416 }
4417 #[cfg(target_arch = "aarch64")]
4418 if sdot_enabled() {
4419 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4420 let out_addr = SendMut(out.as_mut_ptr());
4421 let blocked_ok = blocked_enabled();
4424 let use_i8mm = i8mm_enabled();
4425 if blocked_ok {
4426 let run = |start: usize, end: usize| {
4427 let mut o = start;
4428 while o < end {
4429 if o + 2 <= end {
4430 let r0 = &q[o * cols..(o + 1) * cols];
4431 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4432 let mut bi = 0usize;
4433 while bi + 4 <= acts.len() {
4434 let xs = [
4435 acts[bi].xq.as_slice(),
4436 acts[bi + 1].xq.as_slice(),
4437 acts[bi + 2].xq.as_slice(),
4438 acts[bi + 3].xq.as_slice(),
4439 ];
4440 let d = if use_i8mm {
4441 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4442 } else {
4443 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4444 };
4445 for (r, row) in [r0, r1].into_iter().enumerate() {
4446 for k in 0..4 {
4447 let act = &acts[bi + k];
4448 let mut v = d[r][k] as f32 * act.sx;
4449 for &(j, xv) in &act.outliers {
4450 v += (row[j] as i8) as f32 * xv;
4451 }
4452 unsafe {
4453 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4454 };
4455 }
4456 }
4457 bi += 4;
4458 }
4459 while bi < acts.len() {
4460 for (r, row) in [r0, r1].into_iter().enumerate() {
4461 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4462 unsafe { *out_addr.at(bi * rows + o + r) = v };
4463 }
4464 bi += 1;
4465 }
4466 o += 2;
4467 } else {
4468 let row = &q[o * cols..(o + 1) * cols];
4469 for (bi, act) in acts.iter().enumerate() {
4470 let v = row_dot_sdot(row, act) * row_scale[o];
4471 unsafe { *out_addr.at(bi * rows + o) = v };
4472 }
4473 o += 1;
4474 }
4475 }
4476 };
4477 dispatch_rows(pool, rows, &run);
4478 return;
4479 }
4480 let run = |start: usize, end: usize| {
4481 for o in start..end {
4482 let row = &q[o * cols..(o + 1) * cols];
4483 for (bi, act) in acts.iter().enumerate() {
4484 let v = row_dot_sdot(row, act) * row_scale[o];
4485 unsafe { *out_addr.at(bi * rows + o) = v };
4486 }
4487 }
4488 };
4489 dispatch_rows(pool, rows, &run);
4490 return;
4491 }
4492 #[cfg(target_arch = "x86_64")]
4497 if avx2_a8w8_enabled() {
4498 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4499 let out_addr = SendMut(out.as_mut_ptr());
4500 let blocked_ok = blocked_enabled();
4503 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4504 let run = |start: usize, end: usize| {
4505 let mut o = start;
4506 while o < end {
4507 if o + 2 <= end {
4508 let r0 = &q[o * cols..(o + 1) * cols];
4509 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4510 let mut bi = 0usize;
4511 while bi + 4 <= acts.len() {
4512 let xs = [
4513 acts[bi].xq.as_slice(),
4514 acts[bi + 1].xq.as_slice(),
4515 acts[bi + 2].xq.as_slice(),
4516 acts[bi + 3].xq.as_slice(),
4517 ];
4518 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4519 for (r, row) in [r0, r1].into_iter().enumerate() {
4520 for k in 0..4 {
4521 let act = &acts[bi + k];
4522 let mut v = d[r][k] as f32 * act.sx;
4523 for &(j, xv) in &act.outliers {
4524 v += (row[j] as i8) as f32 * xv;
4525 }
4526 unsafe {
4527 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4528 };
4529 }
4530 }
4531 bi += 4;
4532 }
4533 while bi < acts.len() {
4534 for (r, row) in [r0, r1].into_iter().enumerate() {
4535 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4536 unsafe { *out_addr.at(bi * rows + o + r) = v };
4537 }
4538 bi += 1;
4539 }
4540 o += 2;
4541 } else {
4542 let row = &q[o * cols..(o + 1) * cols];
4543 for (bi, act) in acts.iter().enumerate() {
4544 let v = row_dot_avx2(row, act) * row_scale[o];
4545 unsafe { *out_addr.at(bi * rows + o) = v };
4546 }
4547 o += 1;
4548 }
4549 }
4550 };
4551 dispatch_rows(pool, rows, &run);
4552 return;
4553 }
4554 let run = |start: usize, end: usize| {
4555 for o in start..end {
4556 let row = &q[o * cols..(o + 1) * cols];
4557 for (bi, act) in acts.iter().enumerate() {
4558 let v = row_dot_avx2(row, act) * row_scale[o];
4559 unsafe { *out_addr.at(bi * rows + o) = v };
4560 }
4561 }
4562 };
4563 dispatch_rows(pool, rows, &run);
4564 return;
4565 }
4566 let out_addr = SendMut(out.as_mut_ptr());
4567 let run = |start: usize, end: usize| {
4568 for o in start..end {
4569 let row = &q[o * cols..(o + 1) * cols];
4570 for (bi, x) in pre.iter().enumerate() {
4571 let mut acc = 0f32;
4572 for j in 0..cols {
4573 acc += (row[j] as i8) as f32 * x[j];
4574 }
4575 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4576 }
4577 }
4578 };
4579 dispatch_rows(pool, rows, &run);
4580}
4581
4582fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4585 match pool {
4586 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4587 _ => run(0, rows),
4588 }
4589}
4590
4591fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4593 let groups = rows * cols / GROUP_SIZE;
4594 bytes.split_at(groups * 16)
4595}
4596
4597#[inline]
4602fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4603 #[cfg(target_arch = "aarch64")]
4604 unsafe {
4605 return vbit_fill4_neon(data, buf);
4606 }
4607 #[cfg(target_arch = "x86_64")]
4608 if avx2_enabled() {
4609 return unsafe { vbit_fill4_avx2(data, buf) };
4610 }
4611 #[allow(unreachable_code)]
4612 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4613 let u = unpack8::<4>(&data[blk * 4..]);
4614 for k in 0..8 {
4615 chunk[k] = (u[k] - 7) as i8 as u8;
4616 }
4617 }
4618}
4619
4620#[cfg(target_arch = "aarch64")]
4621#[target_feature(enable = "neon")]
4622unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4623 unsafe {
4626 use core::arch::aarch64::*;
4627 let n = buf.len();
4628 let mask = vdupq_n_u8(0x0F);
4629 let seven = vdupq_n_s8(7);
4630 let mut g = 0usize;
4631 while g * 32 + 32 <= n {
4632 let b = vld1q_u8(data.as_ptr().add(g * 16));
4633 let hi = vshrq_n_u8::<4>(b);
4634 let lo = vandq_u8(b, mask);
4635 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4636 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4637 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4638 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4639 g += 1;
4640 }
4641 }
4642}
4643
4644#[cfg(target_arch = "x86_64")]
4645#[target_feature(enable = "avx2")]
4646unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4647 unsafe {
4649 use core::arch::x86_64::*;
4650 let n = buf.len();
4651 let mask = _mm_set1_epi8(0x0F);
4652 let seven = _mm256_set1_epi8(7);
4653 let mut g = 0usize;
4654 while g * 32 + 32 <= n {
4655 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4656 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4657 let lo = _mm_and_si128(b, mask);
4658 let z = _mm256_sub_epi8(
4659 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4660 seven,
4661 );
4662 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4663 g += 1;
4664 }
4665 }
4666}
4667
4668#[inline(always)]
4673fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4674 let mut acc = 0u64;
4675 for i in 0..B {
4676 acc = (acc << 8) | data[i] as u64;
4677 }
4678 let mask = (1u64 << B) - 1;
4679 let mut out = [0i32; 8];
4680 for (k, o) in out.iter_mut().enumerate() {
4681 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4682 }
4683 out
4684}
4685
4686#[allow(clippy::too_many_arguments)]
4692fn vbitmatvec(
4693 bytes: &[u8],
4694 offsets: &[usize],
4695 x: &[f32],
4696 rows: usize,
4697 cols: usize,
4698 out: &mut [f32],
4699 pool: Option<&Pool>,
4700) {
4701 debug_assert_eq!(out.len(), rows);
4702 debug_assert_eq!(offsets.len(), rows + 1);
4703
4704 if a8w8_enabled() {
4708 let act = split_act(x);
4709 let out_addr = SendMut(out.as_mut_ptr());
4710 let run = move |start: usize, end: usize| {
4711 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4712 };
4713 dispatch_rows(pool, rows, &run);
4714 return;
4715 }
4716
4717 let out_addr = SendMut(out.as_mut_ptr());
4718 let run = move |start: usize, end: usize| {
4719 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4720 };
4721 dispatch_rows(pool, rows, &run);
4722}
4723
4724#[allow(clippy::too_many_arguments)]
4728fn vbit_range_a8w8(
4729 bytes: &[u8],
4730 offsets: &[usize],
4731 x: &[f32],
4732 act: &SplitAct,
4733 rows: usize,
4734 cols: usize,
4735 out: SendMut,
4736 start: usize,
4737 end: usize,
4738) {
4739 let ng = cols / GROUP_SIZE;
4740 let bits = &bytes[..rows];
4741 let sc_off = rows;
4742 let row_dot = |r: usize| -> f32 {
4743 let b = bits[r] as usize;
4744 let l = (1i32 << (b - 1)) - 1;
4745 let mask = (1u64 << b) - 1;
4746 let data = &bytes[offsets[r]..offsets[r + 1]];
4747 if b == 8 {
4748 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4750 let mut dot = 0f32;
4751 for g in 0..ng {
4752 let so = (r * ng + g) * 2;
4753 let sgf = f16_to_f32(u16::from_le_bytes([
4754 bytes[sc_off + so],
4755 bytes[sc_off + so + 1],
4756 ]));
4757 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4758 let mut gd = 0f32;
4759 for &xv in xg.iter() {
4760 if nbits < 8 {
4761 acc = (acc << 8) | data[idx] as u64;
4762 idx += 1;
4763 nbits += 8;
4764 }
4765 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4766 nbits -= 8;
4767 gd += (u - l) as f32 * xv;
4768 }
4769 dot += gd * sgf;
4770 }
4771 return dot;
4772 }
4773 thread_local! {
4777 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4778 const { std::cell::RefCell::new(Vec::new()) };
4779 }
4780 #[inline(always)]
4781 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4782 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4783 let u = unpack8::<B>(&data[blk * B..]);
4784 for k in 0..8 {
4785 chunk[k] = (u[k] - l) as i8 as u8;
4786 }
4787 }
4788 }
4789 let _ = mask;
4790 VBIT_SCRATCH.with(|scratch| {
4791 let mut buf = scratch.borrow_mut();
4792 buf.resize(cols, 0);
4793 match b {
4794 3 => fill::<3>(data, l, &mut buf),
4795 4 => vbit_fill4(data, &mut buf),
4796 5 => fill::<5>(data, l, &mut buf),
4797 6 => fill::<6>(data, l, &mut buf),
4798 _ => unreachable!(),
4799 }
4800 let mut dot = 0f32;
4801 for g in 0..ng {
4802 let so = (r * ng + g) * 2;
4803 let s = f16_to_f32(u16::from_le_bytes([
4804 bytes[sc_off + so],
4805 bytes[sc_off + so + 1],
4806 ]));
4807 let d = dot_i8_i8(
4808 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4809 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4810 ) as f32
4811 * act.sx;
4812 dot += d * s;
4813 }
4814 for &(j, xv) in &act.outliers {
4815 let so = (r * ng + j / GROUP_SIZE) * 2;
4816 let s = f16_to_f32(u16::from_le_bytes([
4817 bytes[sc_off + so],
4818 bytes[sc_off + so + 1],
4819 ]));
4820 dot += (buf[j] as i8) as f32 * s * xv;
4822 }
4823 dot
4824 })
4825 };
4826 for r in start..end {
4827 unsafe { *out.at(r) = row_dot(r) };
4829 }
4830}
4831
4832#[allow(clippy::too_many_arguments)]
4834fn vbit_range_f32(
4835 bytes: &[u8],
4836 offsets: &[usize],
4837 x: &[f32],
4838 rows: usize,
4839 cols: usize,
4840 out: SendMut,
4841 start: usize,
4842 end: usize,
4843) {
4844 let ng = cols / GROUP_SIZE;
4845 let bits = &bytes[..rows];
4846 let sc_off = rows;
4847 #[inline(always)]
4851 fn dot_row<const B: usize>(
4852 data: &[u8],
4853 bytes: &[u8],
4854 sc_off: usize,
4855 r: usize,
4856 ng: usize,
4857 x: &[f32],
4858 ) -> f32 {
4859 let l = ((1i32 << (B - 1)) - 1) as f32;
4860 let gbytes = GROUP_SIZE * B / 8;
4861 let mut dot = 0f32;
4862 for g in 0..ng {
4863 let so = (r * ng + g) * 2;
4864 let s = f16_to_f32(u16::from_le_bytes([
4865 bytes[sc_off + so],
4866 bytes[sc_off + so + 1],
4867 ]));
4868 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4869 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4870 let mut gd = 0f32;
4871 for blk in 0..GROUP_SIZE / 8 {
4872 let u = unpack8::<B>(&gd0[blk * B..]);
4873 let xb = &xg[blk * 8..blk * 8 + 8];
4874 for k in 0..8 {
4875 gd += (u[k] as f32 - l) * xb[k];
4876 }
4877 }
4878 dot += gd * s;
4879 }
4880 dot
4881 }
4882 for r in start..end {
4883 let data = &bytes[offsets[r]..offsets[r + 1]];
4884 let v = match bits[r] {
4885 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4886 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4887 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4888 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4889 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4890 b => unreachable!("vbit bit-width {b} (validated at load)"),
4891 };
4892 unsafe { *out.at(r) = v };
4894 }
4895}
4896
4897#[allow(clippy::too_many_arguments)]
4902fn vbitmatvec2(
4903 bytes: &[u8],
4904 offsets: &[usize],
4905 x1: &[f32],
4906 x2: &[f32],
4907 rows: usize,
4908 cols: usize,
4909 o1: &mut [f32],
4910 o2: &mut [f32],
4911 pool: Option<&Pool>,
4912) {
4913 debug_assert_eq!(o1.len(), rows);
4914 debug_assert_eq!(o2.len(), rows);
4915
4916 if a8w8_enabled() {
4917 let a1 = split_act(x1);
4918 let a2 = split_act(x2);
4919 let p1 = SendMut(o1.as_mut_ptr());
4920 let p2 = SendMut(o2.as_mut_ptr());
4921 let run = move |start: usize, end: usize| {
4922 vbit_range2_a8w8(
4923 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4924 )
4925 };
4926 dispatch_rows(pool, rows, &run);
4927 return;
4928 }
4929
4930 let p1 = SendMut(o1.as_mut_ptr());
4931 let p2 = SendMut(o2.as_mut_ptr());
4932 let run = move |start: usize, end: usize| {
4933 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4934 };
4935 dispatch_rows(pool, rows, &run);
4936}
4937
4938#[allow(clippy::too_many_arguments)]
4942fn vbit_range2_a8w8(
4943 bytes: &[u8],
4944 offsets: &[usize],
4945 x1: &[f32],
4946 x2: &[f32],
4947 a1: &SplitAct,
4948 a2: &SplitAct,
4949 rows: usize,
4950 cols: usize,
4951 p1: SendMut,
4952 p2: SendMut,
4953 start: usize,
4954 end: usize,
4955) {
4956 let ng = cols / GROUP_SIZE;
4957 let bits = &bytes[..rows];
4958 let sc_off = rows;
4959 let row_dots = |r: usize| -> (f32, f32) {
4960 let b = bits[r] as usize;
4961 let l = (1i32 << (b - 1)) - 1;
4962 let data = &bytes[offsets[r]..offsets[r + 1]];
4963 if b == 8 {
4964 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4967 let (mut d1, mut d2) = (0f32, 0f32);
4968 for g in 0..ng {
4969 let so = (r * ng + g) * 2;
4970 let sgf = f16_to_f32(u16::from_le_bytes([
4971 bytes[sc_off + so],
4972 bytes[sc_off + so + 1],
4973 ]));
4974 let (mut g1, mut g2) = (0f32, 0f32);
4975 for k in 0..GROUP_SIZE {
4976 if nbits < 8 {
4977 acc = (acc << 8) | data[idx] as u64;
4978 idx += 1;
4979 nbits += 8;
4980 }
4981 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4982 nbits -= 8;
4983 let w = (u - l) as f32;
4984 g1 += w * x1[g * GROUP_SIZE + k];
4985 g2 += w * x2[g * GROUP_SIZE + k];
4986 }
4987 d1 += g1 * sgf;
4988 d2 += g2 * sgf;
4989 }
4990 return (d1, d2);
4991 }
4992 thread_local! {
4993 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4994 const { std::cell::RefCell::new(Vec::new()) };
4995 }
4996 #[inline(always)]
4997 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4998 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4999 let u = unpack8::<B>(&data[blk * B..]);
5000 for k in 0..8 {
5001 chunk[k] = (u[k] - l) as i8 as u8;
5002 }
5003 }
5004 }
5005 VBIT_SCRATCH2.with(|scratch| {
5006 let mut buf = scratch.borrow_mut();
5007 buf.resize(cols, 0);
5008 match b {
5009 3 => fill::<3>(data, l, &mut buf),
5010 4 => vbit_fill4(data, &mut buf),
5011 5 => fill::<5>(data, l, &mut buf),
5012 6 => fill::<6>(data, l, &mut buf),
5013 _ => unreachable!(),
5014 }
5015 let (mut d1, mut d2) = (0f32, 0f32);
5016 for g in 0..ng {
5017 let so = (r * ng + g) * 2;
5018 let s = f16_to_f32(u16::from_le_bytes([
5019 bytes[sc_off + so],
5020 bytes[sc_off + so + 1],
5021 ]));
5022 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5023 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5024 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5025 d1 += v1 * s;
5026 d2 += v2 * s;
5027 }
5028 for &(j, xv) in &a1.outliers {
5029 let so = (r * ng + j / GROUP_SIZE) * 2;
5030 let s = f16_to_f32(u16::from_le_bytes([
5031 bytes[sc_off + so],
5032 bytes[sc_off + so + 1],
5033 ]));
5034 d1 += (buf[j] as i8) as f32 * s * xv;
5035 }
5036 for &(j, xv) in &a2.outliers {
5037 let so = (r * ng + j / GROUP_SIZE) * 2;
5038 let s = f16_to_f32(u16::from_le_bytes([
5039 bytes[sc_off + so],
5040 bytes[sc_off + so + 1],
5041 ]));
5042 d2 += (buf[j] as i8) as f32 * s * xv;
5043 }
5044 (d1, d2)
5045 })
5046 };
5047 for r in start..end {
5048 let (v1, v2) = row_dots(r);
5049 unsafe {
5051 *p1.at(r) = v1;
5052 *p2.at(r) = v2;
5053 }
5054 }
5055}
5056
5057#[allow(clippy::too_many_arguments)]
5061fn vbit_range2_f32(
5062 bytes: &[u8],
5063 offsets: &[usize],
5064 x1: &[f32],
5065 x2: &[f32],
5066 rows: usize,
5067 cols: usize,
5068 p1: SendMut,
5069 p2: SendMut,
5070 start: usize,
5071 end: usize,
5072) {
5073 let ng = cols / GROUP_SIZE;
5074 let bits = &bytes[..rows];
5075 let sc_off = rows;
5076 #[inline(always)]
5077 #[allow(clippy::too_many_arguments)]
5078 fn dot_row2<const B: usize>(
5079 data: &[u8],
5080 bytes: &[u8],
5081 sc_off: usize,
5082 r: usize,
5083 ng: usize,
5084 x1: &[f32],
5085 x2: &[f32],
5086 ) -> (f32, f32) {
5087 let l = ((1i32 << (B - 1)) - 1) as f32;
5088 let gbytes = GROUP_SIZE * B / 8;
5089 let (mut d1, mut d2) = (0f32, 0f32);
5090 for g in 0..ng {
5091 let so = (r * ng + g) * 2;
5092 let s = f16_to_f32(u16::from_le_bytes([
5093 bytes[sc_off + so],
5094 bytes[sc_off + so + 1],
5095 ]));
5096 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5097 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5098 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5099 let (mut g1, mut g2) = (0f32, 0f32);
5100 for blk in 0..GROUP_SIZE / 8 {
5101 let u = unpack8::<B>(&gd0[blk * B..]);
5102 for k in 0..8 {
5103 let w = u[k] as f32 - l;
5104 g1 += w * x1g[blk * 8 + k];
5105 g2 += w * x2g[blk * 8 + k];
5106 }
5107 }
5108 d1 += g1 * s;
5109 d2 += g2 * s;
5110 }
5111 (d1, d2)
5112 }
5113 for r in start..end {
5114 let data = &bytes[offsets[r]..offsets[r + 1]];
5115 let (v1, v2) = match bits[r] {
5116 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5117 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5118 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5119 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5120 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5121 b => unreachable!("vbit bit-width {b} (validated at load)"),
5122 };
5123 unsafe {
5125 *p1.at(r) = v1;
5126 *p2.at(r) = v2;
5127 }
5128 }
5129}
5130
5131#[inline]
5138#[allow(unreachable_code)]
5139fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5140 #[cfg(target_arch = "aarch64")]
5141 unsafe {
5142 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5143 }
5144 #[cfg(target_arch = "x86_64")]
5145 unsafe {
5146 if vnni_tiles_enabled() {
5147 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5148 }
5149 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5150 }
5151 let mut acc = 0f32;
5152 for gi in 0..gpr {
5153 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5154 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5155 let mut d = 0i32;
5156 for (k, &b) in tile[2..].iter().enumerate() {
5157 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5158 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5159 }
5160 acc += d as f32 * s;
5161 }
5162 acc
5163}
5164
5165#[cfg(target_arch = "aarch64")]
5166#[target_feature(enable = "neon,dotprod")]
5167unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5168 unsafe {
5171 use core::arch::aarch64::*;
5172 use core::arch::asm;
5173 let lomask = vdupq_n_u8(0x0F);
5174 let eight = vdupq_n_s8(8);
5175 let mut acc = 0f32;
5176 for gi in 0..gpr {
5177 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5178 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5179 let b = vld1q_u8(t.add(2));
5180 let lo = vandq_u8(b, lomask);
5181 let hi = vshrq_n_u8::<4>(b);
5182 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5183 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5184 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5185 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5186 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5187 asm!(
5188 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5189 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5190 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5191 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5192 options(pure, nomem, nostack),
5193 );
5194 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5195 }
5196 acc
5197 }
5198}
5199
5200#[cfg(target_arch = "x86_64")]
5201#[target_feature(enable = "avx2")]
5202unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5203 unsafe {
5205 use core::arch::x86_64::*;
5206 let lomask = _mm_set1_epi8(0x0F);
5207 let eight = _mm256_set1_epi8(8);
5208 let ones = _mm256_set1_epi16(1);
5209 let mut acc = 0f32;
5210 for gi in 0..gpr {
5211 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5212 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5213 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5214 let lo = _mm_and_si128(b, lomask);
5215 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5216 let w = _mm256_sub_epi8(
5217 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5218 eight,
5219 );
5220 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5221 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5222 let d = _mm256_madd_epi16(p16, ones);
5223 let hi128 = _mm256_extracti128_si256::<1>(d);
5224 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5225 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5226 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5227 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5228 }
5229 acc
5230 }
5231}
5232
5233#[cfg(target_arch = "x86_64")]
5237#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5238unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5239 unsafe {
5241 use core::arch::x86_64::*;
5242 let lomask = _mm_set1_epi8(0x0F);
5243 let eight = _mm256_set1_epi8(8);
5244 let mut acc = 0f32;
5245 for gi in 0..gpr {
5246 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5247 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5248 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5249 let lo = _mm_and_si128(b, lomask);
5250 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5251 let w = _mm256_sub_epi8(
5252 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5253 eight,
5254 );
5255 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5256 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5257 acc += d as f32 * s;
5258 }
5259 acc
5260 }
5261}
5262
5263#[cfg(target_arch = "x86_64")]
5268#[target_feature(enable = "avx2,fma")]
5273unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5274 unsafe {
5276 use core::arch::x86_64::*;
5277 let lomask = _mm_set1_epi8(0x0F);
5278 let eight = _mm256_set1_epi8(8);
5279 let ones = _mm256_set1_epi16(1);
5280 let mut f0 = _mm256_setzero_ps();
5293 let mut f1 = _mm256_setzero_ps();
5294 let mut f2 = _mm256_setzero_ps();
5295 let mut f3 = _mm256_setzero_ps();
5296 for gi in 0..gpr {
5297 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5298 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5299 let sv = _mm256_set1_ps(s);
5300 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5301 let lo = _mm_and_si128(bb, lomask);
5302 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5303 let w = _mm256_sub_epi8(
5304 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5305 eight,
5306 );
5307 let aw = _mm256_abs_epi8(w);
5308 let off = gi * GROUP_SIZE;
5309 let dot = |xq: &[i8]| {
5310 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5311 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5312 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5313 };
5314 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5315 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5316 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5317 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5318 }
5319 [
5320 hsum256_ps(f0),
5321 hsum256_ps(f1),
5322 hsum256_ps(f2),
5323 hsum256_ps(f3),
5324 ]
5325 }
5326}
5327
5328#[cfg(target_arch = "x86_64")]
5331#[target_feature(enable = "avx2")]
5332#[inline]
5333unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5334 unsafe {
5336 use core::arch::x86_64::*;
5337 let hi = _mm256_extractf128_ps::<1>(v);
5338 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5339 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5340 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5341 _mm_cvtss_f32(s)
5342 }
5343}
5344
5345#[cfg(target_arch = "x86_64")]
5347#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5348unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5349 unsafe {
5351 use core::arch::x86_64::*;
5352 let lomask = _mm_set1_epi8(0x0F);
5353 let eight = _mm256_set1_epi8(8);
5354 let mut f0 = _mm256_setzero_ps();
5357 let mut f1 = _mm256_setzero_ps();
5358 let mut f2 = _mm256_setzero_ps();
5359 let mut f3 = _mm256_setzero_ps();
5360 for gi in 0..gpr {
5361 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5362 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5363 let sv = _mm256_set1_ps(s);
5364 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5365 let lo = _mm_and_si128(bb, lomask);
5366 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5367 let w = _mm256_sub_epi8(
5368 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5369 eight,
5370 );
5371 let aw = _mm256_abs_epi8(w);
5372 let off = gi * GROUP_SIZE;
5373 let dot = |xq: &[i8]| {
5374 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5375 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5376 _mm256_setzero_si256(),
5377 aw,
5378 _mm256_sign_epi8(x, w),
5379 ))
5380 };
5381 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5382 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5383 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5384 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5385 }
5386 let acc = [
5387 hsum256_ps(f0),
5388 hsum256_ps(f1),
5389 hsum256_ps(f2),
5390 hsum256_ps(f3),
5391 ];
5392 acc
5393 }
5394}
5395
5396#[cfg(target_arch = "aarch64")]
5401#[target_feature(enable = "neon,dotprod")]
5402unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5403 unsafe {
5405 use core::arch::aarch64::*;
5406 use core::arch::asm;
5407 let lomask = vdupq_n_u8(0x0F);
5408 let eight = vdupq_n_s8(8);
5409 let mut acc = [0f32; 4];
5410 for gi in 0..gpr {
5411 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5412 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5413 let b = vld1q_u8(t.add(2));
5414 let lo = vandq_u8(b, lomask);
5415 let hi = vshrq_n_u8::<4>(b);
5416 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5417 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5418 for (k, xq) in xs.iter().enumerate() {
5419 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5420 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5421 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5422 asm!(
5423 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5424 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5425 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5426 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5427 options(pure, nomem, nostack),
5428 );
5429 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5430 }
5431 }
5432 acc
5433 }
5434}
5435
5436#[inline]
5438fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5439 let gi = j / GROUP_SIZE;
5440 let k = j % GROUP_SIZE;
5441 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5442 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5443 let byte = tile[2 + k / 2];
5444 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5445 ((nib as i32 - 8) as f32, s)
5446}
5447
5448#[inline]
5451fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5452 let mut acc = 0f32;
5453 for gi in 0..gpr {
5454 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5455 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5456 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5457 let mut ga = 0f32;
5458 for (k, &b) in tile[2..].iter().enumerate() {
5459 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5460 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5461 }
5462 acc += ga * s;
5463 }
5464 acc
5465}
5466
5467struct Q4tpView<'a> {
5471 nib: &'a [u8],
5472 params: &'a [u8],
5473 codes: &'a [u8],
5474 stride: usize,
5475 zero_rung: bool,
5477}
5478
5479impl<'a> Q4tpView<'a> {
5480 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5481 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5482 Self {
5483 nib: &bytes[..params_off],
5484 params: &bytes[params_off..codes_off],
5485 codes: &bytes[codes_off..],
5486 stride,
5487 zero_rung: false,
5488 }
5489 }
5490
5491 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5493 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5494 Self {
5495 nib: &bytes[..params_off],
5496 params: &bytes[params_off..codes_off],
5497 codes: &bytes[codes_off..],
5498 stride,
5499 zero_rung: true,
5500 }
5501 }
5502
5503 #[inline]
5519 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5520 let tab = if self.zero_rung {
5521 q2tp_ladder(self.params, r)
5522 } else {
5523 q4tp_ladder(self.params, r)
5524 };
5525 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5526 let out = &mut out[..gpr];
5527 let mut chunks = out.chunks_exact_mut(8);
5528 let mut ci = 0usize;
5529 for c in &mut chunks {
5530 let w = u64::from(codes[ci])
5531 | u64::from(codes[ci + 1]) << 8
5532 | u64::from(codes[ci + 2]) << 16
5533 | u64::from(codes[ci + 3]) << 24
5534 | u64::from(codes[ci + 4]) << 32;
5535 for (k, o) in c.iter_mut().enumerate() {
5536 *o = tab[((w >> (5 * k)) & 31) as usize];
5537 }
5538 ci += 5;
5539 }
5540 let tail = &codes[ci..];
5543 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5544 *o = tab[q4tp_code(tail, k)];
5545 }
5546 }
5547}
5548
5549#[inline]
5550fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5551 #[cfg(target_arch = "aarch64")]
5552 unsafe {
5553 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5554 }
5555 #[cfg(target_arch = "x86_64")]
5556 unsafe {
5557 if vnni_tiles_enabled() {
5558 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5559 }
5560 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5561 }
5562 #[allow(unreachable_code)]
5563 {
5564 let mut acc = 0f32;
5565 for gi in 0..gpr {
5566 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5567 let s = scales[gi];
5568 let mut d = 0i32;
5569 for (k, &b) in tile.iter().enumerate() {
5570 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5571 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5572 }
5573 acc += d as f32 * s;
5574 }
5575 acc
5576 }
5577}
5578
5579#[cfg(target_arch = "aarch64")]
5582#[target_feature(enable = "neon,dotprod")]
5583unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5584 unsafe {
5587 use core::arch::aarch64::*;
5588 use core::arch::asm;
5589 let lomask = vdupq_n_u8(0x0F);
5590 let eight = vdupq_n_s8(8);
5591 let mut acc = 0f32;
5592 for gi in 0..gpr {
5593 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5594 let s = *scales.get_unchecked(gi);
5595 let b = vld1q_u8(t);
5596 let lo = vandq_u8(b, lomask);
5597 let hi = vshrq_n_u8::<4>(b);
5598 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5599 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5600 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5601 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5602 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5603 asm!(
5604 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5605 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5606 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5607 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5608 options(pure, nomem, nostack),
5609 );
5610 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5611 }
5612 acc
5613 }
5614}
5615
5616#[cfg(target_arch = "x86_64")]
5617#[target_feature(enable = "avx2")]
5618unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5619 unsafe {
5621 use core::arch::x86_64::*;
5622 let lomask = _mm_set1_epi8(0x0F);
5623 let eight = _mm256_set1_epi8(8);
5624 let ones = _mm256_set1_epi16(1);
5625 let mut acc = 0f32;
5626 for gi in 0..gpr {
5627 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5628 let s = *scales.get_unchecked(gi);
5629 let b = _mm_loadu_si128(t as *const __m128i);
5630 let lo = _mm_and_si128(b, lomask);
5631 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5632 let w = _mm256_sub_epi8(
5633 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5634 eight,
5635 );
5636 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5637 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5638 let d = _mm256_madd_epi16(p16, ones);
5639 let hi128 = _mm256_extracti128_si256::<1>(d);
5640 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5641 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5642 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5643 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5644 }
5645 acc
5646 }
5647}
5648
5649#[cfg(target_arch = "x86_64")]
5652#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5653unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5654 unsafe {
5656 use core::arch::x86_64::*;
5657 let lomask = _mm_set1_epi8(0x0F);
5658 let eight = _mm256_set1_epi8(8);
5659 let mut acc = 0f32;
5660 for gi in 0..gpr {
5661 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5662 let s = *scales.get_unchecked(gi);
5663 let b = _mm_loadu_si128(t as *const __m128i);
5664 let lo = _mm_and_si128(b, lomask);
5665 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5666 let w = _mm256_sub_epi8(
5667 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5668 eight,
5669 );
5670 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5671 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5672 }
5673 acc
5674 }
5675}
5676
5677#[inline]
5680fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5681 #[cfg(target_arch = "x86_64")]
5682 if avx2_enabled() {
5683 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5685 }
5686 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5687}
5688
5689#[inline]
5690fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5691 let mut acc = 0f32;
5692 for gi in 0..gpr {
5693 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5694 let s = scales[gi];
5695 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5696 let mut ga = 0f32;
5697 for (k, &b) in tile.iter().enumerate() {
5698 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5699 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5700 }
5701 acc += ga * s;
5702 }
5703 acc
5704}
5705
5706#[cfg(target_arch = "x86_64")]
5710#[target_feature(enable = "avx2")]
5711unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5712 unsafe {
5715 use core::arch::x86_64::*;
5716 let mask = _mm_set1_epi8(15);
5717 let eight = _mm_set1_epi8(8);
5718 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5719 let mut acc = 0.0f32;
5720 for gi in 0..gpr {
5721 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5722 let lo = _mm_and_si128(packed, mask);
5723 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5724 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5725 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5726 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5727 let a = _mm256_mul_ps(
5728 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5729 _mm256_loadu_ps(xp),
5730 );
5731 let b = _mm256_mul_ps(
5732 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5733 _mm256_loadu_ps(xp.add(8)),
5734 );
5735 let c = _mm256_mul_ps(
5736 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5737 _mm256_loadu_ps(xp.add(16)),
5738 );
5739 let d = _mm256_mul_ps(
5740 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5741 _mm256_loadu_ps(xp.add(24)),
5742 );
5743 let mut pairs = [0.0f32; 16];
5744 _mm256_storeu_ps(
5745 pairs.as_mut_ptr(),
5746 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5747 );
5748 _mm256_storeu_ps(
5749 pairs.as_mut_ptr().add(8),
5750 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5751 );
5752 let mut ga = 0.0f32;
5753 for v in pairs {
5754 ga += v;
5755 }
5756 acc += ga * scales[gi];
5757 }
5758 acc
5759 }
5760}
5761
5762#[inline]
5765fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5766 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5767 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5768 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5769 ((n as i32 - 8) as f32, scales[gi])
5770}
5771
5772fn q4tp_matvec(
5774 bytes: &[u8],
5775 x: &[f32],
5776 rows: usize,
5777 cols: usize,
5778 out: &mut [f32],
5779 pool: Option<&Pool>,
5780) {
5781 debug_assert_eq!(out.len(), rows);
5782 let gpr = cols / GROUP_SIZE;
5783 let v = Q4tpView::new(bytes, rows, cols);
5784 let out_addr = SendMut(out.as_mut_ptr());
5785 if a8w8_enabled() {
5786 let act = split_act(x);
5787 let run = |start: usize, end: usize| {
5788 with_krow(gpr, |sc| {
5790 for r in start..end {
5791 v.scales_into(r, gpr, sc);
5792 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5793 for &(j, xv) in &act.outliers {
5794 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5795 acc += w * s * xv;
5796 }
5797 unsafe { *out_addr.at(r) = acc };
5799 }
5800 })
5801 };
5802 dispatch_rows(pool, rows, &run);
5803 return;
5804 }
5805 let run = |start: usize, end: usize| {
5806 with_krow(gpr, |sc| {
5807 for r in start..end {
5808 v.scales_into(r, gpr, sc);
5809 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5811 }
5812 })
5813 };
5814 dispatch_rows(pool, rows, &run);
5815}
5816
5817#[allow(clippy::too_many_arguments)]
5820fn q4tp_matvec2(
5821 bytes: &[u8],
5822 x1: &[f32],
5823 x2: &[f32],
5824 rows: usize,
5825 cols: usize,
5826 o1: &mut [f32],
5827 o2: &mut [f32],
5828 pool: Option<&Pool>,
5829) {
5830 let gpr = cols / GROUP_SIZE;
5831 let v = Q4tpView::new(bytes, rows, cols);
5832 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5833 let run = |start: usize, end: usize| {
5834 let mut sc = vec![0f32; gpr];
5835 for r in start..end {
5836 v.scales_into(r, gpr, &mut sc);
5837 unsafe {
5839 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5840 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5841 }
5842 }
5843 };
5844 dispatch_rows(pool, rows, &run);
5845}
5846
5847#[inline]
5850fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5851 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5852 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5853 let c = (byte >> (2 * (k % 4))) & 3;
5854 (c as f32 - 1.5, scales[gi])
5855}
5856
5857#[cfg(target_arch = "x86_64")]
5858const Q2TP_DECODE_U32: [u32; 256] = {
5859 let mut tab = [0u32; 256];
5860 let mut b = 0usize;
5861 while b < 256 {
5862 tab[b] = ((b as u32) & 3)
5863 | ((((b as u32) >> 2) & 3) << 8)
5864 | ((((b as u32) >> 4) & 3) << 16)
5865 | ((((b as u32) >> 6) & 3) << 24);
5866 b += 1;
5867 }
5868 tab
5869};
5870
5871#[cfg(target_arch = "x86_64")]
5875#[target_feature(enable = "avx2")]
5876unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5877 use core::arch::x86_64::*;
5878 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5879 let codes = _mm256_setr_epi32(
5880 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5881 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5882 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5883 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5884 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5885 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5886 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5887 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5888 );
5889 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5890 let pair = _mm256_maddubs_epi16(codes, xv);
5891 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5892 let sum128 = _mm_add_epi32(
5893 _mm256_castsi256_si128(quad),
5894 _mm256_extracti128_si256(quad, 1),
5895 );
5896 let sum64 = _mm_hadd_epi32(sum128, sum128);
5897 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5898}
5899
5900#[cfg(target_arch = "x86_64")]
5901#[inline]
5902fn q2tp_avx2_enabled() -> bool {
5903 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5904 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5905}
5906
5907#[inline]
5914fn dot_q2tp_row_i8(
5915 chunks: &[u8],
5916 r: usize,
5917 gpr: usize,
5918 xq: &[i8],
5919 gsum: &[i32],
5920 scales: &[f32],
5921) -> f32 {
5922 let mut acc = 0f32;
5923 let base = r * gpr * Q2TP_CHUNK;
5924 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5925 let mut codes = [0i8; GROUP_SIZE];
5926 #[cfg(target_arch = "x86_64")]
5927 let avx2 = q2tp_avx2_enabled();
5932 for gi in 0..gpr {
5933 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5934 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5935 #[cfg(target_arch = "aarch64")]
5936 let dot = unsafe {
5942 use core::arch::aarch64::*;
5943 let b = vld1_u8(ch.as_ptr());
5944 let three = vdup_n_u8(3);
5945 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5946 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5947 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5948 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5949 let x4 = vld4_s8(xg.as_ptr());
5950 let mut acc4 = vdupq_n_s32(0);
5951 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5952 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5953 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5954 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5955 vaddvq_s32(acc4)
5956 };
5957 #[cfg(target_arch = "x86_64")]
5958 let dot: i32 = if avx2 {
5959 unsafe { q2tp_code_dot_avx2(ch, xg) }
5962 } else {
5963 ch.iter()
5964 .enumerate()
5965 .map(|(k, &b)| {
5966 ((b & 3) as i32) * xg[k * 4] as i32
5967 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5968 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5969 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5970 })
5971 .sum()
5972 };
5973 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5974 let dot: i32 = {
5975 for (k, &b) in ch.iter().enumerate() {
5976 codes[k * 4] = (b & 3) as i8;
5977 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5978 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5979 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5980 }
5981 codes
5982 .iter()
5983 .zip(xg)
5984 .map(|(&c, &x)| c as i32 * x as i32)
5985 .sum()
5986 };
5987 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5988 }
5989 acc
5990}
5991
5992fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5996 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5997}
5998
5999#[inline]
6003fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6004 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6005}
6006
6007#[inline]
6008fn q2tp_row_exact_center(
6009 chunks: &[u8],
6010 r: usize,
6011 gpr: usize,
6012 x: &[f32],
6013 scales: &[f32],
6014 center: f32,
6015) -> f32 {
6016 let mut acc = 0f32;
6017 for gi in 0..gpr {
6018 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6019 let s = scales[gi];
6020 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6021 let mut g = 0f32;
6022 for (k, &b) in ch.iter().enumerate() {
6023 g += ((b & 3) as f32 - center) * xb[k * 4]
6024 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6025 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6026 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6027 }
6028 acc += s * g;
6029 }
6030 acc
6031}
6032
6033fn q2tp_matvec(
6034 bytes: &[u8],
6035 x: &[f32],
6036 rows: usize,
6037 cols: usize,
6038 out: &mut [f32],
6039 pool: Option<&Pool>,
6040) {
6041 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6042}
6043
6044fn q2tp_affine_matvec(
6045 bytes: &[u8],
6046 x: &[f32],
6047 rows: usize,
6048 cols: usize,
6049 out: &mut [f32],
6050 pool: Option<&Pool>,
6051) {
6052 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6053}
6054
6055fn q2tp_matvec_mode(
6056 bytes: &[u8],
6057 x: &[f32],
6058 rows: usize,
6059 cols: usize,
6060 out: &mut [f32],
6061 pool: Option<&Pool>,
6062 affine: bool,
6063) {
6064 debug_assert_eq!(out.len(), rows);
6065 let gpr = cols / GROUP_SIZE;
6066 let v = Q4tpView::new_q2(bytes, rows, cols);
6067 let out_addr = SendMut(out.as_mut_ptr());
6068 if !affine && a8w8_enabled() {
6073 let act = split_act(x);
6074 let gsum = q1_group_sums(&act.xq, gpr);
6075 let (act, gsum) = (&act, &gsum);
6076 let run = move |start: usize, end: usize| {
6077 with_krow(gpr, |sc| {
6078 for r in start..end {
6079 v.scales_into(r, gpr, sc);
6080 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6081 for &(j, xv) in &act.outliers {
6082 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6083 acc += w * s * xv;
6084 }
6085 unsafe { *out_addr.at(r) = acc };
6087 }
6088 })
6089 };
6090 dispatch_rows(pool, rows, &run);
6091 return;
6092 }
6093 let run = |start: usize, end: usize| {
6094 with_krow(gpr, |sc| {
6095 for r in start..end {
6096 v.scales_into(r, gpr, sc);
6097 unsafe {
6099 *out_addr.at(r) = if affine {
6100 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6101 } else {
6102 q2tp_row_exact(v.nib, r, gpr, x, sc)
6103 }
6104 };
6105 }
6106 })
6107 };
6108 dispatch_rows(pool, rows, &run);
6109}
6110
6111#[allow(clippy::too_many_arguments)]
6113fn q2tp_matvec2(
6114 bytes: &[u8],
6115 x1: &[f32],
6116 x2: &[f32],
6117 rows: usize,
6118 cols: usize,
6119 o1: &mut [f32],
6120 o2: &mut [f32],
6121 pool: Option<&Pool>,
6122) {
6123 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6124}
6125
6126#[allow(clippy::too_many_arguments)]
6127fn q2tp_affine_matvec2(
6128 bytes: &[u8],
6129 x1: &[f32],
6130 x2: &[f32],
6131 rows: usize,
6132 cols: usize,
6133 o1: &mut [f32],
6134 o2: &mut [f32],
6135 pool: Option<&Pool>,
6136) {
6137 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6138}
6139
6140#[allow(clippy::too_many_arguments)]
6141fn q2tp_matvec2_mode(
6142 bytes: &[u8],
6143 x1: &[f32],
6144 x2: &[f32],
6145 rows: usize,
6146 cols: usize,
6147 o1: &mut [f32],
6148 o2: &mut [f32],
6149 pool: Option<&Pool>,
6150 affine: bool,
6151) {
6152 let gpr = cols / GROUP_SIZE;
6153 let v = Q4tpView::new_q2(bytes, rows, cols);
6154 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6155 let run = |start: usize, end: usize| {
6156 let mut sc = vec![0f32; gpr];
6157 for r in start..end {
6158 v.scales_into(r, gpr, &mut sc);
6159 unsafe {
6161 *p1.at(r) = if affine {
6162 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6163 } else {
6164 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6165 };
6166 *p2.at(r) = if affine {
6167 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6168 } else {
6169 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6170 };
6171 }
6172 }
6173 };
6174 dispatch_rows(pool, rows, &run);
6175}
6176
6177pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6184 let gpr = cols / GROUP_SIZE;
6189 let v = Q4tpView::new_q2(bytes, rows, cols);
6190 with_krow(gpr, |sc| {
6191 for r in 0..rows {
6192 v.scales_into(r, gpr, sc);
6193 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6194 }
6195 });
6196}
6197
6198pub fn q2tp_affine_matvec_for_test(
6201 bytes: &[u8],
6202 x: &[f32],
6203 rows: usize,
6204 cols: usize,
6205 out: &mut [f32],
6206) {
6207 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6208}
6209
6210pub fn q2tp_matmat_for_test(
6211 bytes: &[u8],
6212 xs_all: &[f32],
6213 b: usize,
6214 rows: usize,
6215 cols: usize,
6216 out: &mut [f32],
6217) {
6218 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6219}
6220
6221fn q2tp_matmat(
6222 bytes: &[u8],
6223 xs_all: &[f32],
6224 b: usize,
6225 rows: usize,
6226 cols: usize,
6227 out: &mut [f32],
6228 pool: Option<&Pool>,
6229) {
6230 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6231}
6232
6233fn q2tp_affine_matmat(
6234 bytes: &[u8],
6235 xs_all: &[f32],
6236 b: usize,
6237 rows: usize,
6238 cols: usize,
6239 out: &mut [f32],
6240 pool: Option<&Pool>,
6241) {
6242 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6243}
6244
6245fn q2tp_matmat_mode(
6246 bytes: &[u8],
6247 xs_all: &[f32],
6248 b: usize,
6249 rows: usize,
6250 cols: usize,
6251 out: &mut [f32],
6252 pool: Option<&Pool>,
6253 affine: bool,
6254) {
6255 debug_assert_eq!(out.len(), b * rows);
6256 let gpr = cols / GROUP_SIZE;
6257 let v = Q4tpView::new_q2(bytes, rows, cols);
6258 let out_addr = SendMut(out.as_mut_ptr());
6259 let run = |start: usize, end: usize| {
6260 let mut sc = vec![0f32; gpr];
6261 for r in start..end {
6262 v.scales_into(r, gpr, &mut sc);
6263 for bi in 0..b {
6264 let x = &xs_all[bi * cols..(bi + 1) * cols];
6265 unsafe {
6267 *out_addr.at(bi * rows + r) = if affine {
6268 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6269 } else {
6270 q2tp_row_exact(v.nib, r, gpr, x, &sc)
6271 }
6272 };
6273 }
6274 }
6275 };
6276 dispatch_rows(pool, rows, &run);
6277}
6278
6279#[cfg(target_arch = "aarch64")]
6287#[target_feature(enable = "neon,dotprod")]
6288unsafe fn dot_q4tp_row_1x4_sdot_v1(
6289 nib: &[u8],
6290 r: usize,
6291 gpr: usize,
6292 xs: [&[i8]; 4],
6293 scales: &[f32],
6294) -> [f32; 4] {
6295 unsafe {
6296 use core::arch::aarch64::*;
6297 use core::arch::asm;
6298 let lomask = vdupq_n_u8(0x0F);
6299 let eight = vdupq_n_s8(8);
6300 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6301 for gi in 0..gpr {
6302 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6303 let s = *scales.get_unchecked(gi);
6304 let bb = vld1q_u8(t);
6305 let lo = vandq_u8(bb, lomask);
6306 let hi = vshrq_n_u8::<4>(bb);
6307 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6308 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6309 let mut d = [0f32; 4];
6310 for (k, dk) in d.iter_mut().enumerate() {
6311 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6312 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6313 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6314 asm!(
6315 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6316 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6317 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6318 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6319 options(pure, nomem, nostack),
6320 );
6321 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6322 }
6323 f0 += d[0];
6324 f1 += d[1];
6325 f2 += d[2];
6326 f3 += d[3];
6327 }
6328 [f0, f1, f2, f3]
6329 }
6330}
6331
6332#[allow(dead_code)]
6339static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6340
6341#[cfg(test)]
6344static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6345
6346#[cfg(target_arch = "x86_64")]
6352fn q4tp_blocked_x86() -> bool {
6353 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6354 1 => false,
6355 2 => avx512vnni_enabled(),
6360 _ => blocked_enabled() && avx512vnni_enabled(),
6364 }
6365}
6366
6367#[cfg(target_arch = "aarch64")]
6369#[allow(dead_code)]
6370fn q4tp_v1() -> bool {
6371 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6372 1 => true,
6373 2 => false,
6374 _ => {
6375 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6376 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6377 }
6378 }
6379}
6380
6381#[cfg(target_arch = "x86_64")]
6393#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6394unsafe fn dot_q4tp_2x8_avx512(
6395 nib: &[u8],
6396 r0: usize,
6397 gpr: usize,
6398 xs: [&[i8]; 8],
6399 sc0: &[f32],
6400 sc1: &[f32],
6401) -> [[f32; 8]; 2] {
6402 unsafe {
6405 use core::arch::x86_64::*;
6406 let lomask = _mm256_set1_epi8(0x0F);
6407 let eight = _mm256_set1_epi8(8);
6408 let zero = _mm512_setzero_si512();
6409 let mut v0 = [_mm512_setzero_ps(); 8];
6410 let mut v1 = [_mm512_setzero_ps(); 8];
6411 let pairs = gpr / 2;
6412 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6413 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6414 let bb = _mm256_loadu_si256(t as *const __m256i);
6415 let lo = _mm256_and_si256(bb, lomask);
6416 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6417 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6418 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6419 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6420 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6421 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6422 };
6423 for gp in 0..pairs {
6424 let gi = gp * 2;
6425 let (wa0, neg0) = unpack(r0, gi);
6426 let (wa1, neg1) = unpack(r0 + 1, gi);
6427 let off = gi * GROUP_SIZE;
6428 let sv = |sc: &[f32]| {
6429 _mm512_insertf32x8::<1>(
6430 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6431 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6432 )
6433 };
6434 let s0 = sv(sc0);
6435 let s1 = sv(sc1);
6436 for k in 0..8 {
6437 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6438 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6439 zero,
6440 wa0,
6441 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6442 ));
6443 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6444 zero,
6445 wa1,
6446 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6447 ));
6448 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6449 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6450 }
6451 }
6452 let mut acc = [[0f32; 8]; 2];
6453 for k in 0..8 {
6454 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6455 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6456 }
6457 if gpr % 2 == 1 {
6458 let off = (gpr - 1) * GROUP_SIZE;
6459 for j in off..off + GROUP_SIZE {
6460 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6461 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6462 for k in 0..8 {
6463 let x = *xs[k].get_unchecked(j) as f32;
6464 acc[0][k] += w0 * sa * x;
6465 acc[1][k] += w1 * sb * x;
6466 }
6467 }
6468 }
6469 acc
6470 }
6471}
6472
6473#[cfg(target_arch = "x86_64")]
6478#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6479unsafe fn dot_q4tp_row_1x8_avx512(
6480 nib: &[u8],
6481 r: usize,
6482 gpr: usize,
6483 xs: [&[i8]; 8],
6484 scales: &[f32],
6485) -> [f32; 8] {
6486 unsafe {
6488 use core::arch::x86_64::*;
6489 let lomask = _mm256_set1_epi8(0x0F);
6490 let eight = _mm256_set1_epi8(8);
6491 let zero = _mm512_setzero_si512();
6492 let (mut v0, mut v1, mut v2, mut v3) = (
6493 _mm512_setzero_ps(),
6494 _mm512_setzero_ps(),
6495 _mm512_setzero_ps(),
6496 _mm512_setzero_ps(),
6497 );
6498 let (mut v4, mut v5, mut v6, mut v7) = (
6499 _mm512_setzero_ps(),
6500 _mm512_setzero_ps(),
6501 _mm512_setzero_ps(),
6502 _mm512_setzero_ps(),
6503 );
6504 let pairs = gpr / 2;
6505 for gp in 0..pairs {
6506 let gi = gp * 2;
6507 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6508 let bb = _mm256_loadu_si256(t as *const __m256i);
6509 let lo = _mm256_and_si256(bb, lomask);
6510 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6511 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6516 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6517 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6518 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6519 let wabs = _mm512_abs_epi8(w);
6520 let neg = _mm512_movepi8_mask(w);
6521 let off = gi * GROUP_SIZE;
6522 let sv = _mm512_insertf32x8::<1>(
6523 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6524 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6525 );
6526 let dot = |x: &[i8]| -> __m512 {
6527 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6528 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6529 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6530 };
6531 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6532 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6533 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6534 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6535 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6536 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6537 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6538 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6539 }
6540 let mut acc = [
6541 _mm512_reduce_add_ps(v0),
6542 _mm512_reduce_add_ps(v1),
6543 _mm512_reduce_add_ps(v2),
6544 _mm512_reduce_add_ps(v3),
6545 _mm512_reduce_add_ps(v4),
6546 _mm512_reduce_add_ps(v5),
6547 _mm512_reduce_add_ps(v6),
6548 _mm512_reduce_add_ps(v7),
6549 ];
6550 if gpr % 2 == 1 {
6553 let off = (gpr - 1) * GROUP_SIZE;
6554 for j in off..off + GROUP_SIZE {
6555 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6556 let ws = w * s;
6557 for k in 0..8 {
6558 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6559 }
6560 }
6561 }
6562 acc
6563 }
6564}
6565
6566#[cfg(target_arch = "x86_64")]
6579#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6580unsafe fn dot_q4tp_row_1x4_avx512(
6581 nib: &[u8],
6582 r: usize,
6583 gpr: usize,
6584 xs: [&[i8]; 4],
6585 scales: &[f32],
6586) -> [f32; 4] {
6587 unsafe {
6589 use core::arch::x86_64::*;
6590 let lomask = _mm256_set1_epi8(0x0F);
6591 let eight = _mm256_set1_epi8(8);
6592 let zero = _mm512_setzero_si512();
6593 let (mut v0, mut v1, mut v2, mut v3) = (
6594 _mm512_setzero_ps(),
6595 _mm512_setzero_ps(),
6596 _mm512_setzero_ps(),
6597 _mm512_setzero_ps(),
6598 );
6599 let pairs = gpr / 2;
6600 for gp in 0..pairs {
6601 let gi = gp * 2;
6602 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6603 let bb = _mm256_loadu_si256(t as *const __m256i);
6604 let lo = _mm256_and_si256(bb, lomask);
6605 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6606 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6611 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6612 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6613 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6614 let wabs = _mm512_abs_epi8(w);
6615 let neg = _mm512_movepi8_mask(w);
6616 let off = gi * GROUP_SIZE;
6617 let sv = _mm512_insertf32x8::<1>(
6618 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6619 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6620 );
6621 let dot = |x: &[i8]| -> __m512 {
6622 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6623 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6624 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6625 };
6626 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6627 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6628 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6629 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6630 }
6631 let mut acc = [
6632 _mm512_reduce_add_ps(v0),
6633 _mm512_reduce_add_ps(v1),
6634 _mm512_reduce_add_ps(v2),
6635 _mm512_reduce_add_ps(v3),
6636 ];
6637 if gpr % 2 == 1 {
6640 let off = (gpr - 1) * GROUP_SIZE;
6641 for j in off..off + GROUP_SIZE {
6642 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6643 let ws = w * s;
6644 for k in 0..4 {
6645 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6646 }
6647 }
6648 }
6649 acc
6650 }
6651}
6652
6653#[cfg(target_arch = "aarch64")]
6657#[target_feature(enable = "neon,dotprod")]
6658unsafe fn dot_q4tp_row_1x4_sdot(
6659 nib: &[u8],
6660 r: usize,
6661 gpr: usize,
6662 xs: [&[i8]; 4],
6663 scales: &[f32],
6664) -> [f32; 4] {
6665 unsafe {
6667 use core::arch::aarch64::*;
6668 use core::arch::asm;
6669 let lomask = vdupq_n_u8(0x0F);
6670 let eight = vdupq_n_s8(8);
6671 let (mut v0, mut v1, mut v2, mut v3) = (
6687 vdupq_n_f32(0.0),
6688 vdupq_n_f32(0.0),
6689 vdupq_n_f32(0.0),
6690 vdupq_n_f32(0.0),
6691 );
6692 for gi in 0..gpr {
6693 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6694 let s = *scales.get_unchecked(gi);
6695 let bb = vld1q_u8(t);
6696 let lo = vandq_u8(bb, lomask);
6697 let hi = vshrq_n_u8::<4>(bb);
6698 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6699 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6700 let off = gi * GROUP_SIZE;
6701 let dot4 = |x: &[i8]| -> int32x4_t {
6702 let x0 = vld1q_s8(x.as_ptr().add(off));
6703 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6704 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6705 asm!(
6706 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6707 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6708 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6709 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6710 options(pure, nomem, nostack),
6711 );
6712 vaddq_s32(a0, a1)
6713 };
6714 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6715 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6716 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6717 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6718 }
6719 [
6720 vaddvq_f32(v0),
6721 vaddvq_f32(v1),
6722 vaddvq_f32(v2),
6723 vaddvq_f32(v3),
6724 ]
6725 }
6726}
6727
6728fn q4tp_matmat(
6736 bytes: &[u8],
6737 xs_all: &[f32],
6738 b: usize,
6739 rows: usize,
6740 cols: usize,
6741 out: &mut [f32],
6742 pool: Option<&Pool>,
6743) {
6744 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6745}
6746
6747#[allow(clippy::too_many_arguments)]
6751fn q4tp_matmat_with(
6752 bytes: &[u8],
6753 xs_all: &[f32],
6754 b: usize,
6755 rows: usize,
6756 cols: usize,
6757 out: &mut [f32],
6758 pool: Option<&Pool>,
6759 exact: bool,
6760) {
6761 debug_assert_eq!(out.len(), b * rows);
6762 let gpr = cols / GROUP_SIZE;
6763 let v = Q4tpView::new(bytes, rows, cols);
6764
6765 #[cfg(target_os = "macos")]
6769 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6770 dequant_matmat_accel(
6771 &|r, dst| {
6772 let mut sc = [0f32; 32];
6773 let mut scv;
6774 let s: &[f32] = if gpr <= 32 {
6775 v.scales_into(r, gpr, &mut sc);
6776 &sc[..gpr]
6777 } else {
6778 scv = vec![0f32; gpr];
6779 v.scales_into(r, gpr, &mut scv);
6780 &scv
6781 };
6782 for gi in 0..gpr {
6783 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6784 for (k, &bb) in tile.iter().enumerate() {
6785 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6786 dst[gi * GROUP_SIZE + k * 2 + 1] =
6787 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6788 }
6789 }
6790 },
6791 xs_all,
6792 b,
6793 rows,
6794 cols,
6795 out,
6796 pool,
6797 );
6798 return;
6799 }
6800
6801 let out_addr = SendMut(out.as_mut_ptr());
6802 if a8w8_enabled() {
6803 let acts: Vec<SplitAct> = (0..b)
6804 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6805 .collect();
6806 let acts = &acts;
6807 #[cfg(target_arch = "aarch64")]
6814 let blocked_ok = sdot_enabled() && blocked_enabled();
6815 #[cfg(target_arch = "x86_64")]
6822 let blocked_ok = q4tp_blocked_x86() && !exact;
6823 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6824 let blocked_ok = {
6825 let _ = exact;
6826 false
6827 };
6828 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6838 .ok()
6839 .and_then(|v| v.parse().ok())
6840 .filter(|v| *v > 0)
6841 .unwrap_or(256);
6842 let run = |start: usize, end: usize| {
6843 for abase in (0..acts.len()).step_by(panel_cols) {
6844 let alen = (acts.len() - abase).min(panel_cols);
6845 let mut sc = vec![0f32; gpr];
6846 #[cfg(target_arch = "x86_64")]
6847 let mut r_lo = start;
6848 #[cfg(target_arch = "x86_64")]
6849 if blocked_ok && alen >= 8 {
6850 let mut sc1 = vec![0f32; gpr];
6851 while r_lo + 2 <= end {
6852 v.scales_into(r_lo, gpr, &mut sc);
6853 v.scales_into(r_lo + 1, gpr, &mut sc1);
6854 let mut bi = 0usize;
6855 while bi + 8 <= alen {
6856 let xs = [
6857 acts[abase + bi].xq.as_slice(),
6858 acts[abase + bi + 1].xq.as_slice(),
6859 acts[abase + bi + 2].xq.as_slice(),
6860 acts[abase + bi + 3].xq.as_slice(),
6861 acts[abase + bi + 4].xq.as_slice(),
6862 acts[abase + bi + 5].xq.as_slice(),
6863 acts[abase + bi + 6].xq.as_slice(),
6864 acts[abase + bi + 7].xq.as_slice(),
6865 ];
6866 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6867 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6868 for k in 0..8 {
6869 let act = &acts[abase + bi + k];
6870 let mut acc = dr[k] * act.sx;
6871 for &(j, xv) in &act.outliers {
6872 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6873 acc += w * s * xv;
6874 }
6875 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6877 }
6878 }
6879 bi += 8;
6880 }
6881 for row in [r_lo, r_lo + 1] {
6884 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6885 for b2 in bi..alen {
6886 let act = &acts[abase + b2];
6887 let xs4 = [
6888 act.xq.as_slice(),
6889 act.xq.as_slice(),
6890 act.xq.as_slice(),
6891 act.xq.as_slice(),
6892 ];
6893 let d =
6894 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6895 let mut acc = d[0] * act.sx;
6896 for &(j, xv) in &act.outliers {
6897 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6898 acc += w * s * xv;
6899 }
6900 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6902 }
6903 }
6904 r_lo += 2;
6905 }
6906 }
6907 #[cfg(target_arch = "x86_64")]
6908 let row_start = r_lo;
6909 #[cfg(not(target_arch = "x86_64"))]
6910 let row_start = start;
6911 for r in row_start..end {
6912 v.scales_into(r, gpr, &mut sc);
6913 let mut bi = 0usize;
6914 #[cfg(target_arch = "x86_64")]
6915 if blocked_ok {
6916 while bi + 8 <= alen {
6917 let xs = [
6918 acts[abase + bi].xq.as_slice(),
6919 acts[abase + bi + 1].xq.as_slice(),
6920 acts[abase + bi + 2].xq.as_slice(),
6921 acts[abase + bi + 3].xq.as_slice(),
6922 acts[abase + bi + 4].xq.as_slice(),
6923 acts[abase + bi + 5].xq.as_slice(),
6924 acts[abase + bi + 6].xq.as_slice(),
6925 acts[abase + bi + 7].xq.as_slice(),
6926 ];
6927 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6928 for k in 0..8 {
6929 let act = &acts[abase + bi + k];
6930 let mut acc = d[k] * act.sx;
6931 for &(j, xv) in &act.outliers {
6932 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6933 acc += w * s * xv;
6934 }
6935 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6937 }
6938 bi += 8;
6939 }
6940 while bi + 4 <= alen {
6941 let xs = [
6942 acts[abase + bi].xq.as_slice(),
6943 acts[abase + bi + 1].xq.as_slice(),
6944 acts[abase + bi + 2].xq.as_slice(),
6945 acts[abase + bi + 3].xq.as_slice(),
6946 ];
6947 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6948 for k in 0..4 {
6949 let act = &acts[abase + bi + k];
6950 let mut acc = d[k] * act.sx;
6951 for &(j, xv) in &act.outliers {
6952 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6953 acc += w * s * xv;
6954 }
6955 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6957 }
6958 bi += 4;
6959 }
6960 }
6961 #[cfg(target_arch = "aarch64")]
6962 if blocked_ok {
6963 while bi + 4 <= alen {
6964 let xs = [
6965 acts[abase + bi].xq.as_slice(),
6966 acts[abase + bi + 1].xq.as_slice(),
6967 acts[abase + bi + 2].xq.as_slice(),
6968 acts[abase + bi + 3].xq.as_slice(),
6969 ];
6970 let d = unsafe {
6971 if exact || q4tp_v1() {
6972 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6973 } else {
6974 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6975 }
6976 };
6977 for k in 0..4 {
6978 let act = &acts[abase + bi + k];
6979 let mut acc = d[k] * act.sx;
6980 for &(j, xv) in &act.outliers {
6981 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6982 acc += w * s * xv;
6983 }
6984 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6986 }
6987 bi += 4;
6988 }
6989 }
6990 let _ = blocked_ok;
6991 while bi < alen {
6992 let act = &acts[abase + bi];
6993 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6994 for &(j, xv) in &act.outliers {
6995 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6996 acc += w * s * xv;
6997 }
6998 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7000 bi += 1;
7001 }
7002 }
7003 }
7004 };
7005 dispatch_rows(pool, rows, &run);
7006 return;
7007 }
7008
7009 let run = |start: usize, end: usize| {
7010 let mut sc = vec![0f32; gpr];
7011 for r in start..end {
7012 v.scales_into(r, gpr, &mut sc);
7013 for bi in 0..b {
7014 let x = &xs_all[bi * cols..(bi + 1) * cols];
7015 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7017 }
7018 }
7019 };
7020 dispatch_rows(pool, rows, &run);
7021}
7022
7023fn q4t_matvec(
7025 bytes: &[u8],
7026 x: &[f32],
7027 rows: usize,
7028 cols: usize,
7029 out: &mut [f32],
7030 pool: Option<&Pool>,
7031) {
7032 debug_assert_eq!(out.len(), rows);
7033 let gpr = cols / GROUP_SIZE;
7034 let out_addr = SendMut(out.as_mut_ptr());
7035 if a8w8_enabled() {
7036 let act = split_act(x);
7037 let run = move |start: usize, end: usize| {
7038 for r in start..end {
7039 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7040 for &(j, xv) in &act.outliers {
7041 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7042 acc += w * s * xv;
7043 }
7044 unsafe { *out_addr.at(r) = acc };
7046 }
7047 };
7048 dispatch_rows(pool, rows, &run);
7049 return;
7050 }
7051 let run = move |start: usize, end: usize| {
7052 for r in start..end {
7053 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7055 }
7056 };
7057 dispatch_rows(pool, rows, &run);
7058}
7059
7060#[allow(clippy::too_many_arguments)]
7062fn q4t_matvec2(
7063 bytes: &[u8],
7064 x1: &[f32],
7065 x2: &[f32],
7066 rows: usize,
7067 cols: usize,
7068 o1: &mut [f32],
7069 o2: &mut [f32],
7070 pool: Option<&Pool>,
7071) {
7072 let gpr = cols / GROUP_SIZE;
7073 let p1 = SendMut(o1.as_mut_ptr());
7074 let p2 = SendMut(o2.as_mut_ptr());
7075 if a8w8_enabled() {
7076 let a1 = split_act(x1);
7077 let a2 = split_act(x2);
7078 let run = move |start: usize, end: usize| {
7079 for r in start..end {
7080 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7081 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7082 for &(j, xv) in &a1.outliers {
7083 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7084 v1 += w * s * xv;
7085 }
7086 for &(j, xv) in &a2.outliers {
7087 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7088 v2 += w * s * xv;
7089 }
7090 unsafe {
7092 *p1.at(r) = v1;
7093 *p2.at(r) = v2;
7094 }
7095 }
7096 };
7097 dispatch_rows(pool, rows, &run);
7098 return;
7099 }
7100 let run = move |start: usize, end: usize| {
7101 for r in start..end {
7102 unsafe {
7104 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7105 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7106 }
7107 }
7108 };
7109 dispatch_rows(pool, rows, &run);
7110}
7111
7112#[allow(clippy::too_many_arguments)]
7114#[cfg(target_os = "macos")]
7120fn dequant_matmat_accel(
7121 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7122 xs_all: &[f32],
7123 b: usize,
7124 rows: usize,
7125 cols: usize,
7126 out: &mut [f32],
7127 pool: Option<&Pool>,
7128) {
7129 const TR: usize = 2048;
7130 thread_local! {
7131 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7132 }
7133 WTILE.with(|wt| {
7134 let mut wtile = wt.borrow_mut();
7135 wtile.resize(TR * cols, 0.0);
7136 let mut r0 = 0usize;
7137 while r0 < rows {
7138 let tr = TR.min(rows - r0);
7139 let wt_addr = SendMut(wtile.as_mut_ptr());
7140 let run = |start: usize, end: usize| {
7141 for r in start..end {
7142 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7144 dequant_row(r0 + r, dst);
7145 }
7146 };
7147 dispatch_rows(pool, tr, &run);
7148 unsafe {
7149 accel_blas::cblas_sgemm(
7150 101, 111, 112, b as i32,
7154 tr as i32,
7155 cols as i32,
7156 1.0,
7157 xs_all.as_ptr(),
7158 cols as i32,
7159 wtile.as_ptr(),
7160 cols as i32,
7161 0.0,
7162 out.as_mut_ptr().add(r0),
7163 rows as i32,
7164 );
7165 }
7166 r0 += tr;
7167 }
7168 });
7169}
7170
7171fn q4t_matmat(
7172 bytes: &[u8],
7173 xs_all: &[f32],
7174 b: usize,
7175 rows: usize,
7176 cols: usize,
7177 out: &mut [f32],
7178 pool: Option<&Pool>,
7179) {
7180 debug_assert_eq!(out.len(), b * rows);
7181 let gpr = cols / GROUP_SIZE;
7182 #[cfg(target_os = "macos")]
7186 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7187 dequant_matmat_accel(
7188 &|r, dst| {
7189 for gi in 0..gpr {
7190 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7191 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7192 for (k, &bb) in tile[2..].iter().enumerate() {
7193 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7194 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7195 }
7196 }
7197 },
7198 xs_all,
7199 b,
7200 rows,
7201 cols,
7202 out,
7203 pool,
7204 );
7205 return;
7206 }
7207 let out_addr = SendMut(out.as_mut_ptr());
7208 if a8w8_enabled() {
7209 let acts: Vec<SplitAct> = (0..b)
7210 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7211 .collect();
7212 let acts = &acts;
7213 #[cfg(target_arch = "x86_64")]
7214 let blocked_ok = avx2_enabled() && blocked_enabled();
7215 #[cfg(target_arch = "aarch64")]
7216 let blocked_ok = sdot_enabled() && blocked_enabled();
7217 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7218 let blocked_ok = false;
7219 let run = move |start: usize, end: usize| {
7220 for r in start..end {
7221 let mut bi = 0usize;
7222 #[cfg(target_arch = "aarch64")]
7223 if blocked_ok {
7224 while bi + 4 <= acts.len() {
7225 let xs = [
7226 acts[bi].xq.as_slice(),
7227 acts[bi + 1].xq.as_slice(),
7228 acts[bi + 2].xq.as_slice(),
7229 acts[bi + 3].xq.as_slice(),
7230 ];
7231 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7232 for k in 0..4 {
7233 let act = &acts[bi + k];
7234 let mut acc = d[k] * act.sx;
7235 for &(j, xv) in &act.outliers {
7236 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7237 acc += w * sc * xv;
7238 }
7239 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7241 }
7242 bi += 4;
7243 }
7244 }
7245 #[cfg(target_arch = "x86_64")]
7246 if blocked_ok {
7247 while bi + 4 <= acts.len() {
7248 let xs = [
7249 acts[bi].xq.as_slice(),
7250 acts[bi + 1].xq.as_slice(),
7251 acts[bi + 2].xq.as_slice(),
7252 acts[bi + 3].xq.as_slice(),
7253 ];
7254 let d = unsafe {
7255 if vnni_tiles_enabled() {
7256 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7257 } else {
7258 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7259 }
7260 };
7261 for k in 0..4 {
7262 let act = &acts[bi + k];
7263 let mut acc = d[k] * act.sx;
7264 for &(j, xv) in &act.outliers {
7265 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7266 acc += w * sc * xv;
7267 }
7268 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7270 }
7271 bi += 4;
7272 }
7273 }
7274 let _ = blocked_ok;
7275 while bi < acts.len() {
7276 let act = &acts[bi];
7277 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7278 for &(j, xv) in &act.outliers {
7279 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7280 acc += w * s * xv;
7281 }
7282 unsafe { *out_addr.at(bi * rows + r) = acc };
7284 bi += 1;
7285 }
7286 }
7287 };
7288 dispatch_rows(pool, rows, &run);
7289 return;
7290 }
7291 let run = move |start: usize, end: usize| {
7292 for r in start..end {
7293 for bi in 0..b {
7294 let x = &xs_all[bi * cols..(bi + 1) * cols];
7295 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7297 }
7298 }
7299 };
7300 dispatch_rows(pool, rows, &run);
7301}
7302
7303fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7312 (0..gpr)
7313 .map(|gi| {
7314 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7315 .iter()
7316 .map(|&v| v as i32)
7317 .sum()
7318 })
7319 .collect()
7320}
7321
7322#[inline]
7326#[allow(unreachable_code)]
7327#[cfg(target_arch = "x86_64")]
7332#[target_feature(enable = "avx2")]
7333unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7334 unsafe {
7336 use core::arch::x86_64::*;
7337 let expand = _mm256_setr_epi8(
7339 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,
7340 3, 3, 3,
7341 );
7342 let bitsel = _mm256_setr_epi8(
7343 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7344 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7345 );
7346 let ones8 = _mm256_set1_epi8(1);
7347 let ones16 = _mm256_set1_epi16(1);
7348 let mut acc = 0f32;
7349 for gi in 0..gpr {
7350 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7351 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7352 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7353 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7354 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7355 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7356 let sel = _mm256_and_si256(x, mask);
7357 let p16 = _mm256_maddubs_epi16(ones8, sel);
7359 let d32 = _mm256_madd_epi16(p16, ones16);
7360 let hi128 = _mm256_extracti128_si256::<1>(d32);
7361 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7362 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7363 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7364 let msum = _mm_cvtsi128_si32(s32);
7365 let d = 2 * msum - gsum[gi];
7368 acc += d as f32 * s;
7369 }
7370 acc
7371 }
7372}
7373
7374#[cfg(target_arch = "x86_64")]
7377#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7378unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7379 unsafe {
7381 use core::arch::x86_64::*;
7382 let expand = _mm256_setr_epi8(
7383 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,
7384 3, 3, 3,
7385 );
7386 let bitsel = _mm256_setr_epi8(
7387 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7388 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7389 );
7390 let ones8 = _mm256_set1_epi8(1);
7391 let mut acc = 0f32;
7392 for gi in 0..gpr {
7393 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7394 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7395 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7396 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7397 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7398 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7399 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7400 let d = 2 * msum - gsum[gi];
7401 acc += d as f32 * s;
7402 }
7403 acc
7404 }
7405}
7406
7407#[cfg(target_arch = "x86_64")]
7409#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7410unsafe fn dot_q1_row_1x4_vnni(
7411 bytes: &[u8],
7412 r: usize,
7413 gpr: usize,
7414 xs: [&[i8]; 4],
7415 gsums: [&[i32]; 4],
7416) -> [f32; 4] {
7417 unsafe {
7419 use core::arch::x86_64::*;
7420 let expand = _mm256_setr_epi8(
7421 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,
7422 3, 3, 3,
7423 );
7424 let bitsel = _mm256_setr_epi8(
7425 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7426 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7427 );
7428 let ones8 = _mm256_set1_epi8(1);
7429 let mut acc = [0f32; 4];
7430 for gi in 0..gpr {
7431 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7432 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7433 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7434 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7435 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7436 for (k, xq) in xs.iter().enumerate() {
7437 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7438 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7439 let d = 2 * msum - gsums[k][gi];
7440 acc[k] += d as f32 * s;
7441 }
7442 }
7443 acc
7444 }
7445}
7446
7447#[cfg(target_arch = "x86_64")]
7450#[target_feature(enable = "avx2")]
7451unsafe fn dot_q1_row_1x4_avx2(
7452 bytes: &[u8],
7453 r: usize,
7454 gpr: usize,
7455 xs: [&[i8]; 4],
7456 gsums: [&[i32]; 4],
7457) -> [f32; 4] {
7458 unsafe {
7460 use core::arch::x86_64::*;
7461 let expand = _mm256_setr_epi8(
7462 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,
7463 3, 3, 3,
7464 );
7465 let bitsel = _mm256_setr_epi8(
7466 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7467 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7468 );
7469 let ones8 = _mm256_set1_epi8(1);
7470 let ones16 = _mm256_set1_epi16(1);
7471 let mut acc = [0f32; 4];
7472 for gi in 0..gpr {
7473 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7474 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7475 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7476 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7477 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7478 for (k, xq) in xs.iter().enumerate() {
7479 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7480 let sel = _mm256_and_si256(x, mask);
7481 let p16 = _mm256_maddubs_epi16(ones8, sel);
7482 let d32 = _mm256_madd_epi16(p16, ones16);
7483 let hi128 = _mm256_extracti128_si256::<1>(d32);
7484 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7485 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7486 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7487 let msum = _mm_cvtsi128_si32(s32);
7488 let d = 2 * msum - gsums[k][gi];
7489 acc[k] += d as f32 * s;
7490 }
7491 }
7492 acc
7493 }
7494}
7495
7496#[allow(unreachable_code)]
7497fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7498 #[cfg(target_arch = "aarch64")]
7499 unsafe {
7500 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7501 }
7502 #[cfg(target_arch = "x86_64")]
7503 if avx2_enabled() {
7504 unsafe {
7505 if vnni_tiles_enabled() {
7506 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7507 }
7508 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7509 }
7510 }
7511 let _ = gsum;
7512 let mut acc = 0f32;
7513 for gi in 0..gpr {
7514 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7515 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7516 let mut d = 0i32;
7517 for (j, &b) in tile[2..].iter().enumerate() {
7518 for k in 0..8 {
7519 let w = ((b >> k) & 1) as i32 * 2 - 1;
7520 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7521 }
7522 }
7523 acc += d as f32 * s;
7524 }
7525 acc
7526}
7527
7528#[cfg(target_arch = "aarch64")]
7537#[target_feature(enable = "neon,dotprod")]
7538unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7539 unsafe {
7542 use core::arch::aarch64::*;
7543 use core::arch::asm;
7544 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7545 let m = vld1q_u8(MASKS.as_ptr());
7546 macro_rules! tile_dot {
7548 ($t:expr, $x:expr) => {{
7549 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7550 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7551 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7552 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7553 let x0 = vld1q_s8($x);
7554 let x1 = vld1q_s8($x.add(16));
7555 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7556 asm!(
7557 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7558 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7559 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7560 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7561 options(pure, nomem, nostack),
7562 );
7563 vaddq_s32(a0, a1)
7564 }};
7565 }
7566 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7575 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7576 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7577 const IW11: [u8; 16] = [
7578 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7579 ];
7580 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7581 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7582 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7583 let isc = vld1_u8(ISC.as_ptr());
7584 macro_rules! tile_dot_tbl {
7586 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7587 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7588 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7589 let x0 = vld1q_s8($x);
7590 let x1 = vld1q_s8($x.add(16));
7591 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7592 asm!(
7593 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7594 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7595 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7596 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7597 options(pure, nomem, nostack),
7598 );
7599 vaddq_s32(a0, a1)
7600 }};
7601 }
7602 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7603 let row_base = r * gpr * Q1_TILE;
7604 let abs_end = bytes.len();
7605 let xp = xq.as_ptr();
7606 let gp = gsum.as_ptr();
7607 let mut accv = vdupq_n_f32(0.0);
7608 let mut gi = 0;
7609 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7612 let t0 = base.add(gi * Q1_TILE);
7613 let ld_a = vld1q_u8(t0);
7614 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7615 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7616 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7617 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7618 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7619 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7621 let g = vld1q_s32(gp.add(gi));
7622 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7623 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7624 let scf: float32x4_t;
7625 asm!(
7626 "fcvtl {o:v}.4s, {i:v}.4h",
7627 o = out(vreg) scf, i = in(vreg) sc16,
7628 options(pure, nomem, nostack),
7629 );
7630 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7631 gi += 4;
7632 }
7633 let mut acc = vaddvq_f32(accv);
7634 while gi < gpr {
7635 let t = base.add(gi * Q1_TILE);
7636 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7637 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7638 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7639 gi += 1;
7640 }
7641 acc
7642 }
7643}
7644
7645#[cfg(target_arch = "aarch64")]
7650#[target_feature(enable = "neon,dotprod")]
7651unsafe fn dot_q1_row_1x4_sdot(
7652 bytes: &[u8],
7653 r: usize,
7654 gpr: usize,
7655 xs: [&[i8]; 4],
7656 gs: [&[i32]; 4],
7657) -> [f32; 4] {
7658 unsafe {
7660 use core::arch::aarch64::*;
7661 use core::arch::asm;
7662 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7663 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7664 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7665 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7666 const IW11: [u8; 16] = [
7667 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7668 ];
7669 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7670 let m = vld1q_u8(MASKS.as_ptr());
7671 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7672 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7673 let isc = vld1_u8(ISC.as_ptr());
7674 macro_rules! sdot2 {
7675 ($w0:expr, $w1:expr, $x:expr) => {{
7676 let x0 = vld1q_s8($x);
7677 let x1 = vld1q_s8($x.add(16));
7678 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7679 asm!(
7680 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7681 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7682 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7683 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7684 options(pure, nomem, nostack),
7685 );
7686 vaddq_s32(a0, a1)
7687 }};
7688 }
7689 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7690 let row_base = r * gpr * Q1_TILE;
7691 let abs_end = bytes.len();
7692 let mut accv = [vdupq_n_f32(0.0); 4];
7693 let mut gi = 0;
7694 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7695 let t0 = base.add(gi * Q1_TILE);
7696 let ld_a = vld1q_u8(t0);
7697 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7698 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7700 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7701 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7702 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7703 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7704 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7705 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7706 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7707 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7708 let scf: float32x4_t;
7709 asm!(
7710 "fcvtl {o:v}.4s, {i:v}.4h",
7711 o = out(vreg) scf, i = in(vreg) sc16,
7712 options(pure, nomem, nostack),
7713 );
7714 for k in 0..4 {
7715 let xp = xs[k].as_ptr();
7716 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7717 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7718 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7719 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7720 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7721 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7722 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7723 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7724 }
7725 gi += 4;
7726 }
7727 let mut acc = [
7728 vaddvq_f32(accv[0]),
7729 vaddvq_f32(accv[1]),
7730 vaddvq_f32(accv[2]),
7731 vaddvq_f32(accv[3]),
7732 ];
7733 while gi < gpr {
7734 let t = base.add(gi * Q1_TILE);
7735 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7736 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7737 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7738 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7739 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7740 for k in 0..4 {
7741 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7742 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7743 }
7744 gi += 1;
7745 }
7746 acc
7747 }
7748}
7749
7750#[inline]
7752fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7753 let gi = j / GROUP_SIZE;
7754 let k = j % GROUP_SIZE;
7755 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7756 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7757 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7758 ((bit as i32 * 2 - 1) as f32, s)
7759}
7760
7761#[inline]
7763fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7764 let mut acc = 0f32;
7765 for gi in 0..gpr {
7766 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7767 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7768 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7769 let mut ga = 0f32;
7770 for (j, &b) in tile[2..].iter().enumerate() {
7771 for k in 0..8 {
7772 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7773 }
7774 }
7775 acc += ga * s;
7776 }
7777 acc
7778}
7779
7780#[allow(clippy::too_many_arguments)]
7783fn bin_range_a8w8(
7784 w: BinW,
7785 gpr: usize,
7786 act: &SplitAct,
7787 gsum: &[i32],
7788 out: SendMut,
7789 start: usize,
7790 end: usize,
7791) {
7792 for r in start..end {
7793 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7794 for &(j, xv) in &act.outliers {
7795 let (wv, s) = w.outlier(r, gpr, j);
7796 acc += wv * s * xv;
7797 }
7798 unsafe { *out.at(r) = acc };
7800 }
7801}
7802
7803fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7805 for r in start..end {
7806 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7808 }
7809}
7810
7811#[derive(Clone, Copy)]
7819pub(crate) enum BinW<'a> {
7820 Q1(&'a [u8]),
7821 Q1g {
7822 bytes: &'a [u8],
7823 sc_off: usize,
7825 rows: usize,
7828 },
7829}
7830
7831impl<'a> BinW<'a> {
7832 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
7833 Self::Q1g {
7834 bytes,
7835 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
7836 rows,
7837 }
7838 }
7839
7840 #[inline]
7841 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7842 match self {
7843 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
7844 Self::Q1g {
7845 bytes,
7846 sc_off,
7847 rows,
7848 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
7849 }
7850 }
7851
7852 #[inline]
7853 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
7854 match self {
7855 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
7856 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
7857 }
7858 }
7859
7860 #[inline]
7861 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
7862 match self {
7863 Self::Q1(b) => q1_outlier(b, r, gpr, j),
7864 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
7865 }
7866 }
7867
7868 #[cfg(target_arch = "aarch64")]
7873 #[inline]
7874 unsafe fn dot_1x4_sdot(
7875 self,
7876 r: usize,
7877 gpr: usize,
7878 xs: [&[i8]; 4],
7879 gs: [&[i32]; 4],
7880 ) -> [f32; 4] {
7881 unsafe {
7883 match self {
7884 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
7885 Self::Q1g {
7886 bytes,
7887 sc_off,
7888 rows,
7889 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
7890 }
7891 }
7892 }
7893
7894 #[cfg(target_arch = "x86_64")]
7899 #[inline]
7900 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
7901 unsafe {
7903 match (self, vnni_tiles_enabled()) {
7904 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
7905 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
7906 (Self::Q1g { bytes, sc_off, .. }, true) => {
7907 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
7908 }
7909 (Self::Q1g { bytes, sc_off, .. }, false) => {
7910 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
7911 }
7912 }
7913 }
7914 }
7915}
7916
7917#[inline]
7927fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
7928 u32::from_le_bytes([
7929 bytes[t * 4],
7930 bytes[t * 4 + 1],
7931 bytes[t * 4 + 2],
7932 bytes[t * 4 + 3],
7933 ])
7934}
7935
7936#[inline]
7938fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
7939 let o = sc_off + (t / 4) * 2;
7940 u16::from_le_bytes([bytes[o], bytes[o + 1]])
7941}
7942
7943#[allow(unreachable_code)]
7945#[allow(clippy::too_many_arguments)]
7946fn dot_q1g_row_i8(
7947 bytes: &[u8],
7948 sc_off: usize,
7949 rows: usize,
7950 r: usize,
7951 gpr: usize,
7952 xq: &[i8],
7953 gsum: &[i32],
7954) -> f32 {
7955 #[cfg(target_arch = "aarch64")]
7956 unsafe {
7957 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
7958 }
7959 #[cfg(target_arch = "x86_64")]
7960 if avx2_enabled() {
7961 unsafe {
7962 if vnni_tiles_enabled() {
7963 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
7964 }
7965 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
7966 }
7967 }
7968 let _ = (gsum, rows);
7969 let mut acc = 0f32;
7970 for gi in 0..gpr {
7971 let t = r * gpr + gi;
7972 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
7973 let bits = q1g_bits(bytes, t);
7974 let mut d = 0i32;
7975 for k in 0..GROUP_SIZE {
7976 let w = ((bits >> k) & 1) as i32 * 2 - 1;
7977 d += w * xq[gi * GROUP_SIZE + k] as i32;
7978 }
7979 acc += d as f32 * s;
7980 }
7981 acc
7982}
7983
7984#[cfg(target_arch = "x86_64")]
7989#[target_feature(enable = "avx2")]
7990#[inline]
7991unsafe fn q1g_hsum4(
7992 d0: core::arch::x86_64::__m256i,
7993 d1: core::arch::x86_64::__m256i,
7994 d2: core::arch::x86_64::__m256i,
7995 d3: core::arch::x86_64::__m256i,
7996) -> core::arch::x86_64::__m128i {
7997 unsafe {
7999 use core::arch::x86_64::*;
8000 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8001 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8002 }
8003}
8004
8005#[cfg(target_arch = "x86_64")]
8009#[target_feature(enable = "avx2")]
8010#[inline]
8011unsafe fn q1g_group_acc(
8012 acc: f32,
8013 msum4: core::arch::x86_64::__m128i,
8014 gsum4: &[i32],
8015 s: f32,
8016) -> f32 {
8017 unsafe {
8019 use core::arch::x86_64::*;
8020 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8021 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8022 let mut f = [0f32; 4];
8023 _mm_storeu_ps(f.as_mut_ptr(), d);
8024 let mut acc = acc;
8025 for v in f {
8026 acc += v * s;
8027 }
8028 acc
8029 }
8030}
8031
8032#[cfg(target_arch = "x86_64")]
8037#[target_feature(enable = "avx2")]
8038unsafe fn dot_q1g_row_avx2(
8039 bytes: &[u8],
8040 sc_off: usize,
8041 r: usize,
8042 gpr: usize,
8043 xq: &[i8],
8044 gsum: &[i32],
8045) -> f32 {
8046 unsafe {
8049 use core::arch::x86_64::*;
8050 let expand = _mm256_setr_epi8(
8051 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,
8052 3, 3, 3,
8053 );
8054 let bitsel = _mm256_setr_epi8(
8055 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8056 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8057 );
8058 let ones8 = _mm256_set1_epi8(1);
8059 let ones16 = _mm256_set1_epi16(1);
8060 let bp = bytes.as_ptr().add(r * gpr * 4);
8061 let mut acc = 0f32;
8062 let mut gi = 0;
8063 while gi < gpr {
8064 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8065 let mut d = [_mm256_setzero_si256(); 4];
8066 for (k, dk) in d.iter_mut().enumerate() {
8067 let t = gi + k;
8068 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8069 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8070 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8071 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8072 let sel = _mm256_and_si256(x, mask);
8073 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8074 }
8075 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8076 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8077 gi += 4;
8078 }
8079 acc
8080 }
8081}
8082
8083#[cfg(target_arch = "x86_64")]
8086#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8087unsafe fn dot_q1g_row_vnni(
8088 bytes: &[u8],
8089 sc_off: usize,
8090 r: usize,
8091 gpr: usize,
8092 xq: &[i8],
8093 gsum: &[i32],
8094) -> f32 {
8095 unsafe {
8097 use core::arch::x86_64::*;
8098 let expand = _mm256_setr_epi8(
8099 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,
8100 3, 3, 3,
8101 );
8102 let bitsel = _mm256_setr_epi8(
8103 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8104 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8105 );
8106 let ones8 = _mm256_set1_epi8(1);
8107 let bp = bytes.as_ptr().add(r * gpr * 4);
8108 let mut acc = 0f32;
8109 let mut gi = 0;
8110 while gi < gpr {
8111 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8112 let mut d = [_mm256_setzero_si256(); 4];
8113 for (k, dk) in d.iter_mut().enumerate() {
8114 let t = gi + k;
8115 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8116 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8117 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8118 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8119 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8120 }
8121 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8122 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8123 gi += 4;
8124 }
8125 acc
8126 }
8127}
8128
8129#[cfg(target_arch = "x86_64")]
8133#[target_feature(enable = "avx2")]
8134unsafe fn dot_q1g_row_1x4_avx2(
8135 bytes: &[u8],
8136 sc_off: usize,
8137 r: usize,
8138 gpr: usize,
8139 xs: [&[i8]; 4],
8140 gsums: [&[i32]; 4],
8141) -> [f32; 4] {
8142 unsafe {
8144 use core::arch::x86_64::*;
8145 let expand = _mm256_setr_epi8(
8146 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,
8147 3, 3, 3,
8148 );
8149 let bitsel = _mm256_setr_epi8(
8150 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8151 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8152 );
8153 let ones8 = _mm256_set1_epi8(1);
8154 let ones16 = _mm256_set1_epi16(1);
8155 let bp = bytes.as_ptr().add(r * gpr * 4);
8156 let mut acc = [0f32; 4];
8157 let mut gi = 0;
8158 while gi < gpr {
8159 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8160 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8162 for t in 0..4 {
8163 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8164 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8165 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8166 for (k, xq) in xs.iter().enumerate() {
8167 let x = _mm256_loadu_si256(
8168 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8169 );
8170 let sel = _mm256_and_si256(x, mask);
8171 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8172 }
8173 }
8174 for k in 0..4 {
8175 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8176 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8177 }
8178 gi += 4;
8179 }
8180 acc
8181 }
8182}
8183
8184#[cfg(target_arch = "x86_64")]
8186#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8187unsafe fn dot_q1g_row_1x4_vnni(
8188 bytes: &[u8],
8189 sc_off: usize,
8190 r: usize,
8191 gpr: usize,
8192 xs: [&[i8]; 4],
8193 gsums: [&[i32]; 4],
8194) -> [f32; 4] {
8195 unsafe {
8197 use core::arch::x86_64::*;
8198 let expand = _mm256_setr_epi8(
8199 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,
8200 3, 3, 3,
8201 );
8202 let bitsel = _mm256_setr_epi8(
8203 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8204 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8205 );
8206 let ones8 = _mm256_set1_epi8(1);
8207 let bp = bytes.as_ptr().add(r * gpr * 4);
8208 let mut acc = [0f32; 4];
8209 let mut gi = 0;
8210 while gi < gpr {
8211 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8212 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8213 for t in 0..4 {
8214 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8215 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8216 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8217 for (k, xq) in xs.iter().enumerate() {
8218 let x = _mm256_loadu_si256(
8219 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8220 );
8221 d[k][t] = _mm256_dpbusd_epi32(
8222 _mm256_setzero_si256(),
8223 ones8,
8224 _mm256_and_si256(x, mask),
8225 );
8226 }
8227 }
8228 for k in 0..4 {
8229 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8230 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8231 }
8232 gi += 4;
8233 }
8234 acc
8235 }
8236}
8237
8238#[cfg(target_arch = "aarch64")]
8249#[target_feature(enable = "neon,dotprod")]
8250#[allow(clippy::too_many_arguments)]
8251unsafe fn dot_q1g_row_sdot(
8252 bytes: &[u8],
8253 sc_off: usize,
8254 rows: usize,
8255 r: usize,
8256 gpr: usize,
8257 xq: &[i8],
8258 gsum: &[i32],
8259) -> f32 {
8260 unsafe {
8263 use core::arch::aarch64::*;
8264 use core::arch::asm;
8265 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8266 const IW: [[u8; 16]; 8] = [
8268 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8269 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8270 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8271 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8272 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8273 [
8274 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8275 ],
8276 [
8277 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8278 ],
8279 [
8280 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8281 ],
8282 ];
8283 let m = vld1q_u8(MASKS.as_ptr());
8284 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8285 macro_rules! tile_dot {
8286 ($w0:expr, $w1:expr, $x:expr) => {{
8287 let x0 = vld1q_s8($x);
8288 let x1 = vld1q_s8($x.add(16));
8289 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8290 asm!(
8291 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8292 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8293 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8294 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8295 options(pure, nomem, nostack),
8296 );
8297 vaddq_s32(a0, a1)
8298 }};
8299 }
8300 macro_rules! tbl_mask {
8301 ($ld:expr, $i:expr) => {
8302 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8303 };
8304 }
8305 let bp = bytes.as_ptr().add(r * gpr * 4);
8306 let xp = xq.as_ptr();
8307 let gp = gsum.as_ptr();
8308 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8309 let mut accv = vdupq_n_f32(0.0);
8310 let mut gi = 0;
8311 while gi + 4 <= vec_end {
8312 let ld = vld1q_u8(bp.add(gi * 4));
8313 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8314 let d1 = tile_dot!(
8315 tbl_mask!(ld, 2),
8316 tbl_mask!(ld, 3),
8317 xp.add((gi + 1) * GROUP_SIZE)
8318 );
8319 let d2 = tile_dot!(
8320 tbl_mask!(ld, 4),
8321 tbl_mask!(ld, 5),
8322 xp.add((gi + 2) * GROUP_SIZE)
8323 );
8324 let d3 = tile_dot!(
8325 tbl_mask!(ld, 6),
8326 tbl_mask!(ld, 7),
8327 xp.add((gi + 3) * GROUP_SIZE)
8328 );
8329 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8330 let g = vld1q_s32(gp.add(gi));
8331 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8332 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8333 let scf: float32x4_t;
8334 asm!(
8335 "fcvtl {o:v}.4s, {i:v}.4h",
8336 o = out(vreg) scf, i = in(vreg) h,
8337 options(pure, nomem, nostack),
8338 );
8339 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8340 gi += 4;
8341 }
8342 let mut acc = vaddvq_f32(accv);
8343 while gi < gpr {
8344 let t = bp.add(gi * 4);
8345 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8346 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8347 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8348 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8349 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8350 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8351 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8352 gi += 1;
8353 }
8354 acc
8355 }
8356}
8357
8358#[cfg(target_arch = "aarch64")]
8362#[target_feature(enable = "neon,dotprod")]
8363#[allow(clippy::too_many_arguments)]
8364unsafe fn dot_q1g_row_1x4_sdot(
8365 bytes: &[u8],
8366 sc_off: usize,
8367 rows: usize,
8368 r: usize,
8369 gpr: usize,
8370 xs: [&[i8]; 4],
8371 gs: [&[i32]; 4],
8372) -> [f32; 4] {
8373 unsafe {
8375 use core::arch::aarch64::*;
8376 use core::arch::asm;
8377 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8378 const IW: [[u8; 16]; 8] = [
8379 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8380 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8381 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8382 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8383 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8384 [
8385 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8386 ],
8387 [
8388 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8389 ],
8390 [
8391 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8392 ],
8393 ];
8394 let m = vld1q_u8(MASKS.as_ptr());
8395 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8396 macro_rules! sdot2 {
8397 ($w0:expr, $w1:expr, $x:expr) => {{
8398 let x0 = vld1q_s8($x);
8399 let x1 = vld1q_s8($x.add(16));
8400 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8401 asm!(
8402 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8403 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8404 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8405 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8406 options(pure, nomem, nostack),
8407 );
8408 vaddq_s32(a0, a1)
8409 }};
8410 }
8411 let bp = bytes.as_ptr().add(r * gpr * 4);
8412 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8413 let mut accv = [vdupq_n_f32(0.0); 4];
8414 let mut gi = 0;
8415 while gi + 4 <= vec_end {
8416 let ld = vld1q_u8(bp.add(gi * 4));
8417 let w: [int8x16_t; 8] =
8419 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8420 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8421 let scf: float32x4_t;
8422 asm!(
8423 "fcvtl {o:v}.4s, {i:v}.4h",
8424 o = out(vreg) scf, i = in(vreg) h,
8425 options(pure, nomem, nostack),
8426 );
8427 for k in 0..4 {
8428 let xp = xs[k].as_ptr();
8429 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8430 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8431 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8432 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8433 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8434 let g = vld1q_s32(gs[k].as_ptr().add(gi));
8435 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8436 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8437 }
8438 gi += 4;
8439 }
8440 let mut acc = [
8441 vaddvq_f32(accv[0]),
8442 vaddvq_f32(accv[1]),
8443 vaddvq_f32(accv[2]),
8444 vaddvq_f32(accv[3]),
8445 ];
8446 while gi < gpr {
8447 let t = bp.add(gi * 4);
8448 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8449 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8450 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8451 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8452 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8453 for k in 0..4 {
8454 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8455 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8456 }
8457 gi += 1;
8458 }
8459 acc
8460 }
8461}
8462
8463#[inline]
8465fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8466 let t = r * gpr + j / GROUP_SIZE;
8467 let k = j % GROUP_SIZE;
8468 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8469 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8470 ((bit as i32 * 2 - 1) as f32, s)
8471}
8472
8473#[inline]
8475fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8476 let mut acc = 0f32;
8477 for gi in 0..gpr {
8478 let t = r * gpr + gi;
8479 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8480 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8481 let mut ga = 0f32;
8482 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8483 for k in 0..8 {
8484 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8485 }
8486 }
8487 acc += ga * s;
8488 }
8489 acc
8490}
8491
8492fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8497 let entries = base_len + (rows + 1) * 4;
8498 (base_len, entries, entries <= bytes.len())
8499}
8500
8501#[inline]
8503fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8504 let o = rp_off + r * 4;
8505 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8506}
8507
8508const SIGN5: [[f32; 5]; 256] = {
8512 let mut lut = [[0.0f32; 5]; 256];
8513 let pow3 = [1u16, 3, 9, 27, 81];
8514 let mut byte = 0usize;
8515 while byte < 256 {
8516 let mut i = 0usize;
8517 while i < 5 {
8518 let code = (byte as u16 / pow3[i]) % 3;
8519 lut[byte][i] = if code == 1 {
8520 1.0
8521 } else if code == 2 {
8522 -1.0
8523 } else {
8524 0.0
8525 };
8526 i += 1;
8527 }
8528 byte += 1;
8529 }
8530 lut
8531};
8532
8533const SIGN5_I8: [[i8; 5]; 256] = {
8535 let mut lut = [[0i8; 5]; 256];
8536 let pow3 = [1u16, 3, 9, 27, 81];
8537 let mut byte = 0usize;
8538 while byte < 256 {
8539 let mut i = 0usize;
8540 while i < 5 {
8541 let code = (byte as u16 / pow3[i]) % 3;
8542 lut[byte][i] = if code == 1 {
8543 1
8544 } else if code == 2 {
8545 -1
8546 } else {
8547 0
8548 };
8549 i += 1;
8550 }
8551 byte += 1;
8552 }
8553 lut
8554};
8555
8556const SIGN5_U64: [u64; 256] = {
8562 let mut lut = [0u64; 256];
8563 let pow3 = [1u16, 3, 9, 27, 81];
8564 let mut byte = 0usize;
8565 while byte < 256 {
8566 let mut v = 0u64;
8567 let mut i = 0usize;
8568 while i < 5 {
8569 let code = (byte as u16 / pow3[i]) % 3;
8570 let s: u8 = if code == 1 {
8571 1
8572 } else if code == 2 {
8573 0xFF
8574 } else {
8575 0
8576 };
8577 v |= (s as u64) << (i * 8);
8578 i += 1;
8579 }
8580 lut[byte] = v;
8581 byte += 1;
8582 }
8583 lut
8584};
8585
8586#[inline]
8591fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8592 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8593 let off = (r * gpr + j / GROUP_SIZE) * TILE;
8594 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8595 let within = j % GROUP_SIZE;
8596 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8597}
8598
8599#[cfg(target_arch = "aarch64")]
8602#[target_feature(enable = "neon,dotprod")]
8603#[inline]
8604unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8605 unsafe {
8607 use core::arch::aarch64::*;
8608 use core::arch::asm;
8609 let w0 = vld1q_s8(w);
8610 let w1 = vld1q_s8(w.add(16));
8611 let x0 = vld1q_s8(x);
8612 let x1 = vld1q_s8(x.add(16));
8613 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8614 asm!(
8615 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8616 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8617 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8618 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8619 options(pure, nomem, nostack),
8620 );
8621 vaddvq_s32(vaddq_s32(a0, a1))
8622 }
8623}
8624
8625#[cfg(target_arch = "x86_64")]
8628#[target_feature(enable = "avx2")]
8629#[inline]
8630unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8631 unsafe {
8633 use core::arch::x86_64::*;
8634 let wv = _mm256_loadu_si256(w as *const __m256i);
8635 let xv = _mm256_loadu_si256(x as *const __m256i);
8636 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8637 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8638 let hi128 = _mm256_extracti128_si256::<1>(d);
8639 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8640 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8641 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8642 _mm_cvtsi128_si32(s32)
8643 }
8644}
8645
8646#[inline]
8651fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8652 debug_assert!(dst.len() >= 40);
8653 unsafe {
8656 let p = dst.as_mut_ptr();
8657 for bi in 0..7 {
8658 core::ptr::write_unaligned(
8659 p.add(bi * 5) as *mut u64,
8660 SIGN5_U64[*codes.add(bi) as usize],
8661 );
8662 }
8663 }
8664}
8665
8666#[inline]
8671fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8672 #[cfg(target_arch = "aarch64")]
8673 unsafe {
8674 return sdot32_i8(w, x);
8675 }
8676 #[cfg(target_arch = "x86_64")]
8677 unsafe {
8678 return i8dot32_avx2(w, x);
8679 }
8680 #[allow(unreachable_code)]
8681 unsafe {
8682 let mut s = 0i32;
8683 for k in 0..GROUP_SIZE {
8684 s += *w.add(k) as i32 * *x.add(k) as i32;
8685 }
8686 s
8687 }
8688}
8689
8690#[inline]
8691unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8692 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8693 (
8694 SIGN5_U64[*codes as usize],
8695 SIGN5_U64[*codes.add(1) as usize],
8696 SIGN5_U64[*codes.add(2) as usize],
8697 SIGN5_U64[*codes.add(3) as usize],
8698 SIGN5_U64[*codes.add(4) as usize],
8699 SIGN5_U64[*codes.add(5) as usize],
8700 SIGN5_U64[*codes.add(6) as usize],
8701 )
8702 };
8703
8704 let u0 = s0 | (s1 << 40);
8705 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8706 let u2 = (s3 >> 8) | (s4 << 32);
8707 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8708
8709 (u0, u1, u2, u3)
8710}
8711
8712#[cfg(target_arch = "aarch64")]
8716#[target_feature(enable = "neon,dotprod")]
8717unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8718 use core::arch::aarch64::*;
8719 use core::arch::asm;
8720 unsafe {
8721 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8722 let mut acc = 0f32;
8723 let bytes_ptr = bytes.as_ptr();
8724 let xq_ptr = xq.as_ptr();
8725 let row_off = r * gpr * TILE;
8726
8727 let gpr2 = gpr & !1;
8728 let mut gi = 0;
8729 while gi < gpr2 {
8730 let off0 = row_off + gi * TILE;
8731 let off1 = off0 + TILE;
8732 let s0 = f16_to_f32(u16::from_le_bytes([
8733 *bytes_ptr.add(off0),
8734 *bytes_ptr.add(off0 + 1),
8735 ]));
8736 let s1 = f16_to_f32(u16::from_le_bytes([
8737 *bytes_ptr.add(off1),
8738 *bytes_ptr.add(off1 + 1),
8739 ]));
8740
8741 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8742 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8743
8744 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8745 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8746 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8747 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8748
8749 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8750 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8751 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8752 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8753
8754 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8755 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8756 asm!(
8757 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8758 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8759 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8760 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8761 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8762 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8763 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8764 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8765 options(pure, nomem, nostack),
8766 );
8767 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8768 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8769 acc += d0 as f32 * s0 + d1 as f32 * s1;
8770 gi += 2;
8771 }
8772
8773 if gi < gpr {
8774 let off = row_off + gi * TILE;
8775 let s = f16_to_f32(u16::from_le_bytes([
8776 *bytes_ptr.add(off),
8777 *bytes_ptr.add(off + 1),
8778 ]));
8779 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8780 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8781 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8782 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8783 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8784 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8785 asm!(
8786 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8787 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8788 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8789 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8790 options(pure, nomem, nostack),
8791 );
8792 let d = vaddvq_s32(vaddq_s32(a0, a1));
8793 acc += d as f32 * s;
8794 }
8795 acc
8796 }
8797}
8798
8799#[cfg(target_arch = "x86_64")]
8801#[target_feature(enable = "avx2")]
8802unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8803 use core::arch::x86_64::*;
8804 unsafe {
8805 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8806 let mut acc = 0f32;
8807 let bytes_ptr = bytes.as_ptr();
8808 let xq_ptr = xq.as_ptr();
8809 let row_off = r * gpr * TILE;
8810
8811 let ones = _mm256_set1_epi16(1);
8812 for gi in 0..gpr {
8813 let off = row_off + gi * TILE;
8814 let s = f16_to_f32(u16::from_le_bytes([
8815 *bytes_ptr.add(off),
8816 *bytes_ptr.add(off + 1),
8817 ]));
8818 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8819 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8820 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8821 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8822 let d256 = _mm256_madd_epi16(p16, ones);
8823 let d128 = _mm_add_epi32(
8824 _mm256_castsi256_si128(d256),
8825 _mm256_extracti128_si256(d256, 1),
8826 );
8827 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8828 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8829 acc += d32 as f32 * s;
8830 }
8831 acc
8832 }
8833}
8834
8835#[cfg(target_arch = "x86_64")]
8837#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8838unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8839 use core::arch::x86_64::*;
8840 unsafe {
8842 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8843 let mut acc = 0f32;
8844 let bytes_ptr = bytes.as_ptr();
8845 let xq_ptr = xq.as_ptr();
8846 let row_off = r * gpr * TILE;
8847 for gi in 0..gpr {
8848 let off = row_off + gi * TILE;
8849 let s = f16_to_f32(u16::from_le_bytes([
8850 *bytes_ptr.add(off),
8851 *bytes_ptr.add(off + 1),
8852 ]));
8853 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8854 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8855 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8856 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8857 acc += d as f32 * s;
8858 }
8859 acc
8860 }
8861}
8862
8863#[inline]
8867fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8868 #[cfg(target_arch = "aarch64")]
8869 unsafe {
8870 return q1t_dot_row_sdot(bytes, r, gpr, xq);
8871 }
8872 #[cfg(target_arch = "x86_64")]
8873 unsafe {
8874 if vnni_tiles_enabled() {
8875 return q1t_dot_row_vnni(bytes, r, gpr, xq);
8876 }
8877 return q1t_dot_row_avx2(bytes, r, gpr, xq);
8878 }
8879 #[allow(unreachable_code)]
8880 {
8881 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8882 let mut acc = 0f32;
8883 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
8885 let off = (r * gpr + gi) * TILE;
8886 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8887 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8888 let mut d = 0i32;
8889 for k in 0..GROUP_SIZE {
8890 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8891 }
8892 acc += d as f32 * s;
8893 }
8894 acc
8895 }
8896}
8897
8898fn q1t_row_outlier_correction(
8905 bytes: &[u8],
8906 r: usize,
8907 rp_off: usize,
8908 entries_off: usize,
8909 has_ov: bool,
8910 x: &[f32],
8911) -> f32 {
8912 if !has_ov {
8913 return 0.0;
8914 }
8915 let (c0, c1) = (
8916 q1t_rowptr(bytes, rp_off, r),
8917 q1t_rowptr(bytes, rp_off, r + 1),
8918 );
8919 let mut corr = 0f32;
8920 for p in c0..c1 {
8921 let e = entries_off + p * 4;
8922 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8923 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8924 corr += val * x[col];
8925 }
8926 corr
8927}
8928
8929fn q1t_dequant_row(
8933 bytes: &[u8],
8934 r: usize,
8935 gpr: usize,
8936 rp_off: usize,
8937 entries_off: usize,
8938 has_ov: bool,
8939 buf: &mut [f32],
8940) {
8941 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8942 for g in 0..gpr {
8943 let off = (r * gpr + g) * TILE;
8944 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8945 let codes = &bytes[off + 2..off + TILE];
8946 let bc = g * GROUP_SIZE;
8947 for bi in 0..6 {
8949 let lut = &SIGN5[codes[bi] as usize];
8950 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8951 for i in 0..5 {
8952 d[i] = lut[i] * s;
8953 }
8954 }
8955 let lut = &SIGN5[codes[6] as usize];
8956 buf[bc + 30] = lut[0] * s;
8957 buf[bc + 31] = lut[1] * s;
8958 }
8959 if !has_ov {
8960 return;
8961 }
8962 let (c0, c1) = (
8963 q1t_rowptr(bytes, rp_off, r),
8964 q1t_rowptr(bytes, rp_off, r + 1),
8965 );
8966 for p in c0..c1 {
8967 let e = entries_off + p * 4;
8968 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8969 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8970 }
8971}
8972
8973fn q1t_add_overlay(
8977 bytes: &[u8],
8978 x: &[f32],
8979 rows: usize,
8980 cols: usize,
8981 out: &mut [f32],
8982 pool: Option<&Pool>,
8983) {
8984 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8985 let gpr = cols / GROUP_SIZE;
8986 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8987 if !has_ov {
8988 return;
8989 }
8990 let out_addr = SendMut(out.as_mut_ptr());
8991 let run = move |start: usize, end: usize| {
8992 for r in start..end {
8993 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8994 unsafe { *out_addr.at(r) += corr };
8996 }
8997 };
8998 dispatch_rows(pool, rows, &run);
8999}
9000
9001#[allow(clippy::too_many_arguments)]
9004fn q1t_range_a8w8(
9005 bytes: &[u8],
9006 gpr: usize,
9007 rp_off: usize,
9008 ent_off: usize,
9009 has_ov: bool,
9010 act: &SplitAct,
9011 x: &[f32],
9012 out: SendMut,
9013 start: usize,
9014 end: usize,
9015) {
9016 for r in start..end {
9017 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9018 for &(j, xv) in &act.outliers {
9019 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9020 }
9021 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9022 unsafe { *out.at(r) = acc };
9024 }
9025}
9026
9027#[allow(clippy::too_many_arguments)]
9030fn q1t_range_f32_batch(
9031 bytes: &[u8],
9032 gpr: usize,
9033 rp_off: usize,
9034 ent_off: usize,
9035 has_ov: bool,
9036 x: &[f32],
9037 out: SendMut,
9038 start: usize,
9039 end: usize,
9040) {
9041 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9042 let mut sg = [0f32; GROUP_SIZE];
9043 for r in start..end {
9044 let mut acc = 0f32;
9045 for g in 0..gpr {
9046 let off = (r * gpr + g) * TILE;
9047 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9048 let codes = &bytes[off + 2..off + TILE];
9049 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9050 for bi in 0..6 {
9051 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9052 }
9053 let lut = &SIGN5[codes[6] as usize];
9054 sg[30] = lut[0];
9055 sg[31] = lut[1];
9056 let mut gsum = 0f32;
9057 for k in 0..GROUP_SIZE {
9058 gsum += sg[k] * xg[k];
9059 }
9060 acc += s * gsum;
9061 }
9062 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9063 unsafe { *out.at(r) = acc };
9065 }
9066}
9067
9068fn q1t_matvec(
9072 bytes: &[u8],
9073 x: &[f32],
9074 rows: usize,
9075 cols: usize,
9076 out: &mut [f32],
9077 pool: Option<&Pool>,
9078) {
9079 debug_assert_eq!(out.len(), rows);
9080 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9081 let gpr = cols / GROUP_SIZE;
9082 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9083 let out_addr = SendMut(out.as_mut_ptr());
9084 if a8w8_enabled() {
9088 let act = split_act(x);
9089 let act = &act;
9090 let run = move |start: usize, end: usize| {
9091 for r in start..end {
9092 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9093 for &(j, xv) in &act.outliers {
9094 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9095 }
9096 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9097 unsafe { *out_addr.at(r) = acc };
9099 }
9100 };
9101 dispatch_rows(pool, rows, &run);
9102 return;
9103 }
9104 let run = move |start: usize, end: usize| {
9105 let mut sg = [0f32; GROUP_SIZE];
9109 for r in start..end {
9110 let mut acc = 0f32;
9111 for g in 0..gpr {
9112 let off = (r * gpr + g) * TILE;
9113 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9114 let codes = &bytes[off + 2..off + TILE];
9115 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9116 for bi in 0..6 {
9117 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9118 }
9119 let lut = &SIGN5[codes[6] as usize];
9120 sg[30] = lut[0];
9121 sg[31] = lut[1];
9122 let mut gsum = 0f32;
9123 for k in 0..GROUP_SIZE {
9124 gsum += sg[k] * xg[k];
9125 }
9126 acc += s * gsum;
9127 }
9128 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9129 unsafe { *out_addr.at(r) = acc };
9130 }
9131 };
9132 dispatch_rows(pool, rows, &run);
9133}
9134
9135#[cfg(target_arch = "aarch64")]
9141#[target_feature(enable = "neon,dotprod")]
9142unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9143 use core::arch::aarch64::*;
9144 use core::arch::asm;
9145 unsafe {
9147 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9148 let bytes_ptr = bytes.as_ptr();
9149 let row_off = r * gpr * TILE;
9150 let xp = [xa.as_ptr(), xb.as_ptr()];
9151 let mut acc = [0f32; 2];
9152 macro_rules! sdot2 {
9153 ($w0:expr, $w1:expr, $x:expr) => {{
9154 let x0 = vld1q_s8($x);
9155 let x1 = vld1q_s8($x.add(16));
9156 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9157 asm!(
9158 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9159 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9160 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9161 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9162 options(pure, nomem, nostack),
9163 );
9164 vaddvq_s32(vaddq_s32(a0, a1))
9165 }};
9166 }
9167 let gpr2 = gpr & !1;
9168 let mut gi = 0;
9169 while gi < gpr2 {
9170 let off0 = row_off + gi * TILE;
9171 let off1 = off0 + TILE;
9172 let s0 = f16_to_f32(u16::from_le_bytes([
9173 *bytes_ptr.add(off0),
9174 *bytes_ptr.add(off0 + 1),
9175 ]));
9176 let s1 = f16_to_f32(u16::from_le_bytes([
9177 *bytes_ptr.add(off1),
9178 *bytes_ptr.add(off1 + 1),
9179 ]));
9180 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9181 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9182 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9183 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9184 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9185 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9186 for k in 0..2 {
9187 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9188 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9189 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9190 }
9191 gi += 2;
9192 }
9193 if gi < gpr {
9194 let off = row_off + gi * TILE;
9195 let s = f16_to_f32(u16::from_le_bytes([
9196 *bytes_ptr.add(off),
9197 *bytes_ptr.add(off + 1),
9198 ]));
9199 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9200 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9201 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9202 for k in 0..2 {
9203 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9204 acc[k] += d as f32 * s;
9205 }
9206 }
9207 acc
9208 }
9209}
9210
9211fn q1t_matvec2(
9217 bytes: &[u8],
9218 x1: &[f32],
9219 x2: &[f32],
9220 rows: usize,
9221 cols: usize,
9222 o1: &mut [f32],
9223 o2: &mut [f32],
9224 pool: Option<&Pool>,
9225) {
9226 debug_assert_eq!(o1.len(), rows);
9227 debug_assert_eq!(o2.len(), rows);
9228 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9229 let gpr = cols / GROUP_SIZE;
9230 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9231 let out1 = SendMut(o1.as_mut_ptr());
9232 let out2 = SendMut(o2.as_mut_ptr());
9233 if a8w8_enabled() {
9234 let a1 = split_act(x1);
9235 let a2 = split_act(x2);
9236 let (a1, a2) = (&a1, &a2);
9237 let run = move |start: usize, end: usize| {
9238 for r in start..end {
9239 #[cfg(target_arch = "aarch64")]
9240 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9243 #[cfg(not(target_arch = "aarch64"))]
9244 let ds = [
9245 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9246 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9247 ];
9248 let mut acc1 = ds[0] * a1.sx;
9249 for &(j, xv) in &a1.outliers {
9250 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9251 }
9252 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9253 let mut acc2 = ds[1] * a2.sx;
9254 for &(j, xv) in &a2.outliers {
9255 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9256 }
9257 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9258 unsafe {
9260 *out1.at(r) = acc1;
9261 *out2.at(r) = acc2;
9262 }
9263 }
9264 };
9265 dispatch_rows(pool, rows, &run);
9266 return;
9267 }
9268 let run = move |start: usize, end: usize| {
9269 let mut sg = [0f32; GROUP_SIZE];
9272 for r in start..end {
9273 let mut acc1 = 0f32;
9274 let mut acc2 = 0f32;
9275 for g in 0..gpr {
9276 let off = (r * gpr + g) * TILE;
9277 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9278 let codes = &bytes[off + 2..off + TILE];
9279 for bi in 0..6 {
9280 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9281 }
9282 let lut = &SIGN5[codes[6] as usize];
9283 sg[30] = lut[0];
9284 sg[31] = lut[1];
9285 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9286 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9287 let mut gsum1 = 0f32;
9288 for k in 0..GROUP_SIZE {
9289 gsum1 += sg[k] * xg1[k];
9290 }
9291 acc1 += s * gsum1;
9292 let mut gsum2 = 0f32;
9293 for k in 0..GROUP_SIZE {
9294 gsum2 += sg[k] * xg2[k];
9295 }
9296 acc2 += s * gsum2;
9297 }
9298 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9299 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9300 unsafe {
9302 *out1.at(r) = acc1;
9303 *out2.at(r) = acc2;
9304 }
9305 }
9306 };
9307 dispatch_rows(pool, rows, &run);
9308}
9309
9310fn q1t_matmat(
9313 bytes: &[u8],
9314 xs: &[f32],
9315 b: usize,
9316 rows: usize,
9317 cols: usize,
9318 out: &mut [f32],
9319 pool: Option<&Pool>,
9320) {
9321 debug_assert_eq!(out.len(), b * rows);
9322 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9323 let gpr = cols / GROUP_SIZE;
9324 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9325 let out_addr = SendMut(out.as_mut_ptr());
9326 if a8w8_enabled() {
9330 let acts: Vec<SplitAct> = (0..b)
9331 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9332 .collect();
9333 let acts = &acts;
9334 let run = move |start: usize, end: usize| {
9335 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
9339 for g in 0..gpr {
9340 let off = (r * gpr + g) * TILE;
9341 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9342 q1t_unpack_group_i8(
9343 bytes.as_ptr().wrapping_add(off + 2),
9344 &mut sg[g * GROUP_SIZE..],
9345 );
9346 }
9347 for bi in 0..b {
9348 let act = &acts[bi];
9349 let mut isum = 0f32;
9350 for g in 0..gpr {
9351 let d = q1t_i8dot32(
9352 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9353 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9354 );
9355 isum += d as f32 * sc[g];
9356 }
9357 let mut acc = isum * act.sx;
9358 for &(j, xv) in &act.outliers {
9359 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9360 }
9361 accs[bi] = acc;
9362 }
9363 if has_ov {
9367 let (c0, c1) = (
9368 q1t_rowptr(bytes, rp_off, r),
9369 q1t_rowptr(bytes, rp_off, r + 1),
9370 );
9371 for p in c0..c1 {
9372 let e = ent_off + p * 4;
9373 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9374 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9375 for bi in 0..b {
9376 accs[bi] += val * xs[bi * cols + col];
9377 }
9378 }
9379 }
9380 for bi in 0..b {
9381 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9382 }
9383 }
9384 };
9385 dispatch_rows(pool, rows, &run);
9386 return;
9387 }
9388 let run = move |start: usize, end: usize| {
9389 let mut buf = vec![0f32; cols];
9390 for r in start..end {
9391 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9392 for bi in 0..b {
9393 let xr = &xs[bi * cols..(bi + 1) * cols];
9394 let mut acc = 0f32;
9395 for j in 0..cols {
9396 acc += buf[j] * xr[j];
9397 }
9398 unsafe { *out_addr.at(bi * rows + r) = acc };
9399 }
9400 }
9401 };
9402 dispatch_rows(pool, rows, &run);
9403}
9404
9405#[cfg(test)]
9408fn q1_matvec(
9409 bytes: &[u8],
9410 x: &[f32],
9411 rows: usize,
9412 cols: usize,
9413 out: &mut [f32],
9414 pool: Option<&Pool>,
9415) {
9416 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9417}
9418
9419fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9422 debug_assert_eq!(out.len(), rows);
9423 let gpr = cols / GROUP_SIZE;
9424 let out_addr = SendMut(out.as_mut_ptr());
9425 if a8w8_enabled() {
9426 let act = split_act(x);
9427 let gsum = q1_group_sums(&act.xq, gpr);
9428 let (act, gsum) = (&act, &gsum);
9429 let run =
9430 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9431 dispatch_rows(pool, rows, &run);
9432 return;
9433 }
9434 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9435 dispatch_rows(pool, rows, &run);
9436}
9437
9438#[cfg(test)]
9440#[allow(clippy::too_many_arguments)]
9441fn q1_matvec2(
9442 bytes: &[u8],
9443 x1: &[f32],
9444 x2: &[f32],
9445 rows: usize,
9446 cols: usize,
9447 o1: &mut [f32],
9448 o2: &mut [f32],
9449 pool: Option<&Pool>,
9450) {
9451 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9452}
9453
9454#[allow(clippy::too_many_arguments)]
9456fn bin_matvec2(
9457 w: BinW,
9458 x1: &[f32],
9459 x2: &[f32],
9460 rows: usize,
9461 cols: usize,
9462 o1: &mut [f32],
9463 o2: &mut [f32],
9464 pool: Option<&Pool>,
9465) {
9466 let gpr = cols / GROUP_SIZE;
9467 let p1 = SendMut(o1.as_mut_ptr());
9468 let p2 = SendMut(o2.as_mut_ptr());
9469 if a8w8_enabled() {
9470 let a1 = split_act(x1);
9471 let a2 = split_act(x2);
9472 let g1 = q1_group_sums(&a1.xq, gpr);
9473 let g2 = q1_group_sums(&a2.xq, gpr);
9474 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9475 let run = move |start: usize, end: usize| {
9476 for r in start..end {
9477 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9478 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9479 for &(j, xv) in &a1.outliers {
9480 let (wv, s) = w.outlier(r, gpr, j);
9481 v1 += wv * s * xv;
9482 }
9483 for &(j, xv) in &a2.outliers {
9484 let (wv, s) = w.outlier(r, gpr, j);
9485 v2 += wv * s * xv;
9486 }
9487 unsafe {
9489 *p1.at(r) = v1;
9490 *p2.at(r) = v2;
9491 }
9492 }
9493 };
9494 dispatch_rows(pool, rows, &run);
9495 return;
9496 }
9497 let run = move |start: usize, end: usize| {
9498 for r in start..end {
9499 unsafe {
9501 *p1.at(r) = w.row_exact(r, gpr, x1);
9502 *p2.at(r) = w.row_exact(r, gpr, x2);
9503 }
9504 }
9505 };
9506 dispatch_rows(pool, rows, &run);
9507}
9508
9509#[cfg(test)]
9511#[allow(clippy::too_many_arguments)]
9512fn q1_matmat(
9513 bytes: &[u8],
9514 xs_all: &[f32],
9515 b: usize,
9516 rows: usize,
9517 cols: usize,
9518 out: &mut [f32],
9519 pool: Option<&Pool>,
9520) {
9521 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9522}
9523
9524#[allow(clippy::too_many_arguments)]
9526fn bin_matmat(
9527 w: BinW,
9528 xs_all: &[f32],
9529 b: usize,
9530 rows: usize,
9531 cols: usize,
9532 out: &mut [f32],
9533 pool: Option<&Pool>,
9534) {
9535 debug_assert_eq!(out.len(), b * rows);
9536 let gpr = cols / GROUP_SIZE;
9537 let out_addr = SendMut(out.as_mut_ptr());
9538 if a8w8_enabled() {
9539 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9540 .map(|bi| {
9541 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9542 let gsum = q1_group_sums(&act.xq, gpr);
9543 (act, gsum)
9544 })
9545 .collect();
9546 let acts = &acts;
9547 #[cfg(target_arch = "x86_64")]
9548 let blocked_ok = avx2_enabled() && blocked_enabled();
9549 #[cfg(target_arch = "aarch64")]
9550 let blocked_ok = sdot_enabled() && blocked_enabled();
9551 let run = move |start: usize, end: usize| {
9552 for r in start..end {
9553 let mut bi = 0usize;
9554 #[cfg(target_arch = "aarch64")]
9557 if blocked_ok {
9558 while bi + 4 <= acts.len() {
9559 let xs = [
9560 acts[bi].0.xq.as_slice(),
9561 acts[bi + 1].0.xq.as_slice(),
9562 acts[bi + 2].0.xq.as_slice(),
9563 acts[bi + 3].0.xq.as_slice(),
9564 ];
9565 let gs = [
9566 acts[bi].1.as_slice(),
9567 acts[bi + 1].1.as_slice(),
9568 acts[bi + 2].1.as_slice(),
9569 acts[bi + 3].1.as_slice(),
9570 ];
9571 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9573 for k in 0..4 {
9574 let (act, _) = &acts[bi + k];
9575 let mut acc = d[k] * act.sx;
9576 for &(j, xv) in &act.outliers {
9577 let (wv, sc) = w.outlier(r, gpr, j);
9578 acc += wv * sc * xv;
9579 }
9580 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9582 }
9583 bi += 4;
9584 }
9585 }
9586 #[cfg(target_arch = "x86_64")]
9587 if blocked_ok {
9588 while bi + 4 <= acts.len() {
9589 let xs = [
9590 acts[bi].0.xq.as_slice(),
9591 acts[bi + 1].0.xq.as_slice(),
9592 acts[bi + 2].0.xq.as_slice(),
9593 acts[bi + 3].0.xq.as_slice(),
9594 ];
9595 let gs = [
9596 acts[bi].1.as_slice(),
9597 acts[bi + 1].1.as_slice(),
9598 acts[bi + 2].1.as_slice(),
9599 acts[bi + 3].1.as_slice(),
9600 ];
9601 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9603 for k in 0..4 {
9604 let (act, _) = &acts[bi + k];
9605 let mut acc = d[k] * act.sx;
9606 for &(j, xv) in &act.outliers {
9607 let (wv, sc) = w.outlier(r, gpr, j);
9608 acc += wv * sc * xv;
9609 }
9610 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9612 }
9613 bi += 4;
9614 }
9615 }
9616 while bi < acts.len() {
9617 let (act, gsum) = &acts[bi];
9618 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9619 for &(j, xv) in &act.outliers {
9620 let (wv, s) = w.outlier(r, gpr, j);
9621 acc += wv * s * xv;
9622 }
9623 unsafe { *out_addr.at(bi * rows + r) = acc };
9625 bi += 1;
9626 }
9627 }
9628 };
9629 dispatch_rows(pool, rows, &run);
9630 return;
9631 }
9632 let run = move |start: usize, end: usize| {
9633 for r in start..end {
9634 for bi in 0..b {
9635 let x = &xs_all[bi * cols..(bi + 1) * cols];
9636 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9638 }
9639 }
9640 };
9641 dispatch_rows(pool, rows, &run);
9642}
9643
9644fn q4matvec(
9650 bytes: &[u8],
9651 x: &[f32],
9652 rows: usize,
9653 cols: usize,
9654 out: &mut [f32],
9655 pool: Option<&Pool>,
9656) {
9657 debug_assert_eq!(out.len(), rows);
9658 let (packed, scales) = q4_split(bytes, rows, cols);
9659 let gpr = cols / GROUP_SIZE;
9660 let out_addr = SendMut(out.as_mut_ptr());
9661
9662 if a8w8_enabled() {
9663 let act = split_act(x);
9664 let run = move |start: usize, end: usize| {
9665 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9666 };
9667 dispatch_rows(pool, rows, &run);
9668 return;
9669 }
9670
9671 let run =
9672 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9673 dispatch_rows(pool, rows, &run);
9674}
9675
9676#[inline]
9679#[allow(unreachable_code)]
9680#[cfg(target_arch = "x86_64")]
9685#[target_feature(enable = "avx2")]
9686unsafe fn dot_q4b_row_1x4_avx2(
9687 buf: &[u8],
9688 scales: &[u8],
9689 g0: usize,
9690 gpr: usize,
9691 xs: [&[i8]; 4],
9692) -> [f32; 4] {
9693 unsafe {
9695 use core::arch::x86_64::*;
9696 let ones = _mm256_set1_epi16(1);
9697 let mut acc = [0f32; 4];
9698 for gi in 0..gpr {
9699 let s = f16_to_f32(u16::from_le_bytes([
9700 scales[(g0 + gi) * 2],
9701 scales[(g0 + gi) * 2 + 1],
9702 ]));
9703 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9704 let aw = _mm256_abs_epi8(w);
9705 for (k, xq) in xs.iter().enumerate() {
9706 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9707 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9708 let d = _mm256_madd_epi16(p16, ones);
9709 let hi128 = _mm256_extracti128_si256::<1>(d);
9710 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9711 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9712 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9713 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9714 }
9715 }
9716 acc
9717 }
9718}
9719
9720#[cfg(target_arch = "x86_64")]
9722#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9723unsafe fn dot_q4b_row_1x4_vnni(
9724 buf: &[u8],
9725 scales: &[u8],
9726 g0: usize,
9727 gpr: usize,
9728 xs: [&[i8]; 4],
9729) -> [f32; 4] {
9730 unsafe {
9732 use core::arch::x86_64::*;
9733 let mut acc = [0f32; 4];
9734 for gi in 0..gpr {
9735 let s = f16_to_f32(u16::from_le_bytes([
9736 scales[(g0 + gi) * 2],
9737 scales[(g0 + gi) * 2 + 1],
9738 ]));
9739 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9740 let aw = _mm256_abs_epi8(w);
9741 for (k, xq) in xs.iter().enumerate() {
9742 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9743 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9744 acc[k] += d as f32 * s;
9745 }
9746 }
9747 acc
9748 }
9749}
9750
9751#[cfg(target_arch = "x86_64")]
9757#[target_feature(enable = "avx2")]
9758unsafe fn dot_q4b_row_1x4_sx_avx2(
9759 buf: &[u8],
9760 scales: &[u8],
9761 g0: usize,
9762 gpr: usize,
9763 xs: [&[i8]; 4],
9764 sxs: [f32; 4],
9765) -> [f32; 4] {
9766 unsafe {
9768 use core::arch::x86_64::*;
9769 let ones = _mm256_set1_epi16(1);
9770 let mut acc = [0f32; 4];
9771 for gi in 0..gpr {
9772 let s = f16_to_f32(u16::from_le_bytes([
9773 scales[(g0 + gi) * 2],
9774 scales[(g0 + gi) * 2 + 1],
9775 ]));
9776 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9777 let aw = _mm256_abs_epi8(w);
9778 for (k, xq) in xs.iter().enumerate() {
9779 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9780 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9781 let d = _mm256_madd_epi16(p16, ones);
9782 let hi128 = _mm256_extracti128_si256::<1>(d);
9783 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9784 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9785 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9786 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9787 }
9788 }
9789 acc
9790 }
9791}
9792
9793#[cfg(target_arch = "x86_64")]
9796#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9797unsafe fn dot_q4b_row_1x4_sx_vnni(
9798 buf: &[u8],
9799 scales: &[u8],
9800 g0: usize,
9801 gpr: usize,
9802 xs: [&[i8]; 4],
9803 sxs: [f32; 4],
9804) -> [f32; 4] {
9805 unsafe {
9807 use core::arch::x86_64::*;
9808 let mut acc = [0f32; 4];
9809 for gi in 0..gpr {
9810 let s = f16_to_f32(u16::from_le_bytes([
9811 scales[(g0 + gi) * 2],
9812 scales[(g0 + gi) * 2 + 1],
9813 ]));
9814 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9815 let aw = _mm256_abs_epi8(w);
9816 for (k, xq) in xs.iter().enumerate() {
9817 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9818 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9819 acc[k] += (d as f32 * sxs[k]) * s;
9820 }
9821 }
9822 acc
9823 }
9824}
9825
9826#[allow(unreachable_code)]
9827fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9828 #[cfg(target_arch = "aarch64")]
9829 unsafe {
9830 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9831 }
9832 #[cfg(target_arch = "x86_64")]
9833 unsafe {
9834 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
9835 }
9836 let mut acc = 0f32;
9837 for gi in 0..gpr {
9838 let g = g0 + gi;
9839 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9840 let mut d = 0i32;
9841 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9842 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
9843 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
9844 }
9845 acc += d as f32 * s;
9846 }
9847 acc
9848}
9849
9850#[inline]
9852#[allow(unreachable_code)]
9853fn dot_q4_row_i8_2(
9854 packed: &[u8],
9855 scales: &[u8],
9856 g0: usize,
9857 gpr: usize,
9858 xq1: &[i8],
9859 xq2: &[i8],
9860) -> (f32, f32) {
9861 #[cfg(target_arch = "aarch64")]
9862 unsafe {
9863 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
9864 }
9865 #[cfg(target_arch = "x86_64")]
9866 unsafe {
9867 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
9868 }
9869 (
9870 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
9871 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
9872 )
9873}
9874
9875#[allow(clippy::too_many_arguments)]
9878fn q4_range_a8w8(
9879 packed: &[u8],
9880 scales: &[u8],
9881 gpr: usize,
9882 cols: usize,
9883 act: &SplitAct,
9884 out: SendMut,
9885 start: usize,
9886 end: usize,
9887) {
9888 for r in start..end {
9889 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
9890 for &(j, xv) in &act.outliers {
9892 let flat = r * cols + j;
9893 let byte = packed[flat / 2];
9894 let nib = if flat & 1 == 0 {
9895 byte & 0x0F
9896 } else {
9897 byte >> 4
9898 };
9899 let s = f16_to_f32(u16::from_le_bytes([
9900 scales[(flat / GROUP_SIZE) * 2],
9901 scales[(flat / GROUP_SIZE) * 2 + 1],
9902 ]));
9903 acc += ((nib as i32 - 8) as f32) * s * xv;
9904 }
9905 unsafe { *out.at(r) = acc };
9907 }
9908}
9909
9910#[allow(clippy::too_many_arguments)]
9913fn q4_range2_a8w8(
9914 packed: &[u8],
9915 scales: &[u8],
9916 gpr: usize,
9917 cols: usize,
9918 a1: &SplitAct,
9919 a2: &SplitAct,
9920 p1: SendMut,
9921 p2: SendMut,
9922 start: usize,
9923 end: usize,
9924) {
9925 for r in start..end {
9926 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9927 let mut acc1 = s1 * a1.sx;
9928 let mut acc2 = s2 * a2.sx;
9929 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9931 for &(j, xv) in outliers {
9932 let flat = r * cols + j;
9933 let byte = packed[flat / 2];
9934 let nib = if flat & 1 == 0 {
9935 byte & 0x0F
9936 } else {
9937 byte >> 4
9938 };
9939 let s = f16_to_f32(u16::from_le_bytes([
9940 scales[(flat / GROUP_SIZE) * 2],
9941 scales[(flat / GROUP_SIZE) * 2 + 1],
9942 ]));
9943 *acc += ((nib as i32 - 8) as f32) * s * xv;
9944 }
9945 };
9946 fix(&a1.outliers, &mut acc1);
9947 fix(&a2.outliers, &mut acc2);
9948 unsafe {
9950 *p1.at(r) = acc1;
9951 *p2.at(r) = acc2;
9952 }
9953 }
9954}
9955
9956fn q4_range_f32(
9958 packed: &[u8],
9959 scales: &[u8],
9960 gpr: usize,
9961 x: &[f32],
9962 out: SendMut,
9963 start: usize,
9964 end: usize,
9965) {
9966 for r in start..end {
9967 let mut acc = 0f32;
9968 for gi in 0..gpr {
9969 let g = r * gpr + gi;
9970 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9971 let pk = &packed[g * 16..(g + 1) * 16];
9972 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9973 let mut ga = 0f32;
9974 for (k, &b) in pk.iter().enumerate() {
9975 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
9976 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
9977 }
9978 acc += ga * s;
9979 }
9980 unsafe { *out.at(r) = acc };
9982 }
9983}
9984
9985#[allow(clippy::too_many_arguments)]
9989fn q4matvec2(
9990 bytes: &[u8],
9991 x1: &[f32],
9992 x2: &[f32],
9993 rows: usize,
9994 cols: usize,
9995 o1: &mut [f32],
9996 o2: &mut [f32],
9997 pool: Option<&Pool>,
9998) {
9999 debug_assert_eq!(o1.len(), rows);
10000 debug_assert_eq!(o2.len(), rows);
10001 let (packed, scales) = q4_split(bytes, rows, cols);
10002 let gpr = cols / GROUP_SIZE;
10003
10004 if a8w8_enabled() {
10005 let a1 = split_act(x1);
10006 let a2 = split_act(x2);
10007 let p1 = SendMut(o1.as_mut_ptr());
10008 let p2 = SendMut(o2.as_mut_ptr());
10009 let run = move |start: usize, end: usize| {
10010 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10011 };
10012 dispatch_rows(pool, rows, &run);
10013 return;
10014 }
10015
10016 let p1 = SendMut(o1.as_mut_ptr());
10017 let p2 = SendMut(o2.as_mut_ptr());
10018 let run = move |start: usize, end: usize| {
10019 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10020 };
10021 dispatch_rows(pool, rows, &run);
10022}
10023
10024#[allow(clippy::too_many_arguments)]
10026fn q4_range2_f32(
10027 packed: &[u8],
10028 scales: &[u8],
10029 gpr: usize,
10030 x1: &[f32],
10031 x2: &[f32],
10032 p1: SendMut,
10033 p2: SendMut,
10034 start: usize,
10035 end: usize,
10036) {
10037 for r in start..end {
10038 let (mut acc1, mut acc2) = (0f32, 0f32);
10039 for gi in 0..gpr {
10040 let g = r * gpr + gi;
10041 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10042 let pk = &packed[g * 16..(g + 1) * 16];
10043 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10044 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10045 let (mut g1, mut g2) = (0f32, 0f32);
10046 for (k, &b) in pk.iter().enumerate() {
10047 let wl = (b & 0x0F) as f32 - 8.0;
10048 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10049 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10050 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10051 }
10052 acc1 += g1 * s;
10053 acc2 += g2 * s;
10054 }
10055 unsafe {
10057 *p1.at(r) = acc1;
10058 *p2.at(r) = acc2;
10059 }
10060 }
10061}
10062
10063thread_local! {
10064 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10067 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10068}
10069
10070#[allow(clippy::too_many_arguments)]
10076fn q4matmat(
10077 bytes: &[u8],
10078 xs_all: &[f32],
10079 b: usize,
10080 rows: usize,
10081 cols: usize,
10082 out: &mut [f32],
10083 pool: Option<&Pool>,
10084) {
10085 debug_assert_eq!(xs_all.len(), b * cols);
10086 debug_assert_eq!(out.len(), b * rows);
10087 let (packed, scales) = q4_split(bytes, rows, cols);
10088 let gpr = cols / GROUP_SIZE;
10089 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10090
10091 if a8w8_enabled() {
10092 let acts: Vec<SplitAct> = (0..b)
10093 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10094 .collect();
10095 let acts = &acts;
10096 let out_addr = SendMut(out.as_mut_ptr());
10097 let run = move |start: usize, end: usize| {
10098 ROW_I8.with(|rb| {
10099 let mut buf = rb.borrow_mut();
10100 buf.resize(cols, 0);
10101 for r in start..end {
10102 for gi in 0..gpr {
10106 let g = r * gpr + gi;
10107 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10108 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10109 buf[gi * GROUP_SIZE + k * 2 + 1] =
10110 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10111 }
10112 }
10113 let mut bi = 0usize;
10114 #[cfg(target_arch = "x86_64")]
10115 if avx2_enabled() && blocked_enabled() {
10116 while bi + 4 <= acts.len() {
10117 let xs = [
10118 acts[bi].xq.as_slice(),
10119 acts[bi + 1].xq.as_slice(),
10120 acts[bi + 2].xq.as_slice(),
10121 acts[bi + 3].xq.as_slice(),
10122 ];
10123 let d = unsafe {
10124 if vnni_tiles_enabled() {
10125 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10126 } else {
10127 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10128 }
10129 };
10130 for k in 0..4 {
10131 let act = &acts[bi + k];
10132 let mut acc = d[k] * act.sx;
10133 for &(j, xv) in &act.outliers {
10134 acc += (buf[j] as i8) as f32
10135 * gscale((r * cols + j) / GROUP_SIZE)
10136 * xv;
10137 }
10138 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10140 }
10141 bi += 4;
10142 }
10143 }
10144 while bi < acts.len() {
10145 let act = &acts[bi];
10146 let mut acc = 0f32;
10147 for gi in 0..gpr {
10148 let d = dot_i8_i8(
10149 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10150 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10151 );
10152 acc += d as f32 * gscale(r * gpr + gi);
10153 }
10154 acc *= act.sx;
10155 for &(j, xv) in &act.outliers {
10157 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10158 }
10159 unsafe { *out_addr.at(bi * rows + r) = acc };
10161 bi += 1;
10162 }
10163 }
10164 })
10165 };
10166 dispatch_rows(pool, rows, &run);
10167 return;
10168 }
10169
10170 let out_addr = SendMut(out.as_mut_ptr());
10171 let run = move |start: usize, end: usize| {
10172 ROW_F32.with(|rb| {
10173 let mut buf = rb.borrow_mut();
10174 buf.resize(cols, 0.0);
10175 for r in start..end {
10176 for gi in 0..gpr {
10179 let g = r * gpr + gi;
10180 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10181 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10182 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10183 }
10184 }
10185 for bi in 0..b {
10186 let x = &xs_all[bi * cols..(bi + 1) * cols];
10187 let mut acc = 0f32;
10188 for gi in 0..gpr {
10189 let mut ga = 0f32;
10190 for k in 0..GROUP_SIZE / 2 {
10195 let e = gi * GROUP_SIZE + k * 2;
10196 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10197 }
10198 acc += ga * gscale(r * gpr + gi);
10199 }
10200 unsafe { *out_addr.at(bi * rows + r) = acc };
10202 }
10203 }
10204 })
10205 };
10206 dispatch_rows(pool, rows, &run);
10207}
10208
10209#[allow(clippy::too_many_arguments)]
10214fn vbitmatmat(
10215 bytes: &[u8],
10216 offsets: &[usize],
10217 xs_all: &[f32],
10218 b: usize,
10219 rows: usize,
10220 cols: usize,
10221 out: &mut [f32],
10222 pool: Option<&Pool>,
10223) {
10224 debug_assert_eq!(xs_all.len(), b * cols);
10225 debug_assert_eq!(out.len(), b * rows);
10226 debug_assert_eq!(offsets.len(), rows + 1);
10227 let ng = cols / GROUP_SIZE;
10228 let bits = &bytes[..rows];
10229 let sc_off = rows;
10230 let gscale = |r: usize, g: usize| {
10231 let so = (r * ng + g) * 2;
10232 f16_to_f32(u16::from_le_bytes([
10233 bytes[sc_off + so],
10234 bytes[sc_off + so + 1],
10235 ]))
10236 };
10237
10238 let decode_f32 = |r: usize, dst: &mut [f32]| {
10240 let bw = bits[r] as usize;
10241 let l = ((1i32 << (bw - 1)) - 1) as f32;
10242 let data = &bytes[offsets[r]..offsets[r + 1]];
10243 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10244 for d in dst.iter_mut() {
10245 while nbits < bw {
10246 acc = (acc << 8) | data[idx] as u64;
10247 idx += 1;
10248 nbits += 8;
10249 }
10250 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10251 nbits -= bw;
10252 *d = u - l;
10253 }
10254 };
10255
10256 if a8w8_enabled() {
10257 let acts: Vec<SplitAct> = (0..b)
10258 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10259 .collect();
10260 let acts = &acts;
10261 let out_addr = SendMut(out.as_mut_ptr());
10262 let run = move |start: usize, end: usize| {
10263 for r in start..end {
10264 let bw = bits[r] as usize;
10265 if bw == 8 {
10266 ROW_F32.with(|rb| {
10269 let mut buf = rb.borrow_mut();
10270 buf.resize(cols, 0.0);
10271 decode_f32(r, &mut buf);
10272 for bi in 0..b {
10273 let x = &xs_all[bi * cols..(bi + 1) * cols];
10274 let mut dot = 0f32;
10275 for g in 0..ng {
10276 let mut gd = 0f32;
10277 for k in 0..GROUP_SIZE {
10278 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10279 }
10280 dot += gd * gscale(r, g);
10281 }
10282 unsafe { *out_addr.at(bi * rows + r) = dot };
10284 }
10285 });
10286 continue;
10287 }
10288 let l = (1i32 << (bw - 1)) - 1;
10289 let data = &bytes[offsets[r]..offsets[r + 1]];
10290 ROW_I8.with(|rb| {
10291 let mut buf = rb.borrow_mut();
10292 buf.resize(cols, 0);
10293 #[inline(always)]
10294 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10295 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10296 let u = unpack8::<B>(&data[blk * B..]);
10297 for k in 0..8 {
10298 chunk[k] = (u[k] - l) as i8 as u8;
10299 }
10300 }
10301 }
10302 match bw {
10303 3 => fill::<3>(data, l, &mut buf),
10304 4 => vbit_fill4(data, &mut buf),
10305 5 => fill::<5>(data, l, &mut buf),
10306 6 => fill::<6>(data, l, &mut buf),
10307 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10308 }
10309 let mut bi = 0usize;
10310 #[cfg(target_arch = "x86_64")]
10314 if avx2_enabled() && blocked_enabled() {
10315 while bi + 4 <= acts.len() {
10316 let xs = [
10317 acts[bi].xq.as_slice(),
10318 acts[bi + 1].xq.as_slice(),
10319 acts[bi + 2].xq.as_slice(),
10320 acts[bi + 3].xq.as_slice(),
10321 ];
10322 let sxs = [
10323 acts[bi].sx,
10324 acts[bi + 1].sx,
10325 acts[bi + 2].sx,
10326 acts[bi + 3].sx,
10327 ];
10328 let d = unsafe {
10329 if vnni_tiles_enabled() {
10330 dot_q4b_row_1x4_sx_vnni(
10331 &buf,
10332 &bytes[sc_off..],
10333 r * ng,
10334 ng,
10335 xs,
10336 sxs,
10337 )
10338 } else {
10339 dot_q4b_row_1x4_sx_avx2(
10340 &buf,
10341 &bytes[sc_off..],
10342 r * ng,
10343 ng,
10344 xs,
10345 sxs,
10346 )
10347 }
10348 };
10349 for k in 0..4 {
10350 let act = &acts[bi + k];
10351 let mut dot = d[k];
10352 for &(j, xv) in &act.outliers {
10353 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10354 }
10355 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10357 }
10358 bi += 4;
10359 }
10360 }
10361 while bi < acts.len() {
10362 let act = &acts[bi];
10363 let mut dot = 0f32;
10364 for g in 0..ng {
10365 let d = dot_i8_i8(
10366 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10367 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10368 ) as f32
10369 * act.sx;
10370 dot += d * gscale(r, g);
10371 }
10372 for &(j, xv) in &act.outliers {
10373 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10374 }
10375 unsafe { *out_addr.at(bi * rows + r) = dot };
10377 bi += 1;
10378 }
10379 });
10380 }
10381 };
10382 dispatch_rows(pool, rows, &run);
10383 return;
10384 }
10385
10386 let out_addr = SendMut(out.as_mut_ptr());
10387 let run = move |start: usize, end: usize| {
10388 ROW_F32.with(|rb| {
10389 let mut buf = rb.borrow_mut();
10390 buf.resize(cols, 0.0);
10391 for r in start..end {
10392 decode_f32(r, &mut buf);
10393 for bi in 0..b {
10394 let x = &xs_all[bi * cols..(bi + 1) * cols];
10395 let mut dot = 0f32;
10396 for g in 0..ng {
10397 let mut gd = 0f32;
10398 for k in 0..GROUP_SIZE {
10399 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10400 }
10401 dot += gd * gscale(r, g);
10402 }
10403 unsafe { *out_addr.at(bi * rows + r) = dot };
10405 }
10406 }
10407 })
10408 };
10409 dispatch_rows(pool, rows, &run);
10410}
10411
10412pub(crate) fn gpu_batch_job<'a>(
10416 t: &'a QTensor,
10417 x: &[f32],
10418) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10419 match t {
10420 QTensor::Mapped {
10421 model,
10422 idx,
10423 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10424 rows,
10425 cols,
10426 row_scale,
10427 col_field,
10428 ..
10429 } => Some((
10430 model.clone(),
10431 crate::gpu::BatchJob {
10432 idx: *idx,
10433 rows: *rows,
10434 cols: *cols,
10435 row_scale,
10436 xs: prescale(x, col_field, *dt).into_owned(),
10437 layout: crate::gpu::BatchLayout::Q8,
10438 },
10439 )),
10440 QTensor::Mapped {
10442 model,
10443 idx,
10444 dtype: TensorDtype::Q1,
10445 rows,
10446 cols,
10447 ..
10448 } => Some((
10449 model.clone(),
10450 crate::gpu::BatchJob {
10451 idx: *idx,
10452 rows: *rows,
10453 cols: *cols,
10454 row_scale: &[],
10455 xs: x.to_vec(),
10456 layout: crate::gpu::BatchLayout::Q1,
10457 },
10458 )),
10459 QTensor::Mapped {
10461 model,
10462 idx,
10463 dtype: TensorDtype::Q1G,
10464 rows,
10465 cols,
10466 ..
10467 } => Some((
10468 model.clone(),
10469 crate::gpu::BatchJob {
10470 idx: *idx,
10471 rows: *rows,
10472 cols: *cols,
10473 row_scale: &[],
10474 xs: x.to_vec(),
10475 layout: crate::gpu::BatchLayout::Q1g,
10476 },
10477 )),
10478 QTensor::Mapped {
10483 model,
10484 idx,
10485 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10486 rows,
10487 cols,
10488 ..
10489 } => Some((
10490 model.clone(),
10491 crate::gpu::BatchJob {
10492 idx: *idx,
10493 rows: *rows,
10494 cols: *cols,
10495 row_scale: &[],
10496 xs: x.to_vec(),
10497 layout: if *dt == TensorDtype::Q4Tiled {
10498 crate::gpu::BatchLayout::Q4t
10499 } else {
10500 crate::gpu::BatchLayout::Q4tp
10501 },
10502 },
10503 )),
10504 _ => None,
10505 }
10506}
10507
10508thread_local! {
10509 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10510 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10511}
10512
10513pub(crate) fn prescale<'a>(
10514 x: &'a [f32],
10515 col_field: &[f32],
10516 dtype: TensorDtype,
10517) -> std::borrow::Cow<'a, [f32]> {
10518 if dtype == TensorDtype::Q8_2f {
10519 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10520 } else {
10521 std::borrow::Cow::Borrowed(x)
10522 }
10523}
10524
10525pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10528 x: &[f32],
10529 col_field: &[f32],
10530 dtype: TensorDtype,
10531 buf_id: u8,
10532 f: F,
10533) -> R {
10534 if dtype == TensorDtype::Q8_2f {
10535 if buf_id == 1 {
10536 PRESCALE_BUF1.with(|b| {
10537 let mut buf = b.borrow_mut();
10538 buf.clear();
10539 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10540 f(&buf)
10541 })
10542 } else {
10543 PRESCALE_BUF2.with(|b| {
10544 let mut buf = b.borrow_mut();
10545 buf.clear();
10546 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10547 f(&buf)
10548 })
10549 }
10550 } else {
10551 f(x)
10552 }
10553}
10554
10555#[cfg(target_arch = "x86_64")]
10560pub(crate) fn avx2_enabled() -> bool {
10561 use std::sync::OnceLock;
10562 static ON: OnceLock<bool> = OnceLock::new();
10563 *ON.get_or_init(|| {
10564 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10565 && std::arch::is_x86_feature_detected!("avx2")
10566 && std::arch::is_x86_feature_detected!("fma")
10567 })
10568}
10569
10570#[cfg(target_arch = "x86_64")]
10575fn avx2_a8w8_enabled() -> bool {
10576 if FLOAT_ACTIVATIONS.get() {
10577 return false;
10578 }
10579 use std::sync::OnceLock;
10580 static ON: OnceLock<bool> = OnceLock::new();
10581 *ON.get_or_init(|| {
10582 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10583 })
10584}
10585
10586thread_local! {
10587 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10588}
10589
10590pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10593 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10594 impl Drop for Restore {
10595 fn drop(&mut self) {
10596 FULL_GPU_Q8.set(self.0);
10597 }
10598 }
10599 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10600}
10601
10602thread_local! {
10606 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10607}
10608
10609pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10610 struct Restore(bool);
10611 impl Drop for Restore {
10612 fn drop(&mut self) {
10613 FLOAT_ACTIVATIONS.set(self.0);
10614 }
10615 }
10616 let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10617 f()
10618}
10619
10620static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10632
10633pub(crate) fn row_exact() -> bool {
10634 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10635}
10636
10637fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10638 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10639 impl Drop for Restore<'_> {
10640 fn drop(&mut self) {
10641 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10642 }
10643 }
10644 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10645 let _restore = Restore(active);
10646 f()
10647}
10648
10649pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10651 counted_row_exact_scope(&ROW_EXACT, f)
10652}
10653
10654#[inline]
10658pub(crate) fn a8w8_enabled() -> bool {
10659 #[cfg(target_arch = "aarch64")]
10660 {
10661 sdot_enabled()
10662 }
10663 #[cfg(target_arch = "x86_64")]
10664 {
10665 avx2_a8w8_enabled()
10666 }
10667 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10668 {
10669 false
10670 }
10671}
10672
10673#[inline]
10676#[allow(unreachable_code)]
10677fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10678 #[cfg(target_arch = "aarch64")]
10679 unsafe {
10680 return dot_i8_sdot(w, xq);
10681 }
10682 #[cfg(target_arch = "x86_64")]
10683 unsafe {
10684 if avx512vnni_enabled() {
10685 return dot_i8_i8_vnni(w, xq);
10686 }
10687 return dot_i8_i8_avx2(w, xq);
10688 }
10689 w.iter()
10690 .zip(xq)
10691 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10692 .sum()
10693}
10694
10695#[cfg(target_arch = "x86_64")]
10699fn avx512vnni_enabled() -> bool {
10700 use std::sync::OnceLock;
10701 static ON: OnceLock<bool> = OnceLock::new();
10702 *ON.get_or_init(|| {
10703 std::env::var("CMF_AVX512")
10704 .map(|v| v != "0")
10705 .unwrap_or(true)
10706 && std::arch::is_x86_feature_detected!("avx512f")
10707 && std::arch::is_x86_feature_detected!("avx512bw")
10708 && std::arch::is_x86_feature_detected!("avx512vl")
10709 && std::arch::is_x86_feature_detected!("avx512vnni")
10710 })
10711}
10712
10713#[cfg(target_arch = "x86_64")]
10721fn vnni_tiles_enabled() -> bool {
10722 use std::sync::OnceLock;
10723 static ON: OnceLock<bool> = OnceLock::new();
10724 *ON.get_or_init(|| {
10725 std::env::var("CMF_VNNI_TILES")
10726 .map(|v| v != "0")
10727 .unwrap_or(true)
10728 && avx512vnni_enabled()
10729 })
10730}
10731
10732#[cfg(target_arch = "x86_64")]
10737#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10738#[inline]
10739unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10740 unsafe {
10742 use core::arch::x86_64::*;
10743 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10744 let hi128 = _mm256_extracti128_si256::<1>(d);
10745 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10746 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10747 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10748 _mm_cvtsi128_si32(s32)
10749 }
10750}
10751
10752#[cfg(target_arch = "x86_64")]
10757#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10758unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10759 unsafe {
10761 use core::arch::x86_64::*;
10762 let n = w.len();
10763 let mut j = 0usize;
10764 let mut total: i32;
10765 {
10770 #[inline(always)]
10771 unsafe fn step(
10772 w: *const u8,
10773 x: *const i8,
10774 acc: core::arch::x86_64::__m512i,
10775 ) -> core::arch::x86_64::__m512i {
10776 unsafe {
10777 use core::arch::x86_64::*;
10778 let wv = _mm512_loadu_si512(w as *const _);
10779 let xv = _mm512_loadu_si512(x as *const _);
10780 let aw = _mm512_abs_epi8(wv);
10781 let neg = _mm512_movepi8_mask(wv);
10782 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10783 _mm512_dpbusd_epi32(acc, aw, sx)
10784 }
10785 }
10786 let (mut a0, mut a1, mut a2, mut a3) = (
10787 _mm512_setzero_si512(),
10788 _mm512_setzero_si512(),
10789 _mm512_setzero_si512(),
10790 _mm512_setzero_si512(),
10791 );
10792 while j + 256 <= n {
10793 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10794 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10795 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10796 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10797 j += 256;
10798 }
10799 while j + 64 <= n {
10800 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10801 j += 64;
10802 }
10803 let s01 = _mm512_add_epi32(a0, a1);
10804 let s23 = _mm512_add_epi32(a2, a3);
10805 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10806 }
10807 if j + 32 <= n {
10809 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10810 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10811 let d = _mm256_dpbusd_epi32(
10812 _mm256_setzero_si256(),
10813 _mm256_abs_epi8(wv),
10814 _mm256_sign_epi8(xv, wv),
10815 );
10816 let hi128 = _mm256_extracti128_si256::<1>(d);
10817 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10818 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10819 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10820 total += _mm_cvtsi128_si32(s32);
10821 j += 32;
10822 }
10823 while j < n {
10824 total += (w[j] as i8) as i32 * xq[j] as i32;
10825 j += 1;
10826 }
10827 total
10828 }
10829}
10830
10831#[cfg(target_arch = "x86_64")]
10833#[target_feature(enable = "avx2,fma")]
10834unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
10835 unsafe {
10837 use core::arch::x86_64::*;
10838 let n = x.len();
10839 let wp = w.as_ptr();
10840 let xp = x.as_ptr();
10841 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
10842 let mut j = 0usize;
10843 while j + 16 <= n {
10844 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
10845 let lo = _mm256_cvtepi8_epi32(wb);
10846 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
10847 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
10848 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
10849 j += 16;
10850 }
10851 let acc = _mm256_add_ps(a0, a1);
10852 let hi128 = _mm256_extractf128_ps::<1>(acc);
10853 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
10854 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
10855 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
10856 let mut sum = _mm_cvtss_f32(s32);
10857 while j < n {
10858 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
10859 j += 1;
10860 }
10861 sum
10862 }
10863}
10864
10865#[cfg(target_arch = "x86_64")]
10870#[target_feature(enable = "avx2")]
10871unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
10872 unsafe {
10874 use core::arch::x86_64::*;
10875 let n = w.len();
10876 let ones = _mm256_set1_epi16(1);
10877 let mut acc = _mm256_setzero_si256();
10878 let mut j = 0usize;
10879 while j + 32 <= n {
10880 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10881 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10882 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10883 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
10884 j += 32;
10885 }
10886 let hi128 = _mm256_extracti128_si256::<1>(acc);
10887 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
10888 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10889 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10890 let mut s = _mm_cvtsi128_si32(s32);
10891 while j < n {
10892 s += (w[j] as i8) as i32 * xq[j] as i32;
10893 j += 1;
10894 }
10895 s
10896 }
10897}
10898
10899#[cfg(target_arch = "aarch64")]
10903#[target_feature(enable = "neon,i8mm")]
10904unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10905 unsafe {
10907 use core::arch::aarch64::*;
10908 use core::arch::asm;
10909 let n = w0.len();
10910 let w0p = w0.as_ptr() as *const i8;
10911 let w1p = w1.as_ptr() as *const i8;
10912 let mut acc01 = vdupq_n_s32(0);
10915 let mut acc23 = vdupq_n_s32(0);
10916 let mut i = 0usize;
10917 while i + 8 <= n {
10918 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10919 let xb01 = vcombine_s8(
10920 vld1_s8(xs[0].as_ptr().add(i)),
10921 vld1_s8(xs[1].as_ptr().add(i)),
10922 );
10923 let xb23 = vcombine_s8(
10924 vld1_s8(xs[2].as_ptr().add(i)),
10925 vld1_s8(xs[3].as_ptr().add(i)),
10926 );
10927 asm!(
10928 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10929 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10930 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10931 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10932 options(pure, nomem, nostack),
10933 );
10934 i += 8;
10935 }
10936 let mut out = [[0i32; 4]; 2];
10937 let a01: [i32; 4] = core::mem::transmute(acc01);
10938 let a23: [i32; 4] = core::mem::transmute(acc23);
10939 out[0][0] = a01[0];
10940 out[0][1] = a01[1];
10941 out[1][0] = a01[2];
10942 out[1][1] = a01[3];
10943 out[0][2] = a23[0];
10944 out[0][3] = a23[1];
10945 out[1][2] = a23[2];
10946 out[1][3] = a23[3];
10947 if i < n {
10948 for (k, x) in xs.iter().enumerate() {
10949 for j in i..n {
10950 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10951 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10952 }
10953 }
10954 }
10955 out
10956 }
10957}
10958
10959#[cfg(target_arch = "aarch64")]
10963#[target_feature(enable = "neon,dotprod")]
10964unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10965 unsafe {
10967 use core::arch::aarch64::*;
10968 use core::arch::asm;
10969 let n = w0.len();
10970 let w0p = w0.as_ptr() as *const i8;
10971 let w1p = w1.as_ptr() as *const i8;
10972 let mut acc = [[vdupq_n_s32(0); 4]; 2];
10973 let mut i = 0usize;
10974 while i + 16 <= n {
10975 let wv0 = vld1q_s8(w0p.add(i));
10976 let wv1 = vld1q_s8(w1p.add(i));
10977 for (k, x) in xs.iter().enumerate() {
10978 let xv = vld1q_s8(x.as_ptr().add(i));
10979 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
10980 asm!(
10981 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
10982 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
10983 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10984 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
10985 options(pure, nomem, nostack),
10986 );
10987 acc[0][k] = a0;
10988 acc[1][k] = a1;
10989 }
10990 i += 16;
10991 }
10992 let mut out = [[0i32; 4]; 2];
10993 for r in 0..2 {
10994 for k in 0..4 {
10995 out[r][k] = vaddvq_s32(acc[r][k]);
10996 }
10997 }
10998 if i < n {
10999 for (k, x) in xs.iter().enumerate() {
11000 for j in i..n {
11001 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11002 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11003 }
11004 }
11005 }
11006 out
11007 }
11008}
11009
11010#[cfg(target_arch = "x86_64")]
11016#[target_feature(enable = "avx2")]
11017unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11018 unsafe {
11020 use core::arch::x86_64::*;
11021 let n = w0.len();
11022 let ones = _mm256_set1_epi16(1);
11023 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11024 let mut j = 0usize;
11025 while j + 32 <= n {
11026 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11027 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11028 let aw0 = _mm256_abs_epi8(wv0);
11029 let aw1 = _mm256_abs_epi8(wv1);
11030 for (k, x) in xs.iter().enumerate() {
11031 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11032 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11033 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11034 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11035 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11036 }
11037 j += 32;
11038 }
11039 let mut out = [[0i32; 4]; 2];
11040 for r in 0..2 {
11041 for k in 0..4 {
11042 let a = acc[r][k];
11043 let hi128 = _mm256_extracti128_si256::<1>(a);
11044 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11045 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11046 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11047 out[r][k] = _mm_cvtsi128_si32(s32);
11048 }
11049 }
11050 if j < n {
11051 for (k, x) in xs.iter().enumerate() {
11052 for i in j..n {
11053 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11054 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11055 }
11056 }
11057 }
11058 out
11059 }
11060}
11061
11062#[cfg(target_arch = "x86_64")]
11067#[inline]
11068fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11069 let dot = if avx512vnni_enabled() && row.len() >= 64 {
11070 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11071 } else {
11072 unsafe { dot_i8_i8_avx2(row, &act.xq) }
11073 };
11074 let mut acc = dot as f32 * act.sx;
11075 for &(j, xv) in &act.outliers {
11076 acc += (row[j] as i8) as f32 * xv;
11077 }
11078 acc
11079}
11080
11081#[cfg(target_arch = "x86_64")]
11085#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11086unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11087 unsafe {
11089 use core::arch::x86_64::*;
11090 let n = w.len();
11091 let flip = _mm512_set1_epi8(-128); #[inline(always)]
11093 unsafe fn step(
11094 w: *const u8,
11095 x: *const i8,
11096 flip: core::arch::x86_64::__m512i,
11097 acc: core::arch::x86_64::__m512i,
11098 ) -> core::arch::x86_64::__m512i {
11099 unsafe {
11100 use core::arch::x86_64::*;
11101 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11102 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11103 }
11104 }
11105 let (mut a0, mut a1, mut a2, mut a3) = (
11106 _mm512_setzero_si512(),
11107 _mm512_setzero_si512(),
11108 _mm512_setzero_si512(),
11109 _mm512_setzero_si512(),
11110 );
11111 let mut j = 0usize;
11112 while j + 256 <= n {
11113 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11114 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11115 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11116 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11117 j += 256;
11118 }
11119 while j + 64 <= n {
11120 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11121 j += 64;
11122 }
11123 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11124 _mm512_add_epi32(a0, a1),
11125 _mm512_add_epi32(a2, a3),
11126 ));
11127 while j < n {
11129 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11130 j += 1;
11131 }
11132 total
11133 }
11134}
11135
11136#[cfg(target_arch = "x86_64")]
11142#[target_feature(enable = "avx2")]
11143unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11144 unsafe {
11147 use core::arch::x86_64::*;
11148 let lomask = _mm_set1_epi8(0x0F);
11149 let eight = _mm256_set1_epi8(8);
11150 let ones = _mm256_set1_epi16(1);
11151 let mut acc = 0f32;
11152 for gi in 0..gpr {
11153 let g = g0 + gi;
11154 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11155 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11156 let lo = _mm_and_si128(b, lomask);
11157 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11158 let w = _mm256_sub_epi8(
11159 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11160 eight,
11161 );
11162 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11163 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11164 let d = _mm256_madd_epi16(p16, ones);
11165 let hi128 = _mm256_extracti128_si256::<1>(d);
11166 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11167 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11168 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11169 acc += _mm_cvtsi128_si32(s32) as f32 * s;
11170 }
11171 acc
11172 }
11173}
11174
11175#[cfg(target_arch = "x86_64")]
11178#[target_feature(enable = "avx2")]
11179unsafe fn dot_q4_row_avx2_2(
11180 packed: &[u8],
11181 scales: &[u8],
11182 g0: usize,
11183 gpr: usize,
11184 xq1: &[i8],
11185 xq2: &[i8],
11186) -> (f32, f32) {
11187 unsafe {
11189 use core::arch::x86_64::*;
11190 let lomask = _mm_set1_epi8(0x0F);
11191 let eight = _mm256_set1_epi8(8);
11192 let ones = _mm256_set1_epi16(1);
11193 let (mut acc1, mut acc2) = (0f32, 0f32);
11194 #[inline(always)]
11195 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11196 unsafe {
11197 use core::arch::x86_64::*;
11198 let hi128 = _mm256_extracti128_si256::<1>(d);
11199 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11200 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11201 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11202 _mm_cvtsi128_si32(s32)
11203 }
11204 }
11205 for gi in 0..gpr {
11206 let g = g0 + gi;
11207 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11208 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11209 let lo = _mm_and_si128(b, lomask);
11210 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11211 let w = _mm256_sub_epi8(
11212 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11213 eight,
11214 );
11215 let aw = _mm256_abs_epi8(w);
11216 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11217 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11218 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11219 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11220 acc1 += hsum(d1) as f32 * s;
11221 acc2 += hsum(d2) as f32 * s;
11222 }
11223 (acc1, acc2)
11224 }
11225}
11226
11227#[cfg(target_arch = "x86_64")]
11229fn q8_range_avx2(
11230 q: &[u8],
11231 row_scale: &[f32],
11232 act: &SplitAct,
11233 cols: usize,
11234 out_addr: SendMut,
11235 start: usize,
11236 end: usize,
11237) {
11238 for o in start..end {
11239 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11240 unsafe { *out_addr.at(o) = v };
11242 }
11243}
11244
11245#[cfg(target_arch = "x86_64")]
11247#[allow(clippy::too_many_arguments)]
11248fn q8_range2_avx2(
11249 q: &[u8],
11250 row_scale: &[f32],
11251 a1: &SplitAct,
11252 a2: &SplitAct,
11253 cols: usize,
11254 p1: SendMut,
11255 p2: SendMut,
11256 start: usize,
11257 end: usize,
11258) {
11259 for o in start..end {
11260 let row = &q[o * cols..(o + 1) * cols];
11261 unsafe {
11263 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11264 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11265 }
11266 }
11267}
11268
11269#[cfg(target_arch = "aarch64")]
11280fn i8mm_enabled() -> bool {
11281 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11282 *ON.get_or_init(|| {
11283 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11284 && std::arch::is_aarch64_feature_detected!("i8mm")
11285 })
11286}
11287
11288#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11292fn sdot_enabled() -> bool {
11293 if FLOAT_ACTIVATIONS.get() {
11294 return false;
11295 }
11296 use std::sync::OnceLock;
11297 static ON: OnceLock<bool> = OnceLock::new();
11298 *ON.get_or_init(|| {
11299 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11300 if !want {
11301 return false;
11302 }
11303
11304 #[cfg(target_arch = "aarch64")]
11305 {
11306 if std::arch::is_aarch64_feature_detected!("dotprod") {
11307 return true;
11308 }
11309 #[cfg(target_os = "android")]
11310 {
11311 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11312 if cpuinfo.lines().any(|l| {
11313 (l.starts_with("Features") || l.starts_with("features"))
11314 && l.contains("asimddp")
11315 }) {
11316 return true;
11317 }
11318 }
11319 }
11320 false
11321 }
11322 #[cfg(not(target_arch = "aarch64"))]
11323 {
11324 false
11325 }
11326 })
11327}
11328
11329struct SplitAct {
11334 xq: Vec<i8>,
11335 sx: f32,
11336 outliers: Vec<(usize, f32)>,
11337 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11340 xsum: i32,
11341}
11342
11343thread_local! {
11344 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11347 const { std::cell::RefCell::new(Vec::new()) };
11348}
11349
11350impl Drop for SplitAct {
11351 fn drop(&mut self) {
11352 let buf = std::mem::take(&mut self.xq);
11353 if buf.capacity() > 0 {
11354 XQ_FREE.with(|f| {
11355 let mut f = f.borrow_mut();
11356 if f.len() < 16 {
11357 f.push(buf);
11358 }
11359 });
11360 }
11361 }
11362}
11363
11364thread_local! {
11365 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11372 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11377 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11378}
11379
11380#[inline]
11383fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11384 KROW.with(|s| {
11385 let b = unsafe { &mut *s.get() };
11388 if b.len() < n {
11389 b.resize(n, 0.0);
11390 }
11391 f(&mut b[..n])
11392 })
11393}
11394
11395#[inline(always)]
11406fn q8_round(t: f32) -> i8 {
11407 let t = t.clamp(-127.0, 127.0);
11408 let i = t as i32;
11409 let f = t - i as f32;
11410 let r = if f >= 0.5 {
11411 i + 1
11412 } else if f <= -0.5 {
11413 i - 1
11414 } else {
11415 i
11416 };
11417 r as i8
11418}
11419
11420#[inline]
11421fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11422 KROWS.with(|s| {
11423 let b = unsafe { &mut *s.get() };
11426 for row in b.iter_mut() {
11427 if row.len() < n {
11428 row.resize(n, 0.0);
11429 }
11430 }
11431 let (a, b) = b.split_at_mut(1);
11432 f(&mut a[0][..n], &mut b[0][..n])
11433 })
11434}
11435
11436#[inline]
11437fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11438 if limit > 0.0 {
11439 up = up.clamp(-limit, limit);
11440 gate = gate.min(limit);
11441 }
11442 gate / (1.0 + (-gate).exp()) * up
11443}
11444
11445fn split_act(x: &[f32]) -> SplitAct {
11446 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11447 let n = x.len();
11448 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11449 let thr = 8.0 * rms;
11450 let mut outliers: Vec<(usize, f32)> = Vec::new();
11454 let mut amax = 0f32;
11455 for (j, &v) in x.iter().enumerate() {
11456 let a = v.abs();
11457 if a > thr {
11458 outliers.push((j, v));
11459 } else if a > amax {
11460 amax = a;
11461 }
11462 }
11463 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11464 let inv = 1.0 / sx;
11465 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11466 xq.clear();
11467 xq.reserve(n);
11468 if outliers.is_empty() {
11469 xq.extend(
11470 x.iter()
11471 .map(|&v| q8_round(v * inv)),
11472 );
11473 } else {
11474 xq.extend(x.iter().map(|&v| {
11476 if v.abs() > thr {
11477 0
11478 } else {
11479 q8_round(v * inv)
11480 }
11481 }));
11482 }
11483 let xsum = xq.iter().map(|&v| v as i32).sum();
11484 SplitAct {
11485 xq,
11486 sx,
11487 outliers,
11488 xsum,
11489 }
11490}
11491
11492fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11493 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11494 let n = x.len();
11495 let rms = (x
11496 .iter()
11497 .zip(col)
11498 .map(|(&a, &c)| {
11499 let v = a * c;
11500 (v * v) as f64
11501 })
11502 .sum::<f64>()
11503 / n.max(1) as f64)
11504 .sqrt() as f32;
11505 let thr = 8.0 * rms;
11506
11507 let mut outliers = Vec::new();
11508 let mut amax = 0f32;
11509 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11510 let v = a * c;
11511 let s = v.abs();
11512 if s > thr {
11513 outliers.push((j, v));
11514 } else if s > amax {
11515 amax = s;
11516 }
11517 }
11518
11519 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11520 let inv = 1.0 / sx;
11521 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11522 xq.clear();
11523 xq.reserve(n);
11524 if outliers.is_empty() {
11525 xq.extend(
11526 x.iter()
11527 .zip(col)
11528 .map(|(&a, &c)| q8_round((a * c) * inv)),
11529 );
11530 } else {
11531 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11532 let v = a * c;
11533 if v.abs() > thr {
11534 0
11535 } else {
11536 q8_round(v * inv)
11537 }
11538 }));
11539 }
11540 let xsum = xq.iter().map(|&v| v as i32).sum();
11541 SplitAct {
11542 xq,
11543 sx,
11544 outliers,
11545 xsum,
11546 }
11547}
11548
11549#[cfg(target_arch = "aarch64")]
11552#[target_feature(enable = "neon,dotprod")]
11553unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11554 unsafe {
11556 use core::arch::aarch64::*;
11557 use core::arch::asm;
11558 let wp = w.as_ptr() as *const i8;
11559 let n = w.len();
11560 let (mut a0, mut a1, mut a2, mut a3) = (
11561 vdupq_n_s32(0),
11562 vdupq_n_s32(0),
11563 vdupq_n_s32(0),
11564 vdupq_n_s32(0),
11565 );
11566 let mut i = 0;
11567 while i + 64 <= n {
11568 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11569 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11570 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11571 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11572 asm!(
11573 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11574 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11575 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11576 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11577 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11578 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11579 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11580 options(pure, nomem, nostack),
11581 );
11582 i += 64;
11583 }
11584 while i + 16 <= n {
11585 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11586 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11587 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11588 i += 16;
11589 }
11590 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11591 while i < n {
11592 s += (*wp.add(i)) as i32 * xq[i] as i32;
11593 i += 1;
11594 }
11595 s
11596 }
11597}
11598
11599#[cfg(target_arch = "aarch64")]
11603#[target_feature(enable = "neon,dotprod")]
11604unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11605 unsafe {
11607 use core::arch::aarch64::*;
11608 use core::arch::asm;
11609 let n = xq.len();
11610 let px = xq.as_ptr();
11611 let (p0, p1, p2, p3) = (
11612 w0.as_ptr() as *const i8,
11613 w1.as_ptr() as *const i8,
11614 w2.as_ptr() as *const i8,
11615 w3.as_ptr() as *const i8,
11616 );
11617 let (mut a0, mut a1, mut a2, mut a3) = (
11618 vdupq_n_s32(0),
11619 vdupq_n_s32(0),
11620 vdupq_n_s32(0),
11621 vdupq_n_s32(0),
11622 );
11623 let mut i = 0;
11624 while i + 16 <= n {
11625 let x = vld1q_s8(px.add(i));
11626 let v0 = vld1q_s8(p0.add(i));
11627 let v1 = vld1q_s8(p1.add(i));
11628 let v2 = vld1q_s8(p2.add(i));
11629 let v3 = vld1q_s8(p3.add(i));
11630 asm!(
11631 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11632 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11633 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11634 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11635 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11636 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11637 options(pure, nomem, nostack),
11638 );
11639 i += 16;
11640 }
11641 let mut r = [
11642 vaddvq_s32(a0),
11643 vaddvq_s32(a1),
11644 vaddvq_s32(a2),
11645 vaddvq_s32(a3),
11646 ];
11647 while i < n {
11648 let xi = *px.add(i) as i32;
11649 r[0] += (*p0.add(i)) as i32 * xi;
11650 r[1] += (*p1.add(i)) as i32 * xi;
11651 r[2] += (*p2.add(i)) as i32 * xi;
11652 r[3] += (*p3.add(i)) as i32 * xi;
11653 i += 1;
11654 }
11655 r
11656 }
11657}
11658
11659#[cfg(target_arch = "aarch64")]
11666#[target_feature(enable = "neon,dotprod")]
11667unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11668 unsafe {
11671 use core::arch::aarch64::*;
11672 use core::arch::asm;
11673 let n = xq.len();
11674 let px = xq.as_ptr();
11675 let pg = g.as_ptr() as *const i8;
11676 let (mut a0, mut a1, mut a2, mut a3) = (
11677 vdupq_n_s32(0),
11678 vdupq_n_s32(0),
11679 vdupq_n_s32(0),
11680 vdupq_n_s32(0),
11681 );
11682 let mut i = 0;
11683 while i + 16 <= n {
11684 let x = vld1q_s8(px.add(i));
11685 let base = pg.add(4 * i);
11686 let v0 = vld1q_s8(base);
11687 let v1 = vld1q_s8(base.add(16));
11688 let v2 = vld1q_s8(base.add(32));
11689 let v3 = vld1q_s8(base.add(48));
11690 asm!(
11691 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11692 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11693 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11694 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11695 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11696 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11697 options(pure, nomem, nostack),
11698 );
11699 i += 16;
11700 }
11701 [
11702 vaddvq_s32(a0),
11703 vaddvq_s32(a1),
11704 vaddvq_s32(a2),
11705 vaddvq_s32(a3),
11706 ]
11707 }
11708}
11709
11710#[cfg(target_arch = "aarch64")]
11716fn q8_range_sdot(
11717 q: &[u8],
11718 rep: &[u8],
11719 row_scale: &[f32],
11720 act: &SplitAct,
11721 cols: usize,
11722 out_addr: SendMut,
11723 start: usize,
11724 end: usize,
11725) {
11726 let mut o = start;
11727 if !rep.is_empty() {
11730 while o < end && o % 4 != 0 {
11731 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11732 unsafe { *out_addr.at(o) = v };
11733 o += 1;
11734 }
11735 }
11736 while o + 4 <= end {
11737 let r = if rep.is_empty() {
11738 unsafe {
11739 dot_i8_sdot_4rows(
11740 &q[o * cols..(o + 1) * cols],
11741 &q[(o + 1) * cols..(o + 2) * cols],
11742 &q[(o + 2) * cols..(o + 3) * cols],
11743 &q[(o + 3) * cols..(o + 4) * cols],
11744 &act.xq,
11745 )
11746 }
11747 } else {
11748 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11749 };
11750 for k in 0..4 {
11751 let mut acc = r[k] as f32 * act.sx;
11752 for &(j, xv) in &act.outliers {
11753 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11754 }
11755 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11757 }
11758 o += 4;
11759 }
11760 while o < end {
11761 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11762 unsafe { *out_addr.at(o) = v };
11763 o += 1;
11764 }
11765}
11766
11767#[cfg(target_arch = "aarch64")]
11770#[allow(clippy::too_many_arguments)]
11771fn q8_range2_sdot(
11772 q: &[u8],
11773 row_scale: &[f32],
11774 a1: &SplitAct,
11775 a2: &SplitAct,
11776 cols: usize,
11777 p1: SendMut,
11778 p2: SendMut,
11779 start: usize,
11780 end: usize,
11781) {
11782 for o in start..end {
11783 let row = &q[o * cols..(o + 1) * cols];
11784 unsafe {
11786 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11787 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11788 }
11789 }
11790}
11791
11792#[allow(clippy::too_many_arguments)]
11794fn q8_range2_f32(
11795 q: &[u8],
11796 row_scale: &[f32],
11797 x1: &[f32],
11798 x2: &[f32],
11799 cols: usize,
11800 p1: SendMut,
11801 p2: SendMut,
11802 start: usize,
11803 end: usize,
11804) {
11805 for o in start..end {
11806 let row = &q[o * cols..(o + 1) * cols];
11807 unsafe {
11809 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11810 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11811 }
11812 }
11813}
11814
11815fn q8_range_f32(
11817 q: &[u8],
11818 row_scale: &[f32],
11819 xs: &[f32],
11820 cols: usize,
11821 out_addr: SendMut,
11822 start: usize,
11823 end: usize,
11824) {
11825 for o in start..end {
11826 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11827 unsafe { *out_addr.at(o) = v };
11829 }
11830}
11831
11832#[inline]
11836fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
11837 #[cfg(target_arch = "aarch64")]
11838 return row_dot_sdot(row, act);
11839 #[cfg(target_arch = "x86_64")]
11840 return row_dot_avx2(row, act);
11841 #[allow(unreachable_code)]
11842 q8_row_dot_scalar(row, act)
11843}
11844
11845#[allow(dead_code)]
11846fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
11847 let mut acc = 0i32;
11848 for (k, &b) in row.iter().enumerate() {
11849 acc += (b as i8) as i32 * act.xq[k] as i32;
11850 }
11851 let mut acc = acc as f32 * act.sx;
11852 for &(j, xv) in &act.outliers {
11853 acc += (row[j] as i8) as f32 * xv;
11854 }
11855 acc
11856}
11857
11858#[cfg(target_arch = "aarch64")]
11861#[inline]
11862fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
11863 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
11864 for &(j, xv) in &act.outliers {
11865 acc += (row[j] as i8) as f32 * xv;
11866 }
11867 acc
11868}
11869
11870#[cfg(target_arch = "aarch64")]
11878#[target_feature(enable = "neon,dotprod")]
11879unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11880 unsafe {
11883 use core::arch::aarch64::*;
11884 use core::arch::asm;
11885 let lomask = vdupq_n_u8(0x0F);
11886 let eight = vdupq_n_s8(8);
11887 let mut acc = 0f32;
11888 for gi in 0..gpr {
11889 let g = g0 + gi;
11890 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11891 let b = vld1q_u8(packed.as_ptr().add(g * 16));
11892 let lo = vandq_u8(b, lomask);
11893 let hi = vshrq_n_u8::<4>(b);
11894 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11895 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11896 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
11897 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
11898 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
11899 asm!(
11900 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
11901 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
11902 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11903 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
11904 options(pure, nomem, nostack),
11905 );
11906 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11907 }
11908 acc
11909 }
11910}
11911
11912#[cfg(target_arch = "aarch64")]
11917#[target_feature(enable = "neon,dotprod")]
11918unsafe fn dot_q4_row_sdot2(
11919 packed: &[u8],
11920 scales: &[u8],
11921 g0: usize,
11922 gpr: usize,
11923 xq1: &[i8],
11924 xq2: &[i8],
11925) -> (f32, f32) {
11926 unsafe {
11929 use core::arch::aarch64::*;
11930 use core::arch::asm;
11931 let lomask = vdupq_n_u8(0x0F);
11932 let eight = vdupq_n_s8(8);
11933 let (mut acc1, mut acc2) = (0f32, 0f32);
11934 for gi in 0..gpr {
11935 let g = g0 + gi;
11936 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11937 let b = vld1q_u8(packed.as_ptr().add(g * 16));
11938 let lo = vandq_u8(b, lomask);
11939 let hi = vshrq_n_u8::<4>(b);
11940 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11941 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11942 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11943 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11944 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11945 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11946 let (mut a0, mut a1, mut b0, mut b1) = (
11947 vdupq_n_s32(0),
11948 vdupq_n_s32(0),
11949 vdupq_n_s32(0),
11950 vdupq_n_s32(0),
11951 );
11952 asm!(
11953 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11954 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11955 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11956 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11957 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11958 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11959 e0 = in(vreg) e0, e1 = in(vreg) e1,
11960 x10 = in(vreg) x10, x11 = in(vreg) x11,
11961 x20 = in(vreg) x20, x21 = in(vreg) x21,
11962 options(pure, nomem, nostack),
11963 );
11964 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11965 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11966 }
11967 (acc1, acc2)
11968 }
11969}
11970
11971#[inline]
11976pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
11977 #[cfg(target_arch = "aarch64")]
11978 unsafe {
11979 return axpy_i8_f32_neon(acc, row, w);
11980 }
11981 #[cfg(target_arch = "x86_64")]
11982 if avx2_enabled() {
11983 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
11984 }
11985 #[allow(unreachable_code)]
11986 {
11987 for (a, &b) in acc.iter_mut().zip(row) {
11988 *a += w * b as f32;
11989 }
11990 }
11991}
11992
11993#[cfg(target_arch = "x86_64")]
11995#[target_feature(enable = "avx2,fma")]
11996unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
11997 unsafe {
11999 use core::arch::x86_64::*;
12000 let n = acc.len().min(row.len());
12001 let ap = acc.as_mut_ptr();
12002 let rp = row.as_ptr();
12003 let wv = _mm256_set1_ps(w);
12004 let mut j = 0usize;
12005 while j + 16 <= n {
12006 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12007 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12008 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12009 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12010 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12011 _mm256_storeu_ps(ap.add(j), v0);
12012 _mm256_storeu_ps(ap.add(j + 8), v1);
12013 j += 16;
12014 }
12015 while j < n {
12016 *ap.add(j) += w * (*rp.add(j)) as f32;
12017 j += 1;
12018 }
12019 }
12020}
12021
12022#[cfg(target_arch = "aarch64")]
12023#[target_feature(enable = "neon")]
12024unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12025 unsafe {
12027 use core::arch::aarch64::*;
12028 let n = acc.len().min(row.len());
12029 let ap = acc.as_mut_ptr();
12030 let rp = row.as_ptr();
12031 let wv = vdupq_n_f32(w);
12032 let mut j = 0usize;
12033 while j + 16 <= n {
12034 let rb = vld1q_s8(rp.add(j));
12035 let lo = vmovl_s8(vget_low_s8(rb));
12036 let hi = vmovl_s8(vget_high_s8(rb));
12037 for (off, half) in [(0, lo), (8, hi)] {
12038 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12039 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12040 let o = j + off;
12041 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12042 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12043 }
12044 j += 16;
12045 }
12046 while j < n {
12047 *ap.add(j) += w * (*rp.add(j)) as f32;
12048 j += 1;
12049 }
12050 }
12051}
12052
12053#[inline]
12056pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12057 #[cfg(target_arch = "aarch64")]
12058 unsafe {
12059 return dot_i8_f32_neon(w, x);
12060 }
12061 #[cfg(target_arch = "x86_64")]
12062 if avx2_enabled() {
12063 return unsafe { dot_i8_f32_avx2(w, x) };
12064 }
12065 #[allow(unreachable_code)]
12066 {
12067 let mut sum = 0.0f32;
12068 for (j, &b) in w.iter().enumerate() {
12069 sum += (b as i8) as f32 * x[j];
12070 }
12071 sum
12072 }
12073}
12074
12075#[inline]
12079fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12080 #[cfg(target_arch = "aarch64")]
12081 unsafe {
12082 return dot_i8_col_f32_neon(w, x, col);
12083 }
12084 #[allow(unreachable_code)]
12085 {
12086 let mut sum = 0.0f32;
12087 for (j, &b) in w.iter().enumerate() {
12088 sum += (b as i8) as f32 * x[j] * col[j];
12089 }
12090 sum
12091 }
12092}
12093
12094#[cfg(target_arch = "aarch64")]
12095#[target_feature(enable = "neon")]
12096unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12097 unsafe {
12099 use core::arch::aarch64::*;
12100 let n = x.len();
12101 let wp = w.as_ptr() as *const i8;
12102 let xp = x.as_ptr();
12103 let cp = col.as_ptr();
12104 let (mut a0, mut a1, mut a2, mut a3) = (
12105 vdupq_n_f32(0.0),
12106 vdupq_n_f32(0.0),
12107 vdupq_n_f32(0.0),
12108 vdupq_n_f32(0.0),
12109 );
12110 let mut j = 0usize;
12111 while j + 16 <= n {
12112 let wb = vld1q_s8(wp.add(j));
12113 let lo = vmovl_s8(vget_low_s8(wb));
12114 let hi = vmovl_s8(vget_high_s8(wb));
12115 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12116 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12117 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12118 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12119 a0 = vfmaq_f32(
12120 a0,
12121 w0,
12122 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12123 );
12124 a1 = vfmaq_f32(
12125 a1,
12126 w1,
12127 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12128 );
12129 a2 = vfmaq_f32(
12130 a2,
12131 w2,
12132 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12133 );
12134 a3 = vfmaq_f32(
12135 a3,
12136 w3,
12137 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12138 );
12139 j += 16;
12140 }
12141 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12142 while j < n {
12143 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12144 j += 1;
12145 }
12146 sum
12147 }
12148}
12149
12150#[cfg(target_arch = "aarch64")]
12151#[target_feature(enable = "neon")]
12152unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12153 unsafe {
12155 use core::arch::aarch64::*;
12156 let n = x.len();
12157 let wp = w.as_ptr() as *const i8;
12158 let xp = x.as_ptr();
12159 let (mut a0, mut a1, mut a2, mut a3) = (
12160 vdupq_n_f32(0.0),
12161 vdupq_n_f32(0.0),
12162 vdupq_n_f32(0.0),
12163 vdupq_n_f32(0.0),
12164 );
12165 let mut j = 0usize;
12166 while j + 16 <= n {
12167 let wb = vld1q_s8(wp.add(j));
12168 let lo = vmovl_s8(vget_low_s8(wb));
12169 let hi = vmovl_s8(vget_high_s8(wb));
12170 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12171 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12172 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12173 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12174 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12175 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12176 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12177 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12178 j += 16;
12179 }
12180 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12181 while j < n {
12182 sum += (*wp.add(j)) as f32 * *xp.add(j);
12183 j += 1;
12184 }
12185 sum
12186 }
12187}
12188
12189#[allow(clippy::too_many_arguments)]
12190fn qmatvec(
12191 q: &[u8],
12192 rep: &[u8],
12193 row_scale: &[f32],
12194 x: &[f32],
12195 col_field: &[f32],
12196 dtype: TensorDtype,
12197 rows: usize,
12198 cols: usize,
12199 out: &mut [f32],
12200 pool: Option<&Pool>,
12201) {
12202 debug_assert_eq!(out.len(), rows);
12203 #[cfg(not(target_arch = "aarch64"))]
12204 let _ = rep;
12205
12206 #[cfg(target_arch = "aarch64")]
12207 if sdot_enabled() {
12208 let act = if dtype == TensorDtype::Q8_2f {
12209 split_act_q8_2f(x, col_field)
12210 } else {
12211 split_act(x)
12212 };
12213 let out_addr = SendMut(out.as_mut_ptr());
12214 let run_range = |start: usize, end: usize| {
12215 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12216 };
12217 match pool {
12218 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12219 _ => run_range(0, rows),
12220 }
12221 return;
12222 }
12223 #[cfg(target_arch = "x86_64")]
12226 if avx2_a8w8_enabled() {
12227 let act = if dtype == TensorDtype::Q8_2f {
12228 split_act_q8_2f(x, col_field)
12229 } else {
12230 split_act(x)
12231 };
12232 let out_addr = SendMut(out.as_mut_ptr());
12233 let run_range = |start: usize, end: usize| {
12234 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12235 };
12236 match pool {
12237 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12238 _ => run_range(0, rows),
12239 }
12240 return;
12241 }
12242
12243 prescale_with(x, col_field, dtype, 1, |xs| {
12244 let out_addr = SendMut(out.as_mut_ptr());
12245 let run_range = move |start: usize, end: usize| {
12246 for o in start..end {
12247 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12248 unsafe { *out_addr.at(o) = v };
12250 }
12251 };
12252 match pool {
12253 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12254 _ => run_range(0, rows),
12255 }
12256 });
12257}
12258
12259#[allow(clippy::too_many_arguments)]
12260fn qmatvec2(
12261 q: &[u8],
12262 row_scale: &[f32],
12263 x1: &[f32],
12264 x2: &[f32],
12265 col_field: &[f32],
12266 dtype: TensorDtype,
12267 rows: usize,
12268 cols: usize,
12269 o1: &mut [f32],
12270 o2: &mut [f32],
12271 pool: Option<&Pool>,
12272) {
12273 #[cfg(target_arch = "aarch64")]
12274 if sdot_enabled() {
12275 let a1s = if dtype == TensorDtype::Q8_2f {
12276 split_act_q8_2f(x1, col_field)
12277 } else {
12278 split_act(x1)
12279 };
12280 let a2s = if dtype == TensorDtype::Q8_2f {
12281 split_act_q8_2f(x2, col_field)
12282 } else {
12283 split_act(x2)
12284 };
12285 let p1 = SendMut(o1.as_mut_ptr());
12286 let p2 = SendMut(o2.as_mut_ptr());
12287 let run_range = |start: usize, end: usize| {
12288 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12289 };
12290 match pool {
12291 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12292 _ => run_range(0, rows),
12293 }
12294 return;
12295 }
12296 #[cfg(target_arch = "x86_64")]
12297 if avx2_a8w8_enabled() {
12298 let a1s = if dtype == TensorDtype::Q8_2f {
12299 split_act_q8_2f(x1, col_field)
12300 } else {
12301 split_act(x1)
12302 };
12303 let a2s = if dtype == TensorDtype::Q8_2f {
12304 split_act_q8_2f(x2, col_field)
12305 } else {
12306 split_act(x2)
12307 };
12308 let p1 = SendMut(o1.as_mut_ptr());
12309 let p2 = SendMut(o2.as_mut_ptr());
12310 let run_range = |start: usize, end: usize| {
12311 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12312 };
12313 match pool {
12314 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12315 _ => run_range(0, rows),
12316 }
12317 return;
12318 }
12319
12320 prescale_with(x1, col_field, dtype, 1, |x1s| {
12321 prescale_with(x2, col_field, dtype, 2, |x2s| {
12322 let p1 = SendMut(o1.as_mut_ptr());
12323 let p2 = SendMut(o2.as_mut_ptr());
12324 let run_range = move |start: usize, end: usize| {
12325 for o in start..end {
12326 let row = &q[o * cols..(o + 1) * cols];
12327 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12328 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12329 unsafe {
12331 *p1.at(o) = s1;
12332 *p2.at(o) = s2;
12333 }
12334 }
12335 };
12336 match pool {
12337 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12338 _ => run_range(0, rows),
12339 }
12340 });
12341 });
12342}
12343
12344#[derive(Clone, Copy)]
12345struct SendMut(*mut f32);
12346unsafe impl Send for SendMut {}
12347unsafe impl Sync for SendMut {}
12348
12349impl SendMut {
12350 #[inline]
12351 fn at(self, i: usize) -> *mut f32 {
12352 unsafe { self.0.add(i) }
12353 }
12354}
12355
12356#[cfg(test)]
12357mod tests {
12358 #[test]
12362 fn q8_round_is_round_clamp() {
12363 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12364 let mut probe = vec![
12365 0.0f32,
12366 -0.0,
12367 f32::NAN,
12368 f32::INFINITY,
12369 f32::NEG_INFINITY,
12370 f32::MAX,
12371 f32::MIN,
12372 1e30,
12373 -1e30,
12374 f32::MIN_POSITIVE,
12375 -f32::MIN_POSITIVE,
12376 ];
12377 for k in -300i32..=300 {
12378 let h = k as f32 * 0.5;
12379 let up = f32::from_bits(h.to_bits() + 1);
12380 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12381 for t in [h, up, down] {
12382 probe.push(t);
12383 probe.push(-t);
12384 }
12385 }
12386 let mut t = -140.0f32;
12387 while t < 140.0 {
12388 probe.push(t);
12389 t += 0.000_731;
12390 }
12391 for t in probe {
12392 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12393 }
12394 }
12395
12396 use super::*;
12397
12398 #[test]
12402 fn f32_matmat_equals_matvec_bitwise() {
12403 let (rows, cols, b) = (5usize, 37usize, 7usize);
12404 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12405 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12406 let t = QTensor::from_f32(w, rows, cols);
12407 let mut out = vec![0f32; b * rows];
12408 t.matmat(&xs, b, &mut out, None);
12409 for bi in 0..b {
12410 let mut one = vec![0f32; rows];
12411 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12412 for o in 0..rows {
12413 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12414 }
12415 }
12416 }
12417
12418 #[test]
12419 fn q2tp_i8_dot_matches_exact_on_grid() {
12420 let (rows, cols) = (5, 64);
12424 let gpr = cols / GROUP_SIZE;
12425 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12429 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12430 .collect();
12431 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12432 let x: Vec<f32> = (0..cols)
12433 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12434 .collect();
12435 let act = split_act(&x);
12436 assert!(
12437 act.outliers.is_empty(),
12438 "on-grid input must have no outliers"
12439 );
12440 let gsum = q1_group_sums(&act.xq, gpr);
12441 for r in 0..rows {
12442 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12443 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12444 assert!(
12445 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12446 "row {r}: exact {exact} vs i8 {fast}"
12447 );
12448 }
12449 }
12450
12451 #[cfg(target_arch = "x86_64")]
12452 #[test]
12453 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12454 if !std::arch::is_x86_feature_detected!("avx2") {
12459 return;
12460 }
12461 let mut seed = 0x9e3779b9u32;
12462 let mut next = || {
12463 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12464 seed
12465 };
12466 for _ in 0..20_000 {
12467 let mut ch = [0u8; Q2TP_CHUNK];
12468 let mut x = [0i8; GROUP_SIZE];
12469 for b in &mut ch {
12470 *b = next() as u8;
12471 }
12472 for v in &mut x {
12473 *v = (next() >> 24) as i8;
12474 }
12475 let mut reference = 0i32;
12476 for (k, &b) in ch.iter().enumerate() {
12477 reference += (b & 3) as i32 * x[k * 4] as i32;
12478 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12479 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12480 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12481 }
12482 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12485 assert_eq!(got, reference, "packed q2 lane mismatch");
12486 }
12487 }
12488
12489 #[test]
12490 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12491 let (rows, cols) = (1usize, GROUP_SIZE);
12492 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12493 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12497 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12498 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
12500 let mut raw = vec![0.0f32; rows];
12501 let mut affine = vec![0.0f32; rows];
12502 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12503 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12504 assert_eq!(raw, vec![-24.0]);
12505 assert_eq!(affine, vec![-8.0]);
12506 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12507 }
12508
12509 #[test]
12510 fn q8_row_dot_fast_matches_scalar() {
12511 let cols = 96;
12514 let row: Vec<u8> = (0..cols)
12515 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12516 .collect();
12517 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12518 let act = split_act(&x);
12519 let fast = q8_row_dot(&row, &act);
12520 let scalar = q8_row_dot_scalar(&row, &act);
12521 assert!(
12522 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12523 "fast {fast} vs scalar {scalar}"
12524 );
12525 }
12526
12527 #[test]
12528 fn f32_matvec_matches_matvec_rows_bitexact() {
12529 let (rows, cols) = (300, 40);
12530 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12531 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12532 let qt = QTensor::from_f32(w.clone(), rows, cols);
12533
12534 let mut a = vec![0.0f32; rows];
12535 matvec_rows(None, &w, &x, &mut a);
12536 let mut b = vec![0.0f32; rows];
12537 qt.matvec(&x, &mut b, None);
12538 assert_eq!(a, b);
12539 }
12540
12541 #[test]
12542 fn sdot_kernel_exact_on_grid() {
12543 eprintln!("sdot_enabled = {}", sdot_enabled());
12548 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
12550 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12551 .collect();
12552 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12553 let x: Vec<f32> = (0..cols)
12554 .map(|i| match i % 3 {
12555 0 => 1.0,
12556 1 => -1.0,
12557 _ => 0.0,
12558 })
12559 .collect();
12560 let mut a = vec![0.0f32; rows];
12561 qmatvec(
12562 &w,
12563 &[],
12564 &scales,
12565 &x,
12566 &[],
12567 TensorDtype::Q8Row,
12568 rows,
12569 cols,
12570 &mut a,
12571 None,
12572 );
12573 for o in 0..rows {
12574 let mut acc = 0.0f32;
12575 for j in 0..cols {
12576 acc += (w[o * cols + j] as i8) as f32 * x[j];
12577 }
12578 let expect = acc * scales[o];
12579 assert!(
12580 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12581 "row {o}: {} vs {expect}",
12582 a[o]
12583 );
12584 }
12585 }
12586
12587 #[test]
12588 fn q1_tbl_fast_path_matches_reference() {
12589 let (rows, cols) = (5, 256);
12594 let gpr = cols / GROUP_SIZE;
12595 let mut bytes = Vec::new();
12596 for t in 0..rows * gpr {
12597 let s = 0.007 + (t % 11) as f32 * 0.004;
12598 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12599 for j in 0..4 {
12600 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12601 }
12602 }
12603 let x: Vec<f32> = (0..cols)
12604 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12605 .collect();
12606 let mut w = vec![0.0f32; rows * cols];
12607 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12608 let mut got = vec![0.0f32; rows];
12609 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12610 for o in 0..rows {
12611 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12612 assert!(
12613 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12614 "row {o}: {} vs {expect}",
12615 got[o]
12616 );
12617 }
12618 let b = 5usize;
12621 let mut xs_all = Vec::new();
12622 for bi in 0..b {
12623 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12624 }
12625 let mut mm = vec![0.0f32; b * rows];
12626 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12627 for bi in 0..b {
12628 let mut single = vec![0.0f32; rows];
12629 q1_matvec(
12630 &bytes,
12631 &xs_all[bi * cols..(bi + 1) * cols],
12632 rows,
12633 cols,
12634 &mut single,
12635 None,
12636 );
12637 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12638 }
12639 }
12640
12641 #[test]
12642 fn q1_kernels_match_exact_reference() {
12643 let (rows, cols) = (7, 96);
12645 let gpr = cols / GROUP_SIZE;
12646 let mut bytes = Vec::new();
12647 for t in 0..rows * gpr {
12648 let s = 0.01 + (t % 13) as f32 * 0.003;
12649 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12650 for j in 0..4 {
12651 bytes.push(((t * 31 + j * 97) % 251) as u8);
12652 }
12653 }
12654 let x: Vec<f32> = (0..cols)
12656 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12657 .collect();
12658 let mut w = vec![0.0f32; rows * cols];
12660 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12661 let mut expect = vec![0.0f32; rows];
12662 for o in 0..rows {
12663 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12664 }
12665 let mut got = vec![0.0f32; rows];
12666 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12667 for o in 0..rows {
12668 assert!(
12669 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12670 "row {o}: {} vs {}",
12671 got[o],
12672 expect[o]
12673 );
12674 }
12675 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12677 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12678 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12679 assert_eq!(a1, got);
12680 let mut xs = x.clone();
12681 xs.extend_from_slice(&x2);
12682 let mut mm = vec![0.0f32; 2 * rows];
12683 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12684 assert_eq!(&mm[..rows], got.as_slice());
12685 assert_eq!(&mm[rows..], a2.as_slice());
12686 }
12687
12688 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12691 let tiles = rows * cols / GROUP_SIZE;
12692 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12693 for t in 0..tiles {
12694 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12695 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12696 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12697 q1.extend_from_slice(&bits.to_le_bytes());
12698 }
12699 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12700 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12701 (q1, rp.data)
12702 }
12703
12704 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12707 let mut x: Vec<f32> = (0..cols)
12708 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12709 .collect();
12710 x[3 % cols] = 41.0;
12711 x[(cols / 2 + 5) % cols] = -37.5;
12712 x
12713 }
12714
12715 #[test]
12722 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12723 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12724 let gpr = cols / GROUP_SIZE;
12725 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12726 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12727 let x = q1g_test_x(cols, rows);
12728 let act = split_act(&x);
12729 assert!(!act.outliers.is_empty(), "test wants the outlier path");
12730 let gsum = q1_group_sums(&act.xq, gpr);
12731 for r in 0..rows {
12732 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12733 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12734 let (a, b) = (
12735 wa.dot_i8(r, gpr, &act.xq, &gsum),
12736 wb.dot_i8(r, gpr, &act.xq, &gsum),
12737 );
12738 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12739 for j in [0, 31, 32, 127, cols - 1] {
12740 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12741 assert_eq!(
12742 (a.0.to_bits(), a.1.to_bits()),
12743 (b.0.to_bits(), b.1.to_bits())
12744 );
12745 }
12746 #[cfg(target_arch = "x86_64")]
12747 if avx2_enabled() {
12748 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12749 unsafe {
12751 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12752 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12753 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12754 }
12755 if std::is_x86_feature_detected!("avx512vnni")
12756 && std::is_x86_feature_detected!("avx512bw")
12757 && std::is_x86_feature_detected!("avx512vl")
12758 {
12759 unsafe {
12761 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12762 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12763 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12764 }
12765 }
12766 }
12767 }
12768 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12770 bin_matvec(wa, &x, rows, cols, &mut ga, None);
12771 bin_matvec(wb, &x, rows, cols, &mut gb, None);
12772 assert_eq!(ga, gb, "{rows}x{cols} matvec");
12773 let x2 = q1g_test_x(cols, rows + 1);
12774 let (mut a1, mut a2, mut b1, mut b2) = (
12775 vec![0f32; rows],
12776 vec![0f32; rows],
12777 vec![0f32; rows],
12778 vec![0f32; rows],
12779 );
12780 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12781 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12782 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12783 let b = 5usize;
12784 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12785 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12786 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12787 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12788 assert_eq!(ma, mb, "{rows}x{cols} matmat");
12789 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12792 .map(|k| {
12793 let a = split_act(&xs[k * cols..(k + 1) * cols]);
12794 let g = q1_group_sums(&a.xq, gpr);
12795 (a, g)
12796 })
12797 .collect();
12798 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12799 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12800 for r in 0..rows {
12801 #[cfg(target_arch = "x86_64")]
12802 if avx2_enabled() {
12803 let (a, b) = unsafe {
12805 (
12806 wa.dot_1x4_x86(r, gpr, xq4, gs4),
12807 wb.dot_1x4_x86(r, gpr, xq4, gs4),
12808 )
12809 };
12810 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12811 }
12812 #[cfg(target_arch = "aarch64")]
12813 if sdot_enabled() {
12814 let (a, b) = unsafe {
12816 (
12817 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12818 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12819 )
12820 };
12821 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12822 }
12823 let _ = (r, &xq4, &gs4);
12824 }
12825 }
12826 }
12827
12828 #[test]
12834 #[ignore]
12835 fn q1g_cpu_matvec_speed() {
12836 use crate::pool::Pool;
12837 let threads = std::env::var("CMF_THREADS")
12838 .ok()
12839 .and_then(|v| v.parse().ok())
12840 .unwrap_or(8usize);
12841 let pool = Pool::new(threads);
12842 for &(name, rows, cols, reps) in &[
12843 ("gate/up_proj", 6144usize, 2048usize, 41usize),
12844 ("down_proj", 2048, 6144, 41),
12845 ("lm_head", 151_669, 2048, 9),
12846 ] {
12847 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
12848 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12849 let x = q1g_test_x(cols, 1);
12850 let mut out = vec![0f32; rows];
12851 let mut time = |w: BinW| {
12852 let t0 = std::time::Instant::now();
12853 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
12854 t0.elapsed().as_secs_f64()
12855 };
12856 let _ = (time(wa), time(wb));
12857 let (mut ta, mut tb) = (Vec::new(), Vec::new());
12858 for r in 0..reps {
12859 if r % 2 == 0 {
12860 ta.push(time(wa));
12861 tb.push(time(wb));
12862 } else {
12863 tb.push(time(wb));
12864 ta.push(time(wa));
12865 }
12866 }
12867 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
12868 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
12869 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
12870 eprintln!(
12871 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
12872 a8w8_enabled(),
12873 ma * 1e6,
12874 mb * 1e6,
12875 mb / ma
12876 );
12877 }
12878 }
12879
12880 #[test]
12885 fn q1g_mapped_tensor_paths_equal_q1() {
12886 use crate::pool::Pool;
12887 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12888 let (rows, cols) = (300usize, 256usize);
12889 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
12890 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
12891 let header: CmfHeader = serde_json::from_value(serde_json::json!({
12892 "format": "cmf", "version": cortiq_core::CMF_VERSION,
12893 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
12894 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12895 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
12896 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
12897 "max_position_embeddings": 8},
12898 "quant_type": "VBIT"
12899 }))
12900 .unwrap();
12901 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
12902 name: name.into(),
12903 dtype,
12904 shape: vec![rows, cols],
12905 data: data.clone(),
12906 };
12907 let specs = [
12908 spec("g1", TensorDtype::Q1, &gq1),
12909 spec("u1", TensorDtype::Q1, &uq1),
12910 spec("gg", TensorDtype::Q1G, &gq1g),
12911 spec("ug", TensorDtype::Q1G, &uq1g),
12912 ];
12913 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
12914 std::fs::create_dir_all(&dir).unwrap();
12915 let path = dir.join("m.cmf");
12916 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12917 let model = Arc::new(CmfModel::open(&path).unwrap());
12918 assert!(model.verify().is_empty(), "{:?}", model.verify());
12919 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
12920 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
12921 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
12922 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
12923 let pool = Pool::new(3);
12924 let x = q1g_test_x(cols, 4);
12925 let run = |a: &QTensor| {
12926 let mut o = vec![0f32; rows];
12927 a.matvec(&x, &mut o, Some(&pool));
12928 o
12929 };
12930 assert_eq!(run(&g1), run(&gg), "matvec");
12931 let many = |a: &QTensor, b: &QTensor| {
12932 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
12933 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
12934 (oa, ob)
12935 };
12936 let want = many(&g1, &u1);
12937 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
12938 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
12939 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
12940 let fused = |g: &QTensor, u: &QTensor| {
12941 let mut o = vec![0f32; rows];
12942 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
12943 (ok, o)
12944 };
12945 let (ok_a, fa) = fused(&g1, &u1);
12946 let (ok_b, fb) = fused(&gg, &ug);
12947 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
12948 assert_eq!(fa, fb, "fused gate/up");
12949 let b = 6usize;
12950 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
12951 let mm = |a: &QTensor| {
12952 let mut o = vec![0f32; b * rows];
12953 a.matmat(&xs, b, &mut o, Some(&pool));
12954 o
12955 };
12956 assert_eq!(mm(&g1), mm(&gg), "matmat");
12957 let x2 = q1g_test_x(cols, 77);
12958 let pair = |a: &QTensor| {
12959 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12960 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
12961 (o1, o2)
12962 };
12963 assert_eq!(pair(&g1), pair(&gg), "matvec2");
12964 for r in [0, 1, rows / 2, rows - 1] {
12965 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
12966 g1.row_f32(r, &mut ra);
12967 gg.row_f32(r, &mut rb);
12968 assert_eq!(ra, rb, "embedding row {r}");
12969 }
12970 drop((g1, u1, gg, ug, model));
12971 let _ = std::fs::remove_dir_all(&dir);
12972 }
12973
12974 #[test]
12981 fn blr_mapped_tensor_paths_equal_dense() {
12982 use crate::pool::Pool;
12983 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12984 let (rows, cols, r) = (96usize, 128usize, 16usize);
12985 let mut rng = crate::blr::tests::Rng(41);
12986 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
12987 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
12988 let mut q8: Vec<u8> = (0..rows * cols)
12990 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
12991 .collect();
12992 for o in 0..rows {
12993 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
12994 }
12995 let header: CmfHeader = serde_json::from_value(serde_json::json!({
12996 "format": "cmf", "version": cortiq_core::CMF_VERSION,
12997 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
12998 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12999 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13000 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13001 "max_position_embeddings": 8},
13002 "quant_type": "Q8_2F"
13003 }))
13004 .unwrap();
13005 let specs = [
13006 TensorSpec {
13007 name: "g".into(),
13008 dtype: TensorDtype::BinaryLowRank,
13009 shape: vec![rows, cols],
13010 data: gate.clone(),
13011 },
13012 TensorSpec {
13013 name: "u".into(),
13014 dtype: TensorDtype::BinaryLowRank,
13015 shape: vec![rows, cols],
13016 data: up.clone(),
13017 },
13018 TensorSpec {
13019 name: "q".into(),
13020 dtype: TensorDtype::Q8Row,
13021 shape: vec![rows, cols],
13022 data: q8,
13023 },
13024 ];
13025 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13026 std::fs::create_dir_all(&dir).unwrap();
13027 let path = dir.join("m.cmf");
13028 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13029 let model = Arc::new(CmfModel::open(&path).unwrap());
13030 assert!(model.verify().is_empty(), "{:?}", model.verify());
13031 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13032 let (g, u, q) = (t("g"), t("u"), t("q"));
13033 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13034 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13035 assert!(g.as_f32().is_none() && g.model_idx().is_none());
13036 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13037 assert!(!g.sparse_col_ok());
13038 assert_eq!((g.rows(), g.cols()), (rows, cols));
13039 let deq = |bytes: &[u8]| {
13040 let mut w = vec![0f32; rows * cols];
13041 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13042 w
13043 };
13044 let (wg, wu) = (deq(&gate), deq(&up));
13045 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13046 (0..rows)
13047 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13048 .collect()
13049 };
13050 let close = |a: &[f32], b: &[f32], what: &str| {
13051 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13052 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13053 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13054 };
13055 let pool = Pool::new(3);
13056 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13057 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13058 let mut o = vec![0f32; rows];
13059 g.matvec(&x, &mut o, Some(&pool));
13060 close(&o, &mv(&wg, &x), "matvec");
13061 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13062 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13063 close(&o1, &mv(&wg, &x), "matvec2 a");
13064 close(&o2, &mv(&wg, &x2), "matvec2 b");
13065 for b in [2usize, 5] {
13066 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13067 let mut om = vec![0f32; b * rows];
13068 g.matmat(&xs, b, &mut om, Some(&pool));
13069 for t in 0..b {
13070 close(
13071 &om[t * rows..(t + 1) * rows],
13072 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13073 "matmat",
13074 );
13075 }
13076 }
13077 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13080 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13081 close(&a, &mv(&wg, &x), "matvec_many blr gate");
13082 close(&c, &mv(&wu, &x), "matvec_many blr up");
13083 let mut qo = vec![0f32; rows];
13084 q.matvec(&x, &mut qo, Some(&pool));
13085 assert_eq!(bq, qo, "matvec_many q8_row sibling");
13086 let mut f = vec![0f32; rows];
13087 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13088 let mut row = vec![0f32; cols];
13089 g.row_f32(rows / 3, &mut row);
13090 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13091 let mut scratch = vec![0f32; cols];
13092 let d = g.row_dot(7, &x, &mut scratch);
13093 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13094 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13095 let mut col = vec![0f32; rows];
13096 g.add_col_scaled(cols / 2, 1.5, &mut col);
13097 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13098 close(&col, &want, "add_col_scaled");
13099 drop((g, u, q, model));
13100 let _ = std::fs::remove_dir_all(&dir);
13101 }
13102
13103 #[test]
13104 fn repack_is_bit_identical() {
13105 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
13111 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13112 .collect();
13113 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13114 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13115 let rep = q8_repack_layout(&w, rows, cols);
13116 for g in 0..rows / 4 {
13118 for c in 0..cols / 16 {
13119 for lane in 0..4 {
13120 assert_eq!(
13121 &rep[g * 4 * cols + c * 64 + lane * 16
13122 ..g * 4 * cols + c * 64 + lane * 16 + 16],
13123 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13124 );
13125 }
13126 }
13127 }
13128 let mut a = vec![0.0f32; rows];
13129 qmatvec(
13130 &w,
13131 &[],
13132 &scales,
13133 &x,
13134 &[],
13135 TensorDtype::Q8Row,
13136 rows,
13137 cols,
13138 &mut a,
13139 None,
13140 );
13141 let mut b = vec![0.0f32; rows];
13142 qmatvec(
13143 &w,
13144 &rep,
13145 &scales,
13146 &x,
13147 &[],
13148 TensorDtype::Q8Row,
13149 rows,
13150 cols,
13151 &mut b,
13152 None,
13153 );
13154 assert_eq!(a, b, "full-range repack output diverged");
13155
13156 #[cfg(target_arch = "aarch64")]
13157 if sdot_enabled() {
13158 let act = split_act(&x);
13160 let mut c1 = vec![0.0f32; rows];
13161 let mut c2 = vec![0.0f32; rows];
13162 q8_range_sdot(
13163 &w,
13164 &[],
13165 &scales,
13166 &act,
13167 cols,
13168 SendMut(c1.as_mut_ptr()),
13169 3,
13170 rows - 2,
13171 );
13172 q8_range_sdot(
13173 &w,
13174 &rep,
13175 &scales,
13176 &act,
13177 cols,
13178 SendMut(c2.as_mut_ptr()),
13179 3,
13180 rows - 2,
13181 );
13182 assert_eq!(c1, c2, "unaligned-range repack output diverged");
13183 }
13184 }
13185
13186 #[test]
13187 fn sdot_a8w8_noise_is_bounded() {
13188 let (rows, cols) = (16, 512);
13192 let w: Vec<u8> = (0..rows * cols)
13193 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13194 .collect();
13195 let scales = vec![0.01f32; rows];
13196 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13197 let mut a = vec![0.0f32; rows];
13198 qmatvec(
13199 &w,
13200 &[],
13201 &scales,
13202 &x,
13203 &[],
13204 TensorDtype::Q8Row,
13205 rows,
13206 cols,
13207 &mut a,
13208 None,
13209 );
13210 let (mut num, mut den) = (0f64, 0f64);
13211 for o in 0..rows {
13212 let mut acc = 0.0f32;
13213 for j in 0..cols {
13214 acc += (w[o * cols + j] as i8) as f32 * x[j];
13215 }
13216 let expect = acc * scales[o];
13217 num += ((a[o] - expect) as f64).powi(2);
13218 den += (expect as f64).powi(2);
13219 }
13220 let rel = (num / den.max(1e-12)).sqrt();
13221 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13222 }
13223
13224 #[test]
13225 fn i8_dot_neon_matches_scalar() {
13226 let n = 100;
13227 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13228 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13229 let mut scalar = 0.0f32;
13230 for j in 0..n {
13231 scalar += (w[j] as i8) as f32 * x[j];
13232 }
13233 let fast = dot_i8_f32(&w, &x);
13234 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13235 }
13236
13237 #[test]
13239 fn vbitmatvec_matches_full_dequant() {
13240 let (rows, cols) = (6, 64);
13241 let ng = cols / GROUP_SIZE;
13242 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13244 let mut bytes = bits.clone();
13245 for g in 0..rows * ng {
13246 let s = 0.02 + 0.001 * g as f32;
13247 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13248 }
13249 for r in 0..rows {
13250 let b = bits[r] as usize;
13251 let (mut acc, mut nb) = (0u64, 0usize);
13252 let mut rowbytes = Vec::new();
13253 for i in 0..cols {
13254 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13255 acc = (acc << b) | v;
13256 nb += b;
13257 while nb >= 8 {
13258 nb -= 8;
13259 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13260 }
13261 }
13262 if nb > 0 {
13263 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13264 }
13265 bytes.extend_from_slice(&rowbytes);
13266 }
13267 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13268
13269 let mut reference = vec![0f32; rows * cols];
13270 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13271 let mut expect = vec![0f32; rows];
13272 for r in 0..rows {
13273 expect[r] = reference[r * cols..(r + 1) * cols]
13274 .iter()
13275 .zip(&x)
13276 .map(|(w, xv)| w * xv)
13277 .sum();
13278 }
13279 let mut got = vec![0f32; rows];
13280 let offsets = vbit_row_offsets(&bytes, rows, cols);
13281 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13282 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13286 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13287 for r in 0..rows {
13288 assert!(
13289 (got[r] - expect[r]).abs() < tol * scale,
13290 "row {r}: {} vs {}",
13291 got[r],
13292 expect[r]
13293 );
13294 }
13295 }
13296
13297 #[test]
13302 #[cfg(target_arch = "x86_64")]
13303 fn vbit_matmat_blocked_matches_per_row() {
13304 let (rows, cols, b) = (64usize, 128usize, 9usize);
13305 let ng = cols / GROUP_SIZE;
13306 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13307 let mut bytes = bits.clone();
13308 for g in 0..rows * ng {
13309 let sc = 0.02 + 0.0005 * g as f32;
13310 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13311 }
13312 for r in 0..rows {
13313 let bw = bits[r] as usize;
13314 let (mut acc, mut nb) = (0u64, 0usize);
13315 let mut rowbytes = Vec::new();
13316 for i in 0..cols {
13317 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13318 acc = (acc << bw) | v;
13319 nb += bw;
13320 while nb >= 8 {
13321 nb -= 8;
13322 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13323 }
13324 }
13325 if nb > 0 {
13326 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13327 }
13328 bytes.extend_from_slice(&rowbytes);
13329 }
13330 let x: Vec<f32> = (0..b * cols)
13331 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13332 .collect();
13333 let offsets = vbit_row_offsets(&bytes, rows, cols);
13334 let mut y_a = vec![0f32; b * rows];
13335 let mut y_b = vec![0f32; b * rows];
13336 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13337 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13338 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13339 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13340 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13341 let max_d = y_a
13342 .iter()
13343 .zip(&y_b)
13344 .map(|(p, q)| (p - q).abs())
13345 .fold(0.0f32, f32::max);
13346 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13347 }
13348
13349 #[test]
13357 fn q4t_matmat_blocked_matches_per_row() {
13358 let (rows, cols, b) = (16usize, 64usize, 9usize);
13359 let gpr = cols / GROUP_SIZE;
13360 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13361 for r in 0..rows {
13362 for g in 0..gpr {
13363 let t = (r * gpr + g) * Q4_TILE;
13364 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13365 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13366 for k in 0..16 {
13367 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13368 }
13369 }
13370 }
13371 let x: Vec<f32> = (0..b * cols)
13372 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13373 .collect();
13374 let mut y_blk = vec![0f32; b * rows];
13375 let mut y_row = vec![0f32; b * rows];
13376 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13377 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13378 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13379 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13380 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13381 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13382 }
13383
13384 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13391 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13392 let gpr = cols / GROUP_SIZE;
13393 let stride = q4tp_code_stride(gpr);
13394 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13395 let mut b = vec![0u8; codes_off + rows * stride];
13396 for r in 0..rows {
13397 for g in 0..gpr {
13398 let t = (r * gpr + g) * Q4TP_NIB;
13399 for k in 0..16 {
13400 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13401 }
13402 }
13403 let lo = -6.0 - 0.03 * (r % 17) as f32;
13404 let step = 0.01 + 0.004 * (r % 11) as f32;
13405 let p = params_off + r * 4;
13406 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13407 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13408 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13409 for g in 0..gpr {
13410 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13411 }
13412 }
13413 b
13414 }
13415
13416 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13420 let gpr = cols / GROUP_SIZE;
13421 let v = Q4tpView::new(bytes, rows, cols);
13422 let mut out = vec![0u8; rows * gpr * Q4_TILE];
13423 let mut sc = vec![0f32; gpr];
13424 for r in 0..rows {
13425 v.scales_into(r, gpr, &mut sc);
13426 for g in 0..gpr {
13427 let t = (r * gpr + g) * Q4_TILE;
13428 let s = sc[g];
13429 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13430 let src = (r * gpr + g) * Q4TP_NIB;
13431 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13432 }
13433 }
13434 out
13435 }
13436
13437 #[test]
13443 fn q4tp_exact_path_matches_dequant_reference() {
13444 let (rows, cols) = (256usize, 512usize);
13445 let gpr = cols / GROUP_SIZE;
13446 let bytes = synth_q4tp(rows, cols);
13447 let mut w = vec![0f32; rows * cols];
13448 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13449
13450 let x: Vec<f32> = (0..cols)
13451 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13452 .collect();
13453 let v = Q4tpView::new(&bytes, rows, cols);
13454 let mut sc = vec![0f32; gpr];
13455 for r in 0..rows {
13456 v.scales_into(r, gpr, &mut sc);
13457 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13458 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13459 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13463 assert!(
13464 (got - want).abs() <= 1e-5 * mag,
13465 "row {r}: kernel {got} vs dequant {want}"
13466 );
13467 }
13468 }
13469
13470 #[test]
13476 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13477 let (rows, cols) = (256usize, 512usize);
13478 let bytes = synth_q4tp(rows, cols);
13479 let twin = q4tp_as_q4t(&bytes, rows, cols);
13480 let x: Vec<f32> = (0..cols)
13481 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13482 .collect();
13483
13484 let mut got = vec![0f32; rows];
13485 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13486 let mut want = vec![0f32; rows];
13487 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13488
13489 let mut w = vec![0f32; rows * cols];
13492 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13493 for r in 0..rows {
13494 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13495 assert!(
13496 (got[r] - want[r]).abs() <= 1e-3 * mag,
13497 "row {r}: q4tp {} vs q4t {}",
13498 got[r],
13499 want[r]
13500 );
13501 }
13502 }
13503
13504 #[test]
13509 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13510 let (rows, cols, b) = (256usize, 512usize, 5usize);
13511 let bytes = synth_q4tp(rows, cols);
13512 let twin = q4tp_as_q4t(&bytes, rows, cols);
13513 let xs: Vec<f32> = (0..b * cols)
13514 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13515 .collect();
13516
13517 let mut got = vec![0f32; b * rows];
13518 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13519 let mut want = vec![0f32; b * rows];
13520 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13521
13522 let mut w = vec![0f32; rows * cols];
13523 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13524 for t in 0..b {
13525 for r in 0..rows {
13526 let mag: f32 = (0..cols)
13527 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13528 .sum();
13529 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13530 assert!(
13531 (g - wa).abs() <= 1e-3 * mag,
13532 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13533 );
13534 }
13535 }
13536 }
13537
13538 #[test]
13539 fn q4tp_matvec2_matches_the_single_stream_kernel() {
13540 let (rows, cols) = (128usize, 256usize);
13541 let gpr = cols / GROUP_SIZE;
13542 let bytes = synth_q4tp(rows, cols);
13543 let xs: Vec<f32> = (0..2 * cols)
13544 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13545 .collect();
13546
13547 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13548 q4tp_matvec2(
13549 &bytes,
13550 &xs[..cols],
13551 &xs[cols..],
13552 rows,
13553 cols,
13554 &mut o1,
13555 &mut o2,
13556 None,
13557 );
13558
13559 let v = Q4tpView::new(&bytes, rows, cols);
13562 let mut sc = vec![0f32; gpr];
13563 for r in 0..rows {
13564 v.scales_into(r, gpr, &mut sc);
13565 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13566 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13567 }
13568 }
13569
13570 #[test]
13577 fn q4tp_matvec_keeps_pace_with_q4t() {
13578 let (rows, cols) = (4096usize, 3072usize);
13579 let bytes = synth_q4tp(rows, cols);
13580 let twin = q4tp_as_q4t(&bytes, rows, cols);
13581 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13582 let mut o = vec![0f32; rows];
13583 let n = 12;
13584 let mut best = (f64::MAX, f64::MAX);
13585 for _ in 0..3 {
13588 let t0 = std::time::Instant::now();
13589 for _ in 0..n {
13590 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13591 }
13592 best.0 = best.0.min(t0.elapsed().as_secs_f64());
13593 let t0 = std::time::Instant::now();
13594 for _ in 0..n {
13595 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13596 }
13597 best.1 = best.1.min(t0.elapsed().as_secs_f64());
13598 }
13599 let ratio = best.1 / best.0;
13600 println!(
13601 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13602 best.0 * 1e3 / n as f64,
13603 best.1 * 1e3 / n as f64
13604 );
13605 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13606 }
13607
13608 #[cfg(target_os = "macos")]
13609 #[test]
13610 fn q4t_matmat_accel_matches_dequant_reference() {
13611 if !accel_gemm_enabled() {
13612 return; }
13614 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
13616 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13617 for r in 0..rows {
13618 for g in 0..gpr {
13619 let t = (r * gpr + g) * Q4_TILE;
13620 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13621 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13622 for k in 0..16 {
13623 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13624 }
13625 }
13626 }
13627 let x: Vec<f32> = (0..b * cols)
13628 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13629 .collect();
13630 let mut got = vec![0f32; b * rows];
13631 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13632 let mut w = vec![0f32; rows * cols];
13634 for r in 0..rows {
13635 for g in 0..gpr {
13636 let t = (r * gpr + g) * Q4_TILE;
13637 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13638 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13639 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13640 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13641 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13642 }
13643 }
13644 }
13645 for bi in 0..b {
13646 for r in 0..rows {
13647 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13648 let d = (got[bi * rows + r] - want).abs();
13649 assert!(
13650 d <= want.abs().max(1.0) * 1e-4,
13651 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13652 got[bi * rows + r]
13653 );
13654 }
13655 }
13656 }
13657
13658 #[test]
13659 fn q4matvec_matches_full_dequant() {
13660 let (rows, cols) = (8, 64);
13661 let groups = rows * cols / GROUP_SIZE;
13662 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13664 for i in 0..groups * 16 {
13665 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13666 }
13667 for g in 0..groups {
13668 let s = 0.01 + 0.003 * g as f32;
13669 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13670 }
13671 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13672
13673 let mut reference = vec![0.0f32; rows * cols];
13674 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13675 let mut expect = vec![0.0f32; rows];
13676 for r in 0..rows {
13677 expect[r] = reference[r * cols..(r + 1) * cols]
13678 .iter()
13679 .zip(&x)
13680 .map(|(w, xv)| w * xv)
13681 .sum();
13682 }
13683
13684 let mut got = vec![0.0f32; rows];
13685 q4matvec(&bytes, &x, rows, cols, &mut got, None);
13686 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13690 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13691 for r in 0..rows {
13692 assert!(
13693 (got[r] - expect[r]).abs() < tol * scale,
13694 "row {r}: {} vs {}",
13695 got[r],
13696 expect[r]
13697 );
13698 }
13699 }
13700
13701 #[test]
13704 fn vbitmatvec2_equals_two_singles() {
13705 let (rows, cols) = (6, 64);
13706 let ng = cols / GROUP_SIZE;
13707 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13708 let mut bytes = bits.clone();
13709 for g in 0..rows * ng {
13710 let s = 0.02 + 0.001 * g as f32;
13711 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13712 }
13713 for r in 0..rows {
13714 let b = bits[r] as usize;
13715 let (mut acc, mut nb) = (0u64, 0usize);
13716 let mut rowbytes = Vec::new();
13717 for i in 0..cols {
13718 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13719 acc = (acc << b) | v;
13720 nb += b;
13721 while nb >= 8 {
13722 nb -= 8;
13723 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13724 }
13725 }
13726 if nb > 0 {
13727 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13728 }
13729 bytes.extend_from_slice(&rowbytes);
13730 }
13731 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13732 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13733 let offsets = vbit_row_offsets(&bytes, rows, cols);
13734
13735 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13736 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13737 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13738 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13739 vbitmatvec2(
13740 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13741 );
13742 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13743 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13744 }
13745
13746 #[test]
13748 fn q4matvec2_equals_two_singles() {
13749 let (rows, cols) = (8, 128);
13750 let groups = rows * cols / GROUP_SIZE;
13751 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13752 for i in 0..groups * 16 {
13753 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13754 }
13755 for g in 0..groups {
13756 let s = 0.01 + 0.003 * g as f32;
13757 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13758 }
13759 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13762 x1[9] = 250.0;
13763 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13764
13765 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13766 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13767 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13768 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13769 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13770 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13771 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13772 }
13773
13774 #[test]
13777 fn matvec_many_equals_separate_matvecs() {
13778 use crate::pool::Pool;
13779 let (r1, r2, cols) = (300, 200, 64);
13780 let mk = |salt: usize, rows: usize| {
13781 QTensor::from_f32(
13782 (0..rows * cols)
13783 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13784 .collect(),
13785 rows,
13786 cols,
13787 )
13788 };
13789 let (a, b) = (mk(1, r1), mk(5, r2));
13790 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13791 let pool = Pool::new(3);
13792
13793 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13794 a.matvec(&x, &mut ea, Some(&pool));
13795 b.matvec(&x, &mut eb, Some(&pool));
13796 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13797 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13798 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13799 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13800 }
13801
13802 #[test]
13807 fn q4tp_matvec_many_equals_separate_matvecs() {
13808 use crate::pool::Pool;
13809 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13810
13811 let (r1, r2, cols) = (300usize, 200usize, 64usize);
13812 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13813 "arch_name": "tiny-q4tp",
13814 "hidden_size": cols,
13815 "intermediate_size": cols * 2,
13816 "num_layers": 1,
13817 "num_attention_heads": 2,
13818 "num_kv_heads": 1,
13819 "head_dim": 32,
13820 "vocab_size": r1,
13821 "layer_types": ["FullAttention"],
13822 "rms_norm_eps": 1e-6,
13823 "max_position_embeddings": 8,
13824 "linear_conv_kernel_dim": 0,
13825 "linear_num_key_heads": 0,
13826 "linear_num_value_heads": 0
13827 }))
13828 .unwrap();
13829 let header = CmfHeader {
13830 format: "cmf".into(),
13831 version: CMF_VERSION,
13832 arch,
13833 quant_type: QuantType::Q4Block,
13834 provenance: None,
13835 tokenizer_config: None,
13836 section_hashes: None,
13837 skills: Vec::new(),
13838 shard: None,
13839 calibration: None,
13840 routing: None,
13841 genome: None,
13842 lineage: Vec::new(),
13843 router: None,
13844 segments: Vec::new(),
13845 };
13846 let specs = [
13847 TensorSpec {
13848 name: "q".into(),
13849 dtype: TensorDtype::Q4TiledP,
13850 shape: vec![r1, cols],
13851 data: synth_q4tp(r1, cols),
13852 },
13853 TensorSpec {
13854 name: "kv".into(),
13855 dtype: TensorDtype::Q4TiledP,
13856 shape: vec![r2, cols],
13857 data: synth_q4tp(r2, cols),
13858 },
13859 ];
13860 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
13861 std::fs::create_dir_all(&dir).unwrap();
13862 let path = dir.join("m.cmf");
13863 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13864 let model = Arc::new(CmfModel::open(&path).unwrap());
13865 let (a, b) = (
13866 QTensor::from_model(&model, "q").unwrap(),
13867 QTensor::from_model(&model, "kv").unwrap(),
13868 );
13869 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
13870 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
13871 let x: Vec<f32> = (0..cols)
13872 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
13873 .collect();
13874 let pool = Pool::new(3);
13875 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13876 a.matvec(&x, &mut ea, Some(&pool));
13877 b.matvec(&x, &mut eb, Some(&pool));
13878 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13879 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13880 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
13881 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
13882 let _ = std::fs::remove_dir_all(&dir);
13883 }
13884
13885 #[test]
13891 fn multi_token_moe_rows_equal_single_token_decode() {
13892 use crate::pool::Pool;
13893 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13894
13895 let (h, inter, ne) = (64usize, 128usize, 3usize);
13896 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13897 "arch_name": "tiny-q4tp-moe",
13898 "hidden_size": h,
13899 "intermediate_size": inter,
13900 "num_layers": 1,
13901 "num_attention_heads": 2,
13902 "num_kv_heads": 1,
13903 "head_dim": 32,
13904 "vocab_size": 8,
13905 "layer_types": ["FullAttention"],
13906 "rms_norm_eps": 1e-6,
13907 "max_position_embeddings": 8,
13908 "linear_conv_kernel_dim": 0,
13909 "linear_num_key_heads": 0,
13910 "linear_num_value_heads": 0
13911 }))
13912 .unwrap();
13913 let header = CmfHeader {
13914 format: "cmf".into(),
13915 version: CMF_VERSION,
13916 arch,
13917 quant_type: QuantType::Q4Block,
13918 provenance: None,
13919 tokenizer_config: None,
13920 section_hashes: None,
13921 skills: Vec::new(),
13922 shard: None,
13923 calibration: None,
13924 routing: None,
13925 genome: None,
13926 lineage: Vec::new(),
13927 router: None,
13928 segments: Vec::new(),
13929 };
13930 let mut specs = Vec::new();
13931 for e in 0..ne {
13932 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
13933 .into_iter()
13934 .enumerate()
13935 {
13936 let mut data = synth_q4tp(r, c);
13939 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
13940 .iter_mut()
13941 .enumerate()
13942 {
13943 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
13944 }
13945 specs.push(TensorSpec {
13946 name: format!("{n}{e}"),
13947 dtype: TensorDtype::Q4TiledP,
13948 shape: vec![r, c],
13949 data,
13950 });
13951 }
13952 }
13953 let dir = std::env::temp_dir().join(format!(
13954 "cmf-moe-rows-{}-{}",
13955 std::process::id(),
13956 FLOAT_ACTIVATIONS.get()
13957 ));
13958 std::fs::create_dir_all(&dir).unwrap();
13959 let path = dir.join("m.cmf");
13960 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13961 let model = Arc::new(CmfModel::open(&path).unwrap());
13962 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
13963 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
13964 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
13965 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
13966 let b = 4usize;
13967 let mut xs: Vec<f32> = (0..b * h)
13968 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
13969 .collect();
13970 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
13973 (vec![2, 0], vec![0.6, 0.4]),
13974 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
13975 (vec![1], vec![1.0]),
13976 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
13977 ];
13978 let pool = Pool::new(3);
13979 let mut want = vec![0f32; b * h];
13981 for (tk, (idx, w)) in routes.iter().enumerate() {
13982 let x = &xs[tk * h..(tk + 1) * h];
13983 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
13984 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
13985 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
13986 if FLOAT_ACTIVATIONS.get() {
13987 for (slot, &e) in idx.iter().enumerate() {
13988 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
13989 g[e].matvec(x, &mut gate, Some(&pool));
13990 u[e].matvec(x, &mut up, Some(&pool));
13991 for (v, u) in gate.iter_mut().zip(up) {
13992 *v = (*v / (1.0 + (-*v).exp())) * u;
13993 }
13994 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
13995 }
13996 }
13997 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
13998 assert!(QTensor::moe_down_many(
13999 &downs,
14000 &gs,
14001 w,
14002 &mut want[tk * h..(tk + 1) * h],
14003 Some(&pool)
14004 ));
14005 }
14006 if FLOAT_ACTIVATIONS.get() {
14007 for (tk, (idx, w)) in routes.iter().enumerate() {
14008 let mut scalar = vec![0.0; h];
14009 for (&e, &weight) in idx.iter().zip(w) {
14010 let (mut gate, mut up, mut down) =
14011 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14012 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14013 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14014 for (v, u) in gate.iter_mut().zip(up) {
14015 *v = (*v / (1.0 + (-*v).exp())) * u;
14016 }
14017 d[e].matvec(&gate, &mut down, Some(&pool));
14018 for (v, d) in scalar.iter_mut().zip(down) {
14019 *v += weight * d;
14020 }
14021 }
14022 assert_eq!(
14023 &want[tk * h..(tk + 1) * h],
14024 scalar,
14025 "float many equals scalar experts"
14026 );
14027 }
14028 }
14029 let mut experts: Vec<usize> = Vec::new();
14031 let mut groups: Vec<Vec<usize>> = Vec::new();
14032 for (tk, (idx, _)) in routes.iter().enumerate() {
14033 for &e in idx {
14034 match experts.iter().position(|&x| x == e) {
14035 Some(k) => groups[k].push(tk),
14036 None => {
14037 experts.push(e);
14038 groups.push(vec![tk]);
14039 }
14040 }
14041 }
14042 }
14043 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14044 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14045 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14046 assert!(QTensor::moe_gate_up_rows(
14047 &pairs,
14048 &groups,
14049 &xs,
14050 &mut gs,
14051 Some(&pool)
14052 ));
14053 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14054 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14055 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14056 assert!(QTensor::moe_down_rows(
14057 &downs,
14058 &lens,
14059 &gs,
14060 &mut ds,
14061 Some(&pool)
14062 ));
14063 let slot = |tk: usize, e: usize| {
14064 let k = experts.iter().position(|&x| x == e).unwrap();
14065 groups[..k].iter().map(|g| g.len()).sum::<usize>()
14066 + groups[k].iter().position(|&x| x == tk).unwrap()
14067 };
14068 let mut got = vec![0f32; b * h];
14069 for (tk, (idx, w)) in routes.iter().enumerate() {
14070 for i in 0..h {
14071 let mut acc = 0f32;
14072 for (&e, &we) in idx.iter().zip(w) {
14073 acc += we * ds[slot(tk, e)][i];
14074 }
14075 got[tk * h + i] = acc;
14076 }
14077 }
14078 assert!(want.iter().any(|v| *v != 0.0));
14079 assert_eq!(
14080 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14081 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14082 "multi-token MoE must equal decode bit for bit"
14083 );
14084
14085 let b5 = 5usize;
14088 let x5: Vec<f32> = (0..b5 * h)
14089 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14090 .collect();
14091 let mut mm = vec![0f32; b5 * inter];
14092 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14093 for tk in 0..b5 {
14094 let mut mv = vec![0f32; inter];
14095 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14096 assert_eq!(
14097 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14098 mm[tk * inter..(tk + 1) * inter]
14099 .iter()
14100 .map(|v| v.to_bits())
14101 .collect::<Vec<_>>(),
14102 "row-exact matmat token {tk}"
14103 );
14104 }
14105 let _ = std::fs::remove_dir_all(&dir);
14108 }
14109
14110 #[test]
14118 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14119 use crate::pool::Pool;
14120 use std::sync::atomic::Ordering::Relaxed;
14121 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14122 Q4TP_ALT.store(2, Relaxed);
14125 let pool = Pool::new(3);
14126 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14130 let bytes = synth_q4tp(rows, cols);
14131 let mut xs: Vec<f32> = (0..b * cols)
14132 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14133 .collect();
14134 xs[3] = 7.5; let run = |exact: bool| {
14136 let mut out = vec![0f32; b * rows];
14137 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14138 out
14139 };
14140 let (fast, exact) = (run(false), run(true));
14141 #[cfg(not(target_arch = "aarch64"))]
14142 let _ = fast;
14143 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14144 let mut matvecs = vec![0f32; b * rows];
14145 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14146 q4tp_matvec(
14147 &bytes,
14148 &xs[bi * cols..(bi + 1) * cols],
14149 rows,
14150 cols,
14151 o,
14152 Some(&pool),
14153 );
14154 }
14155 assert!(matvecs.iter().any(|v| *v != 0.0));
14156 assert_eq!(
14157 bits(&exact),
14158 bits(&matvecs),
14159 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14160 );
14161 #[cfg(target_arch = "aarch64")]
14162 {
14163 let gpr = cols / GROUP_SIZE;
14165 let v = Q4tpView::new(&bytes, rows, cols);
14166 let mut old = vec![0f32; b * rows];
14167 let mut sc = vec![0f32; gpr];
14168 let a8w8 = a8w8_enabled();
14169 let blocked = sdot_enabled() && blocked_enabled();
14170 let acts: Vec<SplitAct> = (0..b)
14171 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14172 .collect();
14173 for r in 0..rows {
14174 v.scales_into(r, gpr, &mut sc);
14175 if !a8w8 {
14176 for bi in 0..b {
14177 let x = &xs[bi * cols..(bi + 1) * cols];
14178 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14179 }
14180 continue;
14181 }
14182 let finish = |d: f32, act: &SplitAct| {
14183 let mut acc = d * act.sx;
14184 for &(j, xv) in &act.outliers {
14185 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14186 acc += w * s * xv;
14187 }
14188 acc
14189 };
14190 let mut bi = 0usize;
14191 while blocked && bi + 4 <= b {
14192 let xs4 = [
14193 acts[bi].xq.as_slice(),
14194 acts[bi + 1].xq.as_slice(),
14195 acts[bi + 2].xq.as_slice(),
14196 acts[bi + 3].xq.as_slice(),
14197 ];
14198 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14199 for k in 0..4 {
14200 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14201 }
14202 bi += 4;
14203 }
14204 for (bi, act) in acts.iter().enumerate().skip(bi) {
14205 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14206 old[bi * rows + r] = finish(d, act);
14207 }
14208 }
14209 assert_eq!(
14210 bits(&fast),
14211 bits(&old),
14212 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14213 );
14214 if blocked {
14217 assert_ne!(
14218 bits(&fast),
14219 bits(&matvecs),
14220 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14221 );
14222 }
14223 }
14224 }
14225 Q4TP_ALT.store(0, Relaxed);
14226 }
14227
14228 #[test]
14229 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14230 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14231 let active = AtomicUsize::new(0);
14234 counted_row_exact_scope(&active, || {
14235 assert_eq!(active.load(Ordering::Acquire), 1);
14236 counted_row_exact_scope(&active, || {
14237 assert_eq!(active.load(Ordering::Acquire), 2);
14238 });
14239 assert_eq!(active.load(Ordering::Acquire), 1);
14240 });
14241 assert_eq!(active.load(Ordering::Acquire), 0);
14242
14243 let both_entered = Barrier::new(2);
14244 let release_last = Barrier::new(2);
14245 std::thread::scope(|s| {
14246 let first = s.spawn(|| counted_row_exact_scope(&active, || {
14247 both_entered.wait();
14248 }));
14249 let last = s.spawn(|| counted_row_exact_scope(&active, || {
14250 both_entered.wait();
14251 release_last.wait();
14252 }));
14253 first.join().unwrap();
14254 let after_first = active.load(Ordering::Acquire);
14255 release_last.wait();
14256 last.join().unwrap();
14257 assert_eq!(after_first, 1, "second request must remain exact");
14258 });
14259 assert_eq!(active.load(Ordering::Acquire), 0);
14260 let panic = std::panic::catch_unwind(|| {
14261 counted_row_exact_scope(&active, || panic!("scope unwind"));
14262 });
14263 assert!(panic.is_err());
14264 assert_eq!(active.load(Ordering::Acquire), 0);
14265 }
14266
14267 #[test]
14268 #[cfg(target_arch = "x86_64")]
14269 fn q4tp_float_avx2_is_bitwise_scalar() {
14270 if !avx2_enabled() {
14271 return;
14272 }
14273 for cols in [32, 64, 96, 2048, 4096] {
14274 let rows = 9;
14275 let bytes = synth_q4tp(rows, cols);
14276 let v = Q4tpView::new(&bytes, rows, cols);
14277 let gpr = cols / GROUP_SIZE;
14278 let mut sc = vec![0.0; gpr];
14279 for seed in 1..=5 {
14280 let xs: Vec<f32> = (0..cols)
14281 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14282 .collect();
14283 for r in 0..rows {
14284 v.scales_into(r, gpr, &mut sc);
14285 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14286 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14287 assert_eq!(
14288 scalar.to_bits(),
14289 vector.to_bits(),
14290 "cols={cols} row={r} seed={seed}"
14291 );
14292 }
14293 }
14294 }
14295 }
14296
14297 #[test]
14298 fn multi_token_moe_rows_float_equal_single_token_decode() {
14299 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14300 }
14301
14302 #[test]
14303 fn full_gpu_q8_scope_is_nested_and_thread_local() {
14304 assert!(!FULL_GPU_Q8.get());
14305 let before = gpu_split_frac();
14306 {
14307 let _guard = enter_full_gpu_q8_scope();
14308 assert_eq!(gpu_split_frac(), 1.0);
14309 {
14310 let _nested = enter_full_gpu_q8_scope();
14311 }
14312 assert_eq!(gpu_split_frac(), 1.0);
14313 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14314 }
14315 assert!(!FULL_GPU_Q8.get());
14316 assert_eq!(gpu_split_frac(), before);
14317 }
14318
14319 #[test]
14320 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14321 assert!(!FLOAT_ACTIVATIONS.get());
14322 let before = a8w8_enabled();
14323 float_activations_scope(|| {
14324 assert!(!a8w8_enabled());
14325 float_activations_scope(|| assert!(!a8w8_enabled()));
14326 assert!(FLOAT_ACTIVATIONS.get());
14327 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14328 .join()
14329 .unwrap();
14330 });
14331 assert!(!FLOAT_ACTIVATIONS.get());
14332 assert_eq!(a8w8_enabled(), before);
14333 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14334 assert!(!FLOAT_ACTIVATIONS.get());
14335 }
14336
14337 #[test]
14340 fn batched_matmat_equals_per_position_matvec() {
14341 let (rows, cols, b) = (8, 64, 5);
14342 let groups = rows * cols / GROUP_SIZE;
14344 let mut q4 = Vec::new();
14345 for i in 0..groups * 16 {
14346 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14347 }
14348 for g in 0..groups {
14349 q4.extend_from_slice(
14350 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14351 );
14352 }
14353 let ng = cols / GROUP_SIZE;
14355 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14356 let mut vb = bits.clone();
14357 for g in 0..rows * ng {
14358 vb.extend_from_slice(
14359 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14360 );
14361 }
14362 for r in 0..rows {
14363 let bw = bits[r] as usize;
14364 let (mut acc, mut nb) = (0u64, 0usize);
14365 let mut rowbytes = Vec::new();
14366 for i in 0..cols {
14367 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14368 acc = (acc << bw) | v;
14369 nb += bw;
14370 while nb >= 8 {
14371 nb -= 8;
14372 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14373 }
14374 }
14375 if nb > 0 {
14376 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14377 }
14378 vb.extend_from_slice(&rowbytes);
14379 }
14380 let offsets = vbit_row_offsets(&vb, rows, cols);
14381
14382 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14383
14384 let mut got = vec![0f32; b * rows];
14386 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14387 for bi in 0..b {
14388 let mut expect = vec![0f32; rows];
14389 q4matvec(
14390 &q4,
14391 &xs[bi * cols..(bi + 1) * cols],
14392 rows,
14393 cols,
14394 &mut expect,
14395 None,
14396 );
14397 assert_eq!(
14398 &got[bi * rows..(bi + 1) * rows],
14399 &expect[..],
14400 "q4 batch pos {bi}"
14401 );
14402 }
14403
14404 let mut got = vec![0f32; b * rows];
14406 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14407 for bi in 0..b {
14408 let mut expect = vec![0f32; rows];
14409 vbitmatvec(
14410 &vb,
14411 &offsets,
14412 &xs[bi * cols..(bi + 1) * cols],
14413 rows,
14414 cols,
14415 &mut expect,
14416 None,
14417 );
14418 assert_eq!(
14419 &got[bi * rows..(bi + 1) * rows],
14420 &expect[..],
14421 "vbit batch pos {bi}"
14422 );
14423 }
14424 }
14425
14426 #[test]
14430 fn q4_tiled_matches_q4_block_bitexact() {
14431 let (rows, cols, b) = (8usize, 128usize, 3usize);
14432 let groups = rows * cols / GROUP_SIZE;
14433 let mut split = Vec::with_capacity(groups * 18);
14434 for i in 0..groups * 16 {
14435 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14436 }
14437 for g in 0..groups {
14438 split.extend_from_slice(
14439 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14440 );
14441 }
14442 let (packed, scales) = split.split_at(groups * 16);
14444 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14445 for g in 0..groups {
14446 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14447 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14448 }
14449
14450 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14451 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14453
14454 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14455 q4matvec(&split, &x1, rows, cols, &mut a, None);
14456 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14457 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14458
14459 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14460 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14461 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14462 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14463 assert_eq!(a1, t1);
14464 assert_eq!(a2, t2);
14465
14466 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14467 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14468 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14469 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14470 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14471 }
14472
14473 #[test]
14480 fn q4matvec_sdot_outlier_exact() {
14481 let (rows, cols) = (4, 128);
14482 let groups = rows * cols / GROUP_SIZE;
14483 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14484 for i in 0..groups * 16 {
14485 bytes.push(((i * 11 + 5) % 256) as u8);
14486 }
14487 for g in 0..groups {
14488 let s = 0.02 + 0.002 * g as f32;
14489 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14490 }
14491 let mut x: Vec<f32> = (0..cols)
14492 .map(|i| match i % 3 {
14493 0 => 1.0,
14494 1 => -1.0,
14495 _ => 0.0,
14496 })
14497 .collect();
14498 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
14501 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14502 let mut expect = vec![0.0f32; rows];
14503 for r in 0..rows {
14504 expect[r] = reference[r * cols..(r + 1) * cols]
14505 .iter()
14506 .zip(&x)
14507 .map(|(w, xv)| w * xv)
14508 .sum();
14509 }
14510 let mut got = vec![0.0f32; rows];
14511 q4matvec(&bytes, &x, rows, cols, &mut got, None);
14512 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14513 for r in 0..rows {
14514 assert!(
14515 (got[r] - expect[r]).abs() < 2e-3 * scale,
14516 "row {r}: {} vs {} (outlier term must be exact)",
14517 got[r],
14518 expect[r]
14519 );
14520 }
14521 }
14522
14523 #[test]
14527 fn q1t_matvec_matches_reference() {
14528 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14529 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
14531 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14532 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14534 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14535 let mut bytes = Vec::new();
14536 for r in 0..rows {
14537 for g in 0..gpr {
14538 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14539 let mut c = [0u8; 7];
14540 for k in 0..GROUP_SIZE {
14541 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14543 0
14544 } else {
14545 ((k + r * 3 + g) % 3) as u8 };
14547 cortiq_core::quant::q1t_pack(&mut c, k, code);
14548 }
14549 bytes.extend_from_slice(&c);
14550 }
14551 }
14552 let mut row_ptr = vec![0u32; rows + 1];
14555 for &(idx, _) in &outliers {
14556 row_ptr[idx as usize / cols + 1] += 1;
14557 }
14558 for r in 0..rows {
14559 row_ptr[r + 1] += row_ptr[r];
14560 }
14561 for &p in &row_ptr {
14562 bytes.extend_from_slice(&p.to_le_bytes());
14563 }
14564 for &(idx, v) in &outliers {
14565 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14566 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14567 }
14568
14569 let mut refw = vec![0f32; rows * cols];
14570 dequant_q1t(&bytes, rows, cols, &mut refw);
14571 let x: Vec<f32> = (0..cols)
14574 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14575 .collect();
14576 let mut expect = vec![0f32; rows];
14577 for r in 0..rows {
14578 let mut a = 0.0f32;
14579 for j in 0..cols {
14580 a += refw[r * cols + j] * x[j];
14581 }
14582 expect[r] = a;
14583 }
14584 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14585 let mut got = vec![0f32; rows];
14586 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14587 for r in 0..rows {
14588 assert!(
14589 (got[r] - expect[r]).abs() < tol(expect[r]),
14590 "row {r}: {} vs {}",
14591 got[r],
14592 expect[r]
14593 );
14594 }
14595 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14597 let mut gm = vec![0f32; 2 * rows];
14598 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14599 for r in 0..rows {
14600 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14601 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14602 }
14603 let xb: Vec<f32> = (0..cols)
14607 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14608 .collect();
14609 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14610 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14611 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14612 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14613 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14614 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14615 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14616 }
14617
14618 #[test]
14621 fn q1t_matvec2_odd_gpr_matches_singles() {
14622 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14623 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
14625 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14626 for r in 0..rows {
14627 for g in 0..gpr {
14628 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14629 let mut c = [0u8; 7];
14630 for k in 0..GROUP_SIZE {
14631 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14632 }
14633 bytes.extend_from_slice(&c);
14634 }
14635 }
14636 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14637 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14638 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14639 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14640 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14641 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14642 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14643 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14644 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14645 }
14646
14647 #[test]
14651 #[ignore]
14652 fn q1t_matvec2_speed() {
14653 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14654 use std::time::Instant;
14655 let (rows, cols) = (8192usize, 4096usize);
14656 let gpr = cols / GROUP_SIZE;
14657 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14658 for r in 0..rows {
14659 for g in 0..gpr {
14660 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14661 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14662 let mut c = [0u8; 7];
14663 for k in 0..GROUP_SIZE {
14664 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14665 }
14666 bytes.extend_from_slice(&c);
14667 }
14668 }
14669 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14670 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14671 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14672 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14673 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14675 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14676 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14677 for _ in 0..8 {
14678 let t0 = Instant::now();
14679 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14680 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14681 let t1 = Instant::now();
14682 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14683 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14684 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14685 }
14686 assert_eq!(p1, s1);
14687 assert_eq!(p2, s2);
14688 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14689 }
14690
14691 #[test]
14695 #[ignore]
14696 fn q1t_matvec_speed() {
14697 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14698 use std::time::Instant;
14699 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
14701 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14702 for r in 0..rows {
14703 for g in 0..gpr {
14704 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14705 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14706 let mut c = [0u8; 7];
14707 for k in 0..GROUP_SIZE {
14708 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14709 }
14710 bytes.extend_from_slice(&c);
14711 }
14712 }
14713 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
14715 let mut idx = 0usize;
14716 while idx < n {
14717 row_ptr[idx / cols + 1] += 1;
14718 idx += stride;
14719 }
14720 for r in 0..rows {
14721 row_ptr[r + 1] += row_ptr[r];
14722 }
14723 for &p in &row_ptr {
14724 bytes.extend_from_slice(&p.to_le_bytes());
14725 }
14726 let mut idx = 0usize;
14727 while idx < n {
14728 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14729 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14730 idx += stride;
14731 }
14732 let x: Vec<f32> = (0..cols)
14735 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14736 .collect();
14737 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14738
14739 let slow = |out: &mut [f32]| {
14741 let mut buf = vec![0f32; cols];
14742 for r in 0..rows {
14743 for g in 0..gpr {
14744 let off = (r * gpr + g) * Q1T_TILE;
14745 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14746 let codes = &bytes[off + 2..off + Q1T_TILE];
14747 for k in 0..GROUP_SIZE {
14748 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14749 1 => s,
14750 2 => -s,
14751 _ => 0.0,
14752 };
14753 }
14754 }
14755 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14756 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14757 }
14758 };
14759 let iters = 5;
14760 let mut a = vec![0f32; rows];
14761 slow(&mut a); let t = Instant::now();
14763 for _ in 0..iters {
14764 slow(&mut a);
14765 }
14766 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14767
14768 let mut b = vec![0f32; rows];
14769 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
14771 for _ in 0..iters {
14772 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14773 }
14774 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14775
14776 for r in 0..rows {
14777 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14778 }
14779 println!(
14780 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
14781 slow_ms / fast_ms
14782 );
14783 }
14784}
14785
14786#[cfg(test)]
14787mod gemm_bench {
14788 #[test]
14799 #[ignore]
14800 fn q4tp_matmat_throughput() {
14801 let _alt = super::Q4TP_ALT_TEST_LOCK
14802 .lock()
14803 .unwrap_or_else(|e| e.into_inner());
14804 let b: usize = std::env::var("CMF_BENCH_B")
14808 .ok()
14809 .and_then(|v| v.parse().ok())
14810 .unwrap_or(296);
14811 let (rows, cols) = (9216usize, 2304usize);
14812 let (_, _, _) = (rows, cols, b);
14813 let total =
14814 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14815 .unwrap();
14816 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14821 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
14822 let lo = cortiq_core::quant::f32_to_f16(-4.0);
14823 let step = cortiq_core::quant::f32_to_f16(0.1);
14824 for r in 0..rows {
14825 let o = params_off + r * 4;
14826 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14827 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14828 }
14829 let _ = codes_off;
14830 let xs: Vec<f32> = (0..b * cols)
14831 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
14832 .collect();
14833 let mut out = vec![0f32; b * rows];
14834 let pool = crate::pool::Pool::from_env();
14835 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14841 let reps: usize = std::env::var("CMF_BENCH_REPS")
14842 .ok()
14843 .and_then(|v| v.parse().ok())
14844 .unwrap_or(10);
14845 let mut best = [f64::MAX; 2];
14846 let mut sums = [0f32; 2];
14847 for _ in 0..reps {
14848 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
14849 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
14850 let t = std::time::Instant::now();
14851 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14852 best[k] = best[k].min(t.elapsed().as_secs_f64());
14853 sums[k] = out.iter().take(64).sum::<f32>();
14854 }
14855 }
14856 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
14857 for (k, name) in ["previous", "tuned "].iter().enumerate() {
14858 println!(
14859 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
14860 best[k] * 1e3,
14861 flops / best[k] / 1e9,
14862 sums[k]
14863 );
14864 }
14865 assert!(
14866 (sums[0] - sums[1]).abs() < 1e-2,
14867 "the tuned kernel changed the result: {} vs {}",
14868 sums[0],
14869 sums[1]
14870 );
14871 }
14872
14873 #[test]
14879 fn q4tp_matmat_blocked_matches_scalar() {
14880 use std::sync::atomic::Ordering::Relaxed;
14881 let _alt = super::Q4TP_ALT_TEST_LOCK
14882 .lock()
14883 .unwrap_or_else(|e| e.into_inner());
14884 for &(rows, cols, b) in &[
14892 (64usize, 128usize, 7usize),
14893 (33, 96, 4),
14894 (16, 256, 9),
14895 (192, 2304, 37),
14896 ] {
14897 let total = cortiq_core::quant::expected_nbytes(
14898 cortiq_core::TensorDtype::Q4TiledP,
14899 &[rows, cols],
14900 )
14901 .unwrap();
14902 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14903 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
14904 let lo = cortiq_core::quant::f32_to_f16(-4.0);
14905 let step = cortiq_core::quant::f32_to_f16(0.1);
14906 for r in 0..rows {
14907 let o = params_off + r * 4;
14908 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14909 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14910 }
14911 let xs: Vec<f32> = (0..b * cols)
14912 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
14913 .collect();
14914 let mut got = vec![0f32; b * rows];
14915 let mut want = vec![0f32; b * rows];
14916 let gpr = cols / 32;
14917 let view = super::Q4tpView::new(&bytes, rows, cols);
14918 let pool = crate::pool::Pool::from_env();
14919 super::Q4TP_ALT.store(2, Relaxed);
14920 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
14921 super::Q4TP_ALT.store(1, Relaxed);
14922 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
14923 super::Q4TP_ALT.store(0, Relaxed);
14924 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14931 let (mut worst, mut at) = (0f32, 0usize);
14932 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
14933 if (g - w).abs() > worst {
14934 worst = (g - w).abs();
14935 at = i;
14936 }
14937 }
14938 assert!(
14939 worst <= 1e-4 * scale,
14940 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
14941 (scale {scale:.3e}) at cell {at}: {} vs {}",
14942 got[at],
14943 want[at]
14944 );
14945
14946 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
14953 for bi in 0..b {
14954 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
14955 for r in 0..rows {
14956 let mut sc = vec![0f32; gpr];
14957 view.scales_into(r, gpr, &mut sc);
14958 let mut exact = 0f64;
14959 for j in 0..cols {
14960 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14961 exact += w as f64 * sq as f64 * act.xq[j] as f64;
14962 }
14963 exact *= act.sx as f64;
14964 for &(j, xv) in &act.outliers {
14965 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14966 exact += w as f64 * sq as f64 * xv as f64;
14967 }
14968 let i = bi * rows + r;
14969 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
14970 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
14971 }
14972 }
14973 println!(
14974 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
14975 per-column {e_scalar:.3e}"
14976 );
14977 assert!(
14982 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
14983 "{rows}x{cols} b={b}: error against f64 too large — blocked \
14984 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
14985 );
14986 }
14987 }
14988
14989 #[test]
14995 #[ignore]
14996 fn q4tp_matmat_row_exact_cost() {
14997 let _alt = super::Q4TP_ALT_TEST_LOCK
14998 .lock()
14999 .unwrap_or_else(|e| e.into_inner());
15000 let pool = crate::pool::Pool::from_env();
15001 let reps: usize = std::env::var("CMF_BENCH_REPS")
15002 .ok()
15003 .and_then(|v| v.parse().ok())
15004 .unwrap_or(30);
15005 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15006 let total = cortiq_core::quant::expected_nbytes(
15007 cortiq_core::TensorDtype::Q4TiledP,
15008 &[rows, cols],
15009 )
15010 .unwrap();
15011 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15012 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15013 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15014 let step = cortiq_core::quant::f32_to_f16(0.1);
15015 for r in 0..rows {
15016 let o = params_off + r * 4;
15017 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15018 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15019 }
15020 for &b in &[2usize, 4, 5, 8] {
15021 let xs: Vec<f32> = (0..b * cols)
15022 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15023 .collect();
15024 let mut out = vec![0f32; b * rows];
15025 let mut best = [f64::MAX; 3];
15026 for _ in 0..reps {
15027 for (k, best_k) in best.iter_mut().enumerate() {
15028 let t = std::time::Instant::now();
15029 match k {
15030 0 | 1 => super::q4tp_matmat_with(
15031 &bytes,
15032 &xs,
15033 b,
15034 rows,
15035 cols,
15036 &mut out,
15037 pool.as_deref(),
15038 k == 1,
15039 ),
15040 _ => {
15041 for (bi, o) in out.chunks_mut(rows).enumerate() {
15042 super::q4tp_matvec(
15043 &bytes,
15044 &xs[bi * cols..(bi + 1) * cols],
15045 rows,
15046 cols,
15047 o,
15048 pool.as_deref(),
15049 );
15050 }
15051 }
15052 }
15053 *best_k = best_k.min(t.elapsed().as_secs_f64());
15054 }
15055 }
15056 println!(
15057 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15058 {b} matvecs {:.3} ms",
15059 best[0] * 1e3,
15060 best[1] * 1e3,
15061 best[2] * 1e3
15062 );
15063 }
15064 }
15065 }
15066
15067 #[test]
15071 #[ignore]
15072 fn q4t_matmat_throughput() {
15073 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15074 let total =
15075 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15076 .unwrap();
15077 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15080 let sc = cortiq_core::quant::f32_to_f16(0.02);
15081 for t in bytes.chunks_mut(super::Q4_TILE) {
15082 t[..2].copy_from_slice(&sc.to_le_bytes());
15083 }
15084 let xs: Vec<f32> = (0..b * cols)
15085 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15086 .collect();
15087 let mut out = vec![0f32; b * rows];
15088 let pool = crate::pool::Pool::from_env();
15089 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15090 let reps: usize = std::env::var("CMF_BENCH_REPS")
15091 .ok()
15092 .and_then(|v| v.parse().ok())
15093 .unwrap_or(10);
15094 let mut best = f64::MAX;
15095 for _ in 0..reps {
15096 let t = std::time::Instant::now();
15097 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15098 best = best.min(t.elapsed().as_secs_f64());
15099 }
15100 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15101 println!(
15102 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15103 best * 1e3,
15104 flops / best / 1e9,
15105 out.iter().take(64).sum::<f32>()
15106 );
15107 }
15108}