diff --git a/invokeai/app/invocations/krea2_denoise.py b/invokeai/app/invocations/krea2_denoise.py index 466568d5050..6beec7d270c 100644 --- a/invokeai/app/invocations/krea2_denoise.py +++ b/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) diff --git a/invokeai/app/run_app.py b/invokeai/app/run_app.py index 0f0f18fa8c1..fb2fcea33a3 100644 --- a/invokeai/app/run_app.py +++ b/invokeai/app/run_app.py @@ -85,6 +85,7 @@ def run_app() -> None: check_cudnn, enable_dev_reload, find_open_port, + log_attention_backends, register_mime_types, ) @@ -99,6 +100,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() diff --git a/invokeai/app/util/startup_utils.py b/invokeai/app/util/startup_utils.py index aaa0097b831..248b145cf21 100644 --- a/invokeai/app/util/startup_utils.py +++ b/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() diff --git a/invokeai/backend/krea2/attention.py b/invokeai/backend/krea2/attention.py index b7f96cd93f9..2bbdb63ab7d 100644 --- a/invokeai/backend/krea2/attention.py +++ b/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 diff --git a/tests/app/util/test_attention_backend_probe.py b/tests/app/util/test_attention_backend_probe.py new file mode 100644 index 00000000000..11dba3bce02 --- /dev/null +++ b/tests/app/util/test_attention_backend_probe.py @@ -0,0 +1,78 @@ +"""The startup SDPA availability probe. + +It exists to answer support questions from a log line instead of a probe script, so the properties +that matter are: it never breaks boot, it says nothing where the question does not apply, and what +it does say is specific enough to act on. +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from invokeai.app.util.startup_utils import log_attention_backends, probe_attention_backends + + +class TestProbe: + @pytest.mark.parametrize("device_type", ["cpu", "mps"]) + def test_a_non_cuda_device_is_not_probed(self, device_type): + # There are no fused SDPA backends to report, and can_use_* would raise. + assert probe_attention_backends(torch.device(device_type)) is None + + def test_a_probe_failure_does_not_propagate(self): + # A diagnostic must never be the reason the server does not start. + with patch.object(torch, "empty", side_effect=RuntimeError("no CUDA driver")): + assert probe_attention_backends(torch.device("cuda")) is None + + def test_math_is_reported_as_always_available(self): + # It is the unfused fallback, not a kernel that can be missing -- so it is never probed for. + with ( + patch.object(torch, "empty", return_value=MagicMock()), + patch.object(torch.backends.cuda, "SDPAParams", return_value=MagicMock()), + patch.object(torch.backends.cuda, "can_use_cudnn_attention", return_value=False), + patch.object(torch.backends.cuda, "can_use_flash_attention", return_value=False), + patch.object(torch.backends.cuda, "can_use_efficient_attention", return_value=False), + patch.object(torch.cuda, "empty_cache"), + ): + available = probe_attention_backends(torch.device("cuda")) + assert available == {"cudnn": False, "flash": False, "efficient": False, "math": True} + + def test_the_probe_tensor_is_not_left_in_the_allocator(self): + with ( + patch.object(torch, "empty", return_value=MagicMock()), + patch.object(torch.backends.cuda, "SDPAParams", return_value=MagicMock()), + patch.object(torch.backends.cuda, "can_use_cudnn_attention", return_value=True), + patch.object(torch.backends.cuda, "can_use_flash_attention", return_value=True), + patch.object(torch.backends.cuda, "can_use_efficient_attention", return_value=True), + patch.object(torch.cuda, "empty_cache") as empty_cache, + ): + probe_attention_backends(torch.device("cuda")) + empty_cache.assert_called_once() + + +class TestLogLine: + def test_nothing_is_logged_where_the_question_does_not_apply(self, caplog): + logger = logging.getLogger("test_probe_silent") + with caplog.at_level(logging.INFO, logger=logger.name): + log_attention_backends(logger, torch.device("cpu")) + assert caplog.records == [] + + def test_the_line_names_every_backend_and_its_answer(self, caplog): + logger = logging.getLogger("test_probe_line") + with patch( + "invokeai.app.util.startup_utils.probe_attention_backends", + return_value={"cudnn": True, "flash": False, "efficient": True, "math": True}, + ): + with caplog.at_level(logging.INFO, logger=logger.name): + log_attention_backends(logger, torch.device("cuda")) + + message = caplog.records[0].message + assert "cudnn=yes" in message + assert "flash=no" in message + assert "efficient=yes" in message + # The shape is part of the answer: availability depends on it, so a line without it would be + # unactionable. + assert "head_dim 128" in message + # And the caveat that keeps it from being read as a dispatch table. + assert "no mask" in message diff --git a/tests/backend/krea2/test_sdpa_backend_selection.py b/tests/backend/krea2/test_sdpa_backend_selection.py new file mode 100644 index 00000000000..9bc1422d7f8 --- /dev/null +++ b/tests/backend/krea2/test_sdpa_backend_selection.py @@ -0,0 +1,184 @@ +"""Which SDPA backends the Krea-2 attention processors run under, and the opt-in override. + +The ranking is a measurement, not a preference. On the real Krea-2 attention shape, per call: on a +Windows 4090 (no flash in the build) cuDNN runs at 3.72ms against efficient's 5.92ms; on a Linux +30-series card flash runs at 19.74ms, cuDNN at 21.27ms, efficient at 31.45ms. So flash leads where it +exists and cuDNN takes the blocks it cannot serve -- the masked ones, and every block on a build +without flash. + +These tests pin what makes that reachable at all: the order, and `set_priority=True`. Without the +latter the list only *permits* backends and torch picks by its own order, so the ranking would be a +silent no-op. +""" + +from unittest.mock import MagicMock + +import pytest +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from invokeai.backend.krea2.attention import ( + KREA2_SDPA_BACKEND_ENV_VAR, + Krea2MemoryEfficientAttnProcessor, + Krea2RegionalPromptingState, + build_krea2_attention_processors, + resolve_krea2_sdpa_backends, +) + + +class TestTheDefaultRanking: + def test_flash_leads_and_cudnn_follows(self): + """Flash is the fastest kernel where the build has it -- and is already what runs today, + because torch's own order puts it above efficient. Ranking cuDNN over it would be a small + regression on every flash-capable build.""" + backends = resolve_krea2_sdpa_backends(raw_override=None).backends + assert backends[0] is SDPBackend.FLASH_ATTENTION + assert backends[1] is SDPBackend.CUDNN_ATTENTION + + def test_cudnn_outranks_efficient(self): + """This is where the win actually comes from: flash refuses the additive mask the + regional-prompting blocks pass, so on those blocks it is skipped and cuDNN takes over -- + 1.6x-2.0x over efficient. On a build without flash, that is every block.""" + backends = resolve_krea2_sdpa_backends(raw_override=None).backends + assert backends.index(SDPBackend.CUDNN_ATTENTION) < backends.index(SDPBackend.EFFICIENT_ATTENTION) + + def test_math_is_last(self): + # The unfused fallback: correct everywhere, and ~35x slower at ~35x the memory. + assert resolve_krea2_sdpa_backends(raw_override=None).backends[-1] is SDPBackend.MATH + + def test_priority_is_set(self): + """Load-bearing: without set_priority the list only *permits* backends and torch picks by + its own order, in which cuDNN ranks last and would never be chosen. The ranking would be a + no-op.""" + assert resolve_krea2_sdpa_backends(raw_override=None).set_priority is True + + def test_the_fallbacks_are_all_still_there(self): + # A ranked list, never an exclusive backend: an unavailable entry is skipped by the + # dispatcher, so the list degrades to today's behaviour on its own. + assert set(resolve_krea2_sdpa_backends(raw_override=None).backends) == { + SDPBackend.CUDNN_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.FLASH_ATTENTION, + SDPBackend.MATH, + } + + def test_flash_is_kept_even_where_a_probe_would_call_it_dead(self): + """A sibling plan proposed dropping FLASH where a probe shows it absent. That would be wrong + twice over: on ROCm flash is available and cuDNN is not, and on Linux CUDA it is the fastest + kernel of the four. A dead entry in a ranked list costs nothing; a missing one costs a + platform.""" + assert SDPBackend.FLASH_ATTENTION in resolve_krea2_sdpa_backends(raw_override=None).backends + + def test_no_override_is_recorded_by_default(self): + # `override is None` is what keeps the benchmark instrumentation, and its synchronize() + # calls, entirely off the default path. + assert resolve_krea2_sdpa_backends(raw_override=None).override is None + + @pytest.mark.parametrize("blank", ["", " ", "\t"]) + def test_a_blank_value_is_the_default_not_an_error(self, blank): + assert resolve_krea2_sdpa_backends(raw_override=blank).override is None + + +class TestTheOverride: + @pytest.mark.parametrize( + "value,expected", + [ + ("cudnn", SDPBackend.CUDNN_ATTENTION), + ("efficient", SDPBackend.EFFICIENT_ATTENTION), + ("flash", SDPBackend.FLASH_ATTENTION), + ("math", SDPBackend.MATH), + ], + ) + def test_a_named_backend_is_exclusive_with_no_fallback(self, value, expected): + """The point of the exclusive modes: a run that completes proves that kernel was used, + because an unavailable backend raises visibly instead of degrading to math unnoticed.""" + choice = resolve_krea2_sdpa_backends(raw_override=value) + assert choice.backends == (expected,) + assert choice.set_priority is False + assert choice.override == value + + def test_priority_cudnn_reproduces_the_default_list(self): + # This is the value the sm_86 pre-merge check runs with, so it must be the shipped list. + explicit = resolve_krea2_sdpa_backends(raw_override="priority-cudnn") + default = resolve_krea2_sdpa_backends(raw_override=None) + assert explicit.backends == default.backends + assert explicit.set_priority == default.set_priority + assert explicit.override == "priority-cudnn" + + @pytest.mark.parametrize("value", [" CUDNN ", "Priority-CuDNN", "EFFICIENT"]) + def test_values_are_stripped_and_lowercased(self, value): + assert resolve_krea2_sdpa_backends(raw_override=value).override == value.strip().lower() + + def test_an_unknown_value_raises_and_names_the_valid_ones(self): + with pytest.raises(ValueError) as excinfo: + resolve_krea2_sdpa_backends(raw_override="cudnn-attention") + message = str(excinfo.value) + assert KREA2_SDPA_BACKEND_ENV_VAR in message + for valid in ("cudnn", "efficient", "flash", "math", "priority-cudnn"): + assert valid in message + + def test_the_environment_is_read_when_no_value_is_passed(self, monkeypatch): + monkeypatch.setenv(KREA2_SDPA_BACKEND_ENV_VAR, "math") + assert resolve_krea2_sdpa_backends().backends == (SDPBackend.MATH,) + monkeypatch.delenv(KREA2_SDPA_BACKEND_ENV_VAR) + assert resolve_krea2_sdpa_backends().override is None + + +class TestProcessorsCarryTheChoice: + def test_every_processor_gets_the_same_resolved_choice(self): + """Resolved once per generation and handed down -- not re-read per attention call, of which + there are dozens per step.""" + transformer = MagicMock() + transformer.attn_processors = { + "transformer_blocks.0.attn.processor": object(), + "transformer_blocks.1.attn.processor": object(), + "single_transformer_blocks.0.attn.processor": object(), + } + + processors = build_krea2_attention_processors(transformer, Krea2RegionalPromptingState()) + + choices = {id(p.sdpa_backends) for p in processors.values()} + assert len(choices) == 1 + assert next(iter(processors.values())).sdpa_backends.backends[0] is SDPBackend.FLASH_ATTENTION + + def test_a_standalone_processor_resolves_for_itself(self): + # Constructed directly in tests and by custom code; it must not depend on the builder. + assert Krea2MemoryEfficientAttnProcessor().sdpa_backends.backends[0] is SDPBackend.FLASH_ATTENTION + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required to exercise the SDPA dispatcher") +class TestTheFallbackIsReal: + """The whole design rests on one dispatcher property: an unusable backend in a *ranked* list is + skipped, while an *exclusively* selected one raises. + + This is what makes cuDNN-first safe on hardware where cuDNN cannot serve the call, and it is + worth pinning: if a future torch made a ranked list raise instead, the ranking would turn every + such device from "today's behaviour" into a failed generation, silently as far as our tests go. + """ + + @staticmethod + def _unservable_by_cudnn(): + # fp32 is refused by the fused kernels, which is a portable way to make cuDNN unusable on a + # card where it otherwise works -- i.e. to stand in for the sm_86 report. + t = torch.randn(1, 24, 512, 128, device="cuda", dtype=torch.float32) + params = torch.backends.cuda.SDPAParams(t, t, t, None, 0.0, False, False) + if torch.backends.cuda.can_use_cudnn_attention(params): + pytest.skip("This build serves fp32 with cuDNN, so it cannot stand in for an unusable backend") + return t + + def test_a_ranked_list_completes_where_cudnn_cannot_serve(self): + t = self._unservable_by_cudnn() + choice = resolve_krea2_sdpa_backends(raw_override=None) + with sdpa_kernel(list(choice.backends), set_priority=choice.set_priority): + out = torch.nn.functional.scaled_dot_product_attention(t, t, t) + assert torch.isfinite(out).all() + + def test_the_same_call_raises_when_cudnn_is_selected_exclusively(self): + """The counterpart, and the likely explanation of the sm_86 report: that measurement was + taken with `sdpa_kernel([backend])` -- exclusive, no fallback -- which is exactly the mode + that raises `No available kernel` when the backend is unusable.""" + t = self._unservable_by_cudnn() + choice = resolve_krea2_sdpa_backends(raw_override="cudnn") + with pytest.raises(RuntimeError, match="No available kernel"): + with sdpa_kernel(list(choice.backends), set_priority=choice.set_priority): + torch.nn.functional.scaled_dot_product_attention(t, t, t)