combs_core/
quant.rs

1//! Group-wise 4-bit quantization support (GGUF `q4_0`-style layout).
2//!
3//! Phase 2 ships the dequantize op as portable burn tensor ops (nibble
4//! extraction via `remainder`/`div` — no bitwise ops or custom kernels
5//! required), so it runs on any backend. A fused dequant-matmul CubeCL
6//! kernel (dequantize tiles inside the matmul, avoiding a materialized f32
7//! weight) is future work; until then [`dequantize_q4`] materializes the
8//! f32 weight and callers use the normal matmul path.
9//!
10//! # Packed layout (GGUF `q4_0`)
11//!
12//! Weights are grouped into blocks of `group_size` (32) values. Each block
13//! stores 16 bytes: byte `j`'s **low** nibble is value `j` of the block and
14//! its **high** nibble is value `j + 16`. The dequantized value is
15//! `(nibble - 8) * scale` (symmetric quantization with an implicit zero
16//! point of 8). Scales are stored per block as f32 on device (f16 scales
17//! are widened at load time by the format adapter).
18
19use burn::tensor::{Int, Tensor, backend::Backend};
20
21/// Default quantization group size (GGUF `q4_0`).
22pub const DEFAULT_Q4_GROUP_SIZE: usize = 32;
23
24/// Dequantizes packed 4-bit weights to f32.
25///
26/// - `packed`: `[rows, cols / 2]` int tensor, byte values 0..=255 (GGUF
27///   `q4_0` nibble order: low nibble = first half of the block, high nibble
28///   = second half).
29/// - `scales`: `[rows, cols / group_size]` f32 per-block scales.
30///
31/// Returns the `[rows, cols]` f32 weight matrix.
32///
33/// Panics (via shape assertions) if `cols % group_size != 0`; group sizes
34/// other than 32 keep the low/high split-half convention within each group
35/// (i.e. byte `j` of a group holds values `j` and `j + group_size/2`).
36pub fn dequantize_q4<B: Backend>(
37    packed: Tensor<B, 2, Int>,
38    scales: Tensor<B, 2>,
39    group_size: usize,
40) -> Tensor<B, 2> {
41    let [rows, packed_cols] = packed.dims();
42    let half = group_size / 2;
43    assert_eq!(
44        packed_cols % half,
45        0,
46        "packed width {packed_cols} not a multiple of half group {half}"
47    );
48    let groups_per_row = packed_cols / half;
49    let cols = groups_per_row * group_size;
50    let [s_rows, s_cols] = scales.dims();
51    assert_eq!(
52        (s_rows, s_cols),
53        (rows, groups_per_row),
54        "scales shape [{s_rows}, {s_cols}] does not match [rows, cols/group] = [{rows}, {groups_per_row}]"
55    );
56
57    // Nibble extraction (portable: remainder/div instead of bitwise ops).
58    let nibbles = packed.reshape([rows, groups_per_row, half]);
59    let lo = nibbles.clone().remainder_scalar(16); // values j (first half)
60    let hi = nibbles.div_scalar(16); // values j + half (second half)
61
62    // [rows, groups, group_size] nibbles, split-half order restored.
63    let nibbles = Tensor::cat(vec![lo, hi], 2).float();
64
65    // Symmetric dequant: w = (nibble - 8) * scale.
66    let scales = scales
67        .unsqueeze_dim::<3>(2)
68        .expand([rows, groups_per_row, group_size]);
69    let w = (nibbles - 8.0) * scales;
70    w.reshape([rows, cols])
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use burn::tensor::TensorData;
77
78    type B = burn::backend::NdArray<f32>;
79
80    /// CPU scalar reference for one GGUF q4_0 block.
81    fn reference_block(bytes: &[u8; 16], scale: f32) -> [f32; 32] {
82        let mut out = [0.0f32; 32];
83        for j in 0..16 {
84            out[j] = ((bytes[j] & 0x0F) as f32 - 8.0) * scale;
85            out[j + 16] = ((bytes[j] >> 4) as f32 - 8.0) * scale;
86        }
87        out
88    }
89
90    #[test]
91    fn dequantize_matches_scalar_reference() {
92        let device = Default::default();
93        // 2 rows x 2 groups (64 values per row, 32 packed bytes per row).
94        let packed_bytes: Vec<i32> = (0..64).map(|i| ((i * 37 + 11) % 256) as i32).collect();
95        let scales: Vec<f32> = vec![0.5, -1.25, 2.0, 0.75];
96        let packed = Tensor::<B, 2, Int>::from_data(
97            TensorData::new(packed_bytes.clone(), [2, 32]),
98            &device,
99        );
100        let scales_t = Tensor::<B, 2>::from_data(TensorData::new(scales.clone(), [2, 2]), &device);
101
102        let w = dequantize_q4(packed, scales_t, DEFAULT_Q4_GROUP_SIZE);
103        let got: Vec<f32> = w.into_data().to_vec().unwrap();
104
105        let mut expected = Vec::with_capacity(128);
106        for row in 0..2 {
107            for g in 0..2 {
108                let mut block = [0u8; 16];
109                for j in 0..16 {
110                    block[j] = packed_bytes[row * 32 + g * 16 + j] as u8;
111                }
112                expected.extend_from_slice(&reference_block(&block, scales[row * 2 + g]));
113            }
114        }
115        assert_eq!(got.len(), expected.len());
116        for (g, e) in got.iter().zip(expected.iter()) {
117            assert!((g - e).abs() < 1e-6, "got {g}, expected {e}");
118        }
119    }
120
121    #[test]
122    fn all_eights_pack_to_zero_weights() {
123        let device = Default::default();
124        // Nibble 8 everywhere -> (8 - 8) * scale == 0 regardless of scale.
125        let packed = Tensor::<B, 2, Int>::from_data(
126            TensorData::new(vec![0x88i32; 16], [1, 16]),
127            &device,
128        );
129        let scales = Tensor::<B, 2>::from_data(TensorData::new(vec![3.5f32], [1, 1]), &device);
130        let w = dequantize_q4(packed, scales, DEFAULT_Q4_GROUP_SIZE);
131        let got: Vec<f32> = w.into_data().to_vec().unwrap();
132        assert_eq!(got, vec![0.0; 32]);
133    }
134}
combs_core/quant.rs

combs_core/
quant.rs
