{ai}[foss/2025a] flash-attention v2.8.3 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
• {tools}[GCCcore/14.2.0] pyNVML v13.590.48 #25395

[github-actions]

Updated software flash-attention-2.8.3-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-2.8.3-foss-2025a-CUDA-12.8.0.eb
index 892ee266ef..ff7e1abdd8 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-2.8.3-foss-2025a-CUDA-12.8.0.eb
@@ -1,33 +1,51 @@
 easyblock = 'PythonBundle'
 
 name = 'flash-attention'
-version = '2.6.3'
+version = '2.8.3'
 versionsuffix = '-CUDA-%(cudaver)s'
 
 homepage = 'https://github.com/Dao-AILab/flash-attention'
 description = """Fast and memory-efficient exact attention."""
 
-toolchain = {'name': 'foss', 'version': '2023a'}
+toolchain = {'name': 'foss', 'version': '2025a'}
 
-builddependencies = [('Ninja', '1.11.1')]
+builddependencies = [
+    ('ninja-python', '1.11.1.4'),
+]
 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', f"{versionsuffix}-whl"),
+    ('CUTLASS', '4.3.5', versionsuffix),
+    ('einops', '0.8.1'),
 ]
 
+# default CUDA compute capabilities to use (override via --cuda-compute-capabilities)
+cuda_compute_capabilities = ['6.0', '7.0', '7.5', '8.0', '8.6', '9.0']
+
+# solves "invalid cross-device link" error
+# see https://github.com/Dao-AILab/flash-attention/issues/875
+_fa_fix_linking = "sed -i 's/os.rename/shutil.move/' setup.py && "
+# use CUTLASS from EB
+_fa_fix_cutlass_src = "ln -s $EBROOTCUTLASS/include csrc/cutlass/include && "
+# avoid using pre-built wheels
+_fa_build_env = "export FLASH_ATTENTION_FORCE_BUILD=TRUE && "
+# adjust CUDA compute capabilities
+_fa_build_env += " export FLASH_ATTN_CUDA_ARCHS='%(cuda_cc_cmake)s' && "
+# 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 && ",
         'source_urls': ['https://github.com/Dao-AILab/flash-attention/archive/'],
         'sources': ['v%(version)s.tar.gz'],
-        'checksums': ['136e149165d4c8c67273d16daa957b5cd5e6fc629061ffd39fa5a25224454d6c'],
+        'patches': ['flash-attention-2.8.3_respect-env-cuda-cc.patch'],
+        'checksums': [
+            '61cd5e91507ad7f04dc7c207d8bc8bfb1e43b56b806e51febbc27faeaee208ba',
+            'b2a4e2c2b7c57db19dad01b77ecf4ab821f5ee312a2f10ee2f0dac3a3487b59b',
+        ],
+        'preinstallopts': _fa_fix_linking + _fa_fix_cutlass_src + _fa_build_env,
     }),
 ]
 
@@ -35,4 +53,4 @@ sanity_check_commands = [
     "python -c 'from flash_attn import flash_attn_qkvpacked_func, flash_attn_func'",
 ]
 
-moduleclass = 'lib'
+moduleclass = 'ai'