Skip to main content

ferrotensor_backend/engine/
quant.rs

1//! Dequantisation kernels.
2
3use super::{groups64, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5
6impl ComputeEngine {
7    /// Unpacks bytes to floats: element `i` of `out` is byte `i` of `input`
8    /// (read as packed `u32`s) converted to `f32`, without scaling.
9    pub fn dequantize(&self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor) {
10        let words = input.shape.iter().product::<usize>().div_ceil(4);
11        self.run(
12            ctx,
13            &self.dequant,
14            &[(0, input.buffer()), (1, out.buffer())],
15            groups64(words),
16        );
17    }
18
19    /// 8-bit affine dequantisation with groups of 64 (MLX layout):
20    /// `out[i] = byte(weight, i) * scales[i / 64] + biases[i / 64]`.
21    pub fn affine_dequantize(
22        &self,
23        ctx: &WgpuContext,
24        weight: &wgpu::Buffer,
25        scales: &wgpu::Buffer,
26        biases: &wgpu::Buffer,
27        out: &UnifiedTensor,
28    ) {
29        self.run(
30            ctx,
31            &self.affine_dequant,
32            &[(0, weight), (1, scales), (2, biases), (3, out.buffer())],
33            groups64(out.shape.iter().product()),
34        );
35    }
36}