github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[OPS] enable flash_attention_v2 TMA (#2544)

Browse Source
This commit is contained in:
runseny
2023-10-26 14:31:17 +08:00
committed by GitHub
parent 2323adb387
commit 4c816c2f59
3 changed files with 48 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/test/regression/test_performance.py
View File

@@ -165,7 +165,7 @@ flash_attention_data = {
(4, 48, 4096, 64, False, False, 'forward', 'bfloat16'): 0.266,
(4, 48, 1024, 16, False, False, 'forward', 'float32'): 0.098,
(4, 48, 4096, 64, False, False, 'backward', 'float16'): 0.159,
(4, 48, 4096, 64, False, False, 'backward', 'bfloat16'): 0.136,
(4, 48, 4096, 64, False, False, 'backward', 'bfloat16'): 0.159,
(4, 48, 1024, 16, False, False, 'backward', 'float32'): 0.088,
}
}