From 9187dddfa3d723af8a8a3ac4b47689a56335ebe0 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 16:27:24 +0200 Subject: [PATCH] perf(krea2): keep the K/V heads unexpanded where a fused kernel takes them Krea-2 has 48 query heads over 12 K/V heads. This processor expands the K/V heads with repeat_interleave so that enable_gqa is not needed, because the fused SDPA kernels used not to support it -- and the alternative is the math backend, which materialises the full [heads, seq, seq] score matrix at ~9 GB. That premise has become build-dependent. cuDNN does serve grouped-query attention, mask included: measured on the real shape ([1, 48q/12kv, 4608, 128], bf16) at 3.53 ms masked and 3.64 ms unmasked, against the memory- efficient kernel refusing it outright and flash refusing the additive mask. So the expansion is now decided per call by asking the dispatcher instead of assuming, and skipped where a permitted kernel takes the shape as it is. Measured on an RTX 4090, live allocation at the SDPA call: seq 4608 expanded 216.0 MB -> grouped 135.5 MB 13.09 -> 12.44 ms seq 9216 expanded 432.0 MB -> grouped 271.0 MB 36.53 -> 33.10 ms The expansion allocates two tensors four times larger than the originals, +108 MB at 4608 tokens, every call. What it does *not* do is raise the peak of the whole attention block: that peak is set after SDPA by the output path, when the K/V tensors are already free. So this is less allocator traffic and a few percent of time, not a lower ceiling -- worth stating, because an isolated-kernel measurement (270 MB -> 197 MB) reads like the latter. Two things make the check conservative rather than optimistic, both because guessing wrong costs ~9 GB rather than some speed: - Only backends this call is *permitted* to reach are consulted. Asking about cuDNN while the resolved list excludes it -- which the `efficient` override does -- would answer for a kernel the dispatcher cannot use, and the call would fail with `No available kernel`. An existing CUDA test caught exactly that during development. - Anything unexpected -- a non-CUDA tensor, a torch whose SDPAParams signature differs -- answers no and expands, which is correct everywhere and merely costs memory. The answer depends only on the call shape, dtype, device, mask presence and permitted backends, all fixed for a given block within a generation, so it is cached per processor. Co-Authored-By: Claude Opus 5 (1M context) --- invokeai/backend/krea2/attention.py | 83 ++++++++- .../krea2/test_grouped_query_attention.py | 160 ++++++++++++++++++ 2 files changed, 237 insertions(+), 6 deletions(-) create mode 100644 tests/backend/krea2/test_grouped_query_attention.py diff --git a/invokeai/backend/krea2/attention.py b/invokeai/backend/krea2/attention.py index 2bbdb63ab7d..9a9bac29945 100644 --- a/invokeai/backend/krea2/attention.py +++ b/invokeai/backend/krea2/attention.py @@ -6,9 +6,16 @@ ``[heads, seq, seq]`` score matrix. At 1280x720 (3600 image tokens) that is ~5.7 GB **per attention**, and it grows O(seq^2) — ~40 GB at 2560x1440 — so generation OOMs or the cache offloads the transformer to RAM. -This processor instead expands the K/V heads to match the query heads (``repeat_interleave``) so ``enable_gqa`` -is not needed, and runs under a ranked list of fused SDPA backends (all of which support the additive padding -mask and are O(seq) in memory). Measured: the same 3600-token attention drops from ~5.7 GB to ~0.19 GB. +This processor runs under a ranked list of fused SDPA backends instead, which are O(seq) in memory. +Measured: the same 3600-token attention drops from ~5.7 GB to ~0.19 GB. + +The premise above has since become build-dependent. **cuDNN does serve grouped-query attention**, mask +included, so where it is available the K/V heads can stay at 12 and the expansion is pure waste -- +``repeat_interleave`` allocates two tensors four times larger than the originals, every call. Where no +fused kernel takes GQA (the memory-efficient kernel refuses it outright, flash refuses the additive +mask, and ROCm has no cuDNN at all) the expansion is still the only thing standing between this +processor and the ~5.7 GB math path. So it is now conditional, decided per call shape by asking the +dispatcher rather than by assuming -- see ``_serves_grouped_query_attention``. The math is otherwise identical to ``Krea2AttnProcessor`` (q/k RMSNorm, rotary embeddings, sigmoid output gate). """ @@ -123,6 +130,8 @@ def __init__( self.regional_prompting_state = regional_prompting_state # Resolved once per generation and handed down, not read per attention call. self.sdpa_backends = sdpa_backends if sdpa_backends is not None else resolve_krea2_sdpa_backends() + # Keyed by the call shape; see `_serves_grouped_query_attention`. + self._gqa_support: dict[tuple, bool] = {} def __call__( self, @@ -159,20 +168,82 @@ def __call__( key = key.transpose(1, 2) value = value.transpose(1, 2) - # Expand K/V heads to the query head count so we can drop enable_gqa (which forces the math backend). - if attn.num_heads != attn.num_kv_heads: + # Krea-2 has 48 query heads over 12 K/V heads. Passing them through as they are avoids + # allocating two tensors four times larger, but only some kernels serve that shape -- so ask, + # do not assume. Where the answer is no, expand as before: that is what keeps this off the + # math path, which materialises the full [heads, seq, seq] score matrix. + enable_gqa = attn.num_heads != attn.num_kv_heads and self._serves_grouped_query_attention( + query, key, value, attention_mask + ) + if attn.num_heads != attn.num_kv_heads and not enable_gqa: repeats = attn.num_heads // attn.num_kv_heads key = key.repeat_interleave(repeats, dim=1) value = value.repeat_interleave(repeats, dim=1) with sdpa_kernel(list(self.sdpa_backends.backends), set_priority=self.sdpa_backends.set_priority): - hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, enable_gqa=enable_gqa + ) # [B, H, S, D] -> [B, S, H, D] -> [B, S, H*D], matching Krea2AttnProcessor's output layout. hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) hidden_states = hidden_states * torch.sigmoid(gate) return attn.to_out[0](hidden_states) + def _serves_grouped_query_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + ) -> bool: + """Whether a fused kernel will take the unexpanded K/V heads for this exact call. + + Measured on the Krea-2 shape ([1, 48q/12kv, 4608, 128], bf16): cuDNN serves it, masked and + unmasked alike; the memory-efficient kernel refuses grouped-query attention outright, and + flash refuses the additive mask the regional-prompting blocks pass. On a build without cuDNN + -- ROCm -- a masked call therefore has no fused kernel left, and answering "yes" there would + drop it onto `math` at roughly 9 GB. That is the failure this whole processor exists to + avoid, so the answer is asked of the dispatcher rather than inferred from the platform. + + The answer depends only on the shape, dtype, device, whether a mask is present, and which + backends this processor permits -- all fixed for a given block within a generation -- so it is + cached, and the query costs nothing after the first call of each kind. + """ + if not query.is_cuda: + # SDPAParams is a CUDA-only interface, and no other backend offers a fused GQA path. + return False + + cache_key = ( + query.shape, + key.shape, + query.dtype, + query.device, + None if attention_mask is None else (attention_mask.shape, attention_mask.dtype), + self.sdpa_backends.backends, + ) + cached = self._gqa_support.get(cache_key) + if cached is not None: + return cached + + try: + params = torch.backends.cuda.SDPAParams(query, key, value, attention_mask, 0.0, False, True) + permitted = self.sdpa_backends.backends + # Only the backends this call will actually permit count. Asking about cuDNN while the + # resolved list excludes it -- which the `efficient` override does -- would answer for a + # kernel the dispatcher is not allowed to reach, and the call would fail outright. + supported = bool( + (SDPBackend.CUDNN_ATTENTION in permitted and torch.backends.cuda.can_use_cudnn_attention(params)) + or (SDPBackend.FLASH_ATTENTION in permitted and torch.backends.cuda.can_use_flash_attention(params)) + ) + except Exception: # noqa: BLE001 + # A torch build whose SDPAParams signature differs: fall back to the expansion, which is + # correct everywhere and merely costs memory. + supported = False + + self._gqa_support[cache_key] = supported + return supported + class _Krea2AttentionProcessorContainer(Protocol): @property diff --git a/tests/backend/krea2/test_grouped_query_attention.py b/tests/backend/krea2/test_grouped_query_attention.py new file mode 100644 index 00000000000..9b4cfe4e7dc --- /dev/null +++ b/tests/backend/krea2/test_grouped_query_attention.py @@ -0,0 +1,160 @@ +"""Whether the K/V heads are expanded before SDPA, or passed through as grouped-query attention. + +Krea-2 has 48 query heads over 12 K/V heads. `repeat_interleave` makes them match, which lets any +fused kernel take the call — at the price of allocating two tensors four times larger, every call. +Only some kernels serve the unexpanded shape, so the choice is made per call by asking the +dispatcher. Getting that answer wrong in the optimistic direction is not a slow path but a ~9 GB one: +the math backend materialises the full [heads, seq, seq] score matrix. +""" + +from unittest.mock import MagicMock + +import pytest +import torch +from diffusers.models.transformers.transformer_krea2 import Krea2Attention +from torch.nn.attention import SDPBackend + +from invokeai.backend.krea2 import attention as attention_module +from invokeai.backend.krea2.attention import ( + Krea2MemoryEfficientAttnProcessor, + Krea2RegionalPromptingState, + Krea2SdpaBackends, +) + +HQ, HKV, HEAD_DIM = 48, 12, 128 + + +def _tensors(device: str, seq: int = 64, kv_heads: int = HKV, masked: bool = False): + q = torch.randn(1, HQ, seq, HEAD_DIM, device=device, dtype=torch.bfloat16) + kv = torch.randn(1, kv_heads, seq, HEAD_DIM, device=device, dtype=torch.bfloat16) + mask = torch.zeros(1, 1, seq, seq, device=device, dtype=torch.bfloat16) if masked else None + return q, kv, kv.clone(), mask + + +class TestTheDecisionIsConservative: + def test_a_non_cuda_tensor_never_takes_the_gqa_path(self): + """`SDPAParams` is CUDA-only, and no other backend offers a fused grouped-query path. + Answering yes here would drop the call onto math.""" + proc = Krea2MemoryEfficientAttnProcessor() + assert proc._serves_grouped_query_attention(*_tensors("cpu")) is False + + def test_it_answers_no_when_the_permitted_list_excludes_the_capable_kernels(self): + """The trap: cuDNN may well serve the shape, but if this call is not permitted to reach it — + which `INVOKE_KREA2_SDPA_BACKEND=efficient` does — then answering yes makes the call fail + outright with `No available kernel`, because the memory-efficient kernel refuses GQA.""" + proc = Krea2MemoryEfficientAttnProcessor( + sdpa_backends=Krea2SdpaBackends(backends=(SDPBackend.EFFICIENT_ATTENTION,), set_priority=False) + ) + q, k, v, mask = _tensors("cpu") + assert proc._serves_grouped_query_attention(q, k, v, mask) is False + + def test_a_torch_build_with_a_different_sdpaparams_answers_no(self, monkeypatch): + # Falling back to the expansion is correct everywhere; it merely costs memory. + proc = Krea2MemoryEfficientAttnProcessor() + q, k, v, mask = _tensors("cpu") + monkeypatch.setattr(type(q), "is_cuda", property(lambda self: True), raising=False) + monkeypatch.setattr(torch.backends.cuda, "SDPAParams", MagicMock(side_effect=TypeError("signature changed"))) + assert proc._serves_grouped_query_attention(q, k, v, mask) is False + + +class TestTheAnswerIsCached: + def test_the_dispatcher_is_asked_once_per_call_shape(self, monkeypatch): + """Dozens of attention calls per step share one shape; querying every time would be waste.""" + proc = Krea2MemoryEfficientAttnProcessor() + q, k, v, mask = _tensors("cpu") + monkeypatch.setattr(type(q), "is_cuda", property(lambda self: True), raising=False) + monkeypatch.setattr(torch.backends.cuda, "SDPAParams", MagicMock(return_value=object())) + can_use = MagicMock(return_value=True) + monkeypatch.setattr(torch.backends.cuda, "can_use_cudnn_attention", can_use) + + for _ in range(5): + assert proc._serves_grouped_query_attention(q, k, v, mask) is True + assert can_use.call_count == 1 + + def test_a_different_mask_state_is_a_different_question(self, monkeypatch): + """Flash takes the unmasked call and refuses the masked one, so the two cannot share an + answer.""" + proc = Krea2MemoryEfficientAttnProcessor() + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True), raising=False) + monkeypatch.setattr(torch.backends.cuda, "SDPAParams", MagicMock(return_value=object())) + can_use = MagicMock(return_value=True) + monkeypatch.setattr(torch.backends.cuda, "can_use_cudnn_attention", can_use) + + proc._serves_grouped_query_attention(*_tensors("cpu", masked=False)) + proc._serves_grouped_query_attention(*_tensors("cpu", masked=True)) + assert can_use.call_count == 2 + + +class TestTheExpansionStillHappensWhenItMust: + @pytest.mark.parametrize("masked", [False, True]) + def test_cpu_output_matches_the_stock_processor(self, masked): + """On CPU the answer is always no, so this is the expansion path — and it must stay exactly + what it was before the grouped-query option existed.""" + torch.manual_seed(0) + attn = Krea2Attention(hidden_size=256, num_heads=8, num_kv_heads=2, eps=1e-5).eval() + hidden = torch.randn(1, 16, 256) + state = Krea2RegionalPromptingState(attention_mask=torch.ones(16, 16, dtype=torch.bool)) if masked else None + + with torch.no_grad(): + attn.set_processor(Krea2MemoryEfficientAttnProcessor(regional_prompting_state=state)) + ours = attn(hidden, attention_mask=None, image_rotary_emb=None) + from diffusers.models.transformers.transformer_krea2 import Krea2AttnProcessor + + attn.set_processor(Krea2AttnProcessor()) + stock = attn(hidden, attention_mask=None, image_rotary_emb=None) + + assert torch.allclose(ours, stock, atol=1e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required to reach a fused GQA kernel") +class TestTheTwoPathsAgreeOnCuda: + @pytest.mark.parametrize("masked", [False, True]) + def test_grouped_query_output_matches_the_expanded_one(self, masked, monkeypatch): + """The saving is only worth having if the result is the same. Anything else would be a + silent change to every Krea-2 image.""" + torch.manual_seed(0) + attn = ( + Krea2Attention(hidden_size=HQ * HEAD_DIM, num_heads=HQ, num_kv_heads=HKV, eps=1e-5) + .eval() + .to("cuda", torch.bfloat16) + ) + hidden = torch.randn(1, 256, HQ * HEAD_DIM, device="cuda", dtype=torch.bfloat16) + state = None + if masked: + m = torch.zeros(256, 256, device="cuda", dtype=torch.bool) + m[:128, :128] = True + m[128:, 128:] = True + state = Krea2RegionalPromptingState(attention_mask=m) + + # Record the K/V head count that actually reaches SDPA. Without this the test would pass + # just as happily if both runs expanded, which is exactly what it is meant to rule out. + real_sdpa = torch.nn.functional.scaled_dot_product_attention + seen: list[int] = [] + + def spy(q, k, v, attn_mask=None, enable_gqa=False, **kwargs): + seen.append(k.shape[1]) + return real_sdpa(q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa, **kwargs) + + def run(force_expand: bool) -> torch.Tensor: + proc = Krea2MemoryEfficientAttnProcessor(regional_prompting_state=state) + if force_expand: + proc._serves_grouped_query_attention = lambda *a, **k: False + attn.set_processor(proc) + monkeypatch.setattr(attention_module.F, "scaled_dot_product_attention", spy) + with torch.no_grad(): + out = attn(hidden, attention_mask=None, image_rotary_emb=None).float() + monkeypatch.undo() + return out + + expanded = run(True) + grouped = run(False) + + assert seen[0] == HQ, "the forced run should have expanded the K/V heads" + if seen[1] != HKV: + pytest.skip("no fused kernel on this device serves grouped-query attention for this shape") + + # bf16 kernels differ in accumulation order, so the two are close rather than identical. The + # bound is the one used for the backend comparison elsewhere. + corr = torch.corrcoef(torch.stack([expanded.flatten(), grouped.flatten()]))[0, 1] + assert corr > 0.9999, f"correlation {corr}" + assert (expanded - grouped).abs().max() < 0.05