[Krea-2] `enable_gqa` + `attn_mask` produces high usage of VRAM

### Describe the bug The current implementation uses `enable_gqa` + `attn_mask` which silently falls back to the SDPA math backend. This makes that using a 4-bit quantization with a 5090 OOM on a simple 1920x1080 image and a 16GB GPU to OOM on 832x832, on windows this will make it a lot slower because it will use RAM and not OOM. Every other backend checks this and alerts of this. I won't open a PR because I'm not confident I can follow a review, but claude did make an easy and simple fix to run tests. ### Reproduction ```python import torch import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel q = torch.randn(1, 48, 4608, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn(1, 12, 4608, 128, device="cuda", dtype=torch.bfloat16) v = torch.randn(1, 12, 4608, 128, device="cuda", dtype=torch.bfloat16) mask = torch.ones(1, 1, 1, 4608, device="cuda", dtype=torch.bool) def peak(fn): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() fn() return (torch.cuda.max_memory_allocated() - base) / 2**30 def as_is(): # what `_native_attention` does today return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=True) def expanded(): # expanding K/V instead, as `_xformers_attention` already does. Done inside the call, so the # cost of the repeated K/V tensors is included in the measurement. ke = k.repeat_interleave(4, dim=1) ve = v.repeat_interleave(4, dim=1) return F.scaled_dot_product_attention(q, ke, ve, attn_mask=mask) print(f"enable_gqa=True + mask : {peak(as_is):5.2f} GiB") print(f"expanded K/V + mask : {peak(expanded):5.2f} GiB") # every fused kernel is disqualified, so only `math` survives for backend in (SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION): try: with sdpa_kernel(backend): F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=True) print(f"{backend.name:20s}: ran") except RuntimeError as e: print(f"{backend.name:20s}: {str(e).splitlines()[0]}") ``` Peak memory of the attention call: | resolution | seq | without fix | with fix | | |---|---|---|---|---| | 512x512 | 1536 | 1.12 GiB | 0.05 GiB | 21x | | 768x768 | 2816 | 3.48 GiB | 0.10 GiB | 35x | | 1024x1024 | 4608 | 9.02 GiB | 0.16 GiB | **57x** | Time per attention call: | resolution | seq | without fix | with fix | | |---|---|---|---|---| | 512x512 | 1536 | 13.94 ms | 1.64 ms | 8.5x | | 768x768 | 2816 | 48.05 ms | 5.59 ms | 8.6x | | 1024x1024 | 4608 | 118.22 ms | 13.67 ms | **8.6x** | ### System Info diffusers from main, 5090 on linux and mobile 4090 on windows ### Who can help? _No response_