Flash attention: lines for mask and value (#1072)

Author: louisfd

crates/cubecl-attention/src/base.rs

@@ -87,19 +87,27 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
     out: &TensorHandleRef<R>,
     attention_elems: &AttentionElems,
 ) -> Result<(), AttentionSetupError> {
-    let line_sizes = AvailableLineSizes::from_elem_types(
-        client,
-        query.elem_size,
-        attention_elems.mask.size(),
-        out.elem_size,
-    );
-    let line_sizes = A::filter_line_sizes(line_sizes)
-        .filter_with_tensor(AttentionIdent::Query, query.strides, query.shape)
-        .filter_with_tensor(AttentionIdent::Key, key.strides, key.shape)
-        .filter_with_tensor(AttentionIdent::Value, value.strides, value.shape)
-        .filter_with_tensor(AttentionIdent::Out, out.strides, out.shape)
-        .pick_max()
-        .unwrap();
+    let line_sizes = {
+        let ls = AvailableLineSizes::from_elem_types(
+            client,
+            query.elem_size,
+            attention_elems.mask.size(),
+            out.elem_size,
+        );
+        let ls = A::filter_line_sizes(ls)
+            .filter_with_tensor(AttentionIdent::Query, query.strides, query.shape)
+            .filter_with_tensor(AttentionIdent::Key, key.strides, key.shape)
+            .filter_with_tensor(AttentionIdent::Value, value.strides, value.shape)
+            .filter_with_tensor(AttentionIdent::Out, out.strides, out.shape);
+
+        if let Some(mask) = mask.as_ref() {
+            ls.filter_with_tensor(AttentionIdent::Mask, mask.strides, mask.shape)
+        } else {
+            ls
+        }
+    }
+    .pick_max()
+    .unwrap();
 
     let problem = AttentionProblem {
         batch: query.shape[0],

crates/cubecl-attention/src/components/tile/accelerated/local_tile.rs

@@ -70,6 +70,7 @@ impl<E: Numeric> LocalTile<E> {
     }
 
     pub fn load_from_strided_tile<E2: Numeric>(&mut self, strided_tile: &StridedTile<E2>) {
+        // Assumes line size == 1
         for r in 0..self.layout.unit_size.0 {
             for c in 0..self.layout.unit_size.1 {
                 let (row, col) = self.layout.absolute_pos((r, c));

crates/cubecl-attention/src/components/tile/accelerated/setup.rs

@@ -2,7 +2,9 @@ use cubecl_core::client::ComputeClient;
 use cubecl_matmul::components::ComputeResources;
 
 use crate::components::AttentionElems;
+use crate::components::AttentionIdent;
 use crate::components::AttentionTileSize;
+use crate::components::AvailableLineSizes;
 use crate::components::tile::SharedTileAttentionConfig;
 use crate::components::tile::TileAttentionConfig;
 use crate::components::tile::accelerated::BlackboxAcceleratedTileAttention;
@@ -86,6 +88,11 @@ impl TileAttentionFamily for BlackboxAcceleratedTileAttention {
             selection.reuse_key_value,
         )
     }
+
+    fn filter_line_sizes(available_line_sizes: AvailableLineSizes) -> AvailableLineSizes {
+        // Vectorized mask not supported
+        available_line_sizes.filter(|ls| *ls == 1, AttentionIdent::Mask)
+    }
 }
 
 fn validate(

crates/cubecl-attention/src/components/tile/unit_register/attention.rs

@@ -344,7 +344,7 @@ fn strided_tile_to_unit_tile<E: Numeric, E2: Numeric>(
     let line_size = strided_tile.line_size;
     assert!(unit_tile.layout.num_cols % line_size == 0);
 
-    let col_iterations = comptime!(unit_tile.layout.num_cols / strided_tile.line_size);
+    let col_iterations = comptime!(unit_tile.layout.num_cols / line_size);
 
     for row in 0..unit_tile.layout.num_rows {
         for col in 0..col_iterations {

crates/cubecl-attention/src/kernels/blackbox_accelerated.rs

@@ -1,6 +1,7 @@
 use cubecl_matmul::components::{global::PartitionedStageFamily, stage::StridedStageFamily};
 
 use crate::components::stage::plane::PlanePartitionStageAttentionFamily;
+use crate::components::tile::TileAttentionFamily;
 use crate::components::tile::accelerated::BlackboxAcceleratedTileAttention;
 use crate::{
     components::{
@@ -24,12 +25,14 @@ impl Algorithm for BlackboxAcceleratedAlgorithm {
     type BatchAttention = SimpleBatchAttentionFamily<Self::GlobalAttention>;
 
     fn filter_line_sizes(available_line_sizes: AvailableLineSizes) -> AvailableLineSizes {
-        AvailableLineSizes {
+        let supported = AvailableLineSizes {
             query: available_line_sizes.query,
             key: vec![1],
-            value: vec![1],
-            mask: vec![1],
+            value: available_line_sizes.value,
+            mask: available_line_sizes.mask,
             out: available_line_sizes.out,
-        }
+        };
+
+        Self::TileAttention::filter_line_sizes(supported)
     }
 }

crates/cubecl-attention/src/kernels/unit.rs

@@ -27,8 +27,8 @@ impl Algorithm for UnitAlgorithm {
         AvailableLineSizes {
             query: available_line_sizes.query,
             key: vec![1],
-            value: vec![1],
-            mask: vec![1],
+            value: available_line_sizes.value,
+            mask: available_line_sizes.mask,
             out: available_line_sizes.out,
         }
     }

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

@@ -56,20 +56,27 @@ pub fn test_attention_algorithm<A, P, R>(
     let out = tensor_raw_parts_output::<P, R>(&client, &problem);
 
     let attention_elems = AttentionElems::new::<P::AP>();
-    let line_sizes = AvailableLineSizes::from_elem_types(
-        &client,
-        attention_elems.query_global.size(),
-        attention_elems.mask.size(),
-        attention_elems.out_global.size(),
-    );
-    let line_sizes = A::filter_line_sizes(line_sizes);
-    let line_sizes = line_sizes
-        .filter_with_tensor(AttentionIdent::Query, &query.strides, &query.shape)
-        .filter_with_tensor(AttentionIdent::Key, &key.strides, &key.shape)
-        .filter_with_tensor(AttentionIdent::Value, &value.strides, &value.shape)
-        .filter_with_tensor(AttentionIdent::Out, &out.strides, &out.shape)
-        .pick_max()
-        .unwrap();
+    let line_sizes = {
+        let ls = AvailableLineSizes::from_elem_types(
+            &client,
+            attention_elems.query_global.size(),
+            attention_elems.mask.size(),
+            attention_elems.out_global.size(),
+        );
+        let ls = A::filter_line_sizes(ls)
+            .filter_with_tensor(AttentionIdent::Query, &query.strides, &query.shape)
+            .filter_with_tensor(AttentionIdent::Key, &key.strides, &key.shape)
+            .filter_with_tensor(AttentionIdent::Value, &value.strides, &value.shape)
+            .filter_with_tensor(AttentionIdent::Out, &out.strides, &out.shape);
+
+        if let Some(mask) = mask.as_ref() {
+            ls.filter_with_tensor(AttentionIdent::Mask, &mask.strides, &mask.shape)
+        } else {
+            ls
+        }
+    }
+    .pick_max()
+    .unwrap();
 
     let config = match A::setup(&client, &problem, &selection, &line_sizes, &attention_elems) {
         Ok(config) => config,