perf(krea2): keep the K/V heads unexpanded where a fused kernel takes them

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

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