perf(krea2): rank cuDNN first for SDPA, with a fallback list and a probe

invokeai/app/invocations/krea2_denoise.py

@@ -1,5 +1,7 @@
 import json
 import math
+import statistics
+import time
 from contextlib import ExitStack
 from pathlib import Path
 from typing import Callable, Iterator, Optional
@@ -25,7 +27,11 @@
 from invokeai.app.invocations.model import TransformerField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.krea2.attention import Krea2RegionalPromptingState, build_krea2_attention_processors
+from invokeai.backend.krea2.attention import (
+    Krea2RegionalPromptingState,
+    build_krea2_attention_processors,
+    resolve_krea2_sdpa_backends,
+)
 from invokeai.backend.krea2.regional_prompting import (
     Krea2RegionalPromptingExtension,
     Krea2TextConditioning,
@@ -50,11 +56,64 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Krea2ConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.logging import InvokeAILogger
 
 # Krea-2 latent channels (Qwen-Image VAE z_dim). The packed transformer in_channels is 16 * patch_size**2 = 64.
 KREA2_LATENT_CHANNELS = 16
 
 
+logger = InvokeAILogger.get_logger(__name__)
+
+
+class _Krea2StepBenchmark:
+    """Per-step timing for the Krea-2 denoise loop, for A/B-ing SDPA backends.
+
+    Off unless KREA2_SDPA_BACKEND_ENV_VAR is set, and off on non-CUDA devices. That matters: the
+    `cuda.synchronize()` calls on both sides of a step serialise the loop, so the default path must
+    not reach them at all.
+    """
+
+    def __init__(self, device: torch.device, label: str) -> None:
+        self._device = device
+        self._label = label
+        self._step_ms: list[float] = []
+        self._step_started_at = 0.0
+        # Reset here, immediately before the loop, so the reported peak is the loop's and not the
+        # model load's.
+        torch.cuda.reset_peak_memory_stats(device)
+
+    @classmethod
+    def create(cls, device: torch.device) -> "_Krea2StepBenchmark | None":
+        override = resolve_krea2_sdpa_backends().override
+        if override is None or device.type != "cuda":
+            return None
+        return cls(device, override)
+
+    def start_step(self) -> None:
+        torch.cuda.synchronize(self._device)
+        self._step_started_at = time.perf_counter()
+
+    def end_step(self) -> None:
+        torch.cuda.synchronize(self._device)
+        self._step_ms.append((time.perf_counter() - self._step_started_at) * 1000)
+
+    def log_summary(self) -> None:
+        if not self._step_ms:
+            return
+        peak_gib = torch.cuda.max_memory_allocated(self._device) / 2**30
+        first_ms = self._step_ms[0]
+        # The first step absorbs kernel selection and allocator warmup -- on a cold cuDNN run it has
+        # been seen at 2606 ms against a 1451 ms steady mean -- so it is reported, not averaged in.
+        steady = self._step_ms[1:] or self._step_ms
+        logger.info(
+            f"Krea-2 SDPA benchmark [{self._label}]: {len(self._step_ms)} steps | "
+            f"first {first_ms:.0f} ms | "
+            f"steady mean {statistics.mean(steady):.0f} ms, median {statistics.median(steady):.0f} ms, "
+            f"min {min(steady):.0f} ms, max {max(steady):.0f} ms | "
+            f"peak torch.cuda.max_memory_allocated {peak_gib:.3f} GiB"
+        )
+
+
 @invocation(
     "krea2_denoise",
     title="Denoise - Krea-2",
@@ -459,7 +518,11 @@ def _run_diffusion(self, context: InvocationContext):
             pos_regional_attention_mask = pos_extension.get_attention_mask()
             neg_regional_attention_mask = neg_extension.get_attention_mask() if neg_extension is not None else None
 
+            benchmark = _Krea2StepBenchmark.create(device)
+
             for step_idx, t in enumerate(tqdm(timesteps_sched)):
+                if benchmark is not None:
+                    benchmark.start_step()
                 # The pipeline passes timestep / num_train_timesteps to the transformer.
                 timestep = (t / num_train_timesteps).expand(latents.shape[0]).to(inference_dtype)
 
@@ -504,6 +567,9 @@ def _run_diffusion(self, context: InvocationContext):
                     )
                     latents = pack_latents(latents_4d, 1, KREA2_LATENT_CHANNELS, latent_height, latent_width)
 
+                if benchmark is not None:
+                    benchmark.end_step()
+
                 step_callback(
                     PipelineIntermediateState(
                         step=step_idx + 1,
@@ -514,6 +580,9 @@ def _run_diffusion(self, context: InvocationContext):
                     ),
                 )
 
+        if benchmark is not None:
+            benchmark.log_summary()
+
         # Unpack to 4D then add a frame dim for the Qwen-Image VAE: (B, C, 1, H, W).
         latents = unpack_latents(latents, latent_height, latent_width)
         latents = latents.unsqueeze(2)

invokeai/app/run_app.py

@@ -70,6 +70,7 @@ def run_app() -> None:
         check_cudnn,
         enable_dev_reload,
         find_open_port,
+        log_attention_backends,
         register_mime_types,
     )
 
