Dequantize a packed 4-bit weight tensor with per-group scales and biases. Used for embeddings in some GPTQ models where the format is: - `weight`: uint32, shape `(rows, packed_cols)` where `packed_cols = cols / 8` - `scales`: f16/bf16, shape `(rows, groups)` - `biases`: f16/bf16, shape `(rows, groups)` Formula: `w_dequant[i, j] = w4(i, j) * scale(i, group(j)) + bias(i, group(j))` Output: F32 te
(
packed: &Tensor,
scales: &Tensor,
biases: &Tensor,
group_size: usize,
)
| 149 | /// |
| 150 | /// Output: F32 tensor of shape `(rows, cols)`. |
| 151 | pub fn dequantize_packed_4bit( |
| 152 | packed: &Tensor, |
| 153 | scales: &Tensor, |
| 154 | biases: &Tensor, |
| 155 | group_size: usize, |
| 156 | ) -> candle_core::Result<Tensor> { |
| 157 | // Handle 3D stacked tensors (e.g., [num_experts, rows, packed_cols]) |
| 158 | if packed.rank() == 3 { |
| 159 | let n = packed.dim(0)?; |
| 160 | let slices: Vec<Tensor> = (0..n) |
| 161 | .map(|i| { |
| 162 | let p = packed.get(i)?; |
| 163 | let s = scales.get(i)?; |
| 164 | let b = biases.get(i)?; |
| 165 | dequantize_packed_4bit(&p, &s, &b, group_size) |
| 166 | }) |
| 167 | .collect::<candle_core::Result<_>>()?; |
| 168 | return Tensor::stack(&slices, 0); |
| 169 | } |
| 170 | let (rows, packed_cols) = packed.dims2()?; |
| 171 | let cols = packed_cols * 8; |
| 172 | let (_, groups) = scales.dims2()?; |
| 173 | |
| 174 | // Extract raw data — avoid Tensor intermediates for the hot path |
| 175 | let pw: Vec<u32> = packed.flatten_all()?.to_vec1::<u32>()?; |
| 176 | let sc: Vec<f32> = scales.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?; |
| 177 | let bi: Vec<f32> = biases.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?; |
| 178 | |
| 179 | use rayon::prelude::*; |
| 180 | let mut weight = vec![0f32; rows * cols]; |
| 181 | weight |
| 182 | .par_chunks_mut(cols) |
| 183 | .enumerate() |
| 184 | .for_each(|(i, row)| { |
| 185 | for pc in 0..packed_cols { |
| 186 | let packed_val = pw[i * packed_cols + pc]; |
| 187 | for bit in 0..8u32 { |
| 188 | let j = pc * 8 + bit as usize; |
| 189 | let w4 = ((packed_val >> (bit * 4)) & 0xF) as f32; |
| 190 | let g = j / group_size; |
| 191 | let scale = sc[i * groups + g]; |
| 192 | let bias = bi[i * groups + g]; |
| 193 | row[j] = w4 * scale + bias; |
| 194 | } |
| 195 | } |
| 196 | }); |
| 197 | |
| 198 | Tensor::from_vec(weight, (rows, cols), &Device::Cpu) |
| 199 | } |
| 200 | |
| 201 | /// Custom VarBuilder backend that transparently dequantizes GPTQ weights. |
| 202 | /// |
no test coverage detected