cubecl_std/quant/
base.rs

1use cubecl_common::quant::scheme::{QuantScheme, ScaleDtype};
2use cubecl_core::prelude::Scalar;
3
4/// Run an arbitrary function with the quantization types from the scheme.
5/// Useful when concrete types aren't available.
6pub trait RunWithQuantType {
7    type Output;
8
9    fn execute<Q: Scalar, S: Scalar>(self) -> Self::Output;
10}
11
12/// Panic when the scale bindings and the scheme's levels disagree.
13///
14/// Every level binds a scale buffer of its own, so nothing ties the bindings to the scheme: a
15/// missing level is dropped from the reconstruction and every value comes back short by that
16/// factor, an extra one is a caller quantizing differently than the scheme it passed.
17///
18/// The global level is further constrained by what this reader serves: it binds as f32 rather than
19/// being read as f32 bytes. It has one scale for the whole tensor, so a narrower type saves
20/// nothing and only reintroduces rounding error.
21pub fn check_scale_bindings(scheme: &QuantScheme, bindings: usize) {
22    let levels = scheme.num_levels();
23    assert!(
24        bindings == levels,
25        "a scheme with {levels} scale level(s) takes as many scale bindings, but {bindings} were provided",
26    );
27    check_global_levels(scheme);
28}
29
30/// The global-level half of [`check_scale_bindings`], for a consumer holding global scales already
31/// folded into a register rather than as countable bindings.
32pub fn check_global_levels(scheme: &QuantScheme) {
33    if scheme.block_scale().is_some()
34        && let Some(tensor) = scheme.tensor_scale()
35    {
36        assert!(
37            tensor == ScaleDtype::F32,
38            "an global scale binds as f32, but the scheme stores it as {tensor:?}",
39        );
40    }
41}
42
43/// Panic when the lookup-table binding and the scheme disagree.
44///
45/// The table binds as a buffer of its own, so nothing ties it to the mode: a missing one leaves
46/// [`QuantMode::Lookup`](cubecl_common::quant::scheme::QuantMode) nothing to index, an extra one
47/// is a caller quantizing differently than the scheme it passed. Lookup is also only wired where
48/// the decode goes through the packed-u32 unpack, and only for the integer values whose field is
49/// a plain bit range — a minifloat field carries its own float semantics, which an index does not
50/// have.
51///
52/// The table must hold `2^bits` f32 entries; [`register_table`](crate::quant::view) checks the
53/// binding's length against that, the one host-side site holding both.
54pub fn check_table_bindings(scheme: &QuantScheme, table_provided: bool) {
55    use cubecl_common::quant::scheme::{QuantMode, QuantStore, QuantValue};
56    match (scheme.mode, table_provided) {
57        (QuantMode::Lookup, false) => {
58            panic!(
59                "{:?} takes a lookup table, but none was provided",
60                scheme.mode
61            )
62        }
63        (QuantMode::Lookup, true) => {
64            assert!(
65                matches!(scheme.store, QuantStore::PackedU32(_)),
66                "lookup decode is only wired for packed-u32 storage, got {:?}",
67                scheme.store
68            );
69            assert!(
70                !matches!(
71                    scheme.value,
72                    QuantValue::E5M2 | QuantValue::E4M3 | QuantValue::E2M1
73                ),
74                "a lookup field is an index, so a minifloat value ({:?}) has nothing to mean; \
75                 use the integer value of the same width",
76                scheme.value
77            );
78        }
79        (_, true) => {
80            panic!(
81                "a lookup table was provided, but {:?} does not take one",
82                scheme.mode
83            )
84        }
85        (_, false) => {}
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::{check_scale_bindings, check_table_bindings};
92    use cubecl_common::quant::scheme::{
93        QuantMode, QuantScheme, QuantStore, QuantValue, ScaleDtype,
94    };
95
96    #[test]
97    fn a_one_level_scheme_takes_one_binding() {
98        check_scale_bindings(&QuantScheme::default().per_tensor(ScaleDtype::F32), 1);
99        check_scale_bindings(&QuantScheme::default().per_block([32], ScaleDtype::F32), 1);
100    }
101
102    #[test]
103    fn a_two_level_scheme_takes_two_bindings() {
104        check_scale_bindings(
105            &QuantScheme::default()
106                .per_block([32], ScaleDtype::F32)
107                .per_tensor(ScaleDtype::F32),
108            2,
109        );
110    }
111
112    /// The binding is f32, so a level naming another dtype would have its scale read as f32 bytes.
113    #[test]
114    #[should_panic(expected = "binds as f32, but")]
115    fn a_two_level_scheme_storing_the_tensor_scale_narrower_is_rejected() {
116        check_scale_bindings(
117            &QuantScheme::default()
118                .per_block([32], ScaleDtype::F32)
119                .per_tensor(ScaleDtype::BF16),
120            2,
121        );
122    }
123
124    #[test]
125    #[should_panic(expected = "takes as many scale bindings, but 1 were provided")]
126    fn a_two_level_scheme_with_one_binding_is_rejected() {
127        // Would otherwise dequantize against the block scales alone, dropping the per-tensor factor.
128        check_scale_bindings(
129            &QuantScheme::default()
130                .per_block([32], ScaleDtype::F32)
131                .per_tensor(ScaleDtype::F32),
132            1,
133        );
134    }
135
136    #[test]
137    #[should_panic(expected = "takes as many scale bindings, but 2 were provided")]
138    fn a_one_level_scheme_with_two_bindings_is_rejected() {
139        check_scale_bindings(&QuantScheme::default().per_tensor(ScaleDtype::F32), 2);
140    }
141
142    fn lookup_scheme() -> QuantScheme {
143        QuantScheme::default()
144            .with_value(QuantValue::Q4F)
145            .with_mode(QuantMode::Lookup)
146    }
147
148    #[test]
149    fn a_lookup_scheme_takes_a_table() {
150        check_table_bindings(&lookup_scheme(), true);
151    }
152
153    #[test]
154    fn a_symmetric_scheme_takes_no_table() {
155        check_table_bindings(&QuantScheme::default(), false);
156    }
157
158    #[test]
159    #[should_panic(expected = "takes a lookup table, but none was provided")]
160    fn a_lookup_scheme_without_a_table_is_rejected() {
161        // Would otherwise fall back to the integer cast and reconstruct the index itself.
162        check_table_bindings(&lookup_scheme(), false);
163    }
164
165    #[test]
166    #[should_panic(expected = "does not take one")]
167    fn a_symmetric_scheme_with_a_table_is_rejected() {
168        check_table_bindings(&QuantScheme::default(), true);
169    }
170
171    /// Only the packed-u32 unpack decodes through the table; the native paths cast the storage
172    /// element directly and would silently ignore it.
173    #[test]
174    #[should_panic(expected = "only wired for packed-u32 storage")]
175    fn a_native_lookup_scheme_is_rejected() {
176        let scheme = QuantScheme::default()
177            .with_value(QuantValue::Q8F)
178            .with_store(QuantStore::Native)
179            .with_mode(QuantMode::Lookup);
180        check_table_bindings(&scheme, true);
181    }
182
183    #[test]
184    #[should_panic(expected = "a lookup field is an index")]
185    fn a_minifloat_lookup_scheme_is_rejected() {
186        let scheme = QuantScheme::default()
187            .with_value(QuantValue::E4M3)
188            .with_mode(QuantMode::Lookup);
189        check_table_bindings(&scheme, true);
190    }
191}
cubecl_std/quant/base.rs

cubecl_std/quant/
base.rs