@@ -84,6 +85,7 @@ def run_app() -> None:
     apply_monkeypatches()
     register_mime_types()
     check_cudnn(logger)
+    log_attention_backends(logger, TorchDevice.choose_torch_device())
 
     # Initialize the app and event loop.
     app, loop = get_app()

invokeai/app/util/startup_utils.py

@@ -34,6 +34,67 @@ def check_cudnn(logger: logging.Logger) -> None:
             )
 
 
+# The shape the diffusion transformers actually attend over: one image, many heads, head_dim 128.
+# Availability can depend on the shape, so probing with a toy one would answer a different question.
+_PROBE_HEADS = 24
+_PROBE_SEQ_LEN = 1024
+_PROBE_HEAD_DIM = 128
+
+
+def probe_attention_backends(device: torch.device) -> dict[str, bool] | None:
+    """Which fused SDPA backends this build and device can actually use.
+
+    Availability is not a property of the torch version alone: ROCm builds have no cuDNN attention,
+    Windows CUDA builds usually have no flash, and both depend on the device architecture. Returns
+    None where the question does not apply (no CUDA/ROCm device) or cannot be answered.
+
+    This asks about the *unmasked* case. A backend listed here can still be rejected for a specific
+    call -- flash refuses the additive padding mask the regional-prompting paths pass, for instance --
+    so this is a diagnostic, not a dispatch table.
+    """
+    if device.type != "cuda":
+        return None
+    try:
+        probe = torch.empty(
+            1,
+            _PROBE_HEADS,
+            _PROBE_SEQ_LEN,
+            _PROBE_HEAD_DIM,
+            device=device,
+            dtype=torch.float16,
+        )
+        params = torch.backends.cuda.SDPAParams(probe, probe, probe, None, 0.0, False, False)
+        return {
+            "cudnn": bool(torch.backends.cuda.can_use_cudnn_attention(params)),
+            "flash": bool(torch.backends.cuda.can_use_flash_attention(params)),
+            "efficient": bool(torch.backends.cuda.can_use_efficient_attention(params)),
+            # Always present -- it is the unfused fallback, not a kernel that can be missing.
+            "math": True,
+        }
+    except Exception:
+        # A diagnostic must never be the reason the server does not start.
+        return None
+    finally:
+        # Do not leave the probe tensor sitting in the caching allocator for the first generation.
+        torch.cuda.empty_cache()
+
+
+def log_attention_backends(logger: logging.Logger, device: torch.device) -> None:
+    """Log the SDPA backend availability once at startup.
+
+    The point is support: a question about attention performance can then be answered by reading a
+    log line instead of asking the user to run a probe script.
+    """
+    available = probe_attention_backends(device)
+    if available is None:
+        return
+    summary = " ".join(f"{name}={'yes' if ok else 'no'}" for name, ok in available.items())
+    logger.info(
+        f"SDPA attention backends "
+        f"(fp16, {_PROBE_HEADS} heads, seq {_PROBE_SEQ_LEN}, head_dim {_PROBE_HEAD_DIM}, no mask): {summary}"
+    )
+
+
 def invokeai_source_dir() -> Path:
     # `invokeai.__file__` doesn't always work for editable installs
     this_module_path = Path(__file__).resolve()

