fix: Quant matmul line sizes (#978)

Author: wingertge

crates/cubecl-attention/src/base.rs

@@ -14,7 +14,6 @@ use crate::{
 
 use crate::components::batch::BatchAttentionConfig;
 use crate::components::batch::BatchAttentionFamily;
-use cubecl_core::frontend::CubePrimitive;
 
 pub enum Strategy {
     /// Temporary implementation
@@ -66,9 +65,9 @@ pub fn launch_tmp<R: Runtime, AP: AttentionPrecision>(
     out: &TensorHandleRef<R>,
 ) -> Result<(), AttentionSetupError> {
     let line_sizes = AvailableLineSizes::from_elem_types::<R>(
-        &QG::<AP>::as_type_native_unchecked(),
-        &MSK::<AP>::as_type_native_unchecked(),
-        &OG::<AP>::as_type_native_unchecked(),
+        query.elem_size,
+        size_of::<MSK<AP>>(),
+        out.elem_size,
     );
     let line_sizes = DummyRegisterAlgorithm::filter_line_sizes(line_sizes)
         .filter_with_tensor(AttentionIdent::Query, query.strides, query.shape)

crates/cubecl-attention/src/components/line_size.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 
-use cubecl_core::{LineSizeError, Runtime, ir::StorageType, tensor_line_size_parallel};
+use cubecl_core::{LineSizeError, Runtime, tensor_line_size_parallel};
 
 use crate::components::{AttentionIdent, AttentionSetupError};
 
@@ -29,11 +29,7 @@ pub struct AvailableLineSizes {
 }
 
 impl AvailableLineSizes {
-    pub fn from_elem_types<R: Runtime>(
-        elem_in: &StorageType,
-        elem_mask: &StorageType,
-        elem_out: &StorageType,
-    ) -> Self {
+    pub fn from_elem_types<R: Runtime>(elem_in: usize, elem_mask: usize, elem_out: usize) -> Self {
         let in_available: Vec<u8> = R::io_optimized_line_sizes_unchecked(elem_in).collect();
         let mask_available: Vec<u8> = R::io_optimized_line_sizes_unchecked(elem_mask).collect();
         let out_available = R::io_optimized_line_sizes_unchecked(elem_out).collect();

crates/cubecl-attention/src/tests/attention_test_launcher.rs

@@ -58,9 +58,9 @@ pub fn test_attention_algorithm<A, P, R>(
     let out = tensor_raw_parts_output::<P, R>(&client, &problem);
 
     let line_sizes = AvailableLineSizes::from_elem_types::<R>(
-        &P::EG::as_type_native_unchecked(),
-        &P::EM::as_type_native_unchecked(),
-        &P::EG::as_type_native_unchecked(),
+        size_of::<P::EG>(),
+        size_of::<P::EM>(),
+        size_of::<P::EG>(),
     );
     let line_sizes = A::filter_line_sizes(line_sizes);
     let line_sizes = line_sizes

crates/cubecl-convolution/src/components/stage/reader.rs

@@ -1,8 +1,9 @@
 use cubecl::prelude::*;
 use cubecl_core as cubecl;
 use cubecl_matmul::components::{
-    MatrixLayout, StageIdent,
-    stage::{StageMemoryConfig, StridedStage, TilingLayout},
+    InvalidConfigError, MatrixLayout, StageIdent,
+    global::memory::GlobalMemoryConfig,
+    stage::{StageMemoryConfig, StridedStage, TilingLayout, TilingValidation},
     tile::StridedTile,
 };
 use cubecl_std::tensor::layout::Coords2d;
@@ -39,3 +40,13 @@ impl TilingLayout for BiasTilingLayout {
         )
     }
 }
+
+impl TilingValidation for BiasTilingLayout {
+    fn check(config: GlobalMemoryConfig) -> Result<(), InvalidConfigError> {
+        let stage_width = config.elements_in_stage_col;
+        if config.global_line_size > stage_width {
+            return Err(Box::new("Invalid line size"));
+        }
+        Ok(())
+    }
+}

crates/cubecl-convolution/src/launch.rs

@@ -181,10 +181,10 @@ where
     Input<Alg, MP>: ConcreteInputsFactory,
     Output<Alg, MP>: ConcreteOutputFactory,
 {
-    let line_sizes = AvailableLineSizes::from_types::<R>(
-        &LhsG::<MP>::as_type_native_unchecked(),
-        &RhsG::<MP>::as_type_native_unchecked(),
-        &AccG::<MP>::as_type_native_unchecked(),
+    let line_sizes = AvailableLineSizes::from_type_sizes::<R>(
+        input.data().elem_size,
+        weight.data().elem_size,
+        out.elem_size,
     )
     .filter_lhs_with_tensor(input.data().strides, input.data().shape, problem.lhs_layout)
     .filter_rhs_with_tensor(

crates/cubecl-convolution/src/tests/convolution_test_launcher.rs

@@ -56,7 +56,7 @@ pub fn test_convolution_algorithm<A, Args, P, R>(
     let line_sizes = AvailableLineSizes {
         lhs: vec![1],
         rhs: vec![1],
-        out: R::io_optimized_line_sizes_unchecked(&P::EG::as_type_native_unchecked()).collect(),
+        out: R::io_optimized_line_sizes_unchecked(size_of::<P::EG>()).collect(),
     }
     .filter_lhs_with_tensor(&lhs.strides, &lhs.shape, problem.lhs_layout)
     .filter_rhs_with_tensor(&rhs.strides, &rhs.shape, problem.rhs_layout)

crates/cubecl-core/src/frontend/element/cube_elem.rs

@@ -45,6 +45,11 @@ pub trait CubePrimitive:
         Self::as_type_native().map(|t| t.size_bits())
     }
 
+    /// Only native element types have a size.
+    fn size_bits_unchecked() -> usize {
+        Self::as_type_native_unchecked().size_bits()
+    }
+
     fn from_expand_elem(elem: ExpandElement) -> Self::ExpandType {
         ExpandElementTyped::new(elem)
     }

crates/cubecl-core/src/runtime.rs

@@ -50,8 +50,9 @@ pub trait Runtime: Send + Sync + 'static + core::fmt::Debug {
     /// Returns all line sizes that are useful to perform optimal IO operation on the given element.
     /// Ignores native support, and allows all line sizes. This means the returned size may be
     /// unrolled, and may not support dynamic indexing.
-    fn io_optimized_line_sizes_unchecked(elem: &StorageType) -> impl Iterator<Item = u8> + Clone {
-        let max = LOAD_WIDTH / elem.size_bits();
+    fn io_optimized_line_sizes_unchecked(size: usize) -> impl Iterator<Item = u8> + Clone {
+        let size_bits = size * 8;
+        let max = LOAD_WIDTH / size_bits;
         let max = usize::min(Self::max_global_line_size() as usize, max);
 
         // If the max is 8, we want to test 1, 2, 4, 8 which is log2(8) + 1.

crates/cubecl-cpu/src/runtime.rs

@@ -111,8 +111,9 @@ impl Runtime for CpuRuntime {
         supported.iter().filter(move |v| **v <= max).cloned()
     }
 
-    fn io_optimized_line_sizes_unchecked(elem: &StorageType) -> impl Iterator<Item = u8> + Clone {
-        let max = LOAD_WIDTH / elem.size_bits();
+    fn io_optimized_line_sizes_unchecked(elem_size: usize) -> impl Iterator<Item = u8> + Clone {
+        let elem_size_bits = elem_size * 8;
+        let max = LOAD_WIDTH / elem_size_bits;
         (1..max as u8).rev().filter(|v| v.is_power_of_two())
     }
 

crates/cubecl-matmul/src/components/global/args.rs

@@ -157,10 +157,8 @@ impl<Lhs: Numeric, Rhs: Numeric, Acc: Numeric> ConcreteInputsFactory
                 );
                 let data_view =
                     ViewArg::new::<BatchedGlobalLayout>(data.as_array_arg(line_size), data_layout);
-                let scales_view = ViewArg::new::<BatchedGlobalScaleLayout>(
-                    scale.as_array_arg(line_size),
-                    scales_layout,
-                );
+                let scales_view =
+                    ViewArg::new::<BatchedGlobalScaleLayout>(scale.as_array_arg(1), scales_layout);
                 ViewArg::new_quantized(data_view, scales_view, **scheme)
             }
         };