{ai}[foss/2025a] flash-attention v3.0.0b1 w/ CUDA 12.8.0

[lexming]

(created using eb --new-pr)

Depends on:

• {ai}[foss/2025a] PyTorch v2.9.1, Triton v3.5.1 w/ CUDA 12.8.0 from wheels #25267
• {math}[GCCcore/14.2.0] einops v0.8.1 #25347

[github-actions]

Updated software flash-attention-3.0.0b1-foss-2025a-CUDA-12.8.0.eb

Diff against flash-attention-2.6.3-foss-2023a-CUDA-12.1.1.eb

easybuild/easyconfigs/f/flash-attention/flash-attention-2.6.3-foss-2023a-CUDA-12.1.1.eb

diff --git a/easybuild/easyconfigs/f/flash-attention/flash-attention-2.6.3-foss-2023a-CUDA-12.1.1.eb b/easybuild/easyconfigs/f/flash-attention/flash-attention-3.0.0b1-foss-2025a-CUDA-12.8.0.eb
index 892ee266ef..a68d2de8a5 100644
--- a/easybuild/easyconfigs/f/flash-attention/flash-attention-2.6.3-foss-2023a-CUDA-12.1.1.eb
+++ b/easybuild/easyconfigs/f/flash-attention/flash-attention-3.0.0b1-foss-2025a-CUDA-12.8.0.eb
@@ -1,38 +1,53 @@
 easyblock = 'PythonBundle'
 
 name = 'flash-attention'
-version = '2.6.3'
+version = '3.0.0b1'
+_commit = 'ef9e6a644192eb2b90155abe0372542f6d9a27b6'
 versionsuffix = '-CUDA-%(cudaver)s'
 
 homepage = 'https://github.com/Dao-AILab/flash-attention'
-description = """Fast and memory-efficient exact attention."""
+description = """Fast and Memory-Efficient Exact Attention with IO-Awareness.
+FlashAttention-3 is optimized for Hopper GPUs (e.g. H100)."""
 
-toolchain = {'name': 'foss', 'version': '2023a'}
+toolchain = {'name': 'foss', 'version': '2025a'}
 
-builddependencies = [('Ninja', '1.11.1')]
 dependencies = [
-    ('Python', '3.11.3'),
-    ('Python-bundle-PyPI', '2023.06'),
-    ('CUDA', '12.1.1', '', SYSTEM),
-    ('PyTorch-bundle', '2.1.2', versionsuffix),
-    ('einops', '0.7.0'),
-    ('CUTLASS', '3.4.0', versionsuffix),
+    ('CUDA', '12.8.0', '', SYSTEM),
+    ('Python', '3.13.1'),
+    ('PyTorch-bundle', '2.9.1', versionsuffix),
+    ('CUTLASS', '4.3.5', versionsuffix),
+    ('einops', '0.8.1'),
+    ('ninja-python', '1.11.1.4'),
 ]
 
+# use CUTLASS from EB
+_fa_fix_cutlass_src = "ln -s $EBROOTCUTLASS/include ../csrc/cutlass/include && "
+# avoid downloads
+_fa_build_env = "export FLASH_ATTENTION_OFFLINE_BUILD=TRUE && "
+# avoid using pre-build wheels
+_fa_build_env += "export FLASH_ATTENTION_FORCE_BUILD=TRUE && "
+# adjust parallelization
+_fa_build_env += "export NVCC_THREADS=1 MAX_JOBS=%(parallel)s && "
+
 exts_list = [
     (name, version, {
-        'modulename': 'flash_attn',
-        # solves Invalid cross-device link error
-        # https://github.com/Dao-AILab/flash-attention/issues/875
-        'preinstallopts': "sed -i 's/os.rename/shutil.move/' setup.py && ",
+        'modulename': 'flash_attn_3',
         'source_urls': ['https://github.com/Dao-AILab/flash-attention/archive/'],
-        'sources': ['v%(version)s.tar.gz'],
-        'checksums': ['136e149165d4c8c67273d16daa957b5cd5e6fc629061ffd39fa5a25224454d6c'],
+        'sources': [{'download_filename': f'{_commit}.tar.gz', 'filename': SOURCE_TAR_GZ}],
+        'patches': ['flash-attention-%(version)s_fix-cuda-linking.patch'],
+        'checksums': [
+            {'flash-attention-3.0.0b1.tar.gz':
+             'ad13a51cc6c23df4ef29c86dbe2cd91f443e81984cedda6f3addcbf9c35467f7'},
+            {'flash-attention-3.0.0b1_fix-cuda-linking.patch':
+             '88b216a88cce3d307b7cfc85e201e38b352408e1e0cd3d8b101e09013db476fc'},
+        ],
+        'start_dir': 'hopper',
+        'preinstallopts': _fa_fix_cutlass_src + _fa_build_env,
     }),
 ]
 
 sanity_check_commands = [
-    "python -c 'from flash_attn import flash_attn_qkvpacked_func, flash_attn_func'",
+    "python -c 'import flash_attn_interface'",
 ]
 
-moduleclass = 'lib'
+moduleclass = 'ai'