Flash Attention: test and fix f16 (#1074)

Author: louisfd

crates/cubecl-attention/Cargo.toml

@@ -15,7 +15,16 @@ default = ["std", "cubecl-runtime/default", "cubecl-core/default"]
 export_tests = ["pretty_assertions"]
 std = ["cubecl-runtime/std", "cubecl-core/std"]
 
-attention_tests = []
+attention_tests_f16 = []
+attention_tests_f32 = []
+attention_tests_unit = []
+attention_tests_blackbox_accelerated = []
+attention_tests_all = [
+    "attention_tests_f16",
+    "attention_tests_f32",
+    "attention_tests_unit",
+    "attention_tests_blackbox_accelerated",
+]
 
 [dependencies]
 bytemuck = { workspace = true }

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

@@ -15,11 +15,13 @@ use std::fmt::Debug;
 use std::hash::Hash;
 
 /// Logits below this are considered masked (effectively -inf)
-pub(crate) const LOGIT_MASKED: f32 = -1e5;
+/// Value chosen to fit within f16 range (~-65,504 max)
+pub(crate) const LOGIT_MASKED: f32 = -6e4;
 
 /// Any value smaller than this is considered numerically zero
 /// (used for fully-masked rows or tiny contributions)
-pub(crate) const FULLY_MASKED_ROW_THRESHOLD: f32 = 1e-7;
+/// Value chosen to be above f16 smallest normal (~6.1e-5)
+pub(crate) const FULLY_MASKED_ROW_THRESHOLD: f32 = 1e-4;
 
 #[cube]
 pub trait TileAttention<AP: AttentionPrecision>: Send + Sync + 'static {

crates/cubecl-attention/src/tests/macros/mod.rs

@@ -7,7 +7,7 @@ use crate::{
         AttentionProblem, AttentionSelection, AttentionTilingScheme, batch::HypercubeSelection,
     },
     kernels::Algorithm,
-    tests::attention_test_launcher::test_attention_algorithm,
+    tests::{attention_test_launcher::test_attention_algorithm, test_utils::TestPrecision},
 };
 
 #[derive(Default)]
@@ -38,7 +38,7 @@ pub mod tiling_scheme_ops {
     }
 }
 
-pub fn attention_test_launch<A: Algorithm, R: Runtime>(
+pub fn attention_test_launch<A: Algorithm, P: TestPrecision, R: Runtime>(
     client: ComputeClient<R>,
     tiling_scheme: AttentionTilingScheme,
     problem: AttentionProblem,
@@ -52,16 +52,15 @@ pub fn attention_test_launch<A: Algorithm, R: Runtime>(
         two_rows_in_array_tile: test_options.two_rows_in_array_tile,
     };
 
-    test_attention_algorithm::<A, (f32, f32), R>(client, problem, selection);
-    // test_attention_algorithm::<A, (half::f16, half::f16), R>(client, problem, selection);
+    test_attention_algorithm::<A, P, R>(client, problem, selection);
 }
 
 #[macro_export]
 macro_rules! testgen_attention {
     () => {
         use super::*;
 
-        #[cfg(feature = "attention_tests")]
+        #[cfg(feature = "attention_tests_unit")]
         mod attention_unit {
             type Algorithm = cubecl_attention::kernels::unit::UnitAlgorithm;
             const TILE_SIZE: cubecl_attention::components::AttentionTileSize =
@@ -73,10 +72,10 @@ macro_rules! testgen_attention {
                 };
             const STAGE_Q_BASE: u32 = 32;
 
-            $crate::testgen_attention_suite!();
+            $crate::testgen_attention_precision!();
         }
 
-        #[cfg(feature = "attention_tests")]
+        #[cfg(feature = "attention_tests_blackbox_accelerated")]
         mod attention_blackbox_accelerated {
             type Algorithm =
                 cubecl_attention::kernels::blackbox_accelerated::BlackboxAcceleratedAlgorithm;
@@ -98,7 +97,28 @@ macro_rules! testgen_attention {
                 };
             const STAGE_Q_BASE: u32 = 1;
 
-            $crate::testgen_attention_suite!();
+            $crate::testgen_attention_precision!();
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! testgen_attention_precision {
+    () => {
+        use super::*;
+
+        #[cfg(feature = "attention_tests_f16")]
+        mod f16_ty {
+            use super::*;
+
+            $crate::testgen_attention_suite!((half::f16, half::f16));
+        }
+
+        #[cfg(feature = "attention_tests_f32")]
+        mod f32_ty {
+            use super::*;
+
+            $crate::testgen_attention_suite!((f32, f32));
         }
     };
 }