File size: 3,962 Bytes
19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | //! FP8 E4M3FN encode/decode and the per-output-row FP8 quantiser used for the MTP draft head (quantised at load from its bf16 checkpoint).
//! The served checkpoint itself is stored pre-quantised (NVFP4 experts, FP8 linears and tied table) and is uploaded as is.
pub const FP8_MAX: f32 = 448.0;
pub fn fp8_e4m3_to_f32(b: u8) -> f32 {
let s = if b & 0x80 != 0 { -1.0 } else { 1.0 };
let e = ((b >> 3) & 0xF) as i32;
let m = (b & 7) as f32;
if e == 0xF && (b & 7) == 7 {
return f32::NAN;
}
if e == 0 {
s * m / 8.0 * 2f32.powi(-6)
} else {
s * (1.0 + m / 8.0) * 2f32.powi(e - 7)
}
}
/// round-to-nearest-even encode (torch `.to(float8_e4m3fn)` semantics), saturating at +-448
pub fn f32_to_fp8_e4m3(x: f32) -> u8 {
if x.is_nan() {
return 0x7F;
}
let sign = if x.is_sign_negative() { 0x80u8 } else { 0 };
let a = x.abs().min(FP8_MAX);
let code = if a < 0.015625 {
// subnormal range (step 2^-9); 8 steps reaches the smallest normal 0x08 correctly
(a * 512.0).round_ties_even() as u8
} else {
let bits = a.to_bits();
let mut e = ((bits >> 23) & 0xff) as i32 - 127;
let m = bits & 0x7f_ffff;
let mut q = m >> 20;
let rem = m & 0xf_ffff;
if rem > 0x8_0000 || (rem == 0x8_0000 && q & 1 == 1) {
q += 1;
}
if q == 8 {
q = 0;
e += 1;
}
(((e + 7) as u32) << 3 | q) as u8
};
sign | code.min(0x7E)
}
pub struct Fp8Rows {
pub w: Vec<u8>,
pub scales: Vec<f32>,
}
/// FP8-E4M3 per-output-row quantiser: scale = row_amax/448, w = fp8(x/scale). No calibration.
pub fn quantize_fp8_rows(w: &[f32], rows: usize, cols: usize) -> Fp8Rows {
assert_eq!(w.len(), rows * cols);
let mut q = vec![0u8; rows * cols];
let mut scales = vec![0f32; rows];
for r in 0..rows {
let row = &w[r * cols..(r + 1) * cols];
let amax = row.iter().fold(0.0f32, |m, x| m.max(x.abs())).max(1e-12);
let s = amax / FP8_MAX;
scales[r] = s;
for (o, &x) in q[r * cols..(r + 1) * cols].iter_mut().zip(row) {
*o = f32_to_fp8_e4m3(x / s);
}
}
Fp8Rows { w: q, scales }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fp8_known_values() {
assert_eq!(fp8_e4m3_to_f32(0x00), 0.0);
assert_eq!(fp8_e4m3_to_f32(0x38), 1.0); // e=7,m=0
assert_eq!(fp8_e4m3_to_f32(0x7E), 448.0);
assert_eq!(fp8_e4m3_to_f32(0xB8), -1.0);
assert_eq!(fp8_e4m3_to_f32(0x01), 2f32.powi(-9)); // smallest subnormal
assert!(fp8_e4m3_to_f32(0x7F).is_nan());
for b in 0u8..0x7F {
assert_eq!(f32_to_fp8_e4m3(fp8_e4m3_to_f32(b)), b, "roundtrip {b:#x}");
}
assert_eq!(f32_to_fp8_e4m3(1e9), 0x7E);
}
#[test]
fn fp8_encode_matches_bruteforce_rne() {
// brute-force nearest with ties-to-even-code
let brute = |x: f32| -> u8 {
let a = x.abs().min(FP8_MAX);
let (mut best, mut bd) = (0u8, f32::INFINITY);
for c in 0u8..0x7F {
let d = (fp8_e4m3_to_f32(c) - a).abs();
if d < bd || (d == bd && c % 2 == 0 && best % 2 == 1) {
bd = d;
best = c;
}
}
(if x.is_sign_negative() { 0x80 } else { 0 }) | best
};
let mut s = 1u64;
for i in 0..20000 {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let x = ((s >> 33) as f32 / (1u64 << 31) as f32 - 0.5) * if i % 3 == 0 { 600.0 } else if i % 3 == 1 { 4.0 } else { 0.05 };
assert_eq!(f32_to_fp8_e4m3(x), brute(x), "x={x}");
}
for c in 0u8..0x7E {
let m = (fp8_e4m3_to_f32(c) + fp8_e4m3_to_f32(c + 1)) / 2.0; // exact tie
assert_eq!(f32_to_fp8_e4m3(m), brute(m), "tie {m}");
}
}
}
|