invokeai/backend/krea2/attention.py

@@ -7,12 +7,13 @@
 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 the memory-efficient SDPA backend (which supports the additive padding mask and
-is O(seq) in memory). Measured: the same 3600-token attention drops from ~5.7 GB to ~0.19 GB.
+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.
 
 The math is otherwise identical to ``Krea2AttnProcessor`` (q/k RMSNorm, rotary embeddings, sigmoid output gate).
 """
 
+import os
 import re
 from dataclasses import dataclass
 from typing import Protocol
@@ -22,8 +23,83 @@
 from diffusers.models.embeddings import apply_rotary_emb
 from torch.nn.attention import SDPBackend, sdpa_kernel
 
-# Prefer the memory-efficient kernel; fall back to flash (if the build has it) then math so we never hard-fail.
-_KREA2_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger(__name__)
+
+# Measured on the real Krea-2 attention shape ([1, 48, 4608, 128], bf16), per call:
+#
+#                     RTX 4090 / Windows      RTX 30-series / Linux
+#   flash             not compiled in         19.74 ms
+#   cudnn              3.72 ms                21.27 ms
+#   efficient          5.92 ms                31.45 ms
+#   math              51.23 ms               168.48 ms
+#
+# Two things follow, and the order below encodes both.
+#
+# **Flash first.** Where the build has it, flash is the fastest kernel and is *already* what runs
+# today: without `set_priority` torch picks by its own order, in which flash outranks efficient. On
+# the 30-series card an unprefixed call lands at 20.15 ms, i.e. on flash, not on efficient. Ranking
+# cuDNN above it would therefore be a small regression on every flash-capable build.
+#
+# **cuDNN second, and it is not a formality.** Flash refuses the additive padding mask that the
+# regional-prompting blocks pass, so on exactly those blocks it is skipped and cuDNN takes over --
+# where it beats efficient by 1.6x-2.0x. Windows CUDA builds have no flash at all, so there cuDNN is
+# what every block gets. Both are the cases the win comes from; flash-first does not give either up.
+#
+# Everything here is a fallback, never an exclusive choice: an unavailable backend is skipped by the
+# dispatcher, so the list degrades on its own -- to efficient on ROCm, where cuDNN is absent and
+# flash rejects the mask, and to efficient anywhere neither fused kernel can serve the call.
+_KREA2_SDPA_BACKENDS = [
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+# Opt-in override, for measuring one backend against another and for support questions. Unset -- the
+# only state a user ever sees by default -- is the ranked list above, unchanged.
+KREA2_SDPA_BACKEND_ENV_VAR = "INVOKE_KREA2_SDPA_BACKEND"
+_PRIORITY_CUDNN = "priority-cudnn"
+_EXCLUSIVE_BACKENDS = {
+    "cudnn": SDPBackend.CUDNN_ATTENTION,
+    "efficient": SDPBackend.EFFICIENT_ATTENTION,
+    "flash": SDPBackend.FLASH_ATTENTION,
+    "math": SDPBackend.MATH,
+}
+
+
+@dataclass(frozen=True)
+class Krea2SdpaBackends:
+    """Which SDPA backends a Krea-2 attention call may use, and in what order."""
+
+    backends: tuple[SDPBackend, ...]
+    set_priority: bool
+    override: str | None = None
+
+    def describe(self) -> str:
+        names = ", ".join(b.name for b in self.backends)
+        return f"sdpa_kernel([{names}], set_priority={self.set_priority})"
+
+
+def resolve_krea2_sdpa_backends(raw_override: str | None = None) -> Krea2SdpaBackends:
+    """Resolve the SDPA backend list, honouring KREA2_SDPA_BACKEND_ENV_VAR.
+
+    The exclusive modes are the point of the override: a run that completes proves that kernel was
+    actually used, because an unavailable backend raises visibly instead of quietly degrading to math.
+    """
+    raw = os.environ.get(KREA2_SDPA_BACKEND_ENV_VAR) if raw_override is None else raw_override
+    if raw is None or not raw.strip():
+        return Krea2SdpaBackends(backends=tuple(_KREA2_SDPA_BACKENDS), set_priority=True)
+
+    value = raw.strip().lower()
+    if value == _PRIORITY_CUDNN:
+        return Krea2SdpaBackends(backends=tuple(_KREA2_SDPA_BACKENDS), set_priority=True, override=value)
+    if value in _EXCLUSIVE_BACKENDS:
+        return Krea2SdpaBackends(backends=(_EXCLUSIVE_BACKENDS[value],), set_priority=False, override=value)
+
+    valid = ", ".join([*sorted(_EXCLUSIVE_BACKENDS), _PRIORITY_CUDNN])
+    raise ValueError(f"{KREA2_SDPA_BACKEND_ENV_VAR}={raw!r} is not a valid value. Valid values: {valid}.")
 
 
 @dataclass
@@ -39,8 +115,14 @@ def set_attention_mask(self, attention_mask: torch.Tensor | None) -> None:
 class Krea2MemoryEfficientAttnProcessor:
     """Drop-in replacement for ``Krea2AttnProcessor`` that avoids the ``enable_gqa`` math fallback."""
 
-    def __init__(self, regional_prompting_state: Krea2RegionalPromptingState | None = None) -> None:
+    def __init__(
+        self,
+        regional_prompting_state: Krea2RegionalPromptingState | None = None,
+        sdpa_backends: Krea2SdpaBackends | None = None,
+    ) -> None:
         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()
 
     def __call__(
         self,
@@ -83,7 +165,7 @@ def __call__(
             key = key.repeat_interleave(repeats, dim=1)
             value = value.repeat_interleave(repeats, dim=1)
 
-        with sdpa_kernel(_KREA2_SDPA_BACKENDS):
+        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)
 
         # [B, H, S, D] -> [B, S, H, D] -> [B, S, H*D], matching Krea2AttnProcessor's output layout.
@@ -103,10 +185,20 @@ def build_krea2_attention_processors(
 ) -> dict[str, Krea2MemoryEfficientAttnProcessor]:
     """Build processors that apply regional masks to alternating main transformer blocks only."""
 
+    sdpa_backends = resolve_krea2_sdpa_backends()
+    if sdpa_backends.override is not None:
+        # Once per generation, not once per attention call.
+        logger.info(
+            f"Krea-2 SDPA backend override active: {KREA2_SDPA_BACKEND_ENV_VAR}={sdpa_backends.override} "
+            f"-> {sdpa_backends.describe()}"
+        )
+
     processors: dict[str, Krea2MemoryEfficientAttnProcessor] = {}
     for name in transformer.attn_processors:
         match = re.fullmatch(r"transformer_blocks\.(\d+)\.attn\.processor", name)
         block_index = int(match.group(1)) if match is not None else None
         state = regional_prompting_state if block_index is not None and block_index % 2 == 0 else None
-        processors[name] = Krea2MemoryEfficientAttnProcessor(regional_prompting_state=state)
+        processors[name] = Krea2MemoryEfficientAttnProcessor(
+            regional_prompting_state=state, sdpa_backends=sdpa_backends
+        )
     return processors