wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-10 13:32:36 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Add support for flash attention 3

Browse Source
This commit is contained in:
simonri 2025-11-16 16:36:48 +00:00
parent 7d6103325e
commit 603a721405
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
comfy/ldm/modules/attention.py
View File

@ -32,7 +32,7 @@ except ImportError as e:
FLASH_ATTENTION_IS_AVAILABLE = False FLASH_ATTENTION_IS_AVAILABLE = False
try: try:
from flash_attn import flash_attn_func from flash_attn_interface import flash_attn_func
FLASH_ATTENTION_IS_AVAILABLE = True FLASH_ATTENTION_IS_AVAILABLE = True
except ImportError: except ImportError:
if model_management.flash_attention_enabled(): if model_management.flash_attention_enabled():
@ -565,7 +565,7 @@ try:
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=()) @torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor: dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal) return flash_attn_func(q, k, v, causal=causal)
@flash_attn_wrapper.register_fake @flash_attn_wrapper.register_fake