From 3fed1835643008d21a6896d02201bef82b63ad2f Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 27 Jun 2026 14:51:56 +0300 Subject: [PATCH] update sdnq atten min sizes --- modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/attention.py b/modules/attention.py index ce888d661..a8669d2e6 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -24,7 +24,7 @@ def set_sdnq_attention(): sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention @wraps(sdpa_pre_sdnq_atten) def sdpa_sdnq_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 not is_causal and query.shape[-1] >= 32 and query.shape[-2] >= 256 and query.shape[-3] > 1: + if not is_causal and (query.shape[-1] >= 32 and key.shape[-1] >= 32 and value.shape[-1] >= 32) and (query.shape[-2] >= 512 or key.shape[-2] >= 512) and query.shape[-3] > 1: return sdnq_triton_atten( query=query, key=key, value=value, attn_mask=attn_mask, scale=scale, enable_gqa=enable_gqa,