burn_cubecl/kernel/quantization/
quantize.rs

1use crate::{ops::empty_qtensor_optimized, tensor::CubeTensor};
2use burn_backend::cubecl::dtype_to_elem_type;
3use burn_backend::{TensorMetadata, quantization::QuantScheme};
4
5/// Convert the tensor to a lower precision data type based on the quantization scheme and parameters.
6pub fn quantize(
7    tensor: CubeTensor,
8    scheme: &QuantScheme,
9    scale: CubeTensor,
10    global: Option<CubeTensor>,
11) -> CubeTensor {
12    let output = empty_qtensor_optimized(tensor.shape(), *scheme, &tensor.device);
13    let (out_values, out_params) = output.clone().quantized_handles().unwrap();
14    let out_global = output.global();
15    let dtype = tensor.dtype;
16
17    // Innermost first: the block scales, then the per-tensor scale they are normalized against.
18    let mut scales = vec![scale.binding()];
19    scales.extend(global.map(|global| global.binding()));
20
21    let mut out_scales = vec![out_params.binding()];
22    out_scales.extend(out_global.map(|global| global.binding()));
23
24    cubek::quantization::quantize::launch_ref(
25        &output.client,
26        tensor.binding(),
27        out_values.binding(),
28        &scales,
29        &out_scales,
30        scheme,
31        dtype_to_elem_type(dtype),
32    )
33    .expect("Kernel to never fail");
34
35    output
36}
burn_cubecl/kernel/quantization/quantize.rs

burn_cubecl/kernel/quantization/
quantize.rs
