burn_cubecl/kernel/quantization/
dequantize.rs

1use crate::ops::numeric::empty_device_dtype;
2use crate::tensor::CubeTensor;
3use alloc::{vec, vec::Vec};
4use burn_backend::cubecl::dtype_to_storage_type;
5use burn_backend::{DType, TensorMetadata};
6
7/// Convert the tensor back to a higher precision data type.
8///
9/// # Panics
10///
11/// A storage-tiled tensor: the kernel reads its values as rows.
12pub fn dequantize(tensor: CubeTensor, dtype: DType) -> CubeTensor {
13    let scheme = match tensor.dtype {
14        DType::QFloat(scheme) => scheme,
15        _ => return tensor,
16    };
17    assert!(
18        !tensor.meta.is_tiled(),
19        "dequantize: a storage-tiled quantized tensor is read only by the kernel it was tiled for"
20    );
21    let (tensor, inverse_axes) = match scheme.store {
22        cubecl::quant::scheme::QuantStore::PackedU32(dim)
23        | cubecl::quant::scheme::QuantStore::PackedNative(dim)
24            if dim != 0 =>
25        {
26            let rank = tensor.rank();
27            let packed_axis = rank - dim - 1;
28            let mut axes = (0..rank)
29                .filter(|axis| *axis != packed_axis)
30                .collect::<Vec<_>>();
31            axes.push(packed_axis);
32
33            let mut inverse_axes = vec![0; rank];
34            for (axis, source_axis) in axes.iter().enumerate() {
35                inverse_axes[*source_axis] = axis;
36            }
37
38            let tensor = (packed_axis..rank - 1).fold(tensor, |tensor, axis| {
39                crate::ops::swap_dims(tensor, axis, axis + 1)
40            });
41
42            (tensor, Some(inverse_axes))
43        }
44        _ => (tensor, None),
45    };
46    let scheme = tensor.scheme();
47
48    let output = empty_device_dtype(
49        tensor.client.clone(),
50        tensor.device.clone(),
51        tensor.shape(),
52        dtype,
53    );
54    let global = tensor.global();
55    let (values, params) = tensor.quantized_handles().unwrap();
56
57    // Innermost first: the block scales, then the per-tensor scale they are normalized against.
58    let mut scales = vec![params.binding()];
59    scales.extend(global.map(|g| g.binding()));
60
61    cubek::quantization::dequantize::launch_ref(
62        &output.client,
63        values.binding(),
64        output.clone().binding(),
65        &scales,
66        &scheme,
67        dtype_to_storage_type(dtype),
68    )
69    .expect("Kernel to never fail");
70
71    match inverse_axes {
72        Some(axes) => crate::ops::permute(output, &axes),
73        None => output,
74    }
75}
burn_cubecl/kernel/quantization/dequantize.rs

burn_cubecl/kernel/quantization/
dequantize.rs
