zsfm_checkpoint/
cast.rs

1//! Canonical dtype casting for GGUF conversion: {F32, F16, BF16} sources to any
2//! writable [`GGMLType`], including Q8_0 block quantization. Bit-for-bit the same
3//! math as the per-model converters (which predate this crate and keep their own
4//! verified copies).
5
6use zsfm_gguf::GGMLType;
7
8/// Source dtype of raw checkpoint tensor bytes.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SrcDtype {
11    F32,
12    F16,
13    BF16,
14}
15
16impl SrcDtype {
17    pub fn bytes_per_elem(self) -> usize {
18        match self {
19            SrcDtype::F32 => 4,
20            SrcDtype::F16 | SrcDtype::BF16 => 2,
21        }
22    }
23
24    pub fn name(self) -> &'static str {
25        match self {
26            SrcDtype::F32 => "F32",
27            SrcDtype::F16 => "F16",
28            SrcDtype::BF16 => "BF16",
29        }
30    }
31}
32
33/// Cast raw little-endian tensor bytes from `src` to `dst`.
34///
35/// Same-dtype casts are byte-for-byte passthrough. F32→F16 rounds to nearest
36/// even; F32→BF16 rounds to nearest even with NaN preservation; everything
37/// routed through F32 uses exact widening.
38pub fn cast_data(data: &[u8], src: SrcDtype, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
39    // Passthrough when the representation already matches.
40    match (src, dst) {
41        (SrcDtype::F32, GGMLType::F32)
42        | (SrcDtype::F16, GGMLType::F16)
43        | (SrcDtype::BF16, GGMLType::BF16) => return Ok(data.to_vec()),
44        _ => {}
45    }
46
47    let f32_values = decode_to_f32(data, src)?;
48    match dst {
49        GGMLType::F32 => Ok(f32_to_bytes(&f32_values)),
50        GGMLType::F16 => Ok(f32_values
51            .iter()
52            .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
53            .collect()),
54        GGMLType::BF16 => Ok(f32_values
55            .iter()
56            .flat_map(|&v| f32_to_bf16_bits(v).to_le_bytes())
57            .collect()),
58        GGMLType::Q8_0 => quantize_q8_0(&f32_values),
59    }
60}
61
62/// Round-to-nearest-even F32→BF16 with NaN preservation.
63pub fn f32_to_bf16_bits(v: f32) -> u16 {
64    let bits = v.to_bits();
65    // Preserve NaN payload as quiet NaN.
66    if (bits & 0x7f80_0000) == 0x7f80_0000 && (bits & 0x007f_ffff) != 0 {
67        return 0x7fc0;
68    }
69    // Round to nearest even: add 0x7fff + LSB of result.
70    let rounding = 0x7fff + ((bits >> 16) & 1);
71    ((bits + rounding) >> 16) as u16
72}
73
74/// Decode raw bytes of `src` dtype into f32 values (exact for all three sources).
75pub fn decode_to_f32(data: &[u8], src: SrcDtype) -> anyhow::Result<Vec<f32>> {
76    match src {
77        SrcDtype::F32 => parse_f32_le(data),
78        SrcDtype::F16 => {
79            if data.len() % 2 != 0 {
80                anyhow::bail!("f16 data length not divisible by 2");
81            }
82            Ok(data
83                .chunks_exact(2)
84                .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]])))
85                .collect())
86        }
87        SrcDtype::BF16 => {
88            if data.len() % 2 != 0 {
89                anyhow::bail!("bf16 data length not divisible by 2");
90            }
91            Ok(data
92                .chunks_exact(2)
93                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
94                .collect())
95        }
96    }
97}
98
99pub fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
100    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
101}
102
103/// Q8_0 block-quantization fallback: tensors whose innermost dim or total
104/// element count is not a multiple of 32 stay F32. Centralizes the
105/// `innermost % 32` check duplicated in `generic.rs`, the Python
106/// `write_generic_gguf`, and every per-model `convert.rs`.
107pub fn q8_dst_or_fallback(shape: &[u64], out_dtype: GGMLType) -> (GGMLType, bool) {
108    if out_dtype == GGMLType::Q8_0 {
109        let n_elems: u64 = shape.iter().product();
110        let innermost = shape.last().copied().unwrap_or(1);
111        if innermost % 32 != 0 || n_elems % 32 != 0 {
112            return (GGMLType::F32, true);
113        }
114    }
115    (out_dtype, false)
116}
117
118/// Cast with automatic Q8_0 fallback. Returns `(bytes, actual_dtype, fell_back)`.
119pub fn cast_with_q8_fallback(
120    data: &[u8],
121    src: SrcDtype,
122    shape: &[u64],
123    out_dtype: GGMLType,
124) -> anyhow::Result<(Vec<u8>, GGMLType, bool)> {
125    let (dst, fell_back) = q8_dst_or_fallback(shape, out_dtype);
126    Ok((cast_data(data, src, dst)?, dst, fell_back))
127}
128
129pub fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
130    const BLOCK: usize = 32;
131    if values.len() % BLOCK != 0 {
132        anyhow::bail!(
133            "Q8_0 requires element count divisible by {BLOCK}, got {}",
134            values.len()
135        );
136    }
137    let n_blocks = values.len() / BLOCK;
138    let mut out = vec![0u8; n_blocks * 34];
139    for b in 0..n_blocks {
140        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
141        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
142        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
143        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
144        let base = b * 34;
145        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
146        for i in 0..BLOCK {
147            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
148        }
149    }
150    Ok(out)
151}
152
153fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
154    if data.len() % 4 != 0 {
155        anyhow::bail!("f32 data length not divisible by 4");
156    }
157    Ok(data
158        .chunks_exact(4)
159        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
160        .collect())
161}
162
163pub fn f32_to_f16_bits(v: f32) -> u16 {
164    let bits = v.to_bits();
165    let sign = ((bits >> 16) & 0x8000) as u16;
166    let exp = ((bits >> 23) & 0xFF) as i32;
167    let mantissa = bits & 0x007F_FFFF;
168    if exp == 0xFF {
169        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
170    }
171    let new_exp = exp - 127 + 15;
172    if new_exp >= 31 {
173        return sign | 0x7C00;
174    }
175    if new_exp <= 0 {
176        if new_exp < -10 {
177            return sign;
178        }
179        // Round subnormals to nearest even.
180        let shift = 1 - new_exp;
181        let m = mantissa | 0x0080_0000;
182        let remainder = m & ((1 << shift) - 1);
183        let halfway = 1 << (shift - 1);
184        let mut rounded = m >> shift;
185        if remainder > halfway as u32 || (remainder == halfway as u32 && (rounded & 1) == 1) {
186            rounded += 1;
187        }
188        return sign | (rounded >> 13) as u16;
189    }
190    // Round normal mantissa to nearest even: add 0xFFF + LSB then truncate.
191    let lsb = (mantissa >> 13) & 1;
192    let rounded = mantissa + 0xFFF + lsb;
193    if rounded & 0x0080_0000 != 0 {
194        // Mantissa overflow carries into exponent.
195        let carried_exp = new_exp + 1;
196        if carried_exp >= 31 {
197            return sign | 0x7C00;
198        }
199        return sign | ((carried_exp as u16) << 10);
200    }
201    sign | ((new_exp as u16) << 10) | (rounded >> 13) as u16
202}
203
204pub fn f16_to_f32(bits: u16) -> f32 {
205    let sign = ((bits & 0x8000) as u32) << 16;
206    let exp = ((bits >> 10) & 0x1F) as i32;
207    let mantissa = (bits & 0x03FF) as u32;
208    let f32_bits = if exp == 0 {
209        if mantissa == 0 {
210            sign
211        } else {
212            let mut m = mantissa;
213            let mut e = 0i32;
214            while m & 0x0400 == 0 {
215                m <<= 1;
216                e += 1;
217            }
218            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
219        }
220    } else if exp == 31 {
221        sign | 0x7F80_0000 | (mantissa << 13)
222    } else {
223        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
224    };
225    f32::from_bits(f32_bits)
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn f16_roundtrip_exact_values() {
234        for v in [0.0f32, 1.0, -1.0, 0.5, 65504.0, -65504.0, 0.099975586] {
235            assert_eq!(f16_to_f32(f32_to_f16_bits(v)), v, "roundtrip {v}");
236        }
237    }
238
239    #[test]
240    fn passthrough_same_dtype() {
241        let bytes: Vec<u8> = (0..16).collect();
242        assert_eq!(
243            cast_data(&bytes, SrcDtype::F32, GGMLType::F32).unwrap(),
244            bytes
245        );
246        assert_eq!(
247            cast_data(&bytes, SrcDtype::F16, GGMLType::F16).unwrap(),
248            bytes
249        );
250        assert_eq!(
251            cast_data(&bytes, SrcDtype::BF16, GGMLType::BF16).unwrap(),
252            bytes
253        );
254    }
255
256    #[test]
257    fn f32_to_bf16_rounds() {
258        let v = 1.2345678f32;
259        let out = cast_data(&v.to_le_bytes(), SrcDtype::F32, GGMLType::BF16).unwrap();
260        let bits = u16::from_le_bytes([out[0], out[1]]);
261        assert_eq!(bits, f32_to_bf16_bits(v));
262        // NaN stays NaN.
263        let nan_bits = f32_to_bf16_bits(f32::NAN);
264        assert_eq!(nan_bits & 0x7f80, 0x7f80);
265        assert_ne!(nan_bits & 0x007f, 0);
266    }
267
268    #[test]
269    fn bf16_to_f32_exact_widening() {
270        let bits: u16 = 0x3FA0; // 1.25 in bf16
271        let out = cast_data(&bits.to_le_bytes(), SrcDtype::BF16, GGMLType::F32).unwrap();
272        let v = f32::from_le_bytes([out[0], out[1], out[2], out[3]]);
273        assert_eq!(v, 1.25);
274    }
275
276    #[test]
277    fn q8_0_block_layout() {
278        let vals: Vec<f32> = (0..32).map(|i| i as f32).collect();
279        let out = quantize_q8_0(&vals).unwrap();
280        assert_eq!(out.len(), 34);
281        let d = f16_to_f32(u16::from_le_bytes([out[0], out[1]]));
282        assert!((d - 31.0 / 127.0).abs() < 1e-3);
283        assert_eq!(out[2] as i8, 0); // 0.0 quantizes to 0
284        assert_eq!(out[33] as i8, 127); // amax quantizes to 127
285    }
286
287    #[test]
288    fn q8_0_rejects_partial_block() {
289        assert!(quantize_q8_0(&[1.0f32; 31]).is_err());
290    }
291}
zsfm_checkpoint/cast.rs

zsfm_checkpoint/
cast.rs
