From dc0944ee0faedd1451173846c416c243e3de5b7b Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 21 Aug 2026 20:41:35 +0000 Subject: [PATCH] docs: correct the enable_gqa / cuDNN explanation The Krea 2 PR (#14523) said no fused SDPA kernel takes a mask together with mismatched head counts and that cuDNN raises on `enable_gqa`. Neither is right: torch's cuDNN kernel accepts both, and the `enable_gqa` raises in `attention_dispatch.py` are in the context-parallel ops, not the backends. What actually happens is that torch tries its kernels in a priority order that changes across versions and GPUs, and whenever math comes before cuDNN an `enable_gqa` + mask call lands on math and materializes the full score matrix. Rewrite the GQA section of models.md around that (kernel compatibility table, why to repeat when there is a mask) and fix the Krea 2 processor comment to match. Co-Authored-By: Claude Fable 5 --- .ai/references/models.md | 13 +++++++++---- .../models/transformers/transformer_krea2.py | 7 ++++--- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/.ai/references/models.md b/.ai/references/models.md index 625d9c4891fb..559b445ba111 100644 --- a/.ai/references/models.md +++ b/.ai/references/models.md @@ -87,13 +87,18 @@ if num_key_value_groups > 1: `dim=2` because tensors are `(batch_size, seq_len, num_heads, head_dim)` here. Must be `repeat_interleave`, not `repeat` — the groups are contiguous, and `repeat` gives a silently wrong pairing no shape check catches. -Both compute the same thing, so weigh the two on compatibility and performance and recommend whichever fits the model better. +Both compute the same thing. What differs is which kernel runs. Under the default backend (`native`, a plain `F.scaled_dot_product_attention`), torch has four kernels: -- **Compatibility.** Most backends do not implement `enable_gqa` yet — flash, FA3, sage, cuDNN and the hub kernels raise on it, as does the context-parallel path. Grep `enable_gqa` in `attention_dispatch.py` for the current list rather than trusting this one; it changes as support lands. The flag limits the model to whichever backends still accept it, while repeating works on all of them. +| kernel | mask | mismatched q/kv heads (`enable_gqa`) | +|---|---|---| +| flash | ✗ | ✓ | +| efficient | ✓ | ✗ | +| math | ✓ | ✓ — materializes the full `[batch_size, num_heads, seq_len_q, seq_len_kv]` score matrix | +| cuDNN | ✓ | ✓ | -- **Performance.** Depends on whether the model passes a mask. With a mask, no fused kernel takes a mask *and* mismatched head counts, so SDPA falls back to math and materializes the full `[batch_size, num_heads, seq_len_q, seq_len_kv]` score matrix — no error, no warning, only memory. Without a mask, flash broadcasts inside the kernel and the flag saves the key/value copy. Both effects scale with sequence length and head count, so measure at the model's real shape; `torch.backends.cuda.can_use_flash_attention(params, debug=True)` and `can_use_efficient_attention` print why a kernel was rejected, which is the fastest way to see which one you actually got. +It tries them in a priority order and takes the first that accepts the call; the order changes across torch versions and GPUs. -- **Recommendation.** Repeat by default — it is portable and never pathological. Reach for `enable_gqa=True` only when the model never passes a mask *and* the measured saving justifies the narrower backend support. For scale: on Krea 2 at 1024×1024, masked, the flag cost 9.02 GiB and 26.7 ms per call against 0.16 GiB and 4.1 ms repeated; unmasked at the same shape it saved 0.11 GiB and 0.1 ms. `transformer_cosmos3.py` is the in-repo case where it is defensible — causal, never masked. +If the model passes a mask, repeat the key/value heads: (1) `enable_gqa` is rejected by the context-parallel path, and (2) with a mask it can only land on math or cuDNN — if math comes first, it silently materializes the full score matrix. ## Model class attributes diff --git a/src/diffusers/models/transformers/transformer_krea2.py b/src/diffusers/models/transformers/transformer_krea2.py index 0b2d1d43fadb..55d275e5dca7 100644 --- a/src/diffusers/models/transformers/transformer_krea2.py +++ b/src/diffusers/models/transformers/transformer_krea2.py @@ -74,10 +74,11 @@ def __call__( query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) - # Krea 2 always attends with a text padding mask, and no fused attention kernel handles grouped-query - # attention together with a mask — SDPA would fall back to its math backend and materialize the full + # Krea 2 always attends with a text padding mask. Of torch's SDPA kernels, only math and cuDNN accept a mask + # together with `enable_gqa`, and when math is tried first it materializes the full # [batch_size, num_heads, seq_len, seq_len] attention matrix. Repeat the key/value heads here instead: the - # result is identical, and it keeps every attention backend usable since they all reject `enable_gqa`. + # result is identical, every kernel accepts it, and the context-parallel path (which rejects `enable_gqa`) + # keeps working. num_key_value_groups = attn.num_heads // attn.num_kv_heads if num_key_value_groups > 1: key = key.repeat_interleave(num_key_value_groups, dim=2)