combs_models/
quant_linear.rs

1//! Quantized linear layer: packed 4-bit weights + per-group scales,
2//! dequantized on-device before the matmul.
3//!
4//! The packed weight stays in VRAM in its compact form between forward
5//! passes (a 4x footprint reduction vs f32 for the weight); the current
6//! implementation dequantizes to f32 and runs the standard matmul (see
7//! `combs_core::quant` docs — a fused dequant-matmul kernel is future
8//! work). Format adapters (GGUF in Phase 5) construct these layers via
9//! [`QuantizedLinear::new`].
10
11use burn::tensor::{Int, Tensor, backend::Backend};
12
13use combs_core::quant::{DEFAULT_Q4_GROUP_SIZE, dequantize_q4};
14
15use crate::matmul::safe_matmul;
16use crate::{ModelError, Result};
17
18/// `y = x @ W^T (+ b)` where `W` is stored group-quantized to 4 bits.
19pub struct QuantizedLinear<B: Backend> {
20    /// `[out_features, in_features / 2]` packed nibbles (GGUF q4_0 order).
21    packed: Tensor<B, 2, Int>,
22    /// `[out_features, in_features / group_size]` per-block scales.
23    scales: Tensor<B, 2>,
24    /// Values per quantization block (GGUF q4_0: 32).
25    group_size: usize,
26    /// Optional `[out_features]` bias.
27    bias: Option<Tensor<B, 1>>,
28    in_features: usize,
29    out_features: usize,
30}
31
32impl<B: Backend> QuantizedLinear<B> {
33    /// Builds a layer from on-device packed parts, validating shapes.
34    pub fn new(
35        packed: Tensor<B, 2, Int>,
36        scales: Tensor<B, 2>,
37        group_size: usize,
38        bias: Option<Tensor<B, 1>>,
39    ) -> Result<Self> {
40        let [out_features, packed_cols] = packed.dims();
41        if group_size == 0 || group_size % 2 != 0 {
42            return Err(ModelError::BadShape {
43                tensor: "quantized_weight".into(),
44                expected: vec![DEFAULT_Q4_GROUP_SIZE],
45                got: vec![group_size],
46            });
47        }
48        if packed_cols % (group_size / 2) != 0 {
49            return Err(ModelError::BadShape {
50                tensor: "quantized_weight".into(),
51                expected: vec![group_size / 2],
52                got: vec![packed_cols],
53            });
54        }
55        let in_features = packed_cols * 2;
56        let [s_rows, s_cols] = scales.dims();
57        if s_rows != out_features || s_cols != in_features / group_size {
58            return Err(ModelError::BadShape {
59                tensor: "quantized_scales".into(),
60                expected: vec![out_features, in_features / group_size],
61                got: vec![s_rows, s_cols],
62            });
63        }
64        if let Some(b) = &bias {
65            if b.dims()[0] != out_features {
66                return Err(ModelError::BadShape {
67                    tensor: "quantized_bias".into(),
68                    expected: vec![out_features],
69                    got: vec![b.dims()[0]],
70                });
71            }
72        }
73        Ok(QuantizedLinear {
74            packed,
75            scales,
76            group_size,
77            bias,
78            in_features,
79            out_features,
80        })
81    }
82
83    /// Input feature count (dequantized `in_features`).
84    pub fn in_features(&self) -> usize {
85        self.in_features
86    }
87
88    /// Output feature count.
89    pub fn out_features(&self) -> usize {
90        self.out_features
91    }
92
93    /// Dequantizes the weight to f32 `[out_features, in_features]`.
94    pub fn weight(&self) -> Tensor<B, 2> {
95        dequantize_q4(self.packed.clone(), self.scales.clone(), self.group_size)
96    }
97
98    /// `y = x @ dequant(W)^T (+ b)` for `x: [batch, seq, in_features]`.
99    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
100        let w = self.weight();
101        let out = safe_matmul(x, w.transpose().unsqueeze_dim::<3>(0));
102        match &self.bias {
103            Some(b) => {
104                let [batch, seq, dim] = out.dims();
105                out + b.clone().reshape([1, 1, dim]).expand([batch, seq, dim])
106            }
107            None => out,
108        }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use burn::tensor::TensorData;
116
117    type B = burn::backend::NdArray<f32>;
118
119    /// Packs f32 weights into GGUF q4_0 form (per-block abs-max scale).
120    fn pack_q4(w: &[f32], rows: usize, cols: usize) -> (Vec<i32>, Vec<f32>) {
121        let mut packed = vec![0i32; rows * cols / 2];
122        let mut scales = vec![0f32; rows * (cols / 32)];
123        for r in 0..rows {
124            for g in 0..cols / 32 {
125                let block = &w[r * cols + g * 32..r * cols + g * 32 + 32];
126                let amax = block.iter().fold(0f32, |m, v| m.max(v.abs()));
127                let scale = amax / 7.0;
128                scales[r * (cols / 32) + g] = scale;
129                for j in 0..16 {
130                    let lo = ((block[j] / scale).round() as i32 + 8).clamp(0, 15) as u32;
131                    let hi = ((block[j + 16] / scale).round() as i32 + 8).clamp(0, 15) as u32;
132                    packed[r * cols / 2 + g * 16 + j] = ((hi << 4) | lo) as i32;
133                }
134            }
135        }
136        (packed, scales)
137    }
138
139    #[test]
140    fn forward_matches_dequantized_dense_matmul() {
141        let device = Default::default();
142        let (rows, cols) = (4, 64); // out=4, in=64
143        let w: Vec<f32> = (0..rows * cols)
144            .map(|i| ((i * 13 % 17) as f32 - 8.0) / 4.0)
145            .collect();
146        let (packed, scales) = pack_q4(&w, rows, cols);
147
148        let layer = QuantizedLinear::<B>::new(
149            Tensor::from_data(TensorData::new(packed, [rows, cols / 2]), &device),
150            Tensor::from_data(TensorData::new(scales, [rows, cols / 32]), &device),
151            DEFAULT_Q4_GROUP_SIZE,
152            None,
153        )
154        .unwrap();
155        assert_eq!(layer.in_features(), cols);
156        assert_eq!(layer.out_features(), rows);
157
158        let x_data: Vec<f32> = (0..cols).map(|i| (i as f32) / 16.0).collect();
159        let x = Tensor::<B, 3>::from_data(TensorData::new(x_data.clone(), [1, 1, cols]), &device);
160
161        let got = layer.forward(x.clone());
162        let dense_w = layer.weight();
163        let expect = crate::matmul::safe_matmul(
164            x,
165            dense_w.transpose().unsqueeze_dim::<3>(0),
166        );
167        let got: Vec<f32> = got.into_data().to_vec().unwrap();
168        let expect: Vec<f32> = expect.into_data().to_vec().unwrap();
169        for (g, e) in got.iter().zip(expect.iter()) {
170            assert!((g - e).abs() < 1e-5);
171        }
172
173        // And the quantized weight must be close to the original (q4 error).
174        let deq: Vec<f32> = layer.weight().into_data().to_vec().unwrap();
175        let max_err = deq
176            .iter()
177            .zip(w.iter())
178            .fold(0f32, |m, (d, o)| m.max((d - o).abs()));
179        assert!(max_err < 0.5, "q4 reconstruction error too large: {max_err}");
180    }
181}
combs_models/quant_linear.rs

combs_models/
quant_linear.rs
