cherenkov/qwen4_exp/pack/
affine.rs

1//! Affine Q4 groups: eight low-first nibbles per word, BF16 scale and bias.
2
3use anyhow::{Result, ensure};
4use half::bf16;
5
6pub(super) struct Group {
7    pub weight: [u8; 32],
8    pub scale: [u8; 2],
9    pub bias: [u8; 2],
10}
11
12/// Use MLX's affine endpoint selection and ties-to-even rounding. Quantize
13/// with float parameters, then store them as BF16, as MLX does for BF16 input.
14/// Reference: mlx/backend/cpu/quantized.cpp, quantize<T, U>.
15pub(super) fn quantize(bytes: &[u8]) -> Result<Group> {
16    let mut values = [0.0_f32; 64];
17    let len = bytes.len() / 2;
18
19    ensure!(
20        matches!(len, 8 | 16 | 32 | 64),
21        "unsupported affine group size {len}"
22    );
23
24    for (value, pair) in values.iter_mut().zip(bytes.as_chunks::<2>().0) {
25        *value = bf16::from_bits(u16::from_le_bytes(*pair)).to_f32();
26
27        ensure!(
28            value.is_finite(),
29            "cannot quantize a non-finite BF16 weight"
30        );
31    }
32
33    let values = &values[..len];
34    let min = values.iter().copied().fold(f32::INFINITY, f32::min);
35    let max = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
36    let magnitude = ((max - min) / 15.0).max(1e-7);
37    let (edge, mut scale) = if min.abs() > max.abs() {
38        (min, magnitude)
39    } else {
40        (max, -magnitude)
41    };
42    let zero = (edge / scale).round_ties_even();
43    let mut bias = 0.0;
44
45    if zero != 0.0 {
46        scale = edge / zero;
47        bias = edge;
48    }
49
50    let stored_scale = bf16::from_f32(scale);
51
52    ensure!(
53        scale.is_finite() && scale != 0.0 && stored_scale.is_finite(),
54        "affine scale overflow"
55    );
56
57    let mut group = Group {
58        weight: [0; 32],
59        scale: stored_scale.to_le_bytes(),
60        bias: bf16::from_f32(bias).to_le_bytes(),
61    };
62
63    for (i, &value) in values.iter().enumerate() {
64        let nibble = ((value - bias) / scale).round_ties_even().clamp(0.0, 15.0) as u8;
65        group.weight[i / 2] |= nibble << (4 * (i % 2));
66    }
67
68    Ok(group)
69}
70
71#[cfg(test)]
72#[path = "../../../tests/unit/qwen4_exp/pack/affine.rs"]
73mod tests;
cherenkov/qwen4_exp/pack/affine.rs

cherenkov/qwen4_exp/pack/
affine.rs
