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

[asomoza]

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 1920x1088 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

This will OOM on a 16GB GPU on Linux and use RAM on windows

import torch
from sdnq import SDNQConfig  # noqa: F401

from diffusers import Krea2Pipeline


pipe = Krea2Pipeline.from_pretrained("OzzyGT/Krea_2_Turbo_sdnq_dynamic_4bit", dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()

image = pipe(
    prompt="a photo of a cat holding a chalkboard sign that reads 'diffusers'",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=0.0,
    generator=torch.Generator("cpu").manual_seed(42),
).images[0]

image.save("krea2_no_fix.png")

print(f"peak allocated: {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")
print(f"peak reserved : {torch.cuda.max_memory_reserved() / 2**30:.2f} GiB")

peak allocated: 19.94 GiB
peak reserved: 20.78 GiB
inference time: OOM or 2201s (36.7 min)

Script with fix

This will work on windows with a 16GB GPU

import torch
from sdnq import SDNQConfig  # noqa: F401

from diffusers import Krea2Pipeline
from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry


# --- workaround for #14518: expand K/V so SDPA keeps the memory-efficient kernel ---
_orig_native_attention = _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE]


def _patched_native_attention(query, key, value, attn_mask=None, enable_gqa=False, **kwargs):
    if attn_mask is not None and enable_gqa and query.shape[2] != key.shape[2]:
        rep = query.shape[2] // key.shape[2]
        key = key.repeat_interleave(rep, dim=2)
        value = value.repeat_interleave(rep, dim=2)
        enable_gqa = False
    return _orig_native_attention(query, key, value, attn_mask=attn_mask, enable_gqa=enable_gqa, **kwargs)


_AttentionBackendRegistry._backends[AttentionBackendName.NATIVE] = _patched_native_attention
# --- end workaround ---

pipe = Krea2Pipeline.from_pretrained("OzzyGT/Krea_2_Turbo_sdnq_dynamic_4bit", dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()

image = pipe(
    prompt="a photo of a cat holding a chalkboard sign that reads 'diffusers'",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=0.0,
    generator=torch.Generator("cpu").manual_seed(42),
).images[0]

image.save("krea2.png")

print(f"peak allocated: {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")
print(f"peak reserved : {torch.cuda.max_memory_reserved() / 2**30:.2f} GiB")

peak allocated: 11.52 GiB
peak reserved : 12.94 GiB
inference time: 45s

System Info

diffusers from main, 5090 on linux and mobile 4090 on windows

Who can help?

No response

[furk4neg3]

Hi! I’d like to work on this issue.

I looked through the Krea2 attention path and the shared attention dispatcher, and I think this is better handled centrally in _native_attention rather than with a Krea2-specific workaround.

My proposed approach is to detect the masked GQA case (attn_mask is not None + valid GQA head layout), explicitly repeat the K/V heads to match the query heads, and then call SDPA with enable_gqa=False. This should keep fused SDPA backends eligible and avoid the expensive math fallback.

Unmasked GQA and explicit backend behavior would remain unchanged. There’s also precedent for this pattern in the repo from the Cosmos3 work.

If this direction sounds reasonable, I’d be happy to implement it with focused regression tests.

[yiyixuxu]

hi @furk4neg3

i opened PR here #14523 with the Krea-specific fix

enable_gqa isn't a diffusres concept - we want to align with what it means in pytorch. so would like to avoid patterns where user pass enable_gpa=True and then we silently turn it into False

[furk4neg3]

Thanks, that makes sense.

I had implemented the shared _native_attention approach locally with focused regression tests, but I understand the concern with changing the semantics of enable_gqa inside the dispatcher.

Keeping the dispatcher aligned with PyTorch and handling the K/V expansion explicitly in Krea2 is cleaner. I’ll drop my shared-dispatcher approach.

Thanks for clarifying the intended design.