burn_cubecl/tensor/
quantization.rs

1use crate::CubeDevice;
2use burn_backend::{DType, Shape, TensorMetadata as _, quantization::QParamTensor};
3use burn_std::{Metadata, Strides};
4use cubecl::quant::scheme::{QuantStore, QuantValue};
5use cubecl::{client::Client, server::Handle};
6
7use super::CubeTensor;
8
9/// Runtime parameters for quantization. Can be used to construct a scales handle from the base
10/// tensor handle.
11pub type QParams = burn_backend::quantization::QParams<QParamTensor>;
12
13impl CubeTensor {
14    /// Create a new quantized tensor
15    pub fn new_quantized(
16        client: Client,
17        handle: Handle,
18        shape: Shape,
19        device: CubeDevice,
20        strides: Strides,
21        dtype: DType,
22        qparams: QParams,
23    ) -> Self {
24        CubeTensor {
25            client,
26            handle,
27            meta: Box::new(Metadata::new(shape, strides)),
28            device,
29            dtype,
30            qparams: Some(qparams),
31        }
32    }
33
34    /// Returns the two tensors: (values, params) for a quantized tensor.
35    /// For the values, native types that aren't supported as a normal `DType` will be returned
36    /// as an unsigned integer tensor representing the bits. Should be reconstructed using `from_bits`
37    /// in kernels.
38    ///
39    /// A storage-tiled tensor's values keep its metadata whole: its physical dims already state
40    /// the packing, and only a kernel that reads tiles may bind them — any other refuses them
41    /// as rows. The scales keep their own metadata, tiled or not.
42    pub fn quantized_handles(&self) -> Option<(CubeTensor, CubeTensor)> {
43        let params = self.scales()?;
44        let scheme = match self.dtype {
45            DType::QFloat(sc) => sc,
46            _ => return None,
47        };
48        let values = match scheme.store {
49            QuantStore::Native => match scheme.value {
50                QuantValue::Q8F | QuantValue::Q8S => CubeTensor {
51                    client: self.client.clone(),
52                    handle: self.handle.clone(),
53                    meta: self.meta.clone(),
54                    device: self.device.clone(),
55                    dtype: DType::I8,
56                    qparams: None,
57                },
58                QuantValue::E4M3 | QuantValue::E5M2 => CubeTensor {
59                    client: self.client.clone(),
60                    handle: self.handle.clone(),
61                    meta: self.meta.clone(),
62                    device: self.device.clone(),
63                    dtype: DType::U8,
64                    qparams: None,
65                },
66                QuantValue::Q4F
67                | QuantValue::Q4S
68                | QuantValue::Q2F
69                | QuantValue::Q2S
70                | QuantValue::E2M1 => {
71                    panic!("Can't store native sub-byte values")
72                }
73            },
74            QuantStore::PackedU32(_) if self.meta.is_tiled() => self.stored_values(DType::U32),
75            QuantStore::PackedNative(_) if self.meta.is_tiled() => self.stored_values(DType::U8),
76            QuantStore::PackedU32(packed_dim) => {
77                let packed_dim = self.rank() - packed_dim - 1;
78                let mut shape = self.shape();
79                shape[packed_dim] = shape[packed_dim].div_ceil(scheme.num_quants());
80
81                CubeTensor {
82                    client: self.client.clone(),
83                    handle: self.handle.clone(),
84                    meta: Box::new(Metadata::new(shape, self.meta.strides.clone())),
85                    device: self.device.clone(),
86                    dtype: DType::U32,
87                    qparams: None,
88                }
89            }
90            QuantStore::PackedNative(packed_dim) => match scheme.value {
91                QuantValue::E2M1 => {
92                    let packed_dim = self.rank() - packed_dim - 1;
93                    let mut shape = self.shape();
94                    shape[packed_dim] = shape[packed_dim].div_ceil(scheme.num_quants());
95
96                    CubeTensor {
97                        client: self.client.clone(),
98                        handle: self.handle.clone(),
99                        meta: Box::new(Metadata::new(shape, self.meta.strides.clone())),
100                        device: self.device.clone(),
101                        dtype: DType::U8,
102                        qparams: None,
103                    }
104                }
105                other => panic!("{other:?} doesn't support native packing"),
106            },
107        };
108
109        Some((values, params))
110    }
111
112    /// The packed values of a storage-tiled tensor, as stored.
113    fn stored_values(&self, dtype: DType) -> CubeTensor {
114        CubeTensor {
115            client: self.client.clone(),
116            handle: self.handle.clone(),
117            meta: self.meta.clone(),
118            device: self.device.clone(),
119            dtype,
120            qparams: None,
121        }
122    }
123
124    /// Construct a separate tensor for the quantization scales, if present
125    pub fn scales(&self) -> Option<CubeTensor> {
126        self.param_tensor(|qparams| Some(&qparams.scales))
127    }
128
129    /// Construct a separate tensor for the per-tensor scale, for a two-level scheme.
130    pub fn global(&self) -> Option<CubeTensor> {
131        self.param_tensor(|qparams| qparams.global.as_ref())
132    }
133
134    fn param_tensor(
135        &self,
136        select: impl Fn(&QParams) -> Option<&QParamTensor>,
137    ) -> Option<CubeTensor> {
138        let param = select(self.qparams.as_ref()?)?;
139        let mut handle = self.handle.clone();
140        handle.offset_start = Some(param.offset_start as u64);
141        handle.offset_end = Some(param.offset_end as u64);
142
143        Some(CubeTensor::new(
144            self.client.clone(),
145            handle,
146            param.metadata.clone(),
147            self.device.clone(),
148            param.dtype,
149        ))
150    }
151}
152
153#[cfg(all(
154    test,
155    any(feature = "wgpu", feature = "cpu", feature = "cuda", feature = "hip")
156))]
157mod storage_tiled {
158    use burn_backend::{
159        DType, TensorMetadata,
160        ops::QTensorOps,
161        quantization::{QuantScheme, QuantStore, QuantValue, ScaleDtype},
162    };
163    use burn_std::{FloatDType, Metadata, TensorData};
164    use cubecl::zspace::Tiling;
165
166    use crate::{CubeBackend, CubeDevice, kernel::untile, ops::from_data, tensor::CubeTensor};
167
168    /// A `[64, 128]` weight quantized four values to a word, and the same tensor stated as
169    /// stored in 32x32 tiles: what a load leaves when it lays a weight out for a kernel that
170    /// reads tiles. Only the metadata is restated; no test here reads the tiled bytes.
171    fn quantized_and_tiled() -> (CubeTensor, CubeTensor) {
172        let device = CubeDevice::default();
173        let values = (0..64 * 128)
174            .map(|i| i as f32 / 8192.0 - 0.5)
175            .collect::<Vec<_>>();
176        let scheme = QuantScheme::default()
177            .with_value(QuantValue::Q8S)
178            .with_store(QuantStore::PackedU32(0))
179            .per_block([32], ScaleDtype::F32);
180        let quantized = CubeBackend::quantize_dynamic(
181            from_data(TensorData::new(values, [64, 128]), &device),
182            &scheme,
183        );
184        let stored = Metadata::new([2, 4, 32, 32], [4 * 32 * 32, 32 * 32, 32, 1])
185            .with_tiling(Tiling::new(&[2, 2]).expect("two fragments of each dim"))
186            .expect("the tiling describes the four physical dims");
187        let mut tiled = quantized.clone();
188        tiled.meta = Box::new(stored);
189        (quantized, tiled)
190    }
191
192    /// A tiled weight still presents its logical shape, and its values carry the tiles to the
193    /// kernel that reads them, where an untiled weight's values are rows of words as before.
194    #[test]
195    fn a_tiled_weight_hands_its_tiles_to_its_values() {
196        let (quantized, tiled) = quantized_and_tiled();
197        assert_eq!(tiled.shape().as_slice(), &[64, 128]);
198
199        let (values, scales) = tiled.quantized_handles().unwrap();
200        assert_eq!(values.meta, tiled.meta);
201        assert_eq!(values.dtype, DType::U32);
202        assert_eq!(scales.meta, quantized.scales().unwrap().meta);
203
204        let (rows, _) = quantized.quantized_handles().unwrap();
205        assert!(!rows.meta.is_tiled());
206        assert_eq!(rows.meta.shape().as_slice(), &[64, 32]);
207    }
208
209    /// Reading a tiled weight's values as rows would compute garbage without an error, so every
210    /// path that would is refused by name.
211    #[test]
212    #[should_panic(expected = "dequantize: a storage-tiled quantized tensor")]
213    fn a_tiled_weight_is_not_dequantized() {
214        let (_, tiled) = quantized_and_tiled();
215        CubeBackend::dequantize(tiled, FloatDType::F32);
216    }
217
218    #[test]
219    #[should_panic(expected = "untile: a storage-tiled quantized tensor")]
220    fn a_tiled_weight_is_not_untiled() {
221        let (_, tiled) = quantized_and_tiled();
222        untile(tiled);
223    }
224
225    /// A layout rewrite lays a tensor back first, so it is refused with the same reason.
226    #[test]
227    #[should_panic(expected = "untile: a storage-tiled quantized tensor")]
228    fn a_tiled_weight_is_not_transposed() {
229        let (_, tiled) = quantized_and_tiled();
230        CubeBackend::q_swap_dims(tiled, 0, 1);
231    }
232
233    /// Tiles are laid out per machine at load; what is saved is the rows every machine reads.
234    #[test]
235    #[should_panic(expected = "q_into_data: a storage-tiled quantized tensor is not saved")]
236    fn a_tiled_weight_is_not_saved() {
237        let (_, tiled) = quantized_and_tiled();
238        let _ = burn_std::future::block_on(CubeBackend::q_into_data(tiled));
239    }
240}
burn_cubecl/tensor/quantization.rs

burn_cubecl/tensor/
quantization.rs
