Flash Attention: strengthen test suite (#1077)

Author: louisfd

crates/cubecl-attention/src/base.rs

@@ -127,15 +127,19 @@ pub fn launch_attention<R: Runtime, A: Algorithm>(
         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: 1,
+                head_dim: partition_head_dim,
                 seq_kv: 1,
-                val_dim: 1,
+                val_dim: partition_val_dim,
             },
             stage_size: AttentionStageSize { seq_q: 1 },
         },

crates/cubecl-attention/src/components/global/simple/reader/query.rs

@@ -41,21 +41,20 @@ impl<AP: AttentionPrecision> QueryReader<AP> {
 
         let line_size = self.gmem_config.line_size;
 
+        let tile_head_dim = attention_tile_size.head_dim;
+
         let slice = self
             .query
             .slice(
-                (
-                    row * attention_tile_size.seq_q,
-                    col * attention_tile_size.head_dim,
-                ),
-                (attention_tile_size.seq_q, attention_tile_size.head_dim).runtime(),
+                (row * attention_tile_size.seq_q, col * tile_head_dim),
+                (attention_tile_size.seq_q, tile_head_dim).runtime(),
             )
             .to_linear_slice();
 
         let start = 0;
-        let length = attention_tile_size.seq_q * attention_tile_size.head_dim / line_size;
+        let length = attention_tile_size.seq_q * tile_head_dim / line_size;
         let end = start + length;
-        let stride = partition_head_dim * attention_tile_size.head_dim / line_size;
+        let stride = partition_head_dim * tile_head_dim / line_size;
 
         StridedTile::<QG<AP>>::new_strided(
             slice,

crates/cubecl-attention/src/components/stage/base.rs

@@ -179,14 +179,29 @@ impl<TC: TileAttentionConfig> PartitionAttentionConfig<TC> {
 
 pub fn validate<TC: TileAttentionConfig>(
     config: PartitionAttentionConfig<TC>,
+    problem: &AttentionProblem,
 ) -> Result<PartitionAttentionConfig<TC>, AttentionSetupError> {
     let tile_size = config.shared().tile_config.attention_tile_size();
     let partition_size = config.shared().partition_size;
 
-    if config.shared().reuse_key_value
-        && (tile_size.head_dim != tile_size.val_dim
-            || partition_size.head_dim != partition_size.val_dim)
-    {
+    if partition_size.head_dim * tile_size.head_dim != problem.head_dim as u32 {
+        return Err(AttentionSetupError::InvalidConfig(Box::new(
+            "Tiling scheme's total head dim must equal problem's head dim".to_string(),
+        )));
+    }
+
+    let head_val_different = tile_size.head_dim != tile_size.val_dim
+        || partition_size.head_dim != partition_size.val_dim;
+
+    if head_val_different {
+        return Err(AttentionSetupError::InvalidConfig(Box::new(
+            "Differing head dim and val dim is not yet supported".to_string(),
+        )));
+    }
+
+    // This check is stricter than the previous one, but the other may be removed
+    // eventually while this one will always remain true.
+    if config.shared().reuse_key_value && head_val_different {
         return Err(AttentionSetupError::InvalidConfig(Box::new(
         "When reusing key/value, head_dim must equal val_dim in both tile_size and partition_size."
             .to_string(),

crates/cubecl-attention/src/components/stage/plane/setup.rs

@@ -107,17 +107,20 @@ impl<
             num_stages: 1,
         };
 
-        validate(PartitionAttentionConfig::Plane(PlanePartitionStageConfig {
-            shared: SharedPartitionAttentionConfig {
-                tile_config,
-                partition_size: selection.tiling_scheme.partition_size,
-                stage_size: selection.tiling_scheme.stage_size,
-                reuse_key_value: selection.reuse_key_value,
-                num_planes,
-                key_smem_config,
-                value_smem_config,
-                out_smem_config,
-            },
-        }))
+        validate(
+            PartitionAttentionConfig::Plane(PlanePartitionStageConfig {
+                shared: SharedPartitionAttentionConfig {
+                    tile_config,
+                    partition_size: selection.tiling_scheme.partition_size,
+                    stage_size: selection.tiling_scheme.stage_size,
+                    reuse_key_value: selection.reuse_key_value,
+                    num_planes,
+                    key_smem_config,
+                    value_smem_config,
+                    out_smem_config,
+                },
+            }),
+            problem,
+        )
     }
 }

crates/cubecl-attention/src/components/stage/tile_ops/softmax/base.rs

@@ -16,7 +16,7 @@ use crate::components::tile::TileAttention;
 #[cube]
 /// Applies softmax to a tile with masking and updates the running state.
 ///
-/// Scales by `1 / sqrt(dk)`, applies the mask, computes row-wise max and sum,
+/// Scales by `1 / sqrt(head_dim)`, applies the mask, computes row-wise max and sum,
 /// exponentiates, and updates the softmax state.
 ///
 /// Returns the exponential difference used for normalization.
@@ -26,12 +26,12 @@ pub fn tile_softmax<AP: AttentionPrecision, TA: TileAttention<AP>, R: Reducer>(
     state: &mut RunningState<SM<AP>>,
     max_placeholder: &mut RowWise<SM<AP>>,
     sum_placeholder: &mut RowWise<SM<AP>>,
-    #[comptime] dk: u32,
+    #[comptime] head_dim: u32,
     #[comptime] config: TA::Config,
 ) -> RowWise<SM<AP>> {
     TA::SoftmaxRow::scale_and_mask::<MaskTile<AP, TA>>(
         rowwise_softmax,
-        SM::<AP>::new(comptime!(1.0 / (dk as f32).sqrt())),
+        SM::<AP>::new(comptime!(1.0 / (head_dim as f32).sqrt())),
         mask,
     );
 

crates/cubecl-attention/src/components/stage/unit/setup.rs

@@ -115,17 +115,20 @@ impl<
             num_stages: 1,
         };
 
-        validate(PartitionAttentionConfig::Unit(UnitPartitionStageConfig {
-            shared: SharedPartitionAttentionConfig {
-                tile_config,
-                partition_size: selection.tiling_scheme.partition_size,
-                stage_size: selection.tiling_scheme.stage_size,
-                reuse_key_value: selection.reuse_key_value,
-                num_planes,
-                key_smem_config,
-                value_smem_config,
-                out_smem_config,
-            },
-        }))
+        validate(
+            PartitionAttentionConfig::Unit(UnitPartitionStageConfig {
+                shared: SharedPartitionAttentionConfig {
+                    tile_config,
+                    partition_size: selection.tiling_scheme.partition_size,
+                    stage_size: selection.tiling_scheme.stage_size,
+                    reuse_key_value: selection.reuse_key_value,
+                    num_planes,
+                    key_smem_config,
+                    value_smem_config,
+                    out_smem_config,
+                },
+            }),
+            problem,
+        )
     }
 }