comfyanonymous/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-17 09:08:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Try to fix flash attention related issue on AMD. (#14880)

Browse Source
This commit is contained in:
comfyanonymous
2026-07-11 17:16:40 -07:00
committed by GitHub
parent f3a36e7484
commit 69ea58697b
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
return out
try:
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
@torch.library.custom_op("comfy::flash_attn", mutates_args=())
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale