1use serde::Serialize;
25
26use crate::codegen_sva::fp_hypotheses::{Nan, Semantics, Tininess};
27use crate::model::quant::{f16_to_f32, GgmlType, KVALUES_IQ4NL};
28
29pub fn ftz(x: f32) -> f32 {
33 if x.is_subnormal() {
34 if x.is_sign_negative() {
35 -0.0
36 } else {
37 0.0
38 }
39 } else {
40 x
41 }
42}
43
44pub const NAN: f32 = f32::from_bits(0x7fc0_0000);
46
47fn canonical(x: f32) -> f32 {
49 if x.is_nan() {
50 NAN
51 } else {
52 x
53 }
54}
55
56pub fn fmul(a: f32, b: f32) -> f32 {
57 canonical(ftz(ftz(a) * ftz(b)))
58}
59
60pub fn fadd(a: f32, b: f32) -> f32 {
61 canonical(ftz(ftz(a) + ftz(b)))
62}
63
64pub fn fsub(a: f32, b: f32) -> f32 {
65 canonical(ftz(ftz(a) - ftz(b)))
66}
67
68#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize)]
71pub enum FpUnit {
72 #[default]
74 Soft,
75 Hard,
79}
80
81impl FpUnit {
82 pub const CARD: Semantics = Semantics { tininess: Tininess::AfterRoundingUnbounded, nan: Nan::XorSigned(0x7fc0_0000), ..Semantics::ENGINE };
85
86 pub fn mul(self, a: f32, b: f32) -> f32 {
88 match self {
89 FpUnit::Soft => fmul(a, b),
90 FpUnit::Hard => f32::from_bits(Self::CARD.mul(a.to_bits(), b.to_bits())),
91 }
92 }
93}
94
95pub const LANES: usize = 16;
97
98pub fn tree(mut lanes: Vec<f32>) -> f32 {
100 while lanes.len() > 1 {
101 lanes = lanes.chunks(2).map(|p| if p.len() == 2 { fadd(p[0], p[1]) } else { p[0] }).collect();
102 }
103 lanes.pop().unwrap_or(0.0)
104}
105
106pub fn fdot(fp: FpUnit, a: &[f32], b: &[f32]) -> f32 {
108 assert_eq!(a.len(), b.len(), "a dot of unequal lengths");
109 let mut lanes = vec![0f32; LANES];
110 for (i, (&x, &y)) in a.iter().zip(b).enumerate() {
111 lanes[i % LANES] = fadd(lanes[i % LANES], fp.mul(x, y));
112 }
113 tree(lanes)
114}
115
116pub fn fsum(a: &[f32]) -> f32 {
118 let mut lanes = vec![0f32; LANES];
119 for (i, &x) in a.iter().enumerate() {
120 lanes[i % LANES] = fadd(lanes[i % LANES], x);
121 }
122 tree(lanes)
123}
124
125pub fn fmax(x: &[f32]) -> f32 {
128 let m = x.iter().fold(f32::NEG_INFINITY, |m, &v| if v > m { v } else { m });
129 if m == 0.0 {
130 0.0
131 } else {
132 m
133 }
134}
135
136pub fn absmax(x: &[f32]) -> f32 {
138 x.iter().fold(0f32, |m, &v| {
139 let a = ftz(v).abs();
140 if a > m {
141 a
142 } else {
143 m
144 }
145 })
146}
147
148pub fn argmax(x: &[f32]) -> usize {
152 let mut best: Option<usize> = None;
153 for (i, &v) in x.iter().enumerate() {
154 if !v.is_nan() && best.is_none_or(|b| v > x[b]) {
155 best = Some(i);
156 }
157 }
158 best.unwrap_or(0)
159}
160
161pub fn ldexp(x: f32, n: i32) -> f32 {
163 let x = ftz(x);
164 if x == 0.0 || !x.is_finite() {
165 return x;
166 }
167 let bits = x.to_bits();
168 let e = ((bits >> 23) & 0xff) as i32 + n;
169 if e >= 0xff {
170 return f32::INFINITY.copysign(x);
171 }
172 if e <= 0 {
173 return 0f32.copysign(x);
174 }
175 f32::from_bits((bits & 0x807f_ffff) | ((e as u32) << 23))
176}
177
178pub fn scaled_to_f32(v: i128, exp: i32) -> f32 {
180 if v == 0 {
181 return 0.0;
182 }
183 let neg = v < 0;
184 let mut m = v.unsigned_abs();
185 let mut e = exp;
186 let bits = 128 - m.leading_zeros() as i32;
187 if bits > 24 {
188 let s = (bits - 24) as u32;
189 let rem = m & ((1u128 << s) - 1);
190 let half = 1u128 << (s - 1);
191 m >>= s;
192 e += s as i32;
193 if rem > half || (rem == half && m & 1 == 1) {
194 m += 1;
195 if m == 1 << 24 {
196 m >>= 1;
197 e += 1;
198 }
199 }
200 }
201 let lead = 127 - m.leading_zeros() as i32;
202 let unbiased = e + lead;
203 if unbiased > 127 {
204 return if neg { f32::NEG_INFINITY } else { f32::INFINITY };
205 }
206 if unbiased < -126 {
207 return if neg { -0.0 } else { 0.0 };
208 }
209 let mantissa = ((m << (23 - lead)) as u32) & 0x7f_ffff;
210 f32::from_bits(((neg as u32) << 31) | (((unbiased + 127) as u32) << 23) | mantissa)
211}
212
213pub fn rne(x: f32) -> i64 {
215 x.round_ties_even() as i64
216}
217
218pub fn f32_to_f16(x: f32) -> u16 {
222 let b = x.to_bits();
223 let sign = ((b >> 16) & 0x8000) as u16;
224 let exp = ((b >> 23) & 0xff) as i32;
225 let mant = b & 0x7f_ffff;
226 if exp == 0xff {
227 return sign | 0x7c00 | if mant != 0 { 0x0200 } else { 0 };
228 }
229 if exp == 0 {
230 return sign;
232 }
233 let e = exp - 127;
234 if e > 15 {
235 return sign | 0x7c00;
236 }
237 if e >= -14 {
238 let m = mant >> 13;
239 let rem = mant & 0x1fff;
240 let mut h = (((e + 15) as u32) << 10) | m;
241 if rem > 0x1000 || (rem == 0x1000 && m & 1 == 1) {
242 h += 1;
243 }
244 return sign | h as u16;
245 }
246 let s = (-(e + 1)) as u32;
248 if s > 24 {
249 return sign;
250 }
251 let sig = 0x80_0000 | mant;
252 let n = sig >> s;
253 let rem = sig & ((1u32 << s) - 1);
254 let half = 1u32 << (s - 1);
255 let n = if rem > half || (rem == half && n & 1 == 1) { n + 1 } else { n };
256 sign | n as u16
257}
258
259#[derive(Clone, Copy, Debug, PartialEq, Eq)]
261pub struct Scale {
262 pub sig: i64,
263 pub exp: i32,
264}
265
266impl Scale {
267 pub const ZERO: Scale = Scale { sig: 0, exp: 0 };
268
269 pub fn f16(h: u16) -> Option<Scale> {
271 let neg = h & 0x8000 != 0;
272 let e = ((h >> 10) & 0x1f) as i32;
273 let m = (h & 0x3ff) as i64;
274 let (sig, exp) = match e {
275 0x1f => return None,
276 0 => (m, -24),
277 _ => (1024 + m, e - 25),
278 };
279 Some(Scale { sig: if neg { -sig } else { sig }, exp })
280 }
281}
282
283pub const GROUP: usize = 256;
287
288pub const ACT_BLOCK: usize = 32;
290
291#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
293pub enum Words {
294 One,
296 Two,
299}
300
301#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
303pub enum ActScaling {
304 Sub32,
307 PerBlock,
310}
311
312pub const WIDE_BITS: u32 = 8;
314
315#[derive(Clone, Copy, Debug, PartialEq, Eq)]
318pub struct ActScale {
319 pub m: i64,
320 pub e: i32,
321}
322
323#[derive(Clone, Debug, PartialEq, Eq)]
326pub struct ActRow {
327 pub codes: Vec<i32>,
328 pub scales: Vec<ActScale>,
329 pub scaling: ActScaling,
330 pub block: usize,
331}
332
333impl ActRow {
334 pub fn values(&self) -> Vec<f64> {
336 self.codes
337 .iter()
338 .enumerate()
339 .map(|(i, &c)| {
340 let s = self.scales[i / self.block];
341 c as f64 * s.m as f64 * 2f64.powi(s.e)
342 })
343 .collect()
344 }
345}
346
347fn split(x: f32) -> (i64, i32) {
349 let b = x.to_bits();
350 let e = ((b >> 23) & 0xff) as i32;
351 let m = (b & 0x7f_ffff) as i64;
352 let sig = if e == 0 { m } else { m | 0x80_0000 };
353 let sig = if b >> 31 == 1 { -sig } else { sig };
354 (sig, if e == 0 { -149 } else { e - 150 })
355}
356
357pub fn rne_div(x: f32, m: i64, e: i32) -> i64 {
359 let (sig, exp) = split(ftz(x));
360 if sig == 0 {
361 return 0;
362 }
363 let k = exp - e + 64;
365 if k < 0 {
366 return 0;
367 }
368 if k > 60 + 64 {
369 return if sig < 0 { i64::MIN } else { i64::MAX };
370 }
371 let num = (sig.unsigned_abs() as u128) << k as u32;
372 let q = num / m as u128;
373 let rem = num % m as u128;
374 let whole = (q >> 64) as i64;
375 let frac = q & u64::MAX as u128;
376 let half = 1u128 << 63;
377 let up = frac > half || (frac == half && rem > 0) || (frac == half && rem == 0 && whole & 1 == 1);
378 let n = whole + up as i64;
379 if sig < 0 {
380 -n
381 } else {
382 n
383 }
384}
385
386fn least_scale(amax: f32, bits: u32) -> ActScale {
389 let (sig, exp) = split(amax);
390 let (lo, hi) = (1i64 << (bits - 1), (1i64 << bits) - 1);
391 let fits = |m: i64, e: i32| -> bool {
393 let d = e - exp;
394 if d >= 0 {
395 d > 40 || (127 * m) << d >= sig
396 } else {
397 -d < 60 && 127 * m >= (sig + (1i64 << -d) - 1) >> -d
398 }
399 };
400 let mut e = (amax as f64 / (127.0 * hi as f64)).log2().floor() as i32 - 1;
401 while !fits(hi, e) {
402 e += 1;
403 }
404 while fits(hi, e - 1) {
405 e -= 1;
406 }
407 let m = (lo..=hi).find(|&m| fits(m, e)).expect("the least exponent leaves a significand of its width");
408 ActScale { m, e }
409}
410
411pub fn block_scale(amax: f32) -> ActScale {
413 least_scale(amax, 4)
414}
415
416pub fn wide_scale(amax: f32) -> ActScale {
418 least_scale(amax, WIDE_BITS)
419}
420
421fn quantize_with(x: &[f32], words: Words, block: usize, scaling: ActScaling, scale_of: fn(f32) -> ActScale, zero: ActScale) -> ActRow {
424 let mut codes = Vec::with_capacity(x.len());
425 let mut scales = Vec::with_capacity(x.len() / block);
426 for chunk in x.chunks(block) {
427 let amax = chunk.iter().fold(0f32, |m, &v| m.max(ftz(v).abs()));
428 assert!(amax.is_finite(), "a non-finite activation");
429 if amax == 0.0 {
430 scales.push(zero);
431 codes.extend(std::iter::repeat(0).take(block));
432 continue;
433 }
434 let s = scale_of(amax);
435 match words {
436 Words::One => {
437 scales.push(s);
438 codes.extend(chunk.iter().map(|&v| rne_div(v, s.m, s.e).clamp(-127, 127) as i32));
439 }
440 Words::Two => {
441 scales.push(ActScale { m: s.m, e: s.e - 8 });
442 codes.extend(chunk.iter().map(|&v| rne_div(v, s.m, s.e - 8) as i32));
443 }
444 }
445 }
446 ActRow { codes, scales, scaling, block }
447}
448
449pub fn quantize(x: &[f32], words: Words) -> ActRow {
451 assert!(x.len() % GROUP == 0, "{} elements are not a whole number of groups of {GROUP}", x.len());
452 quantize_with(x, words, ACT_BLOCK, ActScaling::Sub32, block_scale, ActScale { m: 8, e: 0 })
453}
454
455pub fn quantize_blocks(x: &[f32], words: Words, block: usize) -> ActRow {
457 assert!(block > 0 && x.len() % block == 0, "{} elements are not a whole number of blocks of {block}", x.len());
458 quantize_with(x, words, block, ActScaling::PerBlock, wide_scale, ActScale { m: 1 << (WIDE_BITS - 1), e: 0 })
459}
460
461#[derive(Clone, Debug, PartialEq, Eq)]
466pub struct Segment {
467 pub start: usize,
468 pub len: usize,
469 pub a: Scale,
470 pub b: Scale,
471}
472
473#[derive(Clone, Debug, PartialEq, Eq)]
475pub struct CodedRow {
476 pub coef: Vec<i16>,
477 pub offs: Vec<i16>,
478 pub segments: Vec<Segment>,
479}
480
481#[derive(Clone, Debug, PartialEq, Eq)]
483pub struct Unreadable {
484 pub ty: GgmlType,
485 pub block: usize,
486}
487
488fn half_at(b: &[u8], at: usize) -> u16 {
489 u16::from_le_bytes([b[at], b[at + 1]])
490}
491
492fn scale_min_k4(j: usize, q: &[u8]) -> (u8, u8) {
494 if j < 4 {
495 (q[j] & 63, q[j + 4] & 63)
496 } else {
497 ((q[j + 4] & 0xf) | ((q[j - 4] >> 6) << 4), (q[j + 4] >> 4) | ((q[j] >> 6) << 4))
498 }
499}
500
501pub fn reads(ty: GgmlType) -> bool {
503 matches!(
504 ty,
505 GgmlType::Q4_0
506 | GgmlType::Q4_1
507 | GgmlType::Q5_0
508 | GgmlType::Q5_1
509 | GgmlType::Q8_0
510 | GgmlType::Q2K
511 | GgmlType::Q3K
512 | GgmlType::Q4K
513 | GgmlType::Q5K
514 | GgmlType::Q6K
515 | GgmlType::Iq4Nl
516 | GgmlType::Iq4Xs
517 )
518}
519
520impl CodedRow {
521 pub fn decode(ty: GgmlType, bytes: &[u8]) -> Option<Result<CodedRow, Unreadable>> {
523 if !reads(ty) {
524 return None;
525 }
526 let bb = ty.block_bytes();
527 let n = ty.block_len();
528 assert!(bytes.len() % bb == 0, "{} bytes are not whole {} blocks", bytes.len(), ty.name());
529 let blocks = bytes.len() / bb;
530 let mut row = CodedRow { coef: Vec::with_capacity(blocks * n), offs: Vec::with_capacity(blocks * n), segments: Vec::with_capacity(blocks) };
531 for (bi, b) in bytes.chunks_exact(bb).enumerate() {
532 let bad = || Unreadable { ty, block: bi };
533 let start = bi * n;
534 let scale = |at: usize| Scale::f16(half_at(b, at)).ok_or_else(bad);
535 let (a, s) = match ty {
536 GgmlType::Q4_0 => {
537 let qs = &b[2..18];
538 row.coef.extend(qs.iter().map(|&q| (q & 0x0f) as i16 - 8));
539 row.coef.extend(qs.iter().map(|&q| (q >> 4) as i16 - 8));
540 row.offs.extend(std::iter::repeat(0).take(32));
541 (scale(0), Ok(Scale::ZERO))
542 }
543 GgmlType::Q4_1 => {
544 let qs = &b[4..20];
545 row.coef.extend(qs.iter().map(|&q| (q & 0x0f) as i16));
546 row.coef.extend(qs.iter().map(|&q| (q >> 4) as i16));
547 row.offs.extend(std::iter::repeat(1).take(32));
548 (scale(0), scale(2))
549 }
550 GgmlType::Q5_0 | GgmlType::Q5_1 => {
551 let one = ty == GgmlType::Q5_1;
552 let (qh_at, qs_at) = if one { (4, 8) } else { (2, 6) };
553 let qh = u32::from_le_bytes([b[qh_at], b[qh_at + 1], b[qh_at + 2], b[qh_at + 3]]);
554 let qs = &b[qs_at..qs_at + 16];
555 let centre = if one { 0 } else { 16 };
556 let lo: Vec<i16> = qs.iter().enumerate().map(|(j, &q)| ((q & 0x0f) | ((((qh >> j) << 4) & 0x10) as u8)) as i16 - centre).collect();
557 let hi: Vec<i16> = qs.iter().enumerate().map(|(j, &q)| ((q >> 4) | (((qh >> (j + 12)) & 0x10) as u8)) as i16 - centre).collect();
558 row.coef.extend(lo);
559 row.coef.extend(hi);
560 row.offs.extend(std::iter::repeat(one as i16).take(32));
561 (scale(0), if one { scale(2) } else { Ok(Scale::ZERO) })
562 }
563 GgmlType::Q8_0 => {
564 row.coef.extend(b[2..34].iter().map(|&q| (q as i8) as i16));
565 row.offs.extend(std::iter::repeat(0).take(32));
566 (scale(0), Ok(Scale::ZERO))
567 }
568 GgmlType::Q2K => {
569 let scales = &b[0..16];
570 let qs = &b[16..80];
571 let mut is = 0;
572 for nn in (0..256).step_by(128) {
573 let q = &qs[nn / 4..nn / 4 + 32];
574 let mut shift = 0;
575 for _ in 0..4 {
576 for half_block in [0usize, 16] {
577 let sc = scales[is];
578 is += 1;
579 for l in 0..16 {
580 row.coef.push((sc & 0xf) as i16 * ((q[l + half_block] >> shift) & 3) as i16);
581 row.offs.push(-((sc >> 4) as i16));
582 }
583 }
584 shift += 2;
585 }
586 }
587 (scale(80), scale(82))
588 }
589 GgmlType::Q3K => {
590 let hm = &b[0..32];
591 let qs = &b[32..96];
592 let packed = &b[96..108];
593 let (kmask1, kmask2) = (0x0303_0303u32, 0x0f0f_0f0fu32);
594 let mut aux = [0u32; 4];
595 for (k, a) in aux.iter_mut().take(3).enumerate() {
596 *a = u32::from_le_bytes([packed[4 * k], packed[4 * k + 1], packed[4 * k + 2], packed[4 * k + 3]]);
597 }
598 let tmp = aux[2];
599 aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
600 aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
601 aux[0] = (aux[0] & kmask2) | ((tmp & kmask1) << 4);
602 aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
603 let scales: Vec<i16> = aux.iter().flat_map(|a| a.to_le_bytes()).map(|v| (v as i8) as i16).collect();
604 let mut is = 0;
605 let mut m = 1u8;
606 for nn in (0..256).step_by(128) {
607 let q = &qs[nn / 4..nn / 4 + 32];
608 let mut shift = 0;
609 for _ in 0..4 {
610 for half_block in [0usize, 16] {
611 let sc = scales[is] - 32;
612 is += 1;
613 for l in 0..16 {
614 let low = ((q[l + half_block] >> shift) & 3) as i16;
615 let high = if hm[l + half_block] & m != 0 { 0 } else { 4 };
616 row.coef.push(sc * (low - high));
617 row.offs.push(0);
618 }
619 }
620 shift += 2;
621 m <<= 1;
622 }
623 }
624 (scale(108), Ok(Scale::ZERO))
625 }
626 GgmlType::Q4K | GgmlType::Q5K => {
627 let five = ty == GgmlType::Q5K;
628 let sc_bytes = &b[4..16];
629 let (qh, qs) = if five { (&b[16..48], &b[48..176]) } else { (&b[0..0], &b[16..144]) };
630 let (mut u1, mut u2) = (1u8, 2u8);
631 for (pair, j) in (0..256).step_by(64).enumerate() {
632 let q = &qs[j / 2..j / 2 + 32];
633 let (s1, m1) = scale_min_k4(2 * pair, sc_bytes);
634 let (s2, m2) = scale_min_k4(2 * pair + 1, sc_bytes);
635 for l in 0..32 {
636 let hi = if five && qh[l] & u1 != 0 { 16 } else { 0 };
637 row.coef.push(s1 as i16 * ((q[l] & 0xf) as i16 + hi));
638 row.offs.push(-(m1 as i16));
639 }
640 for l in 0..32 {
641 let hi = if five && qh[l] & u2 != 0 { 16 } else { 0 };
642 row.coef.push(s2 as i16 * ((q[l] >> 4) as i16 + hi));
643 row.offs.push(-(m2 as i16));
644 }
645 u1 <<= 2;
646 u2 <<= 2;
647 }
648 (scale(0), scale(2))
649 }
650 GgmlType::Q6K => {
651 let ql_all = &b[0..128];
652 let qh_all = &b[128..192];
653 let sc_all: Vec<i16> = b[192..208].iter().map(|&v| (v as i8) as i16).collect();
654 let mut block = [0i16; 256];
655 for (part, nn) in (0..256).step_by(128).enumerate() {
656 let ql = &ql_all[part * 64..part * 64 + 64];
657 let qh = &qh_all[part * 32..part * 32 + 32];
658 let sc = &sc_all[part * 8..part * 8 + 8];
659 for l in 0..32 {
660 let is = l / 16;
661 let q1 = ((ql[l] & 0xf) | ((qh[l] & 3) << 4)) as i16 - 32;
662 let q2 = ((ql[l + 32] & 0xf) | (((qh[l] >> 2) & 3) << 4)) as i16 - 32;
663 let q3 = ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) as i16 - 32;
664 let q4 = ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) as i16 - 32;
665 block[nn + l] = sc[is] * q1;
666 block[nn + l + 32] = sc[is + 2] * q2;
667 block[nn + l + 64] = sc[is + 4] * q3;
668 block[nn + l + 96] = sc[is + 6] * q4;
669 }
670 }
671 row.coef.extend_from_slice(&block);
672 row.offs.extend(std::iter::repeat(0).take(256));
673 (scale(208), Ok(Scale::ZERO))
674 }
675 GgmlType::Iq4Nl => {
676 let qs = &b[2..18];
677 row.coef.extend(qs.iter().map(|&q| KVALUES_IQ4NL[(q & 0x0f) as usize] as i16));
678 row.coef.extend(qs.iter().map(|&q| KVALUES_IQ4NL[(q >> 4) as usize] as i16));
679 row.offs.extend(std::iter::repeat(0).take(32));
680 (scale(0), Ok(Scale::ZERO))
681 }
682 GgmlType::Iq4Xs => {
683 let scales_h = u16::from_le_bytes([b[2], b[3]]);
684 for ib in 0..8 {
685 let ls = ((b[4 + ib / 2] >> (4 * (ib % 2))) & 0xf) as i16 | ((((scales_h >> (2 * ib)) & 3) as i16) << 4);
686 let q = &b[8 + 16 * ib..8 + 16 * ib + 16];
687 row.coef.extend(q.iter().map(|&v| (ls - 32) * KVALUES_IQ4NL[(v & 0x0f) as usize] as i16));
688 row.coef.extend(q.iter().map(|&v| (ls - 32) * KVALUES_IQ4NL[(v >> 4) as usize] as i16));
689 }
690 row.offs.extend(std::iter::repeat(0).take(256));
691 (scale(0), Ok(Scale::ZERO))
692 }
693 _ => unreachable!("reads() admits only the formats decoded here"),
694 };
695 match (a, s) {
696 (Ok(a), Ok(b)) => row.segments.push(Segment { start, len: n, a, b }),
697 _ => return Some(Err(bad())),
698 }
699 }
700 Some(Ok(row))
701 }
702
703 pub fn exact(&self, i: usize) -> f64 {
705 let s = &self.segments[i / self.segments[0].len];
706 s.a.sig as f64 * 2f64.powi(s.a.exp) * self.coef[i] as f64 + s.b.sig as f64 * 2f64.powi(s.b.exp) * self.offs[i] as f64
707 }
708}
709
710pub fn dot(row: &CodedRow, act: &ActRow) -> f32 {
720 assert_eq!(row.coef.len(), act.codes.len(), "a row and an activation of unequal lengths");
721 match act.scaling {
722 ActScaling::Sub32 => dot_sub32(row, act),
723 ActScaling::PerBlock => dot_per_block(row, act),
724 }
725}
726
727pub const COEF_BOUND: u32 = 1 << 12;
729
730fn int_dot(c: &[i16], a: &[i32]) -> i64 {
734 c.chunks(8).zip(a.chunks(8)).map(|(c, a)| c.iter().zip(a).map(|(&c, &a)| c as i32 * a).sum::<i32>() as i64).sum()
735}
736
737fn dot_per_block(row: &CodedRow, act: &ActRow) -> f32 {
738 let mut acc = 0f32;
739 for (b, s) in row.segments.iter().enumerate() {
740 assert_eq!(act.block, s.len, "an activation scaled per {} elements against a row of {}-element blocks", act.block, s.len);
741 let sc = act.scales[b];
742 let range = s.start..s.start + s.len;
743 let ia = int_dot(&row.coef[range.clone()], &act.codes[range.clone()]);
744 let ib = if s.b.sig == 0 { 0 } else { int_dot(&row.offs[range.clone()], &act.codes[range]) };
746 let mut terms = [(0i128, 0i32); 2];
747 let mut n = 0;
748 for (w, sum) in [(s.a, ia), (s.b, ib)] {
749 if w.sig != 0 && sum != 0 {
750 terms[n] = (w.sig as i128 * sc.m as i128 * sum as i128, w.exp + sc.e);
751 n += 1;
752 }
753 }
754 acc = fadd(acc, exact_sum(&terms[..n]));
756 }
757 acc
758}
759
760fn dot_sub32(row: &CodedRow, act: &ActRow) -> f32 {
761 let seg_len = row.segments.first().map_or(GROUP, |s| s.len);
762 let mut acc = 0f32;
763 let mut terms: Vec<(i128, i32)> = Vec::with_capacity(32);
764 for g in 0..row.coef.len() / GROUP {
765 terms.clear();
766 for b in g * GROUP / ACT_BLOCK..(g + 1) * GROUP / ACT_BLOCK {
767 let s = &row.segments[b * ACT_BLOCK / seg_len];
768 let sc = act.scales[b];
769 let range = b * ACT_BLOCK..(b + 1) * ACT_BLOCK;
770 let ia = int_dot(&row.coef[range.clone()], &act.codes[range.clone()]);
771 let ib = if s.b.sig == 0 { 0 } else { int_dot(&row.offs[range.clone()], &act.codes[range]) };
772 for (w, sum) in [(s.a, ia), (s.b, ib)] {
773 if w.sig != 0 && sum != 0 {
774 terms.push((w.sig as i128 * sc.m as i128 * sum as i128, w.exp + sc.e));
775 }
776 }
777 }
778 acc = fadd(acc, exact_sum(&terms));
780 }
781 acc
782}
783
784pub const SUM_WINDOW: i32 = 64;
786
787pub fn exact_sum(terms: &[(i128, i32)]) -> f32 {
790 let (Some(lo), Some(hi)) = (terms.iter().map(|t| t.1).min(), terms.iter().map(|t| t.1).max()) else {
791 return 0.0;
792 };
793 let least = lo.max(hi - SUM_WINDOW);
794 let sum: i128 = terms.iter().map(|&(v, e)| if e >= least { v << (e - least) } else { v >> (least - e).min(127) }).sum();
796 scaled_to_f32(sum, least)
797}
798
799pub(crate) const EXP2_64: [u32; 64] = [
803 0x3f800000, 0x3f8164d2, 0x3f82cd87, 0x3f843a29, 0x3f85aac3, 0x3f871f62, 0x3f88980f, 0x3f8a14d5,
804 0x3f8b95c2, 0x3f8d1adf, 0x3f8ea43a, 0x3f9031dc, 0x3f91c3d3, 0x3f935a2b, 0x3f94f4f0, 0x3f96942d,
805 0x3f9837f0, 0x3f99e046, 0x3f9b8d3a, 0x3f9d3eda, 0x3f9ef532, 0x3fa0b051, 0x3fa27043, 0x3fa43516,
806 0x3fa5fed7, 0x3fa7cd94, 0x3fa9a15b, 0x3fab7a3a, 0x3fad583f, 0x3faf3b79, 0x3fb123f6, 0x3fb311c4,
807 0x3fb504f3, 0x3fb6fd92, 0x3fb8fbaf, 0x3fbaff5b, 0x3fbd08a4, 0x3fbf179a, 0x3fc12c4d, 0x3fc346cd,
808 0x3fc5672a, 0x3fc78d75, 0x3fc9b9be, 0x3fcbec15, 0x3fce248c, 0x3fd06334, 0x3fd2a81e, 0x3fd4f35b,
809 0x3fd744fd, 0x3fd99d16, 0x3fdbfbb8, 0x3fde60f5, 0x3fe0ccdf, 0x3fe33f89, 0x3fe5b907, 0x3fe8396a,
810 0x3feac0c7, 0x3fed4f30, 0x3fefe4ba, 0x3ff28177, 0x3ff5257d, 0x3ff7d0df, 0x3ffa83b3, 0x3ffd3e0c,
811];
812pub(crate) const RECIP_SEED: [u32; 128] = [
814 0x3f7f00ff, 0x3f7d08e5, 0x3f7b1885, 0x3f792fb2, 0x3f774e40, 0x3f757404, 0x3f73a0d5, 0x3f71d48c,
815 0x3f700f01, 0x3f6e500f, 0x3f6c9791, 0x3f6ae564, 0x3f693965, 0x3f679373, 0x3f65f36d, 0x3f645933,
816 0x3f62c4a7, 0x3f6135aa, 0x3f5fac1f, 0x3f5e27eb, 0x3f5ca8f1, 0x3f5b2f17, 0x3f59ba42, 0x3f584a5a,
817 0x3f56df44, 0x3f5578e9, 0x3f541733, 0x3f52ba08, 0x3f516154, 0x3f500d01, 0x3f4ebcf9, 0x3f4d7127,
818 0x3f4c2978, 0x3f4ae5d8, 0x3f49a634, 0x3f486a79, 0x3f473294, 0x3f45fe74, 0x3f44ce08, 0x3f43a13e,
819 0x3f427806, 0x3f415250, 0x3f40300c, 0x3f3f112b, 0x3f3df59d, 0x3f3cdd53, 0x3f3bc841, 0x3f3ab656,
820 0x3f39a786, 0x3f389bc3, 0x3f379301, 0x3f368d31, 0x3f358a48, 0x3f348a3a, 0x3f338cfa, 0x3f32927c,
821 0x3f319ab6, 0x3f30a59b, 0x3f2fb322, 0x3f2ec33e, 0x3f2dd5e6, 0x3f2ceb10, 0x3f2c02b0, 0x3f2b1cbe,
822 0x3f2a392f, 0x3f2957fb, 0x3f287917, 0x3f279c7b, 0x3f26c21e, 0x3f25e9f7, 0x3f2513fd, 0x3f244029,
823 0x3f236e72, 0x3f229ecf, 0x3f21d13a, 0x3f2105a9, 0x3f203c17, 0x3f1f747a, 0x3f1eaecd, 0x3f1deb07,
824 0x3f1d2922, 0x3f1c6917, 0x3f1baadf, 0x3f1aee73, 0x3f1a33cd, 0x3f197ae7, 0x3f18c3bb, 0x3f180e41,
825 0x3f175a75, 0x3f16a850, 0x3f15f7cc, 0x3f1548e5, 0x3f149b93, 0x3f13efd2, 0x3f13459c, 0x3f129cec,
826 0x3f11f5bd, 0x3f115009, 0x3f10abcc, 0x3f100901, 0x3f0f67a2, 0x3f0ec7ab, 0x3f0e2918, 0x3f0d8be3,
827 0x3f0cf009, 0x3f0c5584, 0x3f0bbc51, 0x3f0b246b, 0x3f0a8dcd, 0x3f09f874, 0x3f09645c, 0x3f08d181,
828 0x3f083fde, 0x3f07af70, 0x3f072033, 0x3f069223, 0x3f06053c, 0x3f05797c, 0x3f04eedd, 0x3f04655e,
829 0x3f03dcf9, 0x3f0355ad, 0x3f02cf75, 0x3f024a4e, 0x3f01c636, 0x3f014328, 0x3f00c122, 0x3f004020,
830];
831pub(crate) const RSQRT_SEED: [u32; 128] = [
834 0x3f7f017e, 0x3f7d0d3e, 0x3f7b2452, 0x3f79464e, 0x3f7772cb, 0x3f75a969, 0x3f73e9cb, 0x3f723399,
835 0x3f708681, 0x3f6ee232, 0x3f6d4660, 0x3f6bb2c5, 0x3f6a271a, 0x3f68a31d, 0x3f672691, 0x3f65b139,
836 0x3f6442db, 0x3f62db41, 0x3f617a36, 0x3f601f87, 0x3f5ecb04, 0x3f5d7c7f, 0x3f5c33cc, 0x3f5af0bf,
837 0x3f59b330, 0x3f587af7, 0x3f5747ee, 0x3f5619f2, 0x3f54f0df, 0x3f53cc93, 0x3f52acee, 0x3f5191d1,
838 0x3f507b1d, 0x3f4f68b6, 0x3f4e5a7f, 0x3f4d505e, 0x3f4c4a38, 0x3f4b47f4, 0x3f4a497a, 0x3f494eb3,
839 0x3f485787, 0x3f4763e2, 0x3f4673ac, 0x3f4586d3, 0x3f449d42, 0x3f43b6e6, 0x3f42d3ad, 0x3f41f383,
840 0x3f411659, 0x3f403c1c, 0x3f3f64bd, 0x3f3e902c, 0x3f3dbe58, 0x3f3cef34, 0x3f3c22b1, 0x3f3b58c0,
841 0x3f3a9154, 0x3f39cc60, 0x3f3909d6, 0x3f3849aa, 0x3f378bd0, 0x3f36d03d, 0x3f3616e4, 0x3f355fba,
842 0x3f3450fc, 0x3f32ef41, 0x3f319589, 0x3f304387, 0x3f2ef8f2, 0x3f2db587, 0x3f2c7904, 0x3f2b432a,
843 0x3f2a13c0, 0x3f28ea8c, 0x3f27c759, 0x3f26a9f4, 0x3f25922c, 0x3f247fd3, 0x3f2372bd, 0x3f226abe,
844 0x3f2167af, 0x3f206968, 0x3f1f6fc4, 0x3f1e7aa0, 0x3f1d89d9, 0x3f1c9d4e, 0x3f1bb4e1, 0x3f1ad073,
845 0x3f19efe6, 0x3f191320, 0x3f183a05, 0x3f17647c, 0x3f16926c, 0x3f15c3bc, 0x3f14f857, 0x3f143026,
846 0x3f136b13, 0x3f12a90b, 0x3f11e9f9, 0x3f112dca, 0x3f10746d, 0x3f0fbdce, 0x3f0f09dd, 0x3f0e5889,
847 0x3f0da9c2, 0x3f0cfd79, 0x3f0c539f, 0x3f0bac25, 0x3f0b06fd, 0x3f0a641a, 0x3f09c36d, 0x3f0924ec,
848 0x3f088889, 0x3f07ee37, 0x3f0755ed, 0x3f06bf9e, 0x3f062b40, 0x3f0598c7, 0x3f05082a, 0x3f04795f,
849 0x3f03ec5c, 0x3f036117, 0x3f02d788, 0x3f024fa6, 0x3f01c967, 0x3f0144c4, 0x3f00c1b4, 0x3f004030,
850];
851pub(crate) const LN_C: [u32; 49] = [
853 0xbe934b11, 0xbe88bc74, 0xbe7cc8e3, 0xbe6881c0, 0xbe549f6a, 0xbe411e0b, 0xbe2dfa03, 0xbe1b2fe6,
854 0xbe08bc74, 0xbded393b, 0xbdc99af3, 0xbda6988b, 0xbd842cc6, 0xbd44a551, 0xbd020aec, 0xbc8102b3,
855 0x00000000, 0x3c7e0546, 0x3cfc14d8, 0x3d3ba2c8, 0x3d785186, 0x3d9a0ebd, 0x3db78694, 0x3dd4936a,
856 0x3df1383b, 0x3e06bbf4, 0x3e14aa98, 0x3e22695b, 0x3e2ff984, 0x3e3d5c48, 0x3e4a92d5, 0x3e579e4a,
857 0x3e647fbe, 0x3e71383b, 0x3e7dc8c3, 0x3e851927, 0x3e8b3ae5, 0x3e914a10, 0x3e974716, 0x3e9d3263,
858 0x3ea30c5e, 0x3ea8d56c, 0x3eae8dee, 0x3eb43641, 0x3eb9cec0, 0x3ebf57c2, 0x3ec4d19c, 0x3eca3ca1,
859 0x3ecf991f,
860];
861pub(crate) const INV_C: [u32; 49] = [
863 0x3faaaaab, 0x3fa72f05, 0x3fa3d70a, 0x3fa0a0a1, 0x3f9d89d9, 0x3f9a90e8, 0x3f97b426, 0x3f94f209,
864 0x3f924925, 0x3f8fb824, 0x3f8d3dcb, 0x3f8ad8f3, 0x3f888889, 0x3f864b8a, 0x3f842108, 0x3f820821,
865 0x3f800000, 0x3f7c0fc1, 0x3f783e10, 0x3f74898d, 0x3f70f0f1, 0x3f6d7304, 0x3f6a0ea1, 0x3f66c2b4,
866 0x3f638e39, 0x3f607038, 0x3f5d67c9, 0x3f5a740e, 0x3f579436, 0x3f54c77b, 0x3f520d21, 0x3f4f6475,
867 0x3f4ccccd, 0x3f4a4588, 0x3f47ce0c, 0x3f4565c8, 0x3f430c31, 0x3f40c0c1, 0x3f3e82fa, 0x3f3c5264,
868 0x3f3a2e8c, 0x3f381703, 0x3f360b61, 0x3f340b41, 0x3f321643, 0x3f302c0b, 0x3f2e4c41, 0x3f2c7692,
869 0x3f2aaaab,
870];
871
872fn table(t: &[u32], i: usize) -> f32 {
873 f32::from_bits(t[i])
874}
875
876pub(crate) const INV_LN2_64: f32 = f32::from_bits(0x42b8_aa3b);
878pub(crate) const LN2_64_HI: f32 = f32::from_bits(0x3c31_4000);
879pub(crate) const LN2_64_LO: f32 = f32::from_bits(0x3748_5fdf);
880pub(crate) const LN2_HI: f32 = f32::from_bits(0x3f31_7000);
882pub(crate) const LN2_LO: f32 = f32::from_bits(0x3805_fdf4);
883pub(crate) const ONE_SIXTH: f32 = 1.0 / 6.0;
884pub(crate) const ONE_THIRD: f32 = 1.0 / 3.0;
885pub(crate) const EXP_TOP: f32 = 88.8;
887pub(crate) const EXP_BOTTOM: f32 = -104.0;
888pub(crate) const SOFTPLUS_TOP: f32 = 16.0;
890
891pub fn exp(fp: FpUnit, x: f32) -> f32 {
894 let x = ftz(x);
895 if x.is_nan() {
896 return x;
897 }
898 if x > EXP_TOP {
899 return f32::INFINITY;
900 }
901 if x < EXP_BOTTOM {
902 return 0.0;
903 }
904 let n = rne(fp.mul(x, INV_LN2_64));
905 let nf = n as f32;
906 let r = fsub(fsub(x, fp.mul(nf, LN2_64_HI)), fp.mul(nf, LN2_64_LO));
907 let p = fadd(1.0, fp.mul(r, fadd(1.0, fp.mul(r, fadd(0.5, fp.mul(r, ONE_SIXTH))))));
908 let j = n.rem_euclid(64) as usize;
909 let m = (n - j as i64) / 64;
910 ldexp(fp.mul(table(&EXP2_64, j), p), m as i32)
911}
912
913pub fn recip(fp: FpUnit, x: f32) -> f32 {
915 let x = ftz(x);
916 if x.is_nan() {
917 return x;
918 }
919 if x == 0.0 {
920 return f32::INFINITY.copysign(x);
921 }
922 if x.is_infinite() {
923 return 0f32.copysign(x);
924 }
925 let bits = x.to_bits();
926 let e = ((bits >> 23) & 0xff) as i32 - 127;
927 let i = ((bits >> 16) & 0x7f) as usize;
928 let seed = ldexp(table(&RECIP_SEED, i), -e).copysign(x);
929 let mut y = seed;
930 for _ in 0..2 {
931 y = fp.mul(y, fsub(2.0, fp.mul(x, y)));
932 }
933 y
934}
935
936pub fn rsqrt(fp: FpUnit, x: f32) -> f32 {
939 let x = ftz(x);
940 if x.is_nan() || x < 0.0 {
941 return f32::NAN;
942 }
943 if x == 0.0 {
944 return f32::INFINITY;
945 }
946 if x.is_infinite() {
947 return 0.0;
948 }
949 let bits = x.to_bits();
950 let e = ((bits >> 23) & 0xff) as i32 - 127;
951 let parity = e.rem_euclid(2);
952 let i = ((bits >> 17) & 0x3f) as usize;
953 let seed = ldexp(table(&RSQRT_SEED, parity as usize * 64 + i), -(e - parity) / 2);
954 let mut y = seed;
955 for _ in 0..2 {
956 y = fp.mul(y, fsub(1.5, fp.mul(0.5, fp.mul(fp.mul(x, y), y))));
957 }
958 y
959}
960
961pub fn ln(fp: FpUnit, x: f32) -> f32 {
964 let x = ftz(x);
965 if x.is_nan() || x < 0.0 {
966 return f32::NAN;
967 }
968 if x == 0.0 {
969 return f32::NEG_INFINITY;
970 }
971 if x.is_infinite() {
972 return x;
973 }
974 let bits = x.to_bits();
975 let mut big_e = ((bits >> 23) & 0xff) as i32 - 127;
976 let mut m = f32::from_bits((bits & 0x007f_ffff) | 0x3f80_0000);
977 if m >= 1.5 {
978 m *= 0.5;
980 big_e += 1;
981 }
982 let j = rne((m - 0.75) * 64.0) as usize;
984 let c = 0.75 + j as f32 / 64.0;
985 let d = fp.mul(fsub(m, c), table(&INV_C, j));
986 let d2 = fp.mul(d, d);
987 let poly = fsub(d, fp.mul(d2, fsub(0.5, fp.mul(d, fsub(ONE_THIRD, fp.mul(d, 0.25))))));
988 let ef = big_e as f32;
989 fadd(fadd(fp.mul(ef, LN2_HI), fadd(table(&LN_C, j), poly)), fp.mul(ef, LN2_LO))
990}
991
992pub fn ln1p(fp: FpUnit, t: f32) -> f32 {
995 let t = ftz(t);
996 let u = fadd(1.0, t);
997 if u == 1.0 {
998 return t;
999 }
1000 fadd(ln(fp, u), fp.mul(fsub(t, fsub(u, 1.0)), recip(fp, u)))
1001}
1002
1003pub fn softplus(fp: FpUnit, u: f32) -> f32 {
1005 let u = ftz(u);
1006 if u > SOFTPLUS_TOP {
1007 return fadd(u, exp(fp, -u));
1008 }
1009 ln1p(fp, exp(fp, u))
1010}
1011
1012pub fn sigmoid(fp: FpUnit, x: f32) -> f32 {
1013 recip(fp, fadd(1.0, exp(fp, -ftz(x))))
1014}
1015
1016pub fn silu(fp: FpUnit, x: f32) -> f32 {
1017 fp.mul(x, sigmoid(fp, x))
1018}
1019
1020pub fn rmsnorm(fp: FpUnit, x: &[f32], w: &[f32], eps: f32) -> Vec<f32> {
1022 let n = x.len() as f32;
1023 let r = rsqrt(fp, fadd(fp.mul(fdot(fp, x, x), recip(fp, n)), eps));
1024 x.iter().zip(w).map(|(&v, &g)| fp.mul(fp.mul(v, r), g)).collect()
1025}
1026
1027pub fn l2norm(fp: FpUnit, x: &[f32], eps: f32) -> Vec<f32> {
1029 let r = rsqrt(fp, fadd(fdot(fp, x, x), eps));
1030 x.iter().map(|&v| fp.mul(v, r)).collect()
1031}
1032
1033pub fn hadamard(x: &mut [f32]) {
1037 let n = x.len();
1038 assert!(n.is_power_of_two(), "a Hadamard transform of {n} elements");
1039 let mut h = 1;
1040 while h < n {
1041 for start in (0..n).step_by(2 * h) {
1042 for j in start..start + h {
1043 let (a, b) = (x[j], x[j + h]);
1044 x[j] = fadd(a, b);
1045 x[j + h] = fsub(a, b);
1046 }
1047 }
1048 h *= 2;
1049 }
1050}
1051
1052pub fn rope(fp: FpUnit, v: &mut [f32], cos: &[f32], sin: &[f32], neox: bool) {
1055 let half = cos.len();
1056 assert!(sin.len() == half && 2 * half <= v.len(), "a rotary table of {half} pairs over {} elements", v.len());
1057 for i in 0..half {
1058 let (a, b) = if neox { (i, i + half) } else { (2 * i, 2 * i + 1) };
1059 let (x0, x1) = (v[a], v[b]);
1060 v[a] = fsub(fp.mul(x0, cos[i]), fp.mul(x1, sin[i]));
1061 v[b] = fadd(fp.mul(x0, sin[i]), fp.mul(x1, cos[i]));
1062 }
1063}
1064
1065pub fn softmax(fp: FpUnit, scores: &[f32]) -> Vec<f32> {
1070 let m = fmax(scores);
1071 let e: Vec<f32> = scores.iter().map(|&s| exp(fp, fsub(s, m))).collect();
1072 let inv = recip(fp, fsum(&e));
1073 e.iter().map(|&v| fp.mul(v, inv)).collect()
1074}
1075
1076pub fn attend_vector(fp: FpUnit, q: &[f32], keys: &[&[f32]], values: &[&[f32]], scale: f32) -> Vec<f32> {
1080 let scores: Vec<f32> = keys.iter().map(|k| fp.mul(fdot(fp, q, k), scale)).collect();
1081 let p = softmax(fp, &scores);
1082 let width = values.first().map_or(0, |v| v.len());
1083 (0..width).map(|e| fdot(fp, &p, &values.iter().map(|v| v[e]).collect::<Vec<_>>())).collect()
1084}
1085
1086pub fn quantize_head(q: &[f32], words: Words, scaling: ActScaling) -> ActRow {
1089 match scaling {
1090 ActScaling::Sub32 => {
1091 let mut padded = q.to_vec();
1092 padded.resize(q.len().div_ceil(GROUP) * GROUP, 0.0);
1093 quantize(&padded, words)
1094 }
1095 ActScaling::PerBlock => quantize_blocks(q, words, q.len()),
1096 }
1097}
1098
1099pub fn int8_score(q: &ActRow, key_scale: u16, key: &[i8]) -> f32 {
1103 let scale = Scale::f16(key_scale).expect("a finite cache scale");
1104 let terms: Vec<(i128, i32)> = key
1105 .chunks(q.block)
1106 .enumerate()
1107 .map(|(b, block)| {
1108 let sum: i64 = block.iter().enumerate().map(|(i, &c)| q.codes[b * q.block + i] as i64 * c as i64).sum();
1109 let sc = q.scales[b];
1110 (scale.sig as i128 * sc.m as i128 * sum as i128, scale.exp + sc.e)
1111 })
1112 .filter(|t| t.0 != 0)
1113 .collect();
1114 exact_sum(&terms)
1115}
1116
1117pub const VALUE_BLOCK: usize = 32;
1119
1120pub fn weighted_int8(fp: FpUnit, p: &[f32], values: &[(u16, &[i8])]) -> Vec<f32> {
1125 let width = values.first().map_or(0, |v| v.1.len());
1126 let mut out = vec![0f32; width];
1127 for (b, block) in p.chunks(VALUE_BLOCK).enumerate() {
1128 let rows = &values[b * VALUE_BLOCK..b * VALUE_BLOCK + block.len()];
1129 let w: Vec<f32> = block.iter().zip(rows).map(|(&pj, (s, _))| fp.mul(pj, f16_to_f32(*s))).collect();
1130 let wmax = w.iter().fold(0f32, |m, &v| m.max(v));
1131 if wmax == 0.0 {
1132 for o in out.iter_mut() {
1133 *o = fadd(*o, 0.0);
1134 }
1135 continue;
1136 }
1137 let mut e = (wmax as f64 / 255.0).log2().ceil() as i32;
1139 while 255.0 * 2f64.powi(e) < wmax as f64 {
1140 e += 1;
1141 }
1142 while 255.0 * 2f64.powi(e - 1) >= wmax as f64 {
1143 e -= 1;
1144 }
1145 let codes: Vec<i64> = w.iter().map(|&v| (v as f64 * 2f64.powi(-e)).round_ties_even() as i64).collect();
1146 for (el, o) in out.iter_mut().enumerate() {
1147 let sum: i64 = codes.iter().zip(rows).map(|(&c, (_, v))| c * v[el] as i64).sum();
1148 *o = fadd(*o, scaled_to_f32(sum as i128, e));
1149 }
1150 }
1151 out
1152}
1153
1154pub fn attend_int8(fp: FpUnit, q: &[f32], keys: &[(u16, &[i8])], values: &[(u16, &[i8])], scale: f32, words: Words, scaling: ActScaling) -> Vec<f32> {
1157 let qa = quantize_head(q, words, scaling);
1158 let scores: Vec<f32> = keys.iter().map(|&(s, k)| fp.mul(int8_score(&qa, s, k), scale)).collect();
1159 weighted_int8(fp, &softmax(fp, &scores), values)
1160}
1161
1162pub fn conv(fp: FpUnit, ring: &[Vec<f32>], x: &[f32], kernel: &[f32], taps: usize) -> Vec<f32> {
1168 assert_eq!(ring.len() + 1, taps, "a ring row for every tap but the token's own");
1169 (0..x.len())
1170 .map(|ch| {
1171 let mut acc = 0f32;
1172 for (j, row) in ring.iter().enumerate() {
1173 acc = fadd(acc, fp.mul(row[ch], kernel[ch * taps + j]));
1174 }
1175 silu(fp, fadd(acc, fp.mul(x[ch], kernel[ch * taps + taps - 1])))
1176 })
1177 .collect()
1178}
1179
1180pub fn gdn_update(fp: FpUnit, row: &mut [f32], k: &[f32], decay: f32, beta: f32, v: f32, store: impl Fn(f32) -> f32) -> f32 {
1183 let u = fp.mul(decay, fdot(fp, row, k));
1184 let coef = fp.mul(beta, fsub(v, u));
1185 for (m, &kc) in row.iter_mut().zip(k) {
1186 *m = store(fadd(fp.mul(decay, *m), fp.mul(coef, kc)));
1187 }
1188 coef
1189}
1190
1191#[allow(clippy::too_many_arguments)]
1195pub fn gdn_row(fp: FpUnit, row: &mut [f32], k: &[f32], q: &[f32], decay: f32, beta: f32, v: f32, kq: f32, scale: f32, store: impl Fn(f32) -> f32) -> f32 {
1196 let sq = fdot(fp, row, q);
1197 let coef = gdn_update(fp, row, k, decay, beta, v, store);
1198 fp.mul(fadd(fp.mul(decay, sq), fp.mul(coef, kq)), scale)
1199}
1200
1201#[derive(Clone, Copy, Debug)]
1204pub struct Logged<'a> {
1205 pub k: &'a [f32],
1206 pub v: &'a [f32],
1207 pub log_decay: f32,
1208 pub beta: f32,
1209}
1210
1211pub fn gdn_fold(fp: FpUnit, q: &[f32], log: &[Logged], width: usize) -> (Vec<f32>, Vec<f32>) {
1215 fold_on(fp, q.to_vec(), vec![0f32; width], log)
1216}
1217
1218pub fn gdn_fold_trailing(fp: FpUnit, q: &[f32], log: &[Logged], width: usize, older: usize) -> [(Vec<f32>, Vec<f32>); 2] {
1223 let newer = gdn_fold(fp, q, &log[older..], width);
1224 let all = fold_on(fp, newer.0.clone(), newer.1.clone(), &log[..older]);
1225 [newer, all]
1226}
1227
1228fn fold_on(fp: FpUnit, mut y: Vec<f32>, mut carried: Vec<f32>, log: &[Logged]) -> (Vec<f32>, Vec<f32>) {
1230 for e in log.iter().rev() {
1231 let decay = exp(fp, e.log_decay);
1232 let bc = fp.mul(e.beta, fdot(fp, e.k, &y));
1233 for (slot, &v) in carried.iter_mut().zip(e.v) {
1234 *slot = fadd(*slot, fp.mul(bc, v));
1235 }
1236 for (yi, &kc) in y.iter_mut().zip(e.k) {
1237 *yi = fp.mul(decay, fsub(*yi, fp.mul(bc, kc)));
1238 }
1239 }
1240 (y, carried)
1241}
1242
1243pub fn gdn_read(fp: FpUnit, row: &[f32], y: &[f32], carried: f32, scale: f32) -> f32 {
1245 fp.mul(fadd(fdot(fp, row, y), carried), scale)
1246}
1247
1248pub fn gdn_flush(fp: FpUnit, row: &mut [f32], log: &[Logged], r: usize, store: impl Fn(f32) -> f32) {
1251 let mut s = row.to_vec();
1252 for e in log {
1253 gdn_update(fp, &mut s, e.k, exp(fp, e.log_decay), e.beta, e.v[r], |x| x);
1254 }
1255 for (m, x) in row.iter_mut().zip(s) {
1256 *m = store(x);
1257 }
1258}
1259
1260pub const INV_127: f32 = f32::from_bits(0x3c01_0204);
1262
1263pub fn int8_codes(fp: FpUnit, x: &[f32], inv: f32) -> Vec<i8> {
1266 if inv.is_infinite() {
1267 return vec![0; x.len()];
1268 }
1269 x.iter().map(|&v| rne(fp.mul(v, inv)).clamp(-127, 127) as i8).collect()
1270}
1271
1272pub fn cache_row(fp: FpUnit, x: &[f32]) -> (u16, Vec<i8>) {
1275 let amax = absmax(x);
1276 assert!(amax.is_finite(), "a non-finite cache row");
1277 let s = f32_to_f16(fp.mul(amax, INV_127));
1278 (s, int8_codes(fp, x, recip(fp, f16_to_f32(s))))
1279}