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    // Caught again on the dequantization path, but reporting it here names the launch that asked
293    // for it rather than a kernel being expanded.
294    assert_level_supported(scheme.level);
295
296    fn run_with_q<F: RunWithQuantType, Q: Scalar>(func: F, scheme: QuantScheme) -> F::Output {
297        match scheme.param {
298            QuantParam::F32 => func.execute::<Q, f32>(),
299            QuantParam::F16 => func.execute::<Q, f16>(),
300            QuantParam::BF16 => func.execute::<Q, bf16>(),
301            QuantParam::UE8M0 => func.execute::<Q, ue8m0>(),
302            QuantParam::UE4M3 => func.execute::<Q, e4m3>(),
303        }
304    }
305
306    let run_q = match scheme.store {
307        QuantStore::Native => match scheme.value {
308            QuantValue::Q8F => run_with_q::<F, i8>,
309            QuantValue::Q8S => run_with_q::<F, i8>,
310            QuantValue::E5M2 => run_with_q::<F, e5m2>,
311            QuantValue::E4M3 => run_with_q::<F, e4m3>,
312            QuantValue::Q4F
313            | QuantValue::Q4S
314            | QuantValue::Q2F
315            | QuantValue::Q2S
316            | QuantValue::E2M1 => {
317                panic!("Sub-byte quantization can't be native")
318            }
319        },
320        QuantStore::PackedU32(_) => run_with_q::<F, u32>,
321        QuantStore::PackedNative(_) => run_with_q::<F, e2m1x2>,
322    };
323    run_q(func, scheme)
324}
325
326/// Dynamically expand based on the quantization scheme. Ugly, but the only way to fully hide the
327/// quantization from the kernel using the view.
328pub(crate) fn expand_dynamic<E: CubePrimitive, C: Coordinates + 'static>(
329    values: &ViewCompilationArg<C>,
330    scales: &ViewCompilationArg<C>,
331    scheme: QuantScheme,
332    builder: &mut KernelBuilder,
333) -> ViewExpand<'static, E, C> {
334    use core::mem::transmute as t;
335
336    // To specify tighter trait bounds
337    fn expand_dynamic_f<F: Numeric, NF: Size, C: Coordinates + 'static>(
338        values: &ViewCompilationArg<C>,
339        scales: &ViewCompilationArg<C>,
340        scheme: QuantScheme,
341        builder: &mut KernelBuilder,
342    ) -> ViewExpand<'static, Vector<F, NF>, C> {
343        let func = ExpandDynamic {
344            values,
345            scales,
346            scheme,
347            builder,
348            _ty: PhantomData::<(F, NF)>,
349        };
350        run_with_quant_type(func, scheme)
351    }
352
353    define_size!(NF);
354
355    let vector_size = E::__expand_vector_size(&builder.scope);
356
357    builder.scope.register_size::<NF>(vector_size);
358
359    #[allow(clippy::missing_transmute_annotations)]
360    unsafe {
361        match E::__expand_as_type(&builder.scope).storage_type() {
362            StorageType::Scalar(ElemType::Float(ty)) => match ty {
363                FloatKind::F16 => t(expand_dynamic_f::<f16, NF, C>(
364                    values, scales, scheme, builder,
365                )),
366                FloatKind::BF16 => t(expand_dynamic_f::<bf16, NF, C>(
367                    values, scales, scheme, builder,
368                )),
369                FloatKind::Flex32 => t(expand_dynamic_f::<flex32, NF, C>(
370                    values, scales, scheme, builder,
371                )),
372                FloatKind::F32 => t(expand_dynamic_f::<f32, NF, C>(
373                    values, scales, scheme, builder,
374                )),
375                FloatKind::TF32 => t(expand_dynamic_f::<tf32, NF, C>(
376                    values, scales, scheme, builder,
377                )),
378                FloatKind::F64 => t(expand_dynamic_f::<f64, NF, C>(
379                    values, scales, scheme, builder,
380                )),
381                FloatKind::E2M1
382                | FloatKind::E2M3
383                | FloatKind::E3M2
384                | FloatKind::E4M3
385                | FloatKind::E5M2
386                | FloatKind::UE8M0 => unreachable!("Minifloats don't implement `Float` ops"),
387            },
388            _ => unreachable!("Quantized view should only be used with floats"),
389        }
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use super::{RunWithQuantType, quant_vector_size_q, run_with_quant_type};
396    use cubecl_common::quant::scheme::{QuantLevel, QuantParam, QuantScheme};
397    use cubecl_core::prelude::Scalar;
398
399    struct Dispatched;
400
401    impl RunWithQuantType for Dispatched {
402        type Output = bool;
403
404        fn execute<Q: Scalar, S: Scalar>(self) -> bool {
405            true
406        }
407    }
408
409    #[test]
410    fn one_level_scheme_dispatches() {
411        assert!(run_with_quant_type(Dispatched, QuantScheme::default()));
412    }
413
414    #[test]
415    #[should_panic(expected = "two-level quantization is not supported")]
416    fn two_level_scheme_is_rejected() {
417        let scheme =
418            QuantScheme::default().with_level(QuantLevel::block_tensor([32], QuantParam::F32));
419        // Would otherwise dequantize against the block scales alone, dropping the per-tensor factor.
420        run_with_quant_type(Dispatched, scheme);
421    }
422
423    #[test]
424    fn vector_size_q_exact_multiple() {
425        assert_eq!(quant_vector_size_q(8, 8), 1);
426        assert_eq!(quant_vector_size_q(16, 8), 2);
427        assert_eq!(quant_vector_size_q(16, 16), 1);
428    }
429
430    #[test]
431    #[should_panic(expected = "positive multiple of num_quants")]
432    fn vector_size_q_non_multiple_panics() {
433        let _ = quant_vector_size_q(8, 16);
434    }
435}
cubecl_std/quant/view.rs

cubecl_std/quant/
view.rs
