enable unit attention (#1079)

Author: louisfd

crates/cubecl-attention/src/base.rs

@@ -4,11 +4,8 @@ use cubecl_std::tensor::TensorHandle;
 
 use crate::{
     components::{
-        AttentionElems, AttentionIdent, AttentionPartitionSize, AttentionProblem,
-        AttentionSelection, AttentionSetupError, AttentionStageSize, AttentionTileSize,
-        AttentionTilingScheme, AvailableLineSizes,
+        AttentionElems, AttentionIdent, AttentionProblem, AttentionSetupError, AvailableLineSizes,
         args::{TensorArgs, TensorInputsLaunch},
-        batch::HypercubeSelection,
     },
     kernels::{Algorithm, blackbox_accelerated::BlackboxAcceleratedAlgorithm, unit::UnitAlgorithm},
 };
@@ -120,48 +117,22 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
         causal: false,
     };
 
-    let tile_size = AttentionTileSize {
-        seq_q: 8,
-        head_dim: 8,
-        seq_kv: 8,
-        val_dim: 8,
-    };
-
-    assert!(problem.head_dim as u32 % tile_size.head_dim == 0);
-    let partition_head_dim = problem.head_dim as u32 / tile_size.head_dim;
-    let partition_val_dim = partition_head_dim;
-
-    let selection = AttentionSelection {
-        hypercube_selection: HypercubeSelection {},
-        tiling_scheme: AttentionTilingScheme {
-            tile_size,
-            partition_size: AttentionPartitionSize {
-                seq_q: 1,
-                head_dim: partition_head_dim,
-                seq_kv: 1,
-                val_dim: partition_val_dim,
-            },
-            stage_size: AttentionStageSize { seq_q: 1 },
-        },
-        plane_dim: 32,
-        reuse_key_value: false,
-        two_rows_in_array_tile: false,
-    };
-
-    let config = BlackboxAcceleratedAlgorithm::setup(
+    let selection = A::selection(
         client,
         &problem,
-        &selection,
+        client.properties().hardware.plane_size_max,
         &line_sizes,
         attention_elems,
     )?;
 
+    let config = A::setup(client, &problem, &selection, &line_sizes, attention_elems)?;
+
     let cube_count_plan = config
         .hypercube_config()
         .cube_count_plan(&problem, &selection);
 
     let result = unsafe {
-        <BlackboxAcceleratedAlgorithm as Algorithm>::BatchAttention::launch_unchecked::<TensorArgs, R>(
+        <A as Algorithm>::BatchAttention::launch_unchecked::<TensorArgs, R>(
             client,
             config.cube_dim(),
             cube_count_plan.resolve(),

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

@@ -27,4 +27,12 @@ pub trait Algorithm {
     ) -> Result<<Self::BatchAttention as BatchAttentionFamily>::Config, AttentionSetupError> {
         Self::BatchAttention::setup(client, problem, selection, line_sizes, dtypes)
     }
+
+    fn selection<R: Runtime>(
+        client: &ComputeClient<R>,
+        problem: &AttentionProblem,
+        plane_dim: u32,
+        line_sizes: &AttentionLineSizes,
+        dtypes: &AttentionElems,
+    ) -> Result<AttentionSelection, AttentionSetupError>;
 }

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

@@ -1,8 +1,15 @@
+use cubecl_core::client::ComputeClient;
 use cubecl_matmul::components::{global::PartitionedStageFamily, stage::StridedStageFamily};
 
+use crate::components::batch::HypercubeSelection;
 use crate::components::stage::plane::PlanePartitionStageAttentionFamily;
 use crate::components::tile::TileAttentionFamily;
 use crate::components::tile::accelerated::BlackboxAcceleratedTileAttention;
+use crate::components::{
+    AttentionElems, AttentionLineSizes, AttentionPartitionSize, AttentionProblem,
+    AttentionSelection, AttentionSetupError, AttentionStageSize, AttentionTileSize,
+    AttentionTilingScheme,
+};
 use crate::{
     components::{
         AvailableLineSizes, batch::simple::SimpleBatchAttentionFamily,
@@ -27,4 +34,48 @@ impl Algorithm for BlackboxAcceleratedAlgorithm {
     fn filter_line_sizes(available_line_sizes: AvailableLineSizes) -> AvailableLineSizes {
         Self::TileAttention::filter_line_sizes(available_line_sizes)
     }
+
+    fn selection<R: cubecl_core::Runtime>(
+        _client: &ComputeClient<R>,
+        problem: &AttentionProblem,
+        plane_dim: u32,
+        _line_sizes: &AttentionLineSizes,
+        _dtypes: &AttentionElems,
+    ) -> Result<AttentionSelection, AttentionSetupError> {
+        #[cfg(target_os = "macos")]
+        let tile_size = AttentionTileSize {
+            seq_q: 8,
+            head_dim: 8,
+            seq_kv: 8,
+            val_dim: 8,
+        };
+        #[cfg(not(target_os = "macos"))]
+        let tile_size = AttentionTileSize {
+            seq_q: 16,
+            head_dim: 16,
+            seq_kv: 16,
+            val_dim: 16,
+        };
+
+        assert!(problem.head_dim as u32 % tile_size.head_dim == 0);
+        let partition_head_dim = problem.head_dim as u32 / tile_size.head_dim;
+        let partition_val_dim = partition_head_dim;
+
+        Ok(AttentionSelection {
+            hypercube_selection: HypercubeSelection {},
+            tiling_scheme: AttentionTilingScheme {
+                tile_size,
+                partition_size: AttentionPartitionSize {
+                    seq_q: 1,
+                    head_dim: partition_head_dim,
+                    seq_kv: 1,
+                    val_dim: partition_val_dim,
+                },
+                stage_size: AttentionStageSize { seq_q: 1 },
+            },
+            plane_dim,
+            reuse_key_value: false,
+            two_rows_in_array_tile: false,
+        })
+    }
 }

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

@@ -1,7 +1,14 @@
+use cubecl_core::client::ComputeClient;
 use cubecl_matmul::components::{global::PartitionedStageFamily, stage::StridedStageFamily};
 
+use crate::components::batch::HypercubeSelection;
 use crate::components::stage::unit::UnitPartitionStageAttentionFamily;
 use crate::components::tile::unit_register::UnitRegisterTileAttention;
+use crate::components::{
+    AttentionElems, AttentionLineSizes, AttentionPartitionSize, AttentionProblem,
+    AttentionSelection, AttentionSetupError, AttentionStageSize, AttentionTileSize,
+    AttentionTilingScheme,
+};
 use crate::{
     components::{
         batch::simple::SimpleBatchAttentionFamily, global::simple::SimpleGlobalAttentionFamily,
@@ -21,4 +28,40 @@ impl Algorithm for UnitAlgorithm {
     >;
     type GlobalAttention = SimpleGlobalAttentionFamily<Self::StageAttention>;
     type BatchAttention = SimpleBatchAttentionFamily<Self::GlobalAttention>;
+
+    fn selection<R: cubecl_core::Runtime>(
+        _client: &ComputeClient<R>,
+        problem: &AttentionProblem,
+        plane_dim: u32,
+        _line_sizes: &AttentionLineSizes,
+        _dtypes: &AttentionElems,
+    ) -> Result<AttentionSelection, AttentionSetupError> {
+        let tile_size = AttentionTileSize {
+            seq_q: 4,
+            head_dim: 4,
+            seq_kv: 4,
+            val_dim: 4,
+        };
+
+        assert!(problem.head_dim as u32 % tile_size.head_dim == 0);
+        let partition_head_dim = problem.head_dim as u32 / tile_size.head_dim;
+        let partition_val_dim = partition_head_dim;
+
+        Ok(AttentionSelection {
+            hypercube_selection: HypercubeSelection {},
+            tiling_scheme: AttentionTilingScheme {
+                tile_size,
+                partition_size: AttentionPartitionSize {
+                    seq_q: 1,
+                    head_dim: partition_head_dim,
+                    seq_kv: 1,
+                    val_dim: partition_val_dim,
+                },
+                stage_size: AttentionStageSize { seq_q: plane_dim },
+            },
+            plane_dim,
+            reuse_key_value: false,
+            two_rows_in_array_tile: false,
+        })
+    }
 }