mirror of
https://github.com/NixOS/nixpkgs.git
synced 2026-07-18 22:51:21 +00:00
f4adbf877d
Diff: https://github.com/Dao-AILab/flash-attention/compare/v2.8.2...v2.8.3 Changelog: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.8.3
63 lines
2.7 KiB
Diff
63 lines
2.7 KiB
Diff
diff --git a/flash_attn/cute/ampere_helpers.py b/flash_attn/cute/ampere_helpers.py
|
|
index 839f407..9f74824 100644
|
|
--- a/flash_attn/cute/ampere_helpers.py
|
|
+++ b/flash_attn/cute/ampere_helpers.py
|
|
@@ -5,6 +5,15 @@ import cutlass
|
|
import cutlass.cute as cute
|
|
|
|
|
|
+# Max dynamic shared memory per SM, in bytes (architectural max minus 1 KB reserved for the driver).
|
|
+# Mirrors the values previously exposed by cutlass.utils.ampere_helpers, which was removed in nvidia-cutlass-dsl 4.3.
|
|
+SMEM_CAPACITY = {
|
|
+ "sm80": (164 - 1) * 1024,
|
|
+ "sm86": (100 - 1) * 1024,
|
|
+ "sm89": (100 - 1) * 1024,
|
|
+}
|
|
+
|
|
+
|
|
def get_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> cute.ComposedLayout:
|
|
dtype_byte = cutlass.const_expr(dtype.width // 8)
|
|
bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)
|
|
diff --git a/flash_attn/cute/flash_bwd.py b/flash_attn/cute/flash_bwd.py
|
|
index 79f5ee8..be008b1 100644
|
|
--- a/flash_attn/cute/flash_bwd.py
|
|
+++ b/flash_attn/cute/flash_bwd.py
|
|
@@ -11,7 +11,6 @@ import cuda.bindings.driver as cuda
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass.cute.nvgpu import cpasync, warp
|
|
-import cutlass.utils.ampere_helpers as sm80_utils_basic
|
|
|
|
from flash_attn.cute import ampere_helpers as sm80_utils
|
|
from flash_attn.cute import utils
|
|
@@ -125,7 +124,7 @@ class FlashAttentionBackwardSm80:
|
|
smem_usage_V = n_block_size * head_dim_v * 2
|
|
smem_usage_QV = (smem_usage_Q + smem_usage_V) if not V_in_regs else max(smem_usage_Q, smem_usage_V)
|
|
smem_usage = smem_usage_QV + smem_usage_dO + smem_usage_K
|
|
- smem_capacity = sm80_utils_basic.SMEM_CAPACITY["sm80"]
|
|
+ smem_capacity = sm80_utils.SMEM_CAPACITY["sm80"]
|
|
if smem_usage > smem_capacity:
|
|
return False
|
|
return True
|
|
diff --git a/flash_attn/cute/flash_fwd.py b/flash_attn/cute/flash_fwd.py
|
|
index ddd5cfc..cf5ac81 100644
|
|
--- a/flash_attn/cute/flash_fwd.py
|
|
+++ b/flash_attn/cute/flash_fwd.py
|
|
@@ -16,7 +16,6 @@ import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass import const_expr
|
|
from cutlass.cute.nvgpu import cpasync, warp, warpgroup
|
|
-import cutlass.utils.ampere_helpers as sm80_utils_basic
|
|
import cutlass.utils.hopper_helpers as sm90_utils_basic
|
|
|
|
from flash_attn.cute import ampere_helpers as sm80_utils
|
|
@@ -127,7 +126,7 @@ class FlashAttentionForwardBase:
|
|
smem_usage_QV = (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V)
|
|
smem_usage = smem_usage_QV + smem_usage_K
|
|
# TODO: sm86 and sm89
|
|
- smem_capacity = sm80_utils_basic.SMEM_CAPACITY["sm80"]
|
|
+ smem_capacity = sm80_utils.SMEM_CAPACITY["sm80"]
|
|
if smem_usage > smem_capacity:
|
|
return False
|
|
# Check if twice the block size is divisible by the number of threads
|