cubecl_std/quant/
dequantize.rs

1use cubecl::prelude::*;
2use cubecl_common::quant::scheme::*;
3use cubecl_common::{e2m1x2, e4m3, e5m2};
4use cubecl_core as cubecl;
5
6/// Dequantize a vector of values, where `vector_size * num_quants` is a power of two.
7/// Unaligned values can't be dequantized in place.
8#[cube]
9pub fn dequantize_aligned<Q: Scalar, S: CubePrimitive, F: Numeric, NQ: Size, NF: Size>(
10    value: Vector<Q, NQ>,
11    scale: S,
12    #[comptime] scheme: QuantScheme,
13) -> Vector<F, NF> {
14    // Every read from a quantized view lands here, so this is where an unsupported level has to be
15    // caught: the static constructors take a scheme without inspecting it.
16    comptime!(crate::quant::assert_level_supported(scheme.level));
17
18    let q_values = match scheme.store {
19        QuantStore::Native | QuantStore::PackedNative(_) => Vector::<F, NF>::cast_from(value),
20        QuantStore::PackedU32(_) => unpack_cast_u32::<F, NQ, NF>(Vector::cast_from(value), scheme),
21    };
22    let scale = Vector::<F, NF>::cast_from(scale);
23
24    match scheme.mode {
25        QuantMode::Symmetric => q_values * scale,
26    }
27}
28
29/// Unpack a set of values from u32, and convert to the specified floating point format.
30#[cube]
31pub fn unpack_cast_u32<F: Numeric, NQ: Size, NF: Size>(
32    value: Vector<u32, NQ>,
33    #[comptime] scheme: QuantScheme,
34) -> Vector<F, NF> {
35    let num_quants = scheme.num_quants();
36    let native_packing = scheme.native_packing();
37    let size_bits = scheme.size_bits_value();
38    let mask = comptime![packing_mask(scheme)];
39    let size!(NP) = native_packing;
40
41    let mut out = Vector::<F, NF>::empty();
42
43    #[unroll]
44    for vector_idx in 0..value.size() {
45        let packed_val = value.extract(vector_idx);
46        let out_offset = vector_idx * num_quants;
47        #[unroll]
48        for packed_idx in range_stepped(0, num_quants, native_packing) {
49            let shift = packed_idx * size_bits;
50            let value = (packed_val >> shift as u32) & mask;
51
52            let float_value = cast_masked::<F, NP>(value, scheme);
53
54            #[unroll]
55            for native_idx in 0..native_packing {
56                let out_offset = out_offset + packed_idx + native_idx;
57                out.insert(out_offset, float_value.extract(native_idx));
58            }
59        }
60    }
61
62    out
63}
64
65/// The mask required for each packed value, taking into account the native packing required for
66/// `e2m1`.
67fn packing_mask(scheme: QuantScheme) -> u32 {
68    let bits = match scheme.value {
69        QuantValue::E2M1 => 8, // Packed conversion
70        other => other.size_bits(),
71    };
72    (1u32 << bits) - 1
73}
74
75/// Cast a masked-out value in the low `n` bits of a `u32` to the specified float type.
76/// Applies sign conversion for integer quantization before casting to the float type,
77/// while minifloats are simply truncated to `u8`, reinterpreted and then cast.
78/// For `e2m1`, casting is done on the packed `e2m1x2` representation.
79///
80/// # Returns
81/// Two floating point numbers for `e2m1`, one for all other formats.
82#[cube]
83fn cast_masked<F: Numeric, N: Size>(value: u32, #[comptime] scheme: QuantScheme) -> Vector<F, N> {
84    match scheme.value {
85        // For minifloat we can assume if they're supported then u8 is supported
86        QuantValue::E5M2 => Vector::<F, N>::cast_from(e5m2::from_bits(value as u8)),
87        QuantValue::E4M3 => Vector::<F, N>::cast_from(e4m3::from_bits(value as u8)),
88        QuantValue::E2M1 => Vector::<F, N>::cast_from(e2m1x2::from_bits(value as u8)),
89        QuantValue::Q8F
90        | QuantValue::Q4F
91        | QuantValue::Q2F
92        | QuantValue::Q8S
93        | QuantValue::Q4S
94        | QuantValue::Q2S => {
95            let size_quant = scheme.size_bits_value() as u32;
96            let sign_bit = 1u32 << (size_quant - 1);
97            let two_pow_n = 1 << size_quant;
98
99            // Branchless two's complement conversion
100            // If raw >= 2^(n-1), then result = raw - 2^n
101            let raw_i32 = value as i32;
102            let is_negative = (value >= sign_bit) as i32; // 1 if negative, 0 if positive
103            let signed_value = raw_i32 - (is_negative * two_pow_n);
104            Vector::<F, N>::cast_from(signed_value)
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use cubecl_core::define_size;
113    use cubecl_core::ir::{ElemType, Scope, UIntKind};
114
115    define_size!(N1);
116
117    /// Expanding is where an unsupported level has to be caught: the static constructors take a
118    /// scheme without inspecting it, so a guard on the dynamic dispatcher alone leaves them open.
119    #[test]
120    #[should_panic(expected = "two-level quantization is not supported")]
121    fn expanding_a_two_level_scheme_panics() {
122        // A root scope carries no typemap; the launcher normally picks the index width.
123        let scope = Scope::root(false);
124        scope.register_size::<N1>(1);
125        scope.register_type::<usize>(ElemType::UInt(UIntKind::U32).into());
126
127        let one = f32::__expand_new(&scope, 1.0);
128        let value = Vector::<f32, N1>::__expand_new(&scope, one);
129        let scheme =
130            QuantScheme::default().with_level(QuantLevel::block_tensor([32], QuantParam::F32));
131
132        dequantize_aligned::expand::<f32, f32, f32, N1, N1>(&scope, value, one, scheme);
133    }
134}
cubecl_std/quant/dequantize.rs

cubecl_std/quant/
dequantize.rs
