From b3f66f3bd5b24aa974204c81764e9b91bed57623 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 18 Aug 2026 02:54:00 +0000 Subject: [PATCH 1/4] Krea 2: repeat key/value heads instead of `enable_gqa` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Krea 2 always attends with a text padding mask, and no fused SDPA kernel takes a mask together with mismatched query/key head counts — flash rejects the mask, the memory-efficient kernel rejects the mismatch. Attention therefore fell back to the math backend, which materializes the full [batch_size, num_heads, seq_len, seq_len] score matrix with no error or warning. Repeating the key/value heads in the processor computes the same thing and keeps the memory-efficient kernel eligible: at 1024x1024 (48/12 heads, 4608 tokens) one attention call goes from 9.02 GiB / 26.7 ms to 0.16 GiB / 4.1 ms. It also unpins the model from the native backend, since cuDNN, flash, FA3, sage and the hub kernels all raise on `enable_gqa`. Fixes #14518 Co-Authored-By: Claude Opus 5 --- src/diffusers/models/transformers/transformer_krea2.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_krea2.py b/src/diffusers/models/transformers/transformer_krea2.py index d1f6cd0ecded..0b2d1d43fadb 100644 --- a/src/diffusers/models/transformers/transformer_krea2.py +++ b/src/diffusers/models/transformers/transformer_krea2.py @@ -74,12 +74,20 @@ 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 + # [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`. + 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) + value = value.repeat_interleave(num_key_value_groups, dim=2) + hidden_states = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, - enable_gqa=attn.num_heads != attn.num_kv_heads, backend=self._attention_backend, parallel_config=self._parallel_config, ) From 4b89f7b65d411477fe27348c1efd48ff50bebec3 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 18 Aug 2026 02:54:06 +0000 Subject: [PATCH 2/4] docs: how to choose between `enable_gqa` and repeating key/value heads Co-Authored-By: Claude Opus 5 --- .ai/models.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/.ai/models.md b/.ai/models.md index b30f43f21db3..3df1e9450ba0 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -75,6 +75,27 @@ What you pass as `attn_mask=` to `dispatch_attention_fn` determines which backen - **Other mask types (structural, BlockMask, etc.)** — if the model requires a different mask pattern, figure out how to support as many backends as possible (e.g. use `window_size` kwarg for sliding window on flash, `BlockMask` for Flex) and document which backends are supported for that model. - **Don't declare `attention_mask` (or `encoder_hidden_states_mask`) in the forward signature if you ignore it.** "For API stability with other transformers" is not a reason; readers assume a declared param is honored, and downstream pipelines will pass padding masks that silently get dropped. Some existing models in the repo carry unused mask params for historical reasons — e.g. `QwenDoubleStreamAttnProcessor2_0.__call__` declares `encoder_hidden_states_mask` but never reads it (the joint mask is routed through `attention_mask` instead), and the block-level forward in `transformer_qwenimage.py` declares it but always receives `None`. This is a legacy behavior and should not be replicated in new models. +### Grouped-query attention + +Fewer key/value heads than query heads can be spelled two ways. Either pass `enable_gqa=True` to `dispatch_attention_fn` and let the backend broadcast, or repeat the key/value heads in the processor after RoPE and pass no flag (`transformer_krea2.py`, `transformer_nucleusmoe_image.py`): + +```python +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) + value = value.repeat_interleave(num_key_value_groups, dim=2) +``` + +`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 say in the PR which you picked and why. + +- **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. + +- **Performance.** Turns 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. + +- **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. + ## Model class attributes Each `ModelMixin` subclass can declare class-level attributes that configure optimization features. Each attribute corresponds to a user-facing API — the attribute controls how that feature behaves for the model. When adding a new transformer, set all that apply — skim `transformer_flux.py`, `transformer_wan.py`, `transformer_qwenimage.py` for examples. From ed14a5d62b78e414f0dc496c503e80c877c23417 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 18 Aug 2026 02:55:12 +0000 Subject: [PATCH 3/4] docs: recommend a choice rather than just reporting one Co-Authored-By: Claude Opus 5 --- .ai/models.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.ai/models.md b/.ai/models.md index 3df1e9450ba0..d39fe03b6352 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -88,7 +88,7 @@ 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 say in the PR which you picked and why. +Both compute the same thing, so weigh the two on compatibility and performance and recommend whichever fits the model better. - **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. From cf84016fd083b7c2952cadc476e67bbf8880c553 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 18 Aug 2026 19:38:33 +0000 Subject: [PATCH 4/4] docs: add enable_gqa reference model, reword performance note Co-Authored-By: Claude Fable 5 --- .ai/models.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.ai/models.md b/.ai/models.md index d39fe03b6352..d273f54fff3b 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -77,7 +77,7 @@ What you pass as `attn_mask=` to `dispatch_attention_fn` determines which backen ### Grouped-query attention -Fewer key/value heads than query heads can be spelled two ways. Either pass `enable_gqa=True` to `dispatch_attention_fn` and let the backend broadcast, or repeat the key/value heads in the processor after RoPE and pass no flag (`transformer_krea2.py`, `transformer_nucleusmoe_image.py`): +Fewer key/value heads than query heads can be spelled two ways. Either pass `enable_gqa=True` to `dispatch_attention_fn` and let the backend broadcast (`transformer_cosmos3.py`), or repeat the key/value heads in the processor after RoPE and pass no flag (`transformer_krea2.py`): ```python num_key_value_groups = attn.num_heads // attn.num_kv_heads @@ -92,7 +92,7 @@ Both compute the same thing, so weigh the two on compatibility and performance a - **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. -- **Performance.** Turns 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. +- **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. - **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.