Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

attention execution guard against cpu tensors

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-05-17 08:44:29 +02:00
parent d8433e53cd
commit 68f00d85ca
2 changed files with 28 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+1
View File
@@ -26,6 +26,7 @@
- `diffusers` patch custom pipelines for `qk_norm`
- *GHSA* fixes, thanks @SSJCorpSec for reporting
- custom `vae` loader
- `attention` execution guard against `cpu` tensors
## Update for 2026-05-13
modules/attention.py
+27 -3
View File
@@ -25,7 +25,15 @@ def set_triton_flash_attention(backend: str):
sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_triton_flash_atten)
def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
use_triton = (
query.shape[-1] <= 128
and attn_mask is None
and query.dtype != torch.float32
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_triton:
if scale is None:
scale = query.shape[-1] ** (-0.5)
head_size_og = query.size(3)
@@ -101,7 +109,15 @@ def set_ck_flash_attention(backend: str, device: torch.device):
sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_flash_atten)
def sdpa_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
use_flash = (
query.shape[-1] <= 128
and attn_mask is None
and query.dtype != torch.float32
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_flash:
is_unsqueezed = False
if query.dim() == 3:
query = query.unsqueeze(0)
@@ -168,7 +184,15 @@ def set_sage_attention(backend: str, device: torch.device):
sdpa_pre_sage_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sage_atten)
def sdpa_sage_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
if (query.shape[-1] in {128, 96, 64}) and (attn_mask is None) and (query.dtype != torch.float32):
use_sage = (
query.shape[-1] in {128, 96, 64}
and attn_mask is None
and query.dtype != torch.float32
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_sage:
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)