cubecl_std/quant/
view.rs

1use std::marker::PhantomData;
2
3use super::*;
4use crate::tensor::{
5    View, ViewExpand, ViewOperations, ViewOperationsExpand,
6    launch::{ViewArg, ViewCompilationArg},
7    layout::Coordinates,
8};
9use cubecl::prelude::*;
10use cubecl_common::{
11    e2m1x2, e4m3, e5m2,
12    quant::scheme::{QuantParam, QuantScheme, QuantStore, QuantValue},
13    ue8m0,
14};
15use cubecl_core::{
16    self as cubecl, define_size,
17    ir::{ElemType, FloatKind, StorageType, VectorSize},
18    prelude::barrier::Barrier,
19    unexpanded,
20};
21use half::{bf16, f16};
22
23/// View that dequantizes after loads. Scales layout should take values coordinates and map them
24/// to the corresponding scale.
25///
26/// # Warning
27/// Assumes only one scale maps to a single load. Adjust vector size of values or block size to ensure
28/// this.
29/// Must ensure `block_size.is_multiple_of(vector_size * scheme.num_quants())`.
30#[expect(dead_code, reason = "only used in expand")]
31#[derive(CubeType, CubeLaunch, Clone)]
32pub struct QuantizedView<
33    'a,
34    Q: Scalar,
35    NQ: Size,
36    S: Scalar,
37    F: Numeric,
38    NF: Size,
39    C: Coordinates + 'static,
40> {
41    values: View<'a, Vector<Q, NQ>, C>,
42    scales: View<'a, S, C>,
43    #[cube(comptime)]
44    scheme: QuantScheme,
45    #[cube(comptime)]
46    _ty: PhantomData<(F, NF)>,
47}
48
49#[cube]
50impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
51    QuantizedView<'a, Q, NQ, S, F, NF, C>
52{
53    pub fn new(
54        values: View<'a, Vector<Q, NQ>, C>,
55        scales: View<'a, S, C>,
56        #[comptime] scheme: QuantScheme,
57    ) -> Self {
58        QuantizedView::<'a, Q, NQ, S, F, NF, C> {
59            values,
60            scales,
61            scheme,
62            _ty: PhantomData,
63        }
64    }
65}
66
67impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
68    QuantizedView<'a, Q, NQ, S, F, NF, C>
69{
70    pub fn view(self) -> View<'a, Vector<F, NF>, C> {
71        unexpanded!()
72    }
73
74    pub fn __expand_view(
75        scope: &Scope,
76        this: QuantizedViewExpand<'a, Q, NQ, S, F, NF, C>,
77    ) -> ViewExpand<'a, Vector<F, NF>, C> {
78        this.__expand_view_method(scope)
79    }
80}
81
82impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
83    QuantizedViewExpand<'a, Q, NQ, S, F, NF, C>
84{
85    pub fn new(
86        values: ViewExpand<'a, Vector<Q, NQ>, C>,
87        scales: ViewExpand<'a, S, C>,
88        scheme: QuantScheme,
89    ) -> Self {
90        QuantizedViewExpand::<'a, Q, NQ, S, F, NF, C> {
91            values,
92            scales,
93            scheme,
94            _ty: PhantomData,
95        }
96    }
97
98    pub fn __expand_view_method(self, scope: &Scope) -> ViewExpand<'a, Vector<F, NF>, C> {
99        ViewExpand::new(scope, self)
100    }
101}
102
103impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static> Vectorized
104    for QuantizedView<'a, Q, NQ, S, F, NF, C>
105{
106}
107impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
108    VectorizedExpand for QuantizedViewExpand<'a, Q, NQ, S, F, NF, C>
109{
110    fn vector_size(&self) -> VectorSize {
111        self.values.vector_size() * self.scheme.num_quants()
112    }
113}
114
115impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
116    ViewOperations<Vector<F, NF>, C> for QuantizedView<'a, Q, NQ, S, F, NF, C>
117{
118}
119
120impl<'a, Q: Scalar, NQ: Size, S: Scalar, F: Numeric, NF: Size, C: Coordinates + 'static>
121    ViewOperationsExpand<Vector<F, NF>, C> for QuantizedViewExpand<'a, Q, NQ, S, F, NF, C>
122{
123    fn __expand_read_method(
124        &self,
125        scope: &Scope,
126        pos: <C>::ExpandType,
127    ) -> NativeExpand<Vector<F, NF>> {
128        let value = self.values.clone().__expand_read_method(scope, pos.clone());
129        let scale = self.scales.clone().__expand_read_method(scope, pos);
130
131        dequantize_aligned::expand::<Q, S, F, NQ, NF>(scope, value, scale, self.scheme)
132    }
133
134    fn __expand_read_checked_method(
135        &self,
136        scope: &Scope,
137        pos: <C>::ExpandType,
138    ) -> NativeExpand<Vector<F, NF>> {
139        let value = self
140            .values
141            .clone()
142            .__expand_read_checked_method(scope, pos.clone());
143        let scale = self
144            .scales
145            .clone()
146            .__expand_read_checked_method(scope, pos.clone());
147
148        dequantize_aligned::expand::<Q, S, F, NQ, NF>(scope, value, scale, self.scheme)
149    }
150
151    fn __expand_read_masked_method(
152        &self,
153        scope: &Scope,
154        pos: <C>::ExpandType,
155        mask_value: NativeExpand<Vector<F, NF>>,
156    ) -> NativeExpand<Vector<F, NF>> {
157        let value = self
158            .values
159            .clone()
160            .__expand_read_checked_method(scope, pos.clone());
161        let scale = self
162            .scales
163            .clone()
164            .__expand_read_checked_method(scope, pos.clone());
165        let in_bounds = self.__expand_is_in_bounds_method(scope, pos);
166
167        let value = dequantize_aligned::expand::<Q, S, F, NQ, NF>(scope, value, scale, self.scheme);
168        select::expand::<Vector<F, NF>>(scope, in_bounds, value, mask_value)
169    }
170
171    fn __expand_read_unchecked_method(
172        &self,
173        scope: &Scope,
174        pos: <C>::ExpandType,
175    ) -> NativeExpand<Vector<F, NF>> {
176        let value = self
177            .values
178            .clone()
179            .__expand_read_unchecked_method(scope, pos.clone());
180        let scale = self
181            .scales
182            .clone()
183            .__expand_read_unchecked_method(scope, pos);
184
185        dequantize_aligned::expand::<Q, S, F, NQ, NF>(scope, value, scale, self.scheme)
186    }
187
188    fn __expand_as_linear_slice_method(
189        &self,
190        _scope: &Scope,
191        _pos: <C>::ExpandType,
192        _end: <C>::ExpandType,
193    ) -> &SliceExpand<Vector<F, NF>> {
194        panic!("Can't create raw slice for quantized view")
195    }
196
197    fn __expand_shape_method(&self, scope: &Scope) -> <C>::ExpandType {
198        self.values.clone().__expand_shape_method(scope)
199    }
200
201    fn __expand_is_in_bounds_method(
202        &self,
203        scope: &Scope,
204        pos: C::ExpandType,
205    ) -> NativeExpand<bool> {
206        self.values.clone().__expand_is_in_bounds_method(scope, pos)
207    }
208
209    fn __expand_tensor_map_load_method(
210        &self,
211        _scope: &Scope,
212        _barrier: &NativeExpand<Barrier>,
213        _shared_memory: &mut SliceExpand<Vector<F, NF>>,
214        _pos: C::ExpandType,
215    ) {
216        panic!("Can't use tensor map functions on quantized view");
217    }
218}
219
220/// Storage (values) vector size: the float vector size divided by `num_quants`. Asserts the float
221/// vector size is a multiple of `num_quants`, so a violation reports clearly here, not a cryptic cast error.
222fn quant_vector_size_q(vector_size: usize, num_quants: usize) -> usize {
223    assert!(
224        vector_size >= num_quants && vector_size.is_multiple_of(num_quants),
225        "quantized view float vector size {vector_size} must be a positive multiple of num_quants {num_quants}"
226    );
227    vector_size / num_quants
228}
229
230struct ExpandDynamic<'a, E: Numeric, N: Size, C: Coordinates + 'static> {
231    values: &'a ViewCompilationArg<C>,
232    scales: &'a ViewCompilationArg<C>,
233    scheme: QuantScheme,
234    builder: &'a mut KernelBuilder,
235    _ty: PhantomData<(E, N)>,
236}
237
238impl<'a, E: Numeric, N: Size, C: Coordinates + 'static> RunWithQuantType
239    for ExpandDynamic<'a, E, N, C>
240{
241    type Output = ViewExpand<'static, Vector<E, N>, C>;
242
243    fn execute<Q: Scalar, S: Scalar>(self) -> Self::Output {
244        define_size!(NQ);
245
246        let vector_size = N::__expand_value(&self.builder.scope);
247        let vector_size_q = quant_vector_size_q(vector_size, self.scheme.num_quants());
248        self.builder.scope.register_size::<NQ>(vector_size_q);
249
250        let values = View::<Vector<Q, NQ>, C>::expand(self.values, self.builder);
251        let scales = View::<S, C>::expand(self.scales, self.builder);
252        let view = QuantizedViewExpand::new(values, scales, self.scheme);
253        ViewExpand::new(&self.builder.scope, view)
254    }
255}
256
257pub(crate) struct RegisterDynamic<'a, E: CubePrimitive, C: Coordinates + 'static, R: Runtime> {
258    pub values: ViewArg<C, R>,
259    pub scales: ViewArg<C, R>,
260    pub scheme: QuantScheme,
261    pub launcher: &'a mut KernelLauncher<R>,
262    pub _ty: PhantomData<E>,
263}
264
265impl<'a, E: CubePrimitive, C: Coordinates + 'static, R: Runtime> RunWithQuantType
266    for RegisterDynamic<'a, E, C, R>
267{
268    type Output = ViewCompilationArg<C>;
269
270    fn execute<Q: Scalar, S: Scalar>(self) -> Self::Output {
271        define_size!(NQ);
272
273        self.launcher.with_scope(|scope| {
274            let vector_size_q =
275                quant_vector_size_q(E::__expand_vector_size(scope), self.scheme.num_quants());
276            scope.register_size::<NQ>(vector_size_q);
277        });
278
279        let values = View::<Vector<Q, NQ>, C>::register(self.values, self.launcher);
280        let scales = View::<S, C>::register(self.scales, self.launcher);
281        ViewCompilationArg::Quantized {
282            values: Box::new(values),
283            scales: Box::new(scales),
284            scheme: self.scheme,
285        }
286    }
287}
288
289/// Run a function with the quantization storage type and scale. Useful when concrete types are
290/// required but aren't available, and only the dynamic schema is known.
291pub fn run_with_quant_type<F: RunWithQuantType>(func: F, scheme: QuantScheme) -> F::Output {
292    fn run_with_q<F: RunWithQuantType, Q: Scalar>(func: F, scheme: QuantScheme) -> F::Output {
293        match scheme.param {
294            QuantParam::F32 => func.execute::<Q, f32>(),
295            QuantParam::F16 => func.execute::<Q, f16>(),
296            QuantParam::BF16 => func.execute::<Q, bf16>(),
297            QuantParam::UE8M0 => func.execute::<Q, ue8m0>(),
298            QuantParam::UE4M3 => func.execute::<Q, e4m3>(),
299        }
300    }
301
302    let run_q = match scheme.store {
303        QuantStore::Native => match scheme.value {
304            QuantValue::Q8F => run_with_q::<F, i8>,
305            QuantValue::Q8S => run_with_q::<F, i8>,
306            QuantValue::E5M2 => run_with_q::<F, e5m2>,
307            QuantValue::E4M3 => run_with_q::<F, e4m3>,
308            QuantValue::Q4F
309            | QuantValue::Q4S
310            | QuantValue::Q2F
311            | QuantValue::Q2S
312            | QuantValue::E2M1 => {
313                panic!("Sub-byte quantization can't be native")
314            }
315        },
316        QuantStore::PackedU32(_) => run_with_q::<F, u32>,
317        QuantStore::PackedNative(_) => run_with_q::<F, e2m1x2>,
318    };
319    run_q(func, scheme)
320}
321
322/// Dynamically expand based on the quantization scheme. Ugly, but the only way to fully hide the
323/// quantization from the kernel using the view.
324pub(crate) fn expand_dynamic<E: CubePrimitive, C: Coordinates + 'static>(
325    values: &ViewCompilationArg<C>,
326    scales: &ViewCompilationArg<C>,
327    scheme: QuantScheme,
328    builder: &mut KernelBuilder,
329) -> ViewExpand<'static, E, C> {
330    use core::mem::transmute as t;
331
332    // To specify tighter trait bounds
333    fn expand_dynamic_f<F: Numeric, NF: Size, C: Coordinates + 'static>(
334        values: &ViewCompilationArg<C>,
335        scales: &ViewCompilationArg<C>,
336        scheme: QuantScheme,
337        builder: &mut KernelBuilder,
338    ) -> ViewExpand<'static, Vector<F, NF>, C> {
339        let func = ExpandDynamic {
340            values,
341            scales,
342            scheme,
343            builder,
344            _ty: PhantomData::<(F, NF)>,
345        };
346        run_with_quant_type(func, scheme)
347    }
348
349    define_size!(NF);
350
351    let vector_size = E::__expand_vector_size(&builder.scope);
352
353    builder.scope.register_size::<NF>(vector_size);
354
355    #[allow(clippy::missing_transmute_annotations)]
356    unsafe {
357        match E::__expand_as_type(&builder.scope).storage_type() {
358            StorageType::Scalar(ElemType::Float(ty)) => match ty {
359                FloatKind::F16 => t(expand_dynamic_f::<f16, NF, C>(
360                    values, scales, scheme, builder,
361                )),
362                FloatKind::BF16 => t(expand_dynamic_f::<bf16, NF, C>(
363                    values, scales, scheme, builder,
364                )),
365                FloatKind::Flex32 => t(expand_dynamic_f::<flex32, NF, C>(
366                    values, scales, scheme, builder,
367                )),
368                FloatKind::F32 => t(expand_dynamic_f::<f32, NF, C>(
369                    values, scales, scheme, builder,
370                )),
371                FloatKind::TF32 => t(expand_dynamic_f::<tf32, NF, C>(
372                    values, scales, scheme, builder,
373                )),
374                FloatKind::F64 => t(expand_dynamic_f::<f64, NF, C>(
375                    values, scales, scheme, builder,
376                )),
377                FloatKind::E2M1
378                | FloatKind::E2M3
379                | FloatKind::E3M2
380                | FloatKind::E4M3
381                | FloatKind::E5M2
382                | FloatKind::UE8M0 => unreachable!("Minifloats don't implement `Float` ops"),
383            },
384            _ => unreachable!("Quantized view should only be used with floats"),
385        }
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::quant_vector_size_q;
392
393    #[test]
394    fn vector_size_q_exact_multiple() {
395        assert_eq!(quant_vector_size_q(8, 8), 1);
396        assert_eq!(quant_vector_size_q(16, 8), 2);
397        assert_eq!(quant_vector_size_q(16, 16), 1);
398    }
399
400    #[test]
401    #[should_panic(expected = "positive multiple of num_quants")]
402    fn vector_size_q_non_multiple_panics() {
403        let _ = quant_vector_size_q(8, 16);
404    }
405}
cubecl_std/quant/view.rs

cubecl_std/quant/
view.rs
