Flash Attention: all lines (#1073)

Author: louisfd

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

@@ -363,10 +363,22 @@ fn strided_tile_to_transposed_unit_tile<E: Numeric, E2: Numeric>(
     strided_tile: &StridedTile<E>,
     unit_tile: &mut UnitTile<E2>,
 ) {
-    for row in 0..unit_tile.layout.num_rows {
-        for col in 0..unit_tile.layout.num_cols {
-            unit_tile.data[row * unit_tile.layout.num_cols + col] =
-                E2::cast_from(strided_tile.get_line(col, row))
+    let line_size = strided_tile.line_size;
+    assert!(unit_tile.layout.num_cols % line_size == 0);
+
+    let input_num_rows = unit_tile.layout.num_cols;
+    let input_num_cols = unit_tile.layout.num_rows;
+    let line_iterations = comptime!(input_num_cols / line_size);
+
+    for input_row in 0..input_num_rows {
+        for input_col_line in 0..line_iterations {
+            let line_read = strided_tile.get_line(input_row, input_col_line);
+
+            #[unroll]
+            for i in 0..line_size {
+                unit_tile.data[(input_col_line + i) * input_num_rows + input_row] =
+                    E2::cast_from(line_read[i]);
+            }
         }
     }
 }

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

@@ -25,14 +25,6 @@ impl Algorithm for BlackboxAcceleratedAlgorithm {
     type BatchAttention = SimpleBatchAttentionFamily<Self::GlobalAttention>;
 
     fn filter_line_sizes(available_line_sizes: AvailableLineSizes) -> AvailableLineSizes {
-        let supported = AvailableLineSizes {
-            query: available_line_sizes.query,
-            key: vec![1],
-            value: available_line_sizes.value,
-            mask: available_line_sizes.mask,
-            out: available_line_sizes.out,
-        };
-
-        Self::TileAttention::filter_line_sizes(supported)
+        Self::TileAttention::filter_line_sizes(available_line_sizes)
     }
 }

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

@@ -4,8 +4,7 @@ use crate::components::stage::unit::UnitPartitionStageAttentionFamily;
 use crate::components::tile::unit_register::UnitRegisterTileAttention;
 use crate::{
     components::{
-        AvailableLineSizes, batch::simple::SimpleBatchAttentionFamily,
-        global::simple::SimpleGlobalAttentionFamily,
+        batch::simple::SimpleBatchAttentionFamily, global::simple::SimpleGlobalAttentionFamily,
     },
     kernels::Algorithm,
 };
@@ -22,14 +21,4 @@ impl Algorithm for UnitAlgorithm {
     >;
     type GlobalAttention = SimpleGlobalAttentionFamily<Self::StageAttention>;
     type BatchAttention = SimpleBatchAttentionFamily<Self::GlobalAttention>;
-
-    fn filter_line_sizes(available_line_sizes: AvailableLineSizes) -> AvailableLineSizes {
-        AvailableLineSizes {
-            query: available_line_sizes.query,
-            key: vec![1],
-            value: available_line_sizes.value,
-            mask: available_line_sizes.mask,
-            out: available_line_sizes.out,
-        }
-    }
 }