Refactor quantization scheme

crates/burn-autodiff/src/ops/qtensor.rs

@@ -4,7 +4,7 @@ use burn_tensor::{
     Device, Shape, TensorData,
     backend::Backend,
     ops::{FloatTensor, IntTensor, QTensorOps, QuantizedTensor},
-    quantization::{QuantizationParametersPrimitive, QuantizationScheme},
+    quantization::{QuantScheme, QuantizationParametersPrimitive},
 };
 
 use crate::{Autodiff, checkpoint::strategy::CheckpointStrategy};
@@ -16,15 +16,15 @@ impl<B: Backend, C: CheckpointStrategy> QTensorOps<Self> for Autodiff<B, C> {
 
     fn quantize(
         _tensor: FloatTensor<Self>,
-        _scheme: &QuantizationScheme,
+        _scheme: &QuantScheme,
         _qparams: QuantizationParametersPrimitive<Self>,
     ) -> QuantizedTensor<Self> {
         todo!() // required for QAT
     }
 
     fn quantize_dynamic(
         _tensor: FloatTensor<Self>,
-        _scheme: &QuantizationScheme,
+        _scheme: &QuantScheme,
     ) -> QuantizedTensor<Self> {
         todo!()
     }

crates/burn-candle/src/ops/qtensor.rs

@@ -4,7 +4,7 @@ use burn_tensor::{
     DType, Device, Shape, TensorData,
     backend::Backend,
     ops::{FloatTensor, IntTensor, QTensorOps, QuantizedTensor},
-    quantization::{QuantizationParametersPrimitive, QuantizationScheme},
+    quantization::{QuantScheme, QuantizationParametersPrimitive},
 };
 
 use crate::{
@@ -19,7 +19,7 @@ impl<F: FloatCandleElement, I: IntCandleElement> QTensorOps<Self> for Candle<F,
 
     fn quantize(
         _tensor: FloatTensor<Self>,
-        _scheme: &QuantizationScheme,
+        _scheme: &QuantScheme,
         _qparams: QuantizationParametersPrimitive<Self>,
     ) -> QuantizedTensor<Self> {
         unimplemented!()

crates/burn-candle/src/tensor.rs

@@ -1,6 +1,6 @@
 use burn_tensor::{
     DType, Element, Shape, TensorData, TensorMetadata,
-    quantization::{QTensorPrimitive, QuantizationScheme},
+    quantization::{QTensorPrimitive, QuantScheme},
 };
 
 use crate::{CandleDevice, element::CandleElement};
@@ -63,11 +63,11 @@ pub struct CandleQTensor {
     // NOTE: candle  does not implement `WithDType` for i8
     pub qtensor: CandleTensor,
     /// The quantization scheme.
-    pub scheme: QuantizationScheme,
+    pub scheme: QuantScheme,
 }
 
 impl QTensorPrimitive for CandleQTensor {
-    fn scheme(&self) -> &QuantizationScheme {
+    fn scheme(&self) -> &QuantScheme {
         &self.scheme
     }
 }

crates/burn-core/src/module/quantize.rs

@@ -1,7 +1,7 @@
 use burn_tensor::{
     Tensor,
     backend::Backend,
-    quantization::{Calibration, QuantizationScheme},
+    quantization::{Calibration, QuantScheme},
 };
 
 use crate::module::{ModuleMapper, ParamId};
@@ -11,7 +11,7 @@ pub struct Quantizer {
     /// The calibration method used in quantization.
     pub calibration: Calibration,
     /// The quantization scheme.
-    pub scheme: QuantizationScheme,
+    pub scheme: QuantScheme,
 }
 
 impl<B: Backend> ModuleMapper<B> for Quantizer {

crates/burn-cubecl-fusion/src/shared/builder.rs

@@ -8,7 +8,7 @@ use burn_ir::{
     BaseOperationIr, BinaryOpIr, FloatOperationIr, NumericOperationIr, OperationIr, ScalarOpIr,
     TensorIr, UnaryOpIr,
 };
-use burn_tensor::Element;
+use burn_tensor::{DType, Element};
 use cubecl::ir::Elem;
 
 /// The base optimization builder that can be used to fuse all elemwise operations.
@@ -212,6 +212,10 @@ impl FuseOptimizationBuilder {
                     return false;
                 }
 
+                if self.input_is_quantized(&desc.input) {
+                    return false;
+                }
+
                 if self.builder.register(|build| {
                     build.input_swap_dims(
                         &desc.input,
@@ -243,6 +247,10 @@ impl FuseOptimizationBuilder {
                     return false;
                 }
 
+                if self.input_is_quantized(&desc.input) {
+                    return false;
+                }
+
                 if self.builder.register(|build| {
                     build.input_reshaped(&desc.input, &desc.out)?;
                     Some(())
@@ -447,6 +455,10 @@ impl FuseOptimizationBuilder {
                     return false;
                 }
 
+                if self.input_is_quantized(&desc.tensor) {
+                    return false;
+                }
+
                 self.builder.register(|build| {
                     let input = build.input_indexed(&desc.tensor)?;
                     let indices = build.input(&desc.indices)?;
@@ -467,6 +479,10 @@ impl FuseOptimizationBuilder {
                     return false;
                 }
 
+                if self.input_is_quantized(&desc.tensor) {
+                    return false;
+                }
+
                 self.builder.register(|build| {
                     let input = build.input_indexed(&desc.tensor)?;
                     let indices = build.input_indexed(&desc.indices)?;
@@ -494,6 +510,10 @@ impl FuseOptimizationBuilder {
             return false;
         }
 
+        if self.input_is_quantized(&desc.lhs) {
+            return false;
+        }
+
         self.builder.register(|build| {
             let lhs = build.input(&desc.lhs)?;
             let rhs = build.input(&desc.rhs)?;
@@ -513,6 +533,10 @@ impl FuseOptimizationBuilder {
             return false;
         }
 
+        if self.input_is_quantized(&desc.input) {
+            return false;
+        }
+
         self.builder.register(|build| {
             let input = build.input(&desc.input)?;
             let out = build.output(&desc.out)?;
@@ -529,6 +553,10 @@ impl FuseOptimizationBuilder {
             return false;
         }
 
+        if self.input_is_quantized(&desc.lhs) {
+            return false;
+        }
+
         self.builder.register(|build| {
             let elem = desc.lhs.dtype;
             let lhs = build.input(&desc.lhs)?;
@@ -541,6 +569,10 @@ impl FuseOptimizationBuilder {
         })
     }
 
+    fn input_is_quantized(&self, input: &TensorIr) -> bool {
+        matches!(input.dtype, DType::QFloat(_scheme))
+    }
+
     fn output_is_compatible(&mut self, out: &TensorIr) -> bool {
         if self.current_output_shape.is_empty() {
             self.current_output_shape.clone_from(&out.shape);

crates/burn-cubecl-fusion/src/shared/ir.rs

@@ -415,7 +415,7 @@ impl From<DType> for FusePrecision {
             DType::U16 => Self::U16,
             DType::U8 => Self::U8,
             DType::Bool => Self::Bool,
-            _ => panic!("Unsupported"),
+            _ => panic!("Unsupported precision for fusion: {value:?}"),
         }
     }
 }

crates/burn-cubecl/src/kernel/matmul/base.rs

@@ -1,6 +1,9 @@
 use super::init_matmul_output;
 use crate::{CubeRuntime, FloatElement, tensor::CubeTensor};
-use burn_tensor::DType;
+use burn_tensor::{
+    DType,
+    quantization::{QTensorPrimitive, QuantAccPrecision},
+};
 use cubecl::linalg::matmul::{components::Quantized, kernels::MatmulLaunchError};
 
 #[cfg(feature = "autotune")]
@@ -65,16 +68,34 @@ pub fn q_matmul<R: CubeRuntime>(
 
     let client = &lhs.client;
 
+    let scheme = *lhs.scheme();
+
     lhs.dtype = DType::I8;
     rhs.dtype = DType::I8;
 
-    cubecl::linalg::matmul::launch_ref::<R, (i8, half::f16, half::f16, half::f16, Quantized)>(
-        &Default::default(),
-        client,
-        &lhs.as_handle_ref(),
-        &rhs.as_handle_ref(),
-        &out.as_handle_ref(),
-    )?;
+    match scheme.acc_precision {
+        QuantAccPrecision::Full => {
+            cubecl::linalg::matmul::launch_ref::<R, (i8, half::f16, f32, half::f16, Quantized)>(
+                &Default::default(),
+                client,
+                &lhs.as_handle_ref(),
+                &rhs.as_handle_ref(),
+                &out.as_handle_ref(),
+            )?;
+        }
+        QuantAccPrecision::Half => {
+            cubecl::linalg::matmul::launch_ref::<
+                R,
+                (i8, half::f16, half::f16, half::f16, Quantized),
+            >(
+                &Default::default(),
+                client,
+                &lhs.as_handle_ref(),
+                &rhs.as_handle_ref(),
+                &out.as_handle_ref(),
+            )?;
+        }
+    }
 
     Ok(out)
 }

crates/burn-cubecl/src/kernel/quantization/dequantize.rs

@@ -1,7 +1,7 @@
 use crate::tensor::CubeTensor;
 use crate::{CubeElement, CubeRuntime};
 use burn_tensor::DType;
-use burn_tensor::quantization::{QuantizationMode, QuantizationScheme, QuantizationType};
+use burn_tensor::quantization::{QuantInputType, QuantLevel, QuantMode, QuantScheme};
 use cubecl::calculate_cube_count_elemwise;
 use cubecl::prelude::*;
 
@@ -39,7 +39,7 @@ fn unpack_i8s(value: u32) -> Line<i32> {
 fn dequantize_per_tensor_symmetric_int8_kernel(
     input: &QTensor,
     output: &mut Tensor<Line<f32>>,
-    #[comptime] scheme: QuantizationScheme,
+    #[comptime] scheme: QuantScheme,
 ) {
     // Last position contains the qparam
     if ABSOLUTE_POS >= input.len() - 1 {
@@ -93,7 +93,12 @@ where
 
     if let DType::QFloat(scheme) = tensor.dtype {
         match scheme {
-            QuantizationScheme::PerTensor(QuantizationMode::Symmetric, QuantizationType::QInt8) => {
+            QuantScheme {
+                level: QuantLevel::Tensor,
+                mode: QuantMode::Symmetric,
+                q_type: QuantInputType::QInt8,
+                ..
+            } => {
                 unsafe {
                     dequantize_per_tensor_symmetric_int8_kernel::launch_unchecked::<R>(
                         &client,

crates/burn-cubecl/src/kernel/quantization/qtensor.rs

@@ -1,13 +1,13 @@
 #![allow(missing_docs)] // cube derive macros
 
-use burn_tensor::quantization::{QuantizationMode, QuantizationScheme};
+use burn_tensor::quantization::{QuantInputType, QuantLevel, QuantMode, QuantScheme};
 use cubecl::prelude::*;
 
 /// Quantization parameters.
 #[derive(CubeLaunch, CubeType)]
 pub struct QParams {
     #[cube(comptime)]
-    scheme: QuantizationScheme,
+    scheme: QuantScheme,
 }
 
 /// Quantized tensor representation.
@@ -16,7 +16,7 @@ pub type QTensor = Array<Line<u32>>;
 #[cube]
 impl QParams {
     /// Create a new quantization parameters instance.
-    pub fn new(scheme: QuantizationScheme) -> Self {
+    pub fn new(#[comptime] scheme: QuantScheme) -> Self {
         QParams { scheme }
     }
 
@@ -25,9 +25,12 @@ impl QParams {
         let len = tensor.len();
         match comptime!(self.scheme) {
             // Symmetric quantization only contains the scaling factor as the last element
-            QuantizationScheme::PerTensor(QuantizationMode::Symmetric, _) => {
-                (f32::reinterpret(tensor[len - 1][tensor.line_size() - 1]), 0)
-            }
+            QuantScheme {
+                level: QuantLevel::Tensor,
+                mode: QuantMode::Symmetric,
+                q_type: QuantInputType::QInt8,
+                ..
+            } => (f32::reinterpret(tensor[len - 1][tensor.line_size() - 1]), 0),
         }
     }
 }