1use std::sync::OnceLock;
25
26pub const Q8_MAX_ABS: i8 = 127;
28
29const TEAM_WORK_THRESHOLD_BYTES: usize = 512 * 1024;
34
35pub fn quantize_row_q8(row: &[f32], output: &mut [i8]) -> f32 {
49 assert_eq!(output.len(), row.len(), "quantize output length mismatch");
50 let mut maximum = 0.0_f32;
51 for (index, &value) in row.iter().enumerate() {
52 assert!(
53 value.is_finite(),
54 "non-finite value {value} at index {index} reached the Q8 quantizer"
55 );
56 maximum = maximum.max(value.abs());
57 }
58 if maximum == 0.0 {
59 output.fill(0);
60 return 1.0;
61 }
62 let scale = maximum / 127.0;
63 for (&value, slot) in row.iter().zip(output.iter_mut()) {
64 let rounded = (value / scale).clamp(-127.0, 127.0).round_ties_even();
65 *slot = rounded as i8;
68 }
69 scale
70}
71
72#[derive(Clone, Debug)]
77pub struct QuantizedMatrix {
78 pub data: Vec<i8>,
80 pub scales: Vec<f32>,
82 pub n: usize,
84 pub k: usize,
86}
87
88impl QuantizedMatrix {
89 #[must_use]
99 pub fn concat_rows(parts: &[&Self]) -> Self {
100 let k = parts.first().expect("at least one part").k;
101 assert!(parts.iter().all(|part| part.k == k), "parts must share k");
102 let n = parts.iter().map(|part| part.n).sum();
103 let mut data = Vec::with_capacity(n * k);
104 let mut scales = Vec::with_capacity(n);
105 for part in parts {
106 data.extend_from_slice(&part.data);
107 scales.extend_from_slice(&part.scales);
108 }
109 Self { data, scales, n, k }
110 }
111
112 #[must_use]
118 pub fn quantize(weight: &[f32], n: usize, k: usize) -> Self {
119 assert_eq!(weight.len(), n * k, "weight must be [n, k]");
120 let mut data = vec![0_i8; n * k];
121 let mut scales = vec![0.0_f32; n];
122 for ((weight_row, data_row), scale) in weight
123 .chunks_exact(k)
124 .zip(data.chunks_exact_mut(k))
125 .zip(scales.iter_mut())
126 {
127 *scale = quantize_row_q8(weight_row, data_row);
128 }
129 Self { data, scales, n, k }
130 }
131}
132
133#[derive(Clone, Copy, Debug, Eq, PartialEq)]
139pub enum Int8Tier {
140 Scalar,
142 Autovec,
147 NeonSdot,
149}
150
151impl Int8Tier {
152 #[must_use]
154 pub const fn as_str(self) -> &'static str {
155 match self {
156 Self::Scalar => "scalar",
157 Self::Autovec => "autovec",
158 Self::NeonSdot => "neon-sdot",
159 }
160 }
161
162 #[must_use]
164 pub fn available() -> Vec<Self> {
165 let mut tiers = vec![Self::Scalar, Self::Autovec];
166 if neon_sdot_available() {
167 tiers.push(Self::NeonSdot);
168 }
169 tiers
170 }
171
172 #[must_use]
182 pub fn dispatch() -> Self {
183 match std::env::var("FTTS_INT8_TIER").as_deref() {
184 Ok("scalar") => Self::Scalar,
185 Ok("autovec") => Self::Autovec,
186 Ok("neon-sdot") if neon_sdot_available() => Self::NeonSdot,
187 _ if neon_sdot_available() => Self::NeonSdot,
188 _ => Self::Scalar,
189 }
190 }
191}
192
193#[derive(Clone, Copy, Debug, Eq, PartialEq)]
203pub enum QuantLinearMode {
204 W8A8(Int8Tier),
206 W8A16,
208}
209
210impl QuantLinearMode {
211 #[must_use]
213 pub const fn as_str(self) -> &'static str {
214 match self {
215 Self::W8A8(_) => "w8a8",
216 Self::W8A16 => "w8a16",
217 }
218 }
219}
220
221pub fn linear_w8a16(
232 x: &[f32],
233 weight: &QuantizedMatrix,
234 bias: Option<&[f32]>,
235 m: usize,
236 out: &mut [f32],
237) {
238 let (n, k) = (weight.n, weight.k);
239 assert_eq!(x.len(), m * k, "x must be [m, k]");
240 assert_eq!(out.len(), m * n, "out must be [m, n]");
241 if let Some(bias) = bias {
242 assert_eq!(bias.len(), n, "bias must be [n]");
243 }
244 for col in 0..n {
245 let w_row = &weight.data[col * k..(col + 1) * k];
246 let w_scale = weight.scales[col];
247 let bias_term = bias.map(|b| b[col]);
248 for row in 0..m {
249 let x_row = &x[row * k..(row + 1) * k];
250 let acc = dot_w8a16(x_row, w_row);
251 let value = acc * w_scale;
252 out[row * n + col] = bias_term.map_or(value, |b| value + b);
253 }
254 }
255}
256
257fn dot_w8a16(x: &[f32], w: &[i8]) -> f32 {
259 const LANES: usize = 8;
260 let mut lanes = [0.0_f32; LANES];
261 let chunks = x.len() / LANES;
262 for chunk in 0..chunks {
263 let base = chunk * LANES;
264 for lane in 0..LANES {
265 lanes[lane] = f32::from(w[base + lane]).mul_add(x[base + lane], lanes[lane]);
266 }
267 }
268 let mut sum: f32 = lanes.iter().sum();
269 for index in chunks * LANES..x.len() {
270 sum = f32::from(w[index]).mul_add(x[index], sum);
271 }
272 sum
273}
274
275#[must_use]
281pub fn quant_mode_from_environment() -> QuantLinearMode {
282 match std::env::var("FTTS_INT8").as_deref() {
283 Ok("w8a16") => QuantLinearMode::W8A16,
284 _ => QuantLinearMode::W8A8(autotuned_plan().decode_gemv),
285 }
286}
287
288pub fn quant_linear(
290 mode: QuantLinearMode,
291 x: &[f32],
292 weight: &QuantizedMatrix,
293 bias: Option<&[f32]>,
294 m: usize,
295 out: &mut [f32],
296) {
297 match mode {
298 QuantLinearMode::W8A8(tier) => linear_q8_dynamic(x, weight, bias, m, out, tier),
299 QuantLinearMode::W8A16 => linear_w8a16(x, weight, bias, m, out),
300 }
301}
302
303#[derive(Clone, Copy, Debug, Eq, PartialEq)]
309pub struct KernelPlanV0 {
310 pub decode_gemv: Int8Tier,
312 pub batch_gemm: Int8Tier,
314}
315
316pub fn autotuned_plan() -> KernelPlanV0 {
322 static PLAN: OnceLock<KernelPlanV0> = OnceLock::new();
323 *PLAN.get_or_init(|| {
324 if std::env::var("FTTS_INT8_TIER").is_ok() {
325 let forced = Int8Tier::dispatch();
326 return KernelPlanV0 {
327 decode_gemv: forced,
328 batch_gemm: forced,
329 };
330 }
331 if let Some(cached) = load_persisted_plan() {
332 return cached;
333 }
334 let plan = KernelPlanV0 {
335 decode_gemv: fastest_tier(&[(1, 1024, 256), (1, 3072, 256)]),
339 batch_gemm: fastest_tier(&[(16, 1024, 128), (16, 3072, 64)]),
341 };
342 persist_plan(plan);
343 plan
344 })
345}
346
347fn plan_cache_path() -> Option<std::path::PathBuf> {
350 std::env::var_os("HOME")
351 .map(|home| std::path::PathBuf::from(home).join(".cache/franken_tts/kernel_plan_v0.txt"))
352}
353
354fn plan_cache_key() -> String {
356 let tiers: Vec<&str> = Int8Tier::available().iter().map(|t| t.as_str()).collect();
357 format!(
358 "v0|crate={}|tiers={}",
359 env!("CARGO_PKG_VERSION"),
360 tiers.join(",")
361 )
362}
363
364fn load_persisted_plan() -> Option<KernelPlanV0> {
365 let text = std::fs::read_to_string(plan_cache_path()?).ok()?;
366 let mut lines = text.lines();
367 if lines.next()? != plan_cache_key() {
368 return None;
369 }
370 let parse = |line: &str| match line {
371 "scalar" => Some(Int8Tier::Scalar),
372 "autovec" => Some(Int8Tier::Autovec),
373 "neon-sdot" if neon_sdot_available() => Some(Int8Tier::NeonSdot),
374 _ => None,
375 };
376 Some(KernelPlanV0 {
377 decode_gemv: parse(lines.next()?)?,
378 batch_gemm: parse(lines.next()?)?,
379 })
380}
381
382fn persist_plan(plan: KernelPlanV0) {
383 let Some(path) = plan_cache_path() else {
384 return;
385 };
386 if let Some(parent) = path.parent() {
387 let _ = std::fs::create_dir_all(parent);
388 }
389 let _ = std::fs::write(
391 path,
392 format!(
393 "{}\n{}\n{}\n",
394 plan_cache_key(),
395 plan.decode_gemv.as_str(),
396 plan.batch_gemm.as_str()
397 ),
398 );
399}
400
401fn fastest_tier(probes: &[(usize, usize, usize)]) -> Int8Tier {
405 use std::time::Instant;
406 let tiers = Int8Tier::available();
407 let mut best = (tiers[0], f64::MAX);
408 for &tier in &tiers {
409 let mut total = 0.0_f64;
410 for &(m, k, n) in probes {
411 let x_q: Vec<i8> = (0..m * k).map(|i| ((i * 37 + 11) % 255) as i8).collect();
412 let x_scales = vec![1.0_f32; m];
413 let weight = QuantizedMatrix {
414 data: (0..n * k).map(|i| ((i * 29 + 5) % 255) as i8).collect(),
415 scales: vec![1.0_f32; n],
416 n,
417 k,
418 };
419 let mut out = vec![0.0_f32; m * n];
420 let mut rounds: Vec<f64> = (0..3)
421 .map(|_| {
422 let start = Instant::now();
423 linear_q8(&x_q, &x_scales, &weight, None, m, &mut out, tier);
424 start.elapsed().as_secs_f64()
425 })
426 .collect();
427 rounds.sort_by(f64::total_cmp);
428 total += rounds[1];
429 }
430 if total < best.1 {
431 best = (tier, total);
432 }
433 }
434 best.0
435}
436
437#[must_use]
439pub fn neon_sdot_available() -> bool {
440 #[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
441 {
442 neon_dotprod::available()
443 }
444 #[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
445 {
446 false
447 }
448}
449
450#[must_use]
456pub fn dot_i32(a: &[i8], b: &[i8], tier: Int8Tier) -> i32 {
457 assert_eq!(a.len(), b.len(), "int8 dot inputs must match");
458 match tier {
459 Int8Tier::Scalar => dot_i32_scalar(a, b),
460 Int8Tier::Autovec => dot_i32_autovec(a, b),
461 Int8Tier::NeonSdot => dot_i32_neon_or_panic(a, b),
462 }
463}
464
465#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
466fn dot_i32_neon_or_panic(a: &[i8], b: &[i8]) -> i32 {
467 assert!(
468 neon_dotprod::available(),
469 "neon-sdot route selected without FEAT_DotProd"
470 );
471 neon_dotprod::dot_i32(a, b)
472}
473
474#[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
475fn dot_i32_neon_or_panic(_a: &[i8], _b: &[i8]) -> i32 {
476 panic!("neon-sdot route selected on a build without the island");
477}
478
479fn dot_i32_scalar(a: &[i8], b: &[i8]) -> i32 {
480 let mut sum = 0_i32;
481 for index in 0..a.len() {
482 sum += i32::from(a[index]) * i32::from(b[index]);
483 }
484 sum
485}
486
487fn dot_i32_autovec(a: &[i8], b: &[i8]) -> i32 {
491 const LANES: usize = 8;
492 let mut lanes = [0_i32; LANES];
493 let chunks = a.len() / LANES;
494 for chunk in 0..chunks {
495 let base = chunk * LANES;
496 for lane in 0..LANES {
497 lanes[lane] += i32::from(a[base + lane]) * i32::from(b[base + lane]);
498 }
499 }
500 let mut sum: i32 = lanes.iter().sum();
501 for index in chunks * LANES..a.len() {
502 sum += i32::from(a[index]) * i32::from(b[index]);
503 }
504 sum
505}
506
507#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
508mod neon_dotprod {
509 use core::arch::aarch64::{vaddq_s32, vaddvq_s32, vdotq_s32, vdupq_n_s32, vld1q_s8};
514
515 #[must_use]
517 pub fn available() -> bool {
518 std::arch::is_aarch64_feature_detected!("dotprod")
519 }
520
521 #[must_use]
528 pub fn dot_i32(a: &[i8], b: &[i8]) -> i32 {
529 debug_assert!(available(), "SDOT island entered without FEAT_DotProd");
530 unsafe { dot_i32_sdot(a, b) }
533 }
534
535 #[target_feature(enable = "neon,dotprod")]
536 unsafe fn dot_i32_sdot(a: &[i8], b: &[i8]) -> i32 {
537 let len = a.len();
538 let a_ptr = a.as_ptr();
539 let b_ptr = b.as_ptr();
540 let mut acc0 = vdupq_n_s32(0);
541 let mut acc1 = vdupq_n_s32(0);
542 let mut acc2 = vdupq_n_s32(0);
543 let mut acc3 = vdupq_n_s32(0);
544 let mut index = 0_usize;
545 while index + 64 <= len {
546 unsafe {
550 acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
551 acc1 = vdotq_s32(
552 acc1,
553 vld1q_s8(a_ptr.add(index + 16)),
554 vld1q_s8(b_ptr.add(index + 16)),
555 );
556 acc2 = vdotq_s32(
557 acc2,
558 vld1q_s8(a_ptr.add(index + 32)),
559 vld1q_s8(b_ptr.add(index + 32)),
560 );
561 acc3 = vdotq_s32(
562 acc3,
563 vld1q_s8(a_ptr.add(index + 48)),
564 vld1q_s8(b_ptr.add(index + 48)),
565 );
566 }
567 index += 64;
568 }
569 while index + 16 <= len {
570 unsafe {
572 acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
573 }
574 index += 16;
575 }
576 let mut sum = vaddvq_s32(vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)));
577 while index < len {
578 sum += i32::from(a[index]) * i32::from(b[index]);
579 index += 1;
580 }
581 sum
582 }
583}
584
585#[allow(clippy::too_many_arguments)]
597pub fn linear_q8(
598 x_q: &[i8],
599 x_scales: &[f32],
600 weight: &QuantizedMatrix,
601 bias: Option<&[f32]>,
602 m: usize,
603 out: &mut [f32],
604 tier: Int8Tier,
605) {
606 let (n, k) = (weight.n, weight.k);
607 assert_eq!(x_q.len(), m * k, "x_q must be [m, k]");
608 assert_eq!(x_scales.len(), m, "x_scales must be [m]");
609 assert_eq!(out.len(), m * n, "out must be [m, n]");
610 if let Some(bias) = bias {
611 assert_eq!(bias.len(), n, "bias must be [n]");
612 }
613 if n * k >= TEAM_WORK_THRESHOLD_BYTES
617 && let Some(team) = crate::team::armed()
618 {
619 team.linear_q8(x_q, x_scales, weight, bias, m, out, tier);
620 return;
621 }
622
623 for col in 0..n {
628 let w_row = &weight.data[col * k..(col + 1) * k];
629 let w_scale = weight.scales[col];
630 let bias_term = bias.map(|b| b[col]);
631 for row in 0..m {
632 let x_row = &x_q[row * k..(row + 1) * k];
633 let acc = dot_i32(x_row, w_row, tier);
634 let value = acc as f32 * (x_scales[row] * w_scale);
635 out[row * n + col] = bias_term.map_or(value, |b| value + b);
636 }
637 }
638}
639
640pub fn linear_q8_dynamic(
649 x: &[f32],
650 weight: &QuantizedMatrix,
651 bias: Option<&[f32]>,
652 m: usize,
653 out: &mut [f32],
654 tier: Int8Tier,
655) {
656 let k = weight.k;
657 assert_eq!(x.len(), m * k, "x must be [m, k]");
658 let mut x_q = vec![0_i8; m * k];
659 let mut x_scales = vec![0.0_f32; m];
660 for ((x_row, q_row), scale) in x
661 .chunks_exact(k)
662 .zip(x_q.chunks_exact_mut(k))
663 .zip(x_scales.iter_mut())
664 {
665 *scale = quantize_row_q8(x_row, q_row);
666 }
667 linear_q8(&x_q, &x_scales, weight, bias, m, out, tier);
668}
669
670#[cfg(test)]
671mod tests {
672 use super::*;
673
674 fn pseudo_random_q8(len: usize, seed: u64) -> Vec<i8> {
676 let mut state = seed;
677 (0..len)
678 .map(|_| {
679 state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
680 let mut z = state;
681 z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
682 z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
683 z ^= z >> 31;
684 ((z % 255) as i32 - 127) as i8
686 })
687 .collect()
688 }
689
690 const MODEL_SHAPES: &[(usize, usize)] = &[
692 (2048, 1024), (1024, 1024), (1024, 2048), (3072, 1024), (1024, 3072), ];
698
699 #[test]
700 fn every_tier_is_exactly_equal_in_i32_at_every_model_shape() {
701 for &(n, k) in MODEL_SHAPES {
702 let a = pseudo_random_q8(k, 0x5eed_0001 ^ (n as u64) << 20 ^ k as u64);
703 let w = pseudo_random_q8(n * k, 0x5eed_0002 ^ (n as u64) << 20 ^ k as u64);
704 for row in [0, n / 2, n - 1] {
705 let w_row = &w[row * k..(row + 1) * k];
706 let reference = dot_i32(&a, w_row, Int8Tier::Scalar);
707 for tier in Int8Tier::available() {
708 assert_eq!(
709 dot_i32(&a, w_row, tier),
710 reference,
711 "tier {} diverged at shape {n}x{k} row {row}",
712 tier.as_str()
713 );
714 }
715 }
716 }
717 }
718
719 #[test]
720 fn every_tier_survives_the_all_extreme_reduction_at_the_binding_census_k() {
721 for k in [2048_usize, 3072, 4608, 7168, 8192] {
723 let a = vec![127_i8; k];
724 let b = vec![127_i8; k];
725 let negative = vec![-127_i8; k];
726 let expected = 127_i64 * 127 * k as i64;
727 for tier in Int8Tier::available() {
728 assert_eq!(
729 i64::from(dot_i32(&a, &b, tier)),
730 expected,
731 "positive all-extreme diverged on {} at K={k}",
732 tier.as_str()
733 );
734 assert_eq!(
735 i64::from(dot_i32(&a, &negative, tier)),
736 -expected,
737 "negative all-extreme diverged on {} at K={k}",
738 tier.as_str()
739 );
740 }
741 }
742 }
743
744 #[test]
745 fn tail_lengths_that_defeat_block_boundaries_stay_exact() {
746 for len in [1_usize, 7, 15, 16, 17, 63, 64, 65, 100, 129] {
748 let a = pseudo_random_q8(len, tail_seed(len));
749 let b = pseudo_random_q8(len, tail_seed(len) ^ 1);
750 let reference = dot_i32(&a, &b, Int8Tier::Scalar);
751 for tier in Int8Tier::available() {
752 assert_eq!(
753 dot_i32(&a, &b, tier),
754 reference,
755 "len={len} {}",
756 tier.as_str()
757 );
758 }
759 }
760 }
761
762 #[test]
763 fn quantizer_matches_the_canonical_converter_semantics() {
764 let row = [
767 -127.0_f32, -126.5, -125.5, -1.5, -0.5, 0.5, 1.5, 125.5, 126.5, 127.0,
768 ];
769 let mut q = [0_i8; 10];
770 let scale = quantize_row_q8(&row, &mut q);
771 assert_eq!(scale.to_bits(), 1.0_f32.to_bits());
772 assert_eq!(q, [-127, -126, -126, -2, 0, 0, 2, 126, 126, 127]);
773
774 let zeros = [0.0_f32; 4];
775 let mut qz = [1_i8; 4];
776 assert_eq!(
777 quantize_row_q8(&zeros, &mut qz).to_bits(),
778 1.0_f32.to_bits()
779 );
780 assert_eq!(qz, [0, 0, 0, 0]);
781
782 let matrix = QuantizedMatrix::quantize(&[2.0, -1.0, 0.0, 3.0], 2, 2);
783 assert_eq!(matrix.scales[0].to_bits(), (2.0_f32 / 127.0).to_bits());
784 assert_eq!(matrix.scales[1].to_bits(), (3.0_f32 / 127.0).to_bits());
785 assert!(matrix.data.iter().all(|&b| b != -128));
786 }
787
788 #[test]
789 fn dynamic_w8a8_linear_tracks_the_f32_reference_within_quant_error() {
790 let (n, k) = (64_usize, 128_usize);
792 let mut weight = vec![0.0_f32; n * k];
793 let mut x = vec![0.0_f32; k];
794 let mut state = 0x1234_5678_u64;
795 let mut next = || {
796 state = state
797 .wrapping_mul(6_364_136_223_846_793_005)
798 .wrapping_add(1);
799 ((state >> 33) as f32 / (1u64 << 31) as f32) - 1.0
800 };
801 for value in weight.iter_mut() {
802 *value = next();
803 }
804 for value in x.iter_mut() {
805 *value = next();
806 }
807 let quantized = QuantizedMatrix::quantize(&weight, n, k);
808 let mut out_q8 = vec![0.0_f32; n];
809 linear_q8_dynamic(&x, &quantized, None, 1, &mut out_q8, Int8Tier::Autovec);
810
811 let mut out_f32 = vec![0.0_f32; n];
812 crate::f32ref::linear(&x, &weight, None, 1, k, n, &mut out_f32);
813
814 let dot = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>();
815 let cosine = dot(&out_q8, &out_f32)
816 / (dot(&out_q8, &out_q8).sqrt() * dot(&out_f32, &out_f32).sqrt());
817 assert!(
818 cosine > 0.999,
819 "W8A8 dequant plumbing is broken: cosine {cosine}"
820 );
821 }
822
823 #[test]
824 fn tiers_produce_bit_identical_f32_output_not_merely_close() {
825 let (n, k) = (256_usize, 1024_usize);
826 let weight: Vec<f32> = pseudo_random_q8(n * k, 77)
827 .iter()
828 .map(|&b| f32::from(b) / 64.0)
829 .collect();
830 let x: Vec<f32> = pseudo_random_q8(k, 78)
831 .iter()
832 .map(|&b| f32::from(b) / 64.0)
833 .collect();
834 let quantized = QuantizedMatrix::quantize(&weight, n, k);
835 let mut reference = vec![0.0_f32; n];
836 linear_q8_dynamic(&x, &quantized, None, 1, &mut reference, Int8Tier::Scalar);
837 for tier in Int8Tier::available() {
838 let mut out = vec![0.0_f32; n];
839 linear_q8_dynamic(&x, &quantized, None, 1, &mut out, tier);
840 for (index, (a, b)) in reference.iter().zip(&out).enumerate() {
841 assert_eq!(
842 a.to_bits(),
843 b.to_bits(),
844 "tier {} f32 output differs at {index}",
845 tier.as_str()
846 );
847 }
848 }
849 }
850
851 fn tail_seed(len: usize) -> u64 {
852 0x7a11_0000 ^ len as u64
853 }
854}