1use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE};
16use cortiq_core::{CmfModel, TensorDtype};
17use std::sync::Arc;
18
19pub enum QTensor {
20 F32 {
21 data: Vec<f32>,
22 rows: usize,
23 cols: usize,
24 },
25 Mapped {
26 model: Arc<CmfModel>,
27 idx: usize,
29 dtype: TensorDtype,
30 rows: usize,
31 cols: usize,
32 row_scale: Vec<f32>,
34 col_field: Vec<f32>,
36 },
37}
38
39impl QTensor {
40 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
41 debug_assert_eq!(data.len(), rows * cols);
42 Self::F32 { data, rows, cols }
43 }
44
45 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
48 let idx = model
49 .tensors
50 .iter()
51 .position(|t| t.name == name)
52 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
53 let entry = &model.tensors[idx];
54 if entry.shape.len() != 2 {
55 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
56 }
57 let (rows, cols) = (entry.shape[0], entry.shape[1]);
58 let bytes = model.entry_bytes(entry);
59
60 match entry.dtype {
61 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
62 let n = rows * cols;
63 let scales_off = n;
64 let row_scale: Vec<f32> = (0..rows)
65 .map(|o| {
66 f16_to_f32(u16::from_le_bytes([
67 bytes[scales_off + o * 2],
68 bytes[scales_off + o * 2 + 1],
69 ]))
70 })
71 .collect();
72 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
73 let col_off = n + rows * 2;
74 (0..cols)
75 .map(|i| {
76 f16_to_f32(u16::from_le_bytes([
77 bytes[col_off + i * 2],
78 bytes[col_off + i * 2 + 1],
79 ]))
80 })
81 .collect()
82 } else {
83 Vec::new()
84 };
85 Ok(Self::Mapped {
86 model: model.clone(),
87 idx,
88 dtype: entry.dtype,
89 rows,
90 cols,
91 row_scale,
92 col_field,
93 })
94 }
95 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
97 model: model.clone(),
98 idx,
99 dtype: entry.dtype,
100 rows,
101 cols,
102 row_scale: Vec::new(),
103 col_field: Vec::new(),
104 }),
105 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
108 model: model.clone(),
109 idx,
110 dtype: entry.dtype,
111 rows,
112 cols,
113 row_scale: Vec::new(),
114 col_field: Vec::new(),
115 }),
116 _ => {
118 let mut data = vec![0.0f32; rows * cols];
119 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
120 Ok(Self::from_f32(data, rows, cols))
121 }
122 }
123 }
124
125 pub fn rows(&self) -> usize {
126 match self {
127 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
128 }
129 }
130
131 pub fn cols(&self) -> usize {
132 match self {
133 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
134 }
135 }
136
137 pub fn as_f32(&self) -> Option<&[f32]> {
140 match self {
141 Self::F32 { data, .. } => Some(data),
142 Self::Mapped { .. } => None,
143 }
144 }
145
146 fn quant_bytes(&self) -> &[u8] {
147 match self {
148 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
149 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
150 }
151 }
152
153 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
155 let cols = self.cols();
156 debug_assert_eq!(dst.len(), cols);
157 match self {
158 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
159 Self::Mapped {
160 dtype,
161 row_scale,
162 col_field,
163 ..
164 } => {
165 if *dtype == TensorDtype::Q4Block {
166 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
167 let gpr = cols / GROUP_SIZE;
168 for gi in 0..gpr {
169 let g = r * gpr + gi;
170 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
171 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
172 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
173 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
174 }
175 }
176 return;
177 }
178 if *dtype == TensorDtype::Vbit {
179 let bytes = self.quant_bytes();
180 let rows = self.rows();
181 let ng = cols / GROUP_SIZE;
182 let bits = &bytes[..rows];
183 let sc_off = rows;
184 let mut off = sc_off + rows * ng * 2;
185 for rr in 0..r {
186 off += (cols * bits[rr] as usize + 7) / 8;
187 }
188 let b = bits[r] as usize;
189 let l = ((1usize << (b - 1)) - 1) as f32;
190 let data = &bytes[off..];
191 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
192 for (i, d) in dst.iter_mut().enumerate() {
193 while nbits < b {
194 acc = (acc << 8) | data[idx] as u64;
195 idx += 1;
196 nbits += 8;
197 }
198 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
199 nbits -= b;
200 let so = (r * ng + i / GROUP_SIZE) * 2;
201 let sv = f16_to_f32(u16::from_le_bytes([
202 bytes[sc_off + so],
203 bytes[sc_off + so + 1],
204 ]));
205 *d = (u - l) * sv;
206 }
207 return;
208 }
209 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
210 let s = row_scale[r];
211 match dtype {
212 TensorDtype::Q8Row => {
213 for (d, &b) in dst.iter_mut().zip(q) {
214 *d = (b as i8) as f32 * s;
215 }
216 }
217 TensorDtype::Q8_2f => {
218 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
219 *d = (b as i8) as f32 * s * col_field[i];
220 }
221 }
222 _ => unreachable!(),
223 }
224 }
225 }
226 }
227
228 pub fn sparse_col_ok(&self) -> bool {
233 match self {
234 Self::F32 { .. } => true,
235 Self::Mapped { dtype, .. } => {
236 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
237 }
238 }
239 }
240
241 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
245 let inter = self.cols();
246 let hidden = self.rows();
247 debug_assert_eq!(out.len(), hidden);
248 match self {
249 Self::F32 { data, .. } => {
250 for (k, o) in out.iter_mut().enumerate() {
251 *o += w * data[k * inter + c];
252 }
253 }
254 Self::Mapped {
255 dtype, row_scale, col_field, ..
256 } => {
257 let q = self.quant_bytes();
258 let colf = if *dtype == TensorDtype::Q8_2f {
259 col_field[c]
260 } else {
261 1.0
262 };
263 let wc = w * colf;
264 for (k, o) in out.iter_mut().enumerate() {
265 let b = q[k * inter + c] as i8 as f32;
266 *o += wc * b * row_scale[k];
267 }
268 }
269 }
270 }
271
272 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
276 let cols = self.cols();
277 match self {
278 Self::F32 { data, .. } => {
279 let row = &data[r * cols..(r + 1) * cols];
280 row.iter().zip(x).map(|(w, v)| w * v).sum()
281 }
282 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
283 TensorDtype::Q8Row => {
284 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
285 let s: f32 = q.iter().zip(x).map(|(&b, v)| (b as i8) as f32 * v).sum();
286 s * row_scale[r]
287 }
288 TensorDtype::Q8_2f => {
289 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
290 let s: f32 = q
291 .iter()
292 .zip(x)
293 .zip(col_field)
294 .map(|((&b, v), c)| (b as i8) as f32 * v * c)
295 .sum();
296 s * row_scale[r]
297 }
298 _ => {
299 self.row_f32(r, scratch);
300 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
301 }
302 },
303 }
304 }
305
306 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
309 match self {
310 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
311 Self::Mapped {
312 model,
313 idx,
314 dtype,
315 rows,
316 cols,
317 row_scale,
318 col_field,
319 } => {
320 let _ = (model, idx);
321 if *dtype == TensorDtype::Q4Block {
322 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
323 return;
324 }
325 if *dtype == TensorDtype::Vbit {
326 vbitmatvec(self.quant_bytes(), x, *rows, *cols, out, pool);
327 return;
328 }
329 let xs = prescale(x, col_field, *dtype);
330 if *rows >= crate::gpu::min_rows()
335 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
336 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
337 && crate::gpu::enabled_here()
338 {
339 let frac = std::env::var("CMF_GPU_SPLIT")
340 .ok()
341 .and_then(|v| v.parse::<f32>().ok())
342 .unwrap_or(0.5)
343 .clamp(0.0, 1.0);
344 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
345 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
346 let bytes = self.quant_bytes();
347 let ok = std::thread::scope(|sc| {
348 let g = sc.spawn(|| {
349 crate::gpu::q8_matvec_range(
350 model,
351 *idx,
352 cpu_rows,
353 &row_scale[cpu_rows..],
354 &xs,
355 *rows - cpu_rows,
356 *cols,
357 out_gpu,
358 )
359 });
360 if cpu_rows > 0 {
361 qmatvec(
362 &bytes[..cpu_rows * *cols],
363 &row_scale[..cpu_rows],
364 &xs,
365 cpu_rows,
366 *cols,
367 out_cpu,
368 pool,
369 );
370 }
371 g.join().unwrap_or(false)
372 });
373 if ok {
374 return;
375 }
376 qmatvec(
378 &bytes[cpu_rows * *cols..(*rows) * *cols],
379 &row_scale[cpu_rows..],
380 &xs,
381 *rows - cpu_rows,
382 *cols,
383 out_gpu,
384 pool,
385 );
386 return;
387 }
388 qmatvec(self.quant_bytes(), row_scale, &xs, *rows, *cols, out, pool);
389 }
390 }
391 }
392
393 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
395 match self {
396 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
397 Self::Mapped {
398 dtype,
399 rows,
400 cols,
401 row_scale,
402 col_field,
403 ..
404 } => {
405 if *dtype == TensorDtype::Q4Block {
406 q4matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
407 q4matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
408 return;
409 }
410 if *dtype == TensorDtype::Vbit {
411 vbitmatvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
412 vbitmatvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
413 return;
414 }
415 let x1s = prescale(x1, col_field, *dtype);
416 let x2s = prescale(x2, col_field, *dtype);
417 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
418 }
419 }
420 }
421}
422
423impl QTensor {
424 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
429 let cols = self.cols();
430 let rows = self.rows();
431 debug_assert_eq!(xs_all.len(), b * cols);
432 debug_assert_eq!(out.len(), b * rows);
433 match self {
434 Self::F32 { data, .. } => {
435 let out_addr = SendMut(out.as_mut_ptr());
436 let run = |start: usize, end: usize| {
437 for o in start..end {
438 let row = &data[o * cols..(o + 1) * cols];
439 for bi in 0..b {
440 let x = &xs_all[bi * cols..(bi + 1) * cols];
441 let mut acc = 0f32;
442 for j in 0..cols {
443 acc += row[j] * x[j];
444 }
445 unsafe { *out_addr.at(bi * rows + o) = acc };
446 }
447 }
448 };
449 dispatch_rows(pool, rows, &run);
450 }
451 Self::Mapped {
452 dtype,
453 row_scale,
454 col_field,
455 ..
456 } => {
457 if matches!(dtype, TensorDtype::Q4Block | TensorDtype::Vbit) {
458 for bi in 0..b {
460 let x = &xs_all[bi * cols..(bi + 1) * cols];
461 let (head, tail) = out[bi * rows..].split_at_mut(rows);
462 let _ = tail;
463 self.matvec(x, head, pool);
464 }
465 return;
466 }
467 let pre: Vec<Vec<f32>> = (0..b)
468 .map(|bi| {
469 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
470 })
471 .collect();
472 if b >= 8
475 && b * rows * cols >= 128_000_000
476 && crate::gpu::enabled_here()
477 {
478 if let Self::Mapped { model, idx, .. } = self {
479 let flat: Vec<f32> =
480 pre.iter().flat_map(|v| v.iter().copied()).collect();
481 if crate::gpu::q8_matmat(
482 model, *idx, row_scale, &flat, b, rows, cols, out)
483 {
484 return;
485 }
486 }
487 }
488 let q = self.quant_bytes();
489 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
490 }
491 }
492 }
493}
494
495fn qmatmat(
498 q: &[u8],
499 row_scale: &[f32],
500 pre: &[Vec<f32>],
501 rows: usize,
502 cols: usize,
503 out: &mut [f32],
504 pool: Option<&Pool>,
505) {
506 let b = pre.len();
507 debug_assert_eq!(out.len(), b * rows);
508 #[cfg(target_arch = "aarch64")]
509 if sdot_enabled() {
510 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
511 let out_addr = SendMut(out.as_mut_ptr());
512 let run = |start: usize, end: usize| {
513 for o in start..end {
514 let row = &q[o * cols..(o + 1) * cols];
515 for (bi, act) in acts.iter().enumerate() {
516 let v = row_dot_sdot(row, act) * row_scale[o];
517 unsafe { *out_addr.at(bi * rows + o) = v };
518 }
519 }
520 };
521 dispatch_rows(pool, rows, &run);
522 return;
523 }
524 let out_addr = SendMut(out.as_mut_ptr());
525 let run = |start: usize, end: usize| {
526 for o in start..end {
527 let row = &q[o * cols..(o + 1) * cols];
528 for (bi, x) in pre.iter().enumerate() {
529 let mut acc = 0f32;
530 for j in 0..cols {
531 acc += (row[j] as i8) as f32 * x[j];
532 }
533 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
534 }
535 }
536 };
537 dispatch_rows(pool, rows, &run);
538}
539
540fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
542 match pool {
543 Some(pool) if rows >= 256 => {
544 pool.run(&|widx, n| {
545 let chunk = rows.div_ceil(n);
546 let start = widx * chunk;
547 let end = (start + chunk).min(rows);
548 if start < end {
549 run(start, end);
550 }
551 });
552 }
553 _ => run(0, rows),
554 }
555}
556
557fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
559 let groups = rows * cols / GROUP_SIZE;
560 bytes.split_at(groups * 16)
561}
562
563#[inline(always)]
568fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
569 let mut acc = 0u64;
570 for i in 0..B {
571 acc = (acc << 8) | data[i] as u64;
572 }
573 let mask = (1u64 << B) - 1;
574 let mut out = [0i32; 8];
575 for (k, o) in out.iter_mut().enumerate() {
576 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
577 }
578 out
579}
580
581fn vbitmatvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
586 debug_assert_eq!(out.len(), rows);
587 let ng = cols / GROUP_SIZE;
588 let bits = &bytes[..rows];
589 let sc_off = rows;
590 let data_off = sc_off + rows * ng * 2;
591 let mut offsets = Vec::with_capacity(rows + 1);
592 let mut off = data_off;
593 for r in 0..rows {
594 offsets.push(off);
595 off += (cols * bits[r] as usize + 7) / 8;
596 }
597 offsets.push(off);
598
599 let offsets = &offsets;
600
601 #[cfg(target_arch = "aarch64")]
605 if sdot_enabled() {
606 let act = split_act(x);
607 let act = &act;
608 let row_dot = move |r: usize| -> f32 {
609 let b = bits[r] as usize;
610 let l = ((1i32 << (b - 1)) - 1) as i32;
611 let mask = (1u64 << b) - 1;
612 let data = &bytes[offsets[r]..offsets[r + 1]];
613 if b == 8 {
614 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
616 let mut dot = 0f32;
617 for g in 0..ng {
618 let so = (r * ng + g) * 2;
619 let sgf = f16_to_f32(u16::from_le_bytes([
620 bytes[sc_off + so],
621 bytes[sc_off + so + 1],
622 ]));
623 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
624 let mut gd = 0f32;
625 for &xv in xg.iter() {
626 if nbits < 8 {
627 acc = (acc << 8) | data[idx] as u64;
628 idx += 1;
629 nbits += 8;
630 }
631 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
632 nbits -= 8;
633 gd += (u - l) as f32 * xv;
634 }
635 dot += gd * sgf;
636 }
637 return dot;
638 }
639 let mut buf = vec![0u8; cols]; #[inline(always)]
641 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
642 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
643 let u = unpack8::<B>(&data[blk * B..]);
644 for k in 0..8 {
645 chunk[k] = (u[k] - l) as i8 as u8;
646 }
647 }
648 }
649 let _ = mask;
650 match b {
651 3 => fill::<3>(data, l, &mut buf),
652 4 => fill::<4>(data, l, &mut buf),
653 5 => fill::<5>(data, l, &mut buf),
654 6 => fill::<6>(data, l, &mut buf),
655 _ => unreachable!(),
656 }
657 let mut dot = 0f32;
658 for g in 0..ng {
659 let so = (r * ng + g) * 2;
660 let s = f16_to_f32(u16::from_le_bytes([
661 bytes[sc_off + so],
662 bytes[sc_off + so + 1],
663 ]));
664 let d = unsafe {
665 dot_i8_sdot(
666 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
667 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
668 )
669 } as f32
670 * act.sx;
671 dot += d * s;
672 }
673 for &(j, xv) in &act.outliers {
674 let so = (r * ng + j / GROUP_SIZE) * 2;
675 let s = f16_to_f32(u16::from_le_bytes([
676 bytes[sc_off + so],
677 bytes[sc_off + so + 1],
678 ]));
679 dot += (buf[j] as i8) as f32 * s * xv;
681 }
682 dot
683 };
684 match pool {
685 Some(pool) if rows >= 256 => {
686 let out_addr = SendMut(out.as_mut_ptr());
687 pool.run(&move |widx, n| {
688 let chunk = rows.div_ceil(n);
689 let start = widx * chunk;
690 let end = (start + chunk).min(rows);
691 for r in start..end {
692 unsafe { *out_addr.at(r) = row_dot(r) };
694 }
695 });
696 }
697 _ => {
698 for (r, dst) in out.iter_mut().enumerate() {
699 *dst = row_dot(r);
700 }
701 }
702 }
703 return;
704 }
705
706 #[inline(always)]
710 fn dot_row<const B: usize>(
711 data: &[u8],
712 bytes: &[u8],
713 sc_off: usize,
714 r: usize,
715 ng: usize,
716 x: &[f32],
717 ) -> f32 {
718 let l = ((1i32 << (B - 1)) - 1) as f32;
719 let gbytes = GROUP_SIZE * B / 8;
720 let mut dot = 0f32;
721 for g in 0..ng {
722 let so = (r * ng + g) * 2;
723 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
724 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
725 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
726 let mut gd = 0f32;
727 for blk in 0..GROUP_SIZE / 8 {
728 let u = unpack8::<B>(&gd0[blk * B..]);
729 let xb = &xg[blk * 8..blk * 8 + 8];
730 for k in 0..8 {
731 gd += (u[k] as f32 - l) * xb[k];
732 }
733 }
734 dot += gd * s;
735 }
736 dot
737 }
738 let row_dot = move |r: usize| -> f32 {
739 let data = &bytes[offsets[r]..offsets[r + 1]];
740 match bits[r] {
741 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
742 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
743 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
744 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
745 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
746 b => unreachable!("vbit bit-width {b} (validated at load)"),
747 }
748 };
749 match pool {
750 Some(pool) if rows >= 256 => {
751 let out_addr = SendMut(out.as_mut_ptr());
752 pool.run(&move |widx, n| {
753 let chunk = rows.div_ceil(n);
754 let start = widx * chunk;
755 let end = (start + chunk).min(rows);
756 for r in start..end {
757 unsafe { *out_addr.at(r) = row_dot(r) };
759 }
760 });
761 }
762 _ => {
763 for (r, dst) in out.iter_mut().enumerate() {
764 *dst = row_dot(r);
765 }
766 }
767 }
768}
769
770fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
774 debug_assert_eq!(out.len(), rows);
775 let (packed, scales) = q4_split(bytes, rows, cols);
776 let gpr = cols / GROUP_SIZE;
777 let row_dot = |r: usize| -> f32 {
778 let mut acc = 0f32;
779 for gi in 0..gpr {
780 let g = r * gpr + gi;
781 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
782 let pk = &packed[g * 16..(g + 1) * 16];
783 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
784 let mut ga = 0f32;
785 for (k, &b) in pk.iter().enumerate() {
786 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
787 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
788 }
789 acc += ga * s;
790 }
791 acc
792 };
793 match pool {
794 Some(pool) if rows >= 256 => {
795 let out_addr = SendMut(out.as_mut_ptr());
796 pool.run(&move |widx, n| {
797 let chunk = rows.div_ceil(n);
798 let start = widx * chunk;
799 let end = (start + chunk).min(rows);
800 for r in start..end {
801 unsafe { *out_addr.at(r) = row_dot(r) };
803 }
804 });
805 }
806 _ => {
807 for (r, dst) in out.iter_mut().enumerate() {
808 *dst = row_dot(r);
809 }
810 }
811 }
812}
813
814pub(crate) fn prescale(x: &[f32], col_field: &[f32], dtype: TensorDtype) -> Vec<f32> {
815 if dtype == TensorDtype::Q8_2f {
816 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
817 } else {
818 x.to_vec()
819 }
820}
821
822fn sdot_enabled() -> bool {
827 use std::sync::OnceLock;
828 static ON: OnceLock<bool> = OnceLock::new();
829 *ON.get_or_init(|| {
830 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
831 #[cfg(target_arch = "aarch64")]
832 {
833 want && std::arch::is_aarch64_feature_detected!("dotprod")
834 }
835 #[cfg(not(target_arch = "aarch64"))]
836 {
837 let _ = want;
838 false
839 }
840 })
841}
842
843struct SplitAct {
848 xq: Vec<i8>,
849 sx: f32,
850 outliers: Vec<(usize, f32)>,
851}
852
853fn split_act(x: &[f32]) -> SplitAct {
854 let n = x.len();
855 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
856 let thr = 8.0 * rms;
857 let outliers: Vec<(usize, f32)> = (0..n)
858 .filter(|&j| x[j].abs() > thr)
859 .map(|j| (j, x[j]))
860 .collect();
861 let mut xb = x.to_vec();
862 for &(j, _) in &outliers {
863 xb[j] = 0.0;
864 }
865 let amax = xb.iter().fold(0f32, |m, &v| m.max(v.abs()));
866 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
867 let inv = 1.0 / sx;
868 let xq = xb
869 .iter()
870 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8)
871 .collect();
872 SplitAct { xq, sx, outliers }
873}
874
875#[cfg(target_arch = "aarch64")]
878#[target_feature(enable = "neon,dotprod")]
879unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
880 unsafe {
882 use core::arch::aarch64::*;
883 use core::arch::asm;
884 let wp = w.as_ptr() as *const i8;
885 let n = w.len();
886 let (mut a0, mut a1, mut a2, mut a3) =
887 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
888 let mut i = 0;
889 while i + 64 <= n {
890 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
891 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
892 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
893 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
894 asm!(
895 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
896 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
897 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
898 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
899 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
900 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
901 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
902 options(pure, nomem, nostack),
903 );
904 i += 64;
905 }
906 while i + 16 <= n {
907 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
908 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
909 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
910 i += 16;
911 }
912 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
913 while i < n {
914 s += (*wp.add(i)) as i32 * xq[i] as i32;
915 i += 1;
916 }
917 s
918}
919}
920
921#[cfg(target_arch = "aarch64")]
925#[target_feature(enable = "neon,dotprod")]
926unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
927 unsafe {
929 use core::arch::aarch64::*;
930 use core::arch::asm;
931 let n = xq.len();
932 let px = xq.as_ptr();
933 let (p0, p1, p2, p3) = (
934 w0.as_ptr() as *const i8,
935 w1.as_ptr() as *const i8,
936 w2.as_ptr() as *const i8,
937 w3.as_ptr() as *const i8,
938 );
939 let (mut a0, mut a1, mut a2, mut a3) =
940 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
941 let mut i = 0;
942 while i + 16 <= n {
943 let x = vld1q_s8(px.add(i));
944 let v0 = vld1q_s8(p0.add(i));
945 let v1 = vld1q_s8(p1.add(i));
946 let v2 = vld1q_s8(p2.add(i));
947 let v3 = vld1q_s8(p3.add(i));
948 asm!(
949 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
950 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
951 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
952 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
953 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
954 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
955 options(pure, nomem, nostack),
956 );
957 i += 16;
958 }
959 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
960 while i < n {
961 let xi = *px.add(i) as i32;
962 r[0] += (*p0.add(i)) as i32 * xi;
963 r[1] += (*p1.add(i)) as i32 * xi;
964 r[2] += (*p2.add(i)) as i32 * xi;
965 r[3] += (*p3.add(i)) as i32 * xi;
966 i += 1;
967 }
968 r
969}
970}
971
972#[cfg(target_arch = "aarch64")]
975#[inline]
976fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
977 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
978 for &(j, xv) in &act.outliers {
979 acc += (row[j] as i8) as f32 * xv;
980 }
981 acc
982}
983
984#[inline]
989fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
990 #[cfg(target_arch = "aarch64")]
991 unsafe {
992 return dot_i8_f32_neon(w, x);
993 }
994 #[allow(unreachable_code)]
995 {
996 let mut sum = 0.0f32;
997 for (j, &b) in w.iter().enumerate() {
998 sum += (b as i8) as f32 * x[j];
999 }
1000 sum
1001 }
1002}
1003
1004#[cfg(target_arch = "aarch64")]
1005#[target_feature(enable = "neon")]
1006unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
1007 unsafe {
1009 use core::arch::aarch64::*;
1010 let n = x.len();
1011 let wp = w.as_ptr() as *const i8;
1012 let xp = x.as_ptr();
1013 let (mut a0, mut a1, mut a2, mut a3) =
1014 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1015 let mut j = 0usize;
1016 while j + 16 <= n {
1017 let wb = vld1q_s8(wp.add(j));
1018 let lo = vmovl_s8(vget_low_s8(wb));
1019 let hi = vmovl_s8(vget_high_s8(wb));
1020 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
1021 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
1022 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
1023 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
1024 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
1025 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
1026 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
1027 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
1028 j += 16;
1029 }
1030 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
1031 while j < n {
1032 sum += (*wp.add(j)) as f32 * *xp.add(j);
1033 j += 1;
1034 }
1035 sum
1036}
1037}
1038
1039#[allow(clippy::too_many_arguments)]
1040fn qmatvec(
1041 q: &[u8],
1042 row_scale: &[f32],
1043 xs: &[f32],
1044 rows: usize,
1045 cols: usize,
1046 out: &mut [f32],
1047 pool: Option<&Pool>,
1048) {
1049 debug_assert_eq!(out.len(), rows);
1050
1051 #[cfg(target_arch = "aarch64")]
1052 if sdot_enabled() {
1053 let act = split_act(xs);
1054 let out_addr = SendMut(out.as_mut_ptr());
1055 let run_range = |start: usize, end: usize| {
1056 let mut o = start;
1057 while o + 4 <= end {
1058 let r = unsafe {
1059 dot_i8_sdot_4rows(
1060 &q[o * cols..(o + 1) * cols],
1061 &q[(o + 1) * cols..(o + 2) * cols],
1062 &q[(o + 2) * cols..(o + 3) * cols],
1063 &q[(o + 3) * cols..(o + 4) * cols],
1064 &act.xq,
1065 )
1066 };
1067 for k in 0..4 {
1068 let mut acc = r[k] as f32 * act.sx;
1069 for &(j, xv) in &act.outliers {
1070 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
1071 }
1072 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
1074 }
1075 o += 4;
1076 }
1077 while o < end {
1078 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], &act) * row_scale[o];
1079 unsafe { *out_addr.at(o) = v };
1080 o += 1;
1081 }
1082 };
1083 match pool {
1084 Some(pool) if rows >= 256 => {
1085 pool.run(&|widx, n| {
1086 let chunk = rows.div_ceil(n);
1087 let start = widx * chunk;
1088 run_range(start, (start + chunk).min(rows));
1089 });
1090 }
1091 _ => run_range(0, rows),
1092 }
1093 return;
1094 }
1095
1096 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
1097 match pool {
1098 Some(pool) if rows >= 256 => {
1099 let out_addr = SendMut(out.as_mut_ptr());
1100 pool.run(&move |widx, n| {
1101 let chunk = rows.div_ceil(n);
1102 let start = widx * chunk;
1103 let end = (start + chunk).min(rows);
1104 for o in start..end {
1105 unsafe { *out_addr.at(o) = row_dot(o) };
1107 }
1108 });
1109 }
1110 _ => {
1111 for (o, dst) in out.iter_mut().enumerate() {
1112 *dst = row_dot(o);
1113 }
1114 }
1115 }
1116}
1117
1118#[allow(clippy::too_many_arguments)]
1119fn qmatvec2(
1120 q: &[u8],
1121 row_scale: &[f32],
1122 x1: &[f32],
1123 x2: &[f32],
1124 rows: usize,
1125 cols: usize,
1126 o1: &mut [f32],
1127 o2: &mut [f32],
1128 pool: Option<&Pool>,
1129) {
1130 #[cfg(target_arch = "aarch64")]
1131 if sdot_enabled() {
1132 let a1s = split_act(x1);
1133 let a2s = split_act(x2);
1134 let p1 = SendMut(o1.as_mut_ptr());
1135 let p2 = SendMut(o2.as_mut_ptr());
1136 let run_range = |start: usize, end: usize| {
1137 for o in start..end {
1138 let row = &q[o * cols..(o + 1) * cols];
1139 unsafe {
1141 *p1.at(o) = row_dot_sdot(row, &a1s) * row_scale[o];
1142 *p2.at(o) = row_dot_sdot(row, &a2s) * row_scale[o];
1143 }
1144 }
1145 };
1146 match pool {
1147 Some(pool) if rows >= 256 => {
1148 pool.run(&|widx, n| {
1149 let chunk = rows.div_ceil(n);
1150 let start = widx * chunk;
1151 run_range(start, (start + chunk).min(rows));
1152 });
1153 }
1154 _ => run_range(0, rows),
1155 }
1156 return;
1157 }
1158
1159 let row_dots = |o: usize| -> (f32, f32) {
1160 let row = &q[o * cols..(o + 1) * cols];
1161 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
1162 };
1163 match pool {
1164 Some(pool) if rows >= 256 => {
1165 let p1 = SendMut(o1.as_mut_ptr());
1166 let p2 = SendMut(o2.as_mut_ptr());
1167 pool.run(&move |widx, n| {
1168 let chunk = rows.div_ceil(n);
1169 let start = widx * chunk;
1170 let end = (start + chunk).min(rows);
1171 for o in start..end {
1172 let (s1, s2) = row_dots(o);
1173 unsafe {
1175 *p1.at(o) = s1;
1176 *p2.at(o) = s2;
1177 }
1178 }
1179 });
1180 }
1181 _ => {
1182 for o in 0..rows {
1183 let (s1, s2) = row_dots(o);
1184 o1[o] = s1;
1185 o2[o] = s2;
1186 }
1187 }
1188 }
1189}
1190
1191#[derive(Clone, Copy)]
1192struct SendMut(*mut f32);
1193unsafe impl Send for SendMut {}
1194unsafe impl Sync for SendMut {}
1195
1196impl SendMut {
1197 #[inline]
1198 fn at(self, i: usize) -> *mut f32 {
1199 unsafe { self.0.add(i) }
1200 }
1201}
1202
1203#[cfg(test)]
1204mod tests {
1205 use super::*;
1206
1207 #[test]
1208 fn f32_matvec_matches_matvec_rows_bitexact() {
1209 let (rows, cols) = (300, 40);
1210 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
1211 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
1212 let qt = QTensor::from_f32(w.clone(), rows, cols);
1213
1214 let mut a = vec![0.0f32; rows];
1215 matvec_rows(None, &w, &x, &mut a);
1216 let mut b = vec![0.0f32; rows];
1217 qt.matvec(&x, &mut b, None);
1218 assert_eq!(a, b);
1219 }
1220
1221 #[test]
1222 fn sdot_kernel_exact_on_grid() {
1223 eprintln!("sdot_enabled = {}", sdot_enabled());
1228 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
1230 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
1231 .collect();
1232 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
1233 let x: Vec<f32> = (0..cols)
1234 .map(|i| match i % 3 {
1235 0 => 1.0,
1236 1 => -1.0,
1237 _ => 0.0,
1238 })
1239 .collect();
1240 let mut a = vec![0.0f32; rows];
1241 qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
1242 for o in 0..rows {
1243 let mut acc = 0.0f32;
1244 for j in 0..cols {
1245 acc += (w[o * cols + j] as i8) as f32 * x[j];
1246 }
1247 let expect = acc * scales[o];
1248 assert!(
1249 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
1250 "row {o}: {} vs {expect}",
1251 a[o]
1252 );
1253 }
1254 }
1255
1256 #[test]
1257 fn sdot_a8w8_noise_is_bounded() {
1258 let (rows, cols) = (16, 512);
1262 let w: Vec<u8> = (0..rows * cols)
1263 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
1264 .collect();
1265 let scales = vec![0.01f32; rows];
1266 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
1267 let mut a = vec![0.0f32; rows];
1268 qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
1269 let (mut num, mut den) = (0f64, 0f64);
1270 for o in 0..rows {
1271 let mut acc = 0.0f32;
1272 for j in 0..cols {
1273 acc += (w[o * cols + j] as i8) as f32 * x[j];
1274 }
1275 let expect = acc * scales[o];
1276 num += ((a[o] - expect) as f64).powi(2);
1277 den += (expect as f64).powi(2);
1278 }
1279 let rel = (num / den.max(1e-12)).sqrt();
1280 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
1281 }
1282
1283 #[test]
1284 fn i8_dot_neon_matches_scalar() {
1285 let n = 100;
1286 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
1287 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
1288 let mut scalar = 0.0f32;
1289 for j in 0..n {
1290 scalar += (w[j] as i8) as f32 * x[j];
1291 }
1292 let fast = dot_i8_f32(&w, &x);
1293 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
1294 }
1295
1296 #[test]
1298 fn vbitmatvec_matches_full_dequant() {
1299 let (rows, cols) = (6, 64);
1300 let ng = cols / GROUP_SIZE;
1301 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
1303 let mut bytes = bits.clone();
1304 for g in 0..rows * ng {
1305 let s = 0.02 + 0.001 * g as f32;
1306 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
1307 }
1308 for r in 0..rows {
1309 let b = bits[r] as usize;
1310 let (mut acc, mut nb) = (0u64, 0usize);
1311 let mut rowbytes = Vec::new();
1312 for i in 0..cols {
1313 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
1314 acc = (acc << b) | v;
1315 nb += b;
1316 while nb >= 8 {
1317 nb -= 8;
1318 rowbytes.push(((acc >> nb) & 0xFF) as u8);
1319 }
1320 }
1321 if nb > 0 {
1322 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
1323 }
1324 bytes.extend_from_slice(&rowbytes);
1325 }
1326 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
1327
1328 let mut reference = vec![0f32; rows * cols];
1329 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
1330 let mut expect = vec![0f32; rows];
1331 for r in 0..rows {
1332 expect[r] = reference[r * cols..(r + 1) * cols]
1333 .iter()
1334 .zip(&x)
1335 .map(|(w, xv)| w * xv)
1336 .sum();
1337 }
1338 let mut got = vec![0f32; rows];
1339 vbitmatvec(&bytes, &x, rows, cols, &mut got, None);
1340 let tol = if sdot_enabled() { 6e-2 } else { 1e-4 };
1344 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
1345 for r in 0..rows {
1346 assert!(
1347 (got[r] - expect[r]).abs() < tol * scale,
1348 "row {r}: {} vs {}",
1349 got[r],
1350 expect[r]
1351 );
1352 }
1353 }
1354
1355 #[test]
1358 fn q4matvec_matches_full_dequant() {
1359 let (rows, cols) = (8, 64);
1360 let groups = rows * cols / GROUP_SIZE;
1361 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
1363 for i in 0..groups * 16 {
1364 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
1365 }
1366 for g in 0..groups {
1367 let s = 0.01 + 0.003 * g as f32;
1368 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
1369 }
1370 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
1371
1372 let mut reference = vec![0.0f32; rows * cols];
1373 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
1374 let mut expect = vec![0.0f32; rows];
1375 for r in 0..rows {
1376 expect[r] = reference[r * cols..(r + 1) * cols]
1377 .iter()
1378 .zip(&x)
1379 .map(|(w, xv)| w * xv)
1380 .sum();
1381 }
1382
1383 let mut got = vec![0.0f32; rows];
1384 q4matvec(&bytes, &x, rows, cols, &mut got, None);
1385 for r in 0..rows {
1386 assert!(
1387 (got[r] - expect[r]).abs() < 1e-4 * expect[r].abs().max(1.0),
1388 "row {r}: {} vs {}",
1389 got[r],
1390 expect[r]
1391 );
1392 }
1393 }
1394}