From 68f00d85ca30e0ae657b2c50427894eff6aab350 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 17 May 2026 08:44:29 +0200 Subject: [PATCH] attention execution guard against cpu tensors Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + modules/attention.py | 30 +++++++++++++++++++++++++++--- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c0faa985..31afb545b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/modules/attention.py b/modules/attention.py index 7b3ac8e13..2287ca90e 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -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)