replace enable_gpa with repeat_interleave for krea2

[yiyixuxu]

fix #14518

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Thanks! Just flagging the following table:

Claude derived it by using the following code and running it against this branch and main:

from diffusers import Krea2Transformer2DModel

device, dtype = "cuda", torch.bfloat16

model = Krea2Transformer2DModel.from_pretrained(
    "krea/Krea-2-Raw", subfolder="transformer", torch_dtype=dtype
).to(device)
model.eval()

# 1024x1024 image -> 128x128 latents -> 64x64 grid of patch-2 tokens;
# 512 text positions with only the first 77 valid (padding mask exercised).
torch.manual_seed(0)
batch, text_seq_len, grid = 1, 512, 64
image_seq_len = grid * grid

hidden_states = torch.randn(batch, image_seq_len, 64, device=device, dtype=dtype)
encoder_hidden_states = torch.randn(batch, text_seq_len, 12, 2560, device=device, dtype=dtype)
timestep = torch.tensor([0.7], device=device)
encoder_attention_mask = torch.zeros(batch, text_seq_len, dtype=torch.bool, device=device)
encoder_attention_mask[:, :77] = True

text_ids = torch.zeros(text_seq_len, 3, device=device)
image_ids = torch.zeros(grid, grid, 3, device=device)
image_ids[..., 1] = torch.arange(grid, device=device)[:, None]
image_ids[..., 2] = torch.arange(grid, device=device)[None, :]
position_ids = torch.cat([text_ids, image_ids.reshape(-1, 3)], dim=0)

with torch.no_grad():
    out = model(
        hidden_states=hidden_states,
        encoder_hidden_states=encoder_hidden_states,
        timestep=timestep,
        position_ids=position_ids,
        encoder_attention_mask=encoder_attention_mask,
    ).sample

I think this is acceptable because now the kernel invocations internally have changed.

[sayakpaul]

Either pass enable_gqa=True to dispatch_attention_fn

Can we also supplement a model file here for the reference?

[sayakpaul]

Turns on whether the model passes a mask.

Do we mean performance is turned on?

[asomoza]

thanks! this fixes the issue.

Not sure if it's worth to mention, but this changes the resulting image a little:

main	PR

But still, it doesn't make it worse or better, and the change is minimal and mostly imperceptible if you're not comparing both.

[dxqb]

commit message:

It also unpins the model from the native backend, since cuDNN, flash, FA3, sage
and the hub kernels all raise on enable_gqa.

the commit message isn't entirely correct: cuDNN worked with gqa and mask, and still does. I checked if this is a regression for cuDNN but it performs the same without gqa and the interleaving this PR introduces.
So not an issue, but cuDNN can still outperform SDPA for Krea2:

batch size 2 with mask

  │                    │ interleave + cuDNN │ enable_gqa + cuDNN │  interleave + SDPA   │
  ├────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┤
  │ attention backend  │ cuDNN                      │ cuDNN                      │ mem-efficient (60/64 calls) │
  ├────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┤
  │ attention subtotal │ 238.5 ms                   │ 231.2 ms                   │ 441.6 ms                    │
  ├────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┤
  │ total kernel       │ 1844.4 ms                  │ 1838.8 ms                  │ 1972.2 ms                   │
  ├────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┤
  │ reserved VRAM      │ 12.89 GiB                  │ 12.90 GiB                  │ 13.46 GiB                   │
  └────────────────────┴────────────────────────────┴────────────────────────────┴─────────────────────────────┘

[sayakpaul]

Yes, cuDNN is very underrated IMO.

[yiyixuxu]

thanks @dxqb for pointing out the mistake, i updated our docs and comments here #14561