From 1904479a4f1576e7769b17e80582f8a98b85c458 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 19 Aug 2026 21:36:24 +0200 Subject: [PATCH 1/5] fix(flux2): estimate working memory for denoise and both VAE directions The FLUX.2 path called model_on_device() with no working_mem_bytes anywhere, so the model cache reserved only the default device_working_mem_gb and filled the rest of the card with the model. Reference images make that fatal rather than merely tight: their latents are concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of a 1024x1024 generation -- 6.5GB of activations against a 3GB reservation. Measured on CUDA in bf16 as peak reserved memory: transformer activations scale linearly at ~0.39 MB/token (no O(seq^2) term, SDPA) and are independent of block count; the FLUX.2 VAE costs ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, so a 1024x1024 decode peaks at ~4.3GB. Add Flux2DenoiseInvocation._estimate_working_memory() and estimate_vae_working_memory_flux2(), and pass them at every load site so the cache evicts enough to make room instead of hitting the shortfall as an OOM. Closes #9500 --- invokeai/app/invocations/flux2_denoise.py | 61 +++- invokeai/app/invocations/flux2_vae_decode.py | 10 +- invokeai/app/invocations/flux2_vae_encode.py | 9 +- invokeai/backend/flux2/ref_image_extension.py | 21 +- invokeai/backend/util/vae_working_memory.py | 37 +++ .../invocations/test_flux2_working_memory.py | 287 ++++++++++++++++++ 6 files changed, 418 insertions(+), 7 deletions(-) create mode 100644 tests/app/invocations/test_flux2_working_memory.py diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 7a5ce158c98..25c8b5d8f30 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -458,10 +458,33 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: bn_std=bn_std, ) + # Estimate the peak activation memory the transformer forward will need and ask the model cache + # to keep that much VRAM free. Without this hint the cache reserves only the small default + # working memory and fills the rest of the card with the model, so anything beyond a plain + # low-resolution generation OOMs. Reference images are the dominant term: their latents are + # concatenated onto the image stream, so three 1024x1024 references quadruple the sequence + # (and with it the activation footprint) of a 1024x1024 generation. + ref_image_seq_len = ref_image_extension.ref_image_latents.shape[1] if ref_image_extension is not None else 0 + estimated_working_memory = self._estimate_working_memory( + image_seq_len=packed_h * packed_w, + ref_image_seq_len=ref_image_seq_len, + text_seq_len=max(txt.shape[1], neg_txt.shape[1] if neg_txt is not None else 0), + num_loras=len(self.transformer.loras), + # The mask itself is already allocated; only the additive bias built per forward is new. + # It is skipped entirely when reference images are present (see below). + regional_attention_bias_bytes=( + regional_extension.restricted_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() + if regional_extension.restricted_attn_mask is not None and ref_image_seq_len == 0 + else 0 + ), + ) + with ExitStack() as exit_stack: # Load the transformer model (cached_weights, transformer) = exit_stack.enter_context( - context.models.load(self.transformer.transformer).model_on_device() + context.models.load(self.transformer.transformer).model_on_device( + working_mem_bytes=estimated_working_memory + ) ) config = transformer_config @@ -578,6 +601,42 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) mask = mask.to(device=latents.device, dtype=latents.dtype) return mask.expand_as(latents) + def _estimate_working_memory( + self, + image_seq_len: int, + ref_image_seq_len: int, + text_seq_len: int, + num_loras: int, + regional_attention_bias_bytes: int = 0, + ) -> int: + """Estimate peak transformer activation memory (bytes) so the model cache reserves enough headroom. + + FLUX.2 attention runs through SDPA without materializing the O(seq^2) score matrix, so the + activation footprint scales *linearly* with the total attended sequence -- text tokens, image + tokens, and reference-image tokens alike. Measured on the Klein 9B geometry in bf16 as peak + reserved memory, that slope is ~0.39 MB per token and holds from 1.5k to 28k tokens; it is + also independent of the block count (a no-grad forward frees each block's intermediates), so + the constant applies to both the 4B and 9B variants. + + The reference-image term is what makes this estimate necessary rather than merely nice to + have: a 1024x1024 generation is 4096 image tokens (~1.7GB), but attaching three 1024x1024 + references adds 12288 more for ~6.5GB, and a 1328px tile with three 1328px references reaches + ~10.9GB -- against a default ``device_working_mem_gb`` of 3. + + A fixed base covers resolution-independent overhead (transient fp8/GGUF -> bf16 weight casts + during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra + activation branch per patched layer, so we add a per-LoRA margin. + """ + GB = 1024**3 + MB = 1024**2 + per_token_bytes = int(0.4 * MB) + estimated = (image_seq_len + ref_image_seq_len + text_seq_len) * per_token_bytes + estimated += int(1.0 * GB) + estimated += regional_attention_bias_bytes + if num_loras > 0: + estimated += int(0.5 * num_loras * GB) + return estimated + def _load_text_conditioning( self, context: InvocationContext, diff --git a/invokeai/app/invocations/flux2_vae_decode.py b/invokeai/app/invocations/flux2_vae_decode.py index 58297ec1ae4..d4cad4ad75d 100644 --- a/invokeai/app/invocations/flux2_vae_decode.py +++ b/invokeai/app/invocations/flux2_vae_decode.py @@ -21,6 +21,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 @invocation( @@ -49,7 +50,14 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima Input latents should already be in the correct space after BN denormalization was applied in the denoiser. The VAE expects (B, 32, H, W) format. """ - with vae_info.model_on_device() as (_, vae): + # Decoding at FLUX.2 resolutions costs multiple GB of activations (~4.3GB at 1024x1024), + # far above the default working memory the cache would otherwise reserve. Tell it up front so + # it offloads enough of the (possibly still resident) transformer to leave room. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=latents, vae=vae_info.model + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would diff --git a/invokeai/app/invocations/flux2_vae_encode.py b/invokeai/app/invocations/flux2_vae_encode.py index 1b43483a408..9d92b3819a4 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -19,6 +19,7 @@ from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 @invocation( @@ -46,7 +47,13 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc The VAE encodes to 32-channel latent space. Output latents shape: (B, 32, H/8, W/8). """ - with vae_info.model_on_device() as (_, vae): + # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to + # free room rather than discovering the shortfall as an OOM. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="encode", image_tensor=image_tensor, vae=vae_info.model + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype device = TorchDevice.choose_torch_device() image_tensor = image_tensor.to(device=device, dtype=vae_dtype) diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index 368f3c4452f..ccdb1390e28 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -21,12 +21,16 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux2.sampling_utils import pack_flux2 from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 # Maximum pixel counts for reference images (matches BFL FLUX.2 sampling.py) # Single reference image: 2024² pixels, Multiple: 1024² pixels MAX_PIXELS_SINGLE_REF = 2024**2 # ~4.1M pixels MAX_PIXELS_MULTI_REF = 1024**2 # ~1M pixels +# Tile size (in pixels) forced on the VAE for reference-image encoding, see _prepare_ref_images(). +REF_ENCODE_TILE_SIZE = 512 + def resize_image_to_max_pixels(image: Image.Image, max_pixels: int) -> Image.Image: """Resize image to fit within max_pixels while preserving aspect ratio. @@ -203,8 +207,17 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: image_tensor = image_tensor * 2.0 - 1.0 image_tensor = image_tensor.unsqueeze(0) # Add batch dimension - # Encode using FLUX.2 VAE - with vae_info.model_on_device() as (_, vae): + # Encode using FLUX.2 VAE. The encode below forces REF_ENCODE_TILE_SIZE tiling, so the + # peak is bounded by one tile; tell the cache that up front so it frees the room instead + # of hitting the shortfall as an OOM. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + tile_size=REF_ENCODE_TILE_SIZE, + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) @@ -219,8 +232,8 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: downsample = 2 ** (len(vae.config.block_out_channels) - 1) prev_tiling = (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size) vae.use_tiling = True - vae.tile_sample_min_size = 512 - vae.tile_latent_min_size = 512 // downsample + vae.tile_sample_min_size = REF_ENCODE_TILE_SIZE + vae.tile_latent_min_size = REF_ENCODE_TILE_SIZE // downsample try: # FLUX.2 VAE uses diffusers API latent_dist = vae.encode(image_tensor, return_dict=False)[0] diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index bd780d4c0b3..8ccddbf99d0 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -2,6 +2,7 @@ import torch from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny @@ -98,6 +99,42 @@ def estimate_vae_working_memory_flux( return int(working_memory) +def estimate_vae_working_memory_flux2( + operation: Literal["encode", "decode"], + image_tensor: torch.Tensor, + vae: AutoencoderKLFlux2, + tile_size: int | None = None, +) -> int: + """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. + + Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE -- + ``AutoencoderKLFlux2``'s mid-block attention runs through SDPA, so no O(area^2) term appears. + Measured on CUDA/bf16 as peak *reserved* memory (the conservative quantity, including allocator + overhead), the implied constants are ~2170 (decode) and ~1070 (encode) bytes per pixel per + element byte, flat across 512-1536px; the constants below round those up and match the FLUX.1 + ones. For reference, decoding 1024x1024 peaks at ~4.3GB and 1536x1536 at ~9.6GB -- far above the + default ``device_working_mem_gb``, which is why this estimate must be passed to the model cache. + + When tiling is enabled the peak is bounded by a single tile instead of the full image (measured + ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap). + """ + element_size = next(vae.parameters()).element_size() + + # Encoding uses ~50% the working memory of decoding. + scaling_constant = 2200 if operation == "decode" else 1100 + + if tile_size is not None: + # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. + working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 + else: + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] + out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] + working_memory = out_h * out_w * element_size * scaling_constant + + return int(working_memory) + + def estimate_vae_working_memory_anima( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py new file mode 100644 index 00000000000..816a7f71f37 --- /dev/null +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -0,0 +1,287 @@ +"""FLUX.2 working-memory estimates: the transformer denoise and both VAE directions. + +The FLUX.2 path originally called `model_on_device()` with no `working_mem_bytes` anywhere, so the +model cache reserved only the small default `device_working_mem_gb` and filled the rest of the card +with the model. Reference images make that fatal rather than merely tight: their latents are +concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of +a 1024x1024 generation. See https://github.com/invoke-ai/InvokeAI/issues/9500. + +The `MEASURED_*` tables below are peak *reserved* memory measured on CUDA in bf16 (the conservative +quantity, including allocator overhead). Every estimate must stay an upper bound on them. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 + +from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation +from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation +from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 + +MB = 1024**2 +GB = 1024**3 + + +def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, regional_bias=0): + return Flux2DenoiseInvocation._estimate_working_memory( + MagicMock(spec=Flux2DenoiseInvocation), + image_seq_len=image_seq_len, + ref_image_seq_len=ref_image_seq_len, + text_seq_len=text_seq_len, + num_loras=num_loras, + regional_attention_bias_bytes=regional_bias, + ) + + +class TestFlux2DenoiseWorkingMemoryEstimate: + # (image tokens, reference tokens, measured peak reserved MB) on the Klein 9B geometry. + # Token grids are pixels/16, so a 1024px square is 4096 tokens. + MEASURED_DENOISE = [ + (1024, 0, 448), # 512px + (4096, 0, 1702), # 1024px + (4096, 4096, 3324), # 1024px + one 1024px reference + (4096, 8192, 4840), # + two references + (4096, 12288, 6538), # + three references (the tiled-refiner case from #9500) + (6889, 12288, 7528), # 1328px tile + three 1024px references + (6889, 20667, 10954), # 1328px tile + three 1328px references + (16384, 0, 6538), # 2048px, no references + ] + + @pytest.mark.parametrize("image_seq_len, ref_image_seq_len, measured_mb", MEASURED_DENOISE) + def test_estimate_is_an_upper_bound_on_measured_peak(self, image_seq_len, ref_image_seq_len, measured_mb): + """The cache treats the estimate as the amount it must keep free, so under-estimating OOMs.""" + assert _estimate(image_seq_len, ref_image_seq_len) >= measured_mb * MB + + @pytest.mark.parametrize("image_seq_len, ref_image_seq_len, measured_mb", MEASURED_DENOISE) + def test_estimate_does_not_wildly_over_reserve(self, image_seq_len, ref_image_seq_len, measured_mb): + """Over-estimating is not free: the cache offloads the transformer to RAM to honor the + reservation, and a model running over PCIe is indistinguishable from a hang.""" + assert _estimate(image_seq_len, ref_image_seq_len) <= measured_mb * MB + 2 * GB + + def test_reference_image_tokens_are_counted(self): + """The regression this whole module exists for: reference tokens are attended like image + tokens and cost the same per token, so they must enter the estimate.""" + without_refs = _estimate(image_seq_len=4096) + with_refs = _estimate(image_seq_len=4096, ref_image_seq_len=12288) + assert with_refs - without_refs == 12288 * int(0.4 * MB) + + def test_estimate_is_linear_in_total_sequence(self): + """Attention runs through SDPA, so there is no O(seq^2) term to model -- image, reference and + text tokens are interchangeable at the same per-token cost.""" + assert _estimate(image_seq_len=8192) == _estimate(image_seq_len=4096, ref_image_seq_len=4096) + assert _estimate(image_seq_len=4096, text_seq_len=1024) - _estimate(image_seq_len=4096, text_seq_len=512) == ( + 512 * int(0.4 * MB) + ) + + def test_lora_margin_is_added_per_lora(self): + """Sidecar-patched LoRAs add an activation branch per patched layer.""" + base = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, num_loras=1) - base == int(0.5 * GB) + assert _estimate(image_seq_len=4096, num_loras=3) - base == int(1.5 * GB) + + def test_regional_attention_bias_is_added(self): + base = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, regional_bias=123 * MB) - base == 123 * MB + + +class TestFlux2VaeWorkingMemoryEstimate: + # (operation, pixel size, measured peak reserved MB), bf16, untiled. + MEASURED_VAE = [ + ("decode", 512, 1086), + ("decode", 768, 2414), + ("decode", 1024, 4260), + ("decode", 1328, 7146), + ("decode", 1536, 9578), + ("encode", 512, 536), + ("encode", 1024, 2122), + ("encode", 1328, 3022), + ] + + def _mock_bf16_vae(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) # element_size == 2 + return vae + + def _tensor_for(self, operation, px): + # decode receives 32-channel latents at pixels/8; encode receives a pixel image. + return torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) + + @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) + def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae() + ) + assert estimate >= measured_mb * MB + + @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) + def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae() + ) + assert estimate == 1024 * 1024 * 2 * expected_constant + + def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): + """Reference-image encoding forces 512px tiling precisely so the peak stops following the + reference resolution -- measured flat at ~0.55GB from 1024px up to the 2024px reference cap.""" + estimates = [ + estimate_vae_working_memory_flux2( + operation="encode", + image_tensor=torch.zeros(1, 3, px, px), + vae=self._mock_bf16_vae(), + tile_size=512, + ) + for px in (1024, 1328, 2024) + ] + assert len(set(estimates)) == 1 + assert estimates[0] == int(512 * 512 * 2 * 1100 * 1.25) + assert estimates[0] >= 558 * MB # measured tiled peak at 2024px + # The whole point of tiling: it must shrink the reservation, not just bound the VAE. + untiled = estimate_vae_working_memory_flux2( + operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae() + ) + assert estimates[0] < untiled / 4 + + +class TestFlux2VaeInvocationsRequestWorkingMemory: + """The estimate is worthless unless it reaches `model_on_device()`.""" + + def _mock_vae_info(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + + vae_info = MagicMock() + vae_info.model = vae + vae_info.compute_device = torch.device("cpu") + cm = MagicMock() + cm.__enter__ = MagicMock(return_value=(None, vae)) + cm.__exit__ = MagicMock(return_value=None) + vae_info.model_on_device = MagicMock(return_value=cm) + return vae_info + + def test_decode_requests_working_memory(self): + vae_info = self._mock_vae_info() + context = MagicMock() + context.models.load.return_value = vae_info + context.tensors.load.return_value = torch.zeros(1, 32, 128, 128) + + expected = 10 * GB + with patch( + "invokeai.app.invocations.flux2_vae_decode.estimate_vae_working_memory_flux2", return_value=expected + ) as estimate: + invocation = Flux2VaeDecodeInvocation.model_construct( + latents=MagicMock(latents_name="latents"), vae=MagicMock(vae=MagicMock()) + ) + try: + invocation.invoke(context) + except Exception: + # The mocked decode math fails downstream; we only care that the cache was asked to + # reserve the estimate before the device context was entered. + pass + + estimate.assert_called_once() + assert estimate.call_args.kwargs["operation"] == "decode" + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + + def test_encode_requests_working_memory(self): + vae_info = self._mock_vae_info() + context = MagicMock() + context.models.load.return_value = vae_info + + expected = 4 * GB + with ( + patch( + "invokeai.app.invocations.flux2_vae_encode.estimate_vae_working_memory_flux2", return_value=expected + ) as estimate, + patch( + "invokeai.app.invocations.flux2_vae_encode.image_resized_to_grid_as_tensor", + return_value=torch.zeros(3, 1024, 1024), + ), + ): + invocation = Flux2VaeEncodeInvocation.model_construct( + image=MagicMock(image_name="image"), vae=MagicMock(vae=MagicMock()) + ) + try: + invocation.invoke(context) + except Exception: + pass + + estimate.assert_called_once() + assert estimate.call_args.kwargs["operation"] == "encode" + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + + +class _StopBeforeLoad(Exception): + """Raised in place of entering the transformer's device context, to end _run_diffusion early.""" + + +class TestFlux2DenoiseRequestsWorkingMemory: + """The denoise node must hand its estimate to the cache, and that estimate must grow with the + attached reference images -- the combination that #9500 was missing.""" + + def _run(self, num_ref_tokens: int): + """Drive `_run_diffusion` up to the transformer load and return the requested working memory.""" + from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType + from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + ConditioningFieldData, + FLUXConditioningInfo, + ) + + transformer_info = MagicMock() + transformer_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) + + context = MagicMock() + context.models.load.return_value = transformer_info + context.models.get_config.return_value = MagicMock( + base=BaseModelType.Flux2, type=ModelType.Main, format=ModelFormat.Checkpoint + ) + context.conditioning.load.return_value = ConditioningFieldData( + conditionings=[FLUXConditioningInfo(clip_embeds=torch.zeros(1, 768), t5_embeds=torch.zeros(1, 512, 12288))] + ) + + ref_extension = MagicMock() + ref_extension.ref_image_latents = torch.zeros(1, num_ref_tokens, 128) + + invocation = Flux2DenoiseInvocation.model_construct( + latents=None, + noise=None, + denoise_mask=None, + denoising_start=0.0, + denoising_end=1.0, + add_noise=True, + transformer=MagicMock(transformer=MagicMock(), loras=[]), + positive_text_conditioning=MagicMock(conditioning_name="pos", mask=None), + negative_text_conditioning=None, + guidance=4.0, + cfg_scale=1.0, + width=1024, + height=1024, + num_steps=4, + scheduler="euler", + seed=0, + vae=MagicMock(vae=MagicMock()), + kontext_conditioning=MagicMock() if num_ref_tokens else None, + ) + + with ( + patch.object(Flux2DenoiseInvocation, "_get_bn_stats", return_value=None), + patch("invokeai.backend.util.devices.TorchDevice.choose_torch_device", return_value=torch.device("cpu")), + patch("invokeai.app.invocations.flux2_denoise.Flux2RefImageExtension", return_value=ref_extension), + pytest.raises(_StopBeforeLoad), + ): + invocation._run_diffusion(context) + + transformer_info.model_on_device.assert_called_once() + return transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"] + + def test_estimate_reaches_the_model_cache(self): + """Without this the cache reserves only the default `device_working_mem_gb`.""" + assert self._run(num_ref_tokens=0) == _estimate(image_seq_len=64 * 64, text_seq_len=512) + + def test_reference_images_raise_the_reservation(self): + """Three 1024x1024 references add 12288 tokens to a 1024x1024 generation's 4096.""" + without_refs = self._run(num_ref_tokens=0) + with_refs = self._run(num_ref_tokens=12288) + assert with_refs - without_refs == 12288 * int(0.4 * MB) From b457f60cae4a38527ec6ebe0e37237d62c238d8e Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 24 Aug 2026 22:40:46 +0200 Subject: [PATCH 2/5] fix(flux2): budget SDPA's materialized score matrix where it is real The FLUX.2 working-memory estimates were linear in the sequence length, which holds only while SDPA picks a fused kernel. That is a property of the torch build, not of FLUX.2: ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks, so both the VAE's 512-wide mid-block head and the dense S x S bias regional prompting attaches fall through to the math fallback and materialize the score matrix -- ~17GB for a 1536px decode, and heads x S^2 for a masked forward. Rather than assume either way, ask torch: sdpa_score_matrix_bytes() queries can_use_flash/efficient/cudnn_attention for the real head dim, dtype and mask, and adds 13 bytes per score element only when no fused kernel is eligible. Measured on CUDA with SDPBackend.MATH forced: 12.9 bytes/element at 4k tokens, 10.3 at 8k, 9.7 at 16k, identical for bf16, fp16 and fp32 because the fallback's softmax intermediates are always fp32. On CUDA every shape reports fused, so the term is zero and the existing calibration is untouched. Non-CUDA devices keep the fused assumption -- torch exposes no equivalent query there, and guessing would reserve double-digit GB on no evidence. --- invokeai/app/invocations/flux2_denoise.py | 41 +++- invokeai/app/invocations/flux2_vae_decode.py | 2 +- invokeai/app/invocations/flux2_vae_encode.py | 5 +- invokeai/backend/flux2/ref_image_extension.py | 1 + invokeai/backend/util/attention.py | 80 +++++++ invokeai/backend/util/vae_working_memory.py | 39 ++- .../invocations/test_flux2_working_memory.py | 222 +++++++++++++++++- 7 files changed, 374 insertions(+), 16 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 25c8b5d8f30..6dcb412e39e 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -49,8 +49,16 @@ from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo +from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice +# FLUX.2 attention geometry. The head dim is 128 across every variant; the head count follows the +# hidden size (Klein 4B: 24, Klein 9B: 32, FLUX.2 dev: 48). Only the head dim decides which SDPA +# kernel is eligible; the head count scales the `math` fallback's score matrix, and since the +# working-memory estimate is computed before the transformer is loaded, we use the largest. +FLUX2_ATTENTION_HEAD_DIM = 128 +FLUX2_MAX_ATTENTION_HEADS = 48 + @invocation( "flux2_denoise", @@ -465,18 +473,23 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # concatenated onto the image stream, so three 1024x1024 references quadruple the sequence # (and with it the activation footprint) of a 1024x1024 generation. ref_image_seq_len = ref_image_extension.ref_image_latents.shape[1] if ref_image_extension is not None else 0 + # The additive bias is skipped entirely when reference images are present (see below), so the + # mask only costs anything -- storage, and possibly a materialized score matrix -- without them. + regional_attn_mask = regional_extension.restricted_attn_mask if ref_image_seq_len == 0 else None estimated_working_memory = self._estimate_working_memory( image_seq_len=packed_h * packed_w, ref_image_seq_len=ref_image_seq_len, text_seq_len=max(txt.shape[1], neg_txt.shape[1] if neg_txt is not None else 0), num_loras=len(self.transformer.loras), # The mask itself is already allocated; only the additive bias built per forward is new. - # It is skipped entirely when reference images are present (see below). regional_attention_bias_bytes=( - regional_extension.restricted_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() - if regional_extension.restricted_attn_mask is not None and ref_image_seq_len == 0 + regional_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() + if regional_attn_mask is not None else 0 ), + has_regional_attention_mask=regional_attn_mask is not None, + device=device, + dtype=inference_dtype, ) with ExitStack() as exit_stack: @@ -608,6 +621,9 @@ def _estimate_working_memory( text_seq_len: int, num_loras: int, regional_attention_bias_bytes: int = 0, + has_regional_attention_mask: bool = False, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, ) -> int: """Estimate peak transformer activation memory (bytes) so the model cache reserves enough headroom. @@ -626,13 +642,30 @@ def _estimate_working_memory( A fixed base covers resolution-independent overhead (transient fp8/GGUF -> bf16 weight casts during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra activation branch per patched layer, so we add a per-LoRA margin. + + The linear model holds only while SDPA picks a fused kernel. Regional prompting is where that + stops being a given: it hands the transformer a dense additive ``S x S`` bias, which flash + attention never accepts and which ROCm's memory-efficient kernel rejects as well, leaving the + ``math`` fallback and its materialized ``heads x S x S`` score matrix. We ask torch which path + this build will take for these shapes and add the score matrix only when it is really there -- + on CUDA the memory-efficient kernel takes the bias and the term is zero (verified: peak stays + linear with the bias attached). """ GB = 1024**3 MB = 1024**2 per_token_bytes = int(0.4 * MB) - estimated = (image_seq_len + ref_image_seq_len + text_seq_len) * per_token_bytes + total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len + estimated = total_seq_len * per_token_bytes estimated += int(1.0 * GB) estimated += regional_attention_bias_bytes + estimated += sdpa_score_matrix_bytes( + device=device if device is not None else TorchDevice.choose_torch_device(), + dtype=dtype, + num_heads=FLUX2_MAX_ATTENTION_HEADS, + head_dim=FLUX2_ATTENTION_HEAD_DIM, + seq_len=total_seq_len, + has_attn_mask=has_regional_attention_mask, + ) if num_loras > 0: estimated += int(0.5 * num_loras * GB) return estimated diff --git a/invokeai/app/invocations/flux2_vae_decode.py b/invokeai/app/invocations/flux2_vae_decode.py index d4cad4ad75d..f0852f1880f 100644 --- a/invokeai/app/invocations/flux2_vae_decode.py +++ b/invokeai/app/invocations/flux2_vae_decode.py @@ -54,7 +54,7 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima # far above the default working memory the cache would otherwise reserve. Tell it up front so # it offloads enough of the (possibly still resident) transformer to leave room. estimated_working_memory = estimate_vae_working_memory_flux2( - operation="decode", image_tensor=latents, vae=vae_info.model + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/app/invocations/flux2_vae_encode.py b/invokeai/app/invocations/flux2_vae_encode.py index 9d92b3819a4..2da6f38b517 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -50,7 +50,10 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to # free room rather than discovering the shortfall as an OOM. estimated_working_memory = estimate_vae_working_memory_flux2( - operation="encode", image_tensor=image_tensor, vae=vae_info.model + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + device=TorchDevice.choose_torch_device(), ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index ccdb1390e28..9184b15c1d2 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -215,6 +215,7 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: image_tensor=image_tensor, vae=vae_info.model, tile_size=REF_ENCODE_TILE_SIZE, + device=TorchDevice.choose_torch_device(), ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 1df0f99280b..dd75a710446 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,6 +4,8 @@ for attention mechanism. """ +from functools import lru_cache + import psutil import torch @@ -35,3 +37,81 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: return "max" else: return "balanced" + + +# SDPA computes attention one of two ways: a fused kernel (flash / memory-efficient / cuDNN) that +# never materializes the O(S^2) score matrix, or the `math` fallback, which does. torch picks per +# call from the dtype, the head dim and whether an attention mask was passed -- and the answer +# differs between builds. CUDA's memory-efficient kernel accepts head dims well past 128 and +# arbitrary additive masks; ROCm's fused kernels reject both and drop to `math`. A working-memory +# estimate that assumes the fused path is therefore only correct on the build it was measured on, +# which is why the helper below asks torch instead of assuming. + +# Peak *reserved* bytes per element of the materialized score matrix, measured on CUDA with +# `SDPBackend.MATH` forced, each point in a fresh process: 12.9 bytes/element at 4k tokens, 10.3 at +# 8k, 9.7 at 16k -- and the same figures for bf16, fp16 and fp32 inputs, because the fallback's +# softmax intermediates are fp32 regardless. So this is an absolute byte count, not a multiple of +# the element size. 13 is an upper bound on every measured point from 4k tokens up; below that it +# can fall a couple of MB short of the allocator's rounding, which is noise next to the GB-scale +# linear terms this is added to. +SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 + + +@lru_cache(maxsize=None) +def _sdpa_has_fused_kernel( + device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool +) -> bool: + """Ask torch whether any non-materializing SDPA kernel is eligible for these attention shapes. + + Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence + length, so a tiny probe answers for the real forward. Falls back to ``True`` (the status quo + assumption) whenever torch gives us nothing to go on -- over-reserving many GB on a guess would + push the model out of VRAM and be worse than the shortfall we are trying to avoid. + """ + if device_type != "cuda": + # `can_use_*_attention` is CUDA/ROCm-only. MPS, XPU and CPU all ship fused SDPA kernels, so + # keep the fused assumption there rather than guessing at their dispatch rules. + return True + + try: + from torch.backends.cuda import SDPAParams, can_use_efficient_attention, can_use_flash_attention + + device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) + q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) + mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None + try: + params = SDPAParams(q, q, q, mask, 0.0, False, False) + except TypeError: + # torch < 2.5: no `enable_gqa` field. + params = SDPAParams(q, q, q, mask, 0.0, False) + + checks = [can_use_flash_attention, can_use_efficient_attention] + can_use_cudnn_attention = getattr(torch.backends.cuda, "can_use_cudnn_attention", None) + if can_use_cudnn_attention is not None: + checks.append(can_use_cudnn_attention) + return any(check(params, False) for check in checks) + except Exception: + return True + + +def sdpa_score_matrix_bytes( + *, + device: torch.device, + dtype: torch.dtype, + num_heads: int, + head_dim: int, + seq_len: int, + has_attn_mask: bool = False, +) -> int: + """Bytes SDPA spends on a materialized score matrix for one attention call, 0 if fused. + + Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On + CUDA it is almost always 0; on a build whose fused kernels reject the shapes (notably ROCm, + which caps the head dim at 128 and does not take arbitrary additive masks) it is the dominant + term -- a 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~17GB of them. + """ + if seq_len <= 0 or num_heads <= 0: + return 0 + if _sdpa_has_fused_kernel(device.type, device.index, dtype, head_dim, has_attn_mask): + return 0 + return num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 8ccddbf99d0..849efefbad7 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -9,6 +9,8 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.backend.flux.modules.autoencoder import AutoEncoder +from invokeai.backend.util.attention import sdpa_score_matrix_bytes +from invokeai.backend.util.devices import TorchDevice _WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT = 2900 _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B = 6500 @@ -99,26 +101,45 @@ def estimate_vae_working_memory_flux( return int(working_memory) +# The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the +# decoder, on the 8x-downsampled grid. Both are single-head, with the head dim set to the block +# width: 512 for the stock VAE and 384 for the small-decoder variant. The distinction does not +# matter here -- what matters is that both sit far above the 128 head dim ROCm's fused SDPA kernels +# accept, so only the value's side of that limit is load-bearing. +_FLUX2_VAE_MID_BLOCK_HEADS = 1 +_FLUX2_VAE_MID_BLOCK_HEAD_DIM = 512 +_FLUX2_VAE_SPATIAL_COMPRESSION = 8 + + def estimate_vae_working_memory_flux2( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLFlux2, tile_size: int | None = None, + device: torch.device | None = None, ) -> int: """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. - Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE -- - ``AutoencoderKLFlux2``'s mid-block attention runs through SDPA, so no O(area^2) term appears. + Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE. Measured on CUDA/bf16 as peak *reserved* memory (the conservative quantity, including allocator overhead), the implied constants are ~2170 (decode) and ~1070 (encode) bytes per pixel per element byte, flat across 512-1536px; the constants below round those up and match the FLUX.1 ones. For reference, decoding 1024x1024 peaks at ~4.3GB and 1536x1536 at ~9.6GB -- far above the default ``device_working_mem_gb``, which is why this estimate must be passed to the model cache. + That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a + fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the + 512-wide head, and measured peak stays linear from 512 to 1536px). A build whose fused kernels + reject the head dim -- ROCm caps it at 128 -- drops to SDPA's ``math`` fallback and materializes + a (pixels/8)^2 score matrix on top of the linear term: ~3.5GB at 1024px and ~17GB at 1536px. We + ask torch which path applies rather than assuming, so the estimate is right on both. + When tiling is enabled the peak is bounded by a single tile instead of the full image (measured - ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap). + ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score + matrix, if one is materialized at all, is bounded by the tile too. """ - element_size = next(vae.parameters()).element_size() + param = next(vae.parameters()) + element_size = param.element_size() # Encoding uses ~50% the working memory of decoding. scaling_constant = 2200 if operation == "decode" else 1100 @@ -126,11 +147,21 @@ def estimate_vae_working_memory_flux2( if tile_size is not None: # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 + mid_block_seq_len = (tile_size // _FLUX2_VAE_SPATIAL_COMPRESSION) ** 2 else: latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] working_memory = out_h * out_w * element_size * scaling_constant + mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) + + working_memory += sdpa_score_matrix_bytes( + device=device if device is not None else TorchDevice.choose_torch_device(), + dtype=param.dtype, + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS, + head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, + seq_len=mid_block_seq_len, + ) return int(working_memory) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 816a7f71f37..aac5409930d 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -14,18 +14,34 @@ import pytest import torch +import torch.nn.functional as F from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 -from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation +from invokeai.app.invocations.flux2_denoise import FLUX2_MAX_ATTENTION_HEADS, Flux2DenoiseInvocation from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation +from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 MB = 1024**2 GB = 1024**3 - -def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, regional_bias=0): +# The measured tables in this module were all taken on CUDA, where SDPA runs a fused kernel and no +# score matrix is materialized. `_sdpa_has_fused_kernel` reports non-CUDA devices as fused, so +# passing a CPU device reproduces that regime without needing a GPU on the test runner. The +# materializing regime gets its own class below. +FUSED = torch.device("cpu") + + +def _estimate( + image_seq_len, + ref_image_seq_len=0, + text_seq_len=512, + num_loras=0, + regional_bias=0, + has_regional_mask=False, + device=FUSED, +): return Flux2DenoiseInvocation._estimate_working_memory( MagicMock(spec=Flux2DenoiseInvocation), image_seq_len=image_seq_len, @@ -33,6 +49,8 @@ def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, text_seq_len=text_seq_len, num_loras=num_loras, regional_attention_bias_bytes=regional_bias, + has_regional_attention_mask=has_regional_mask, + device=device, ) @@ -112,14 +130,14 @@ def _tensor_for(self, operation, px): @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae() + operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae(), device=FUSED ) assert estimate >= measured_mb * MB @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae() + operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae(), device=FUSED ) assert estimate == 1024 * 1024 * 2 * expected_constant @@ -132,6 +150,7 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): image_tensor=torch.zeros(1, 3, px, px), vae=self._mock_bf16_vae(), tile_size=512, + device=FUSED, ) for px in (1024, 1328, 2024) ] @@ -140,7 +159,7 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): assert estimates[0] >= 558 * MB # measured tiled peak at 2024px # The whole point of tiling: it must shrink the reservation, not just bound the VAE. untiled = estimate_vae_working_memory_flux2( - operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae() + operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae(), device=FUSED ) assert estimates[0] < untiled / 4 @@ -285,3 +304,194 @@ def test_reference_images_raise_the_reservation(self): without_refs = self._run(num_ref_tokens=0) with_refs = self._run(num_ref_tokens=12288) assert with_refs - without_refs == 12288 * int(0.4 * MB) + + +def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): + """Stand in for `_sdpa_has_fused_kernel` on a build with ROCm's fused-kernel rules. + + ROCm's fused SDPA kernels cap the head dim at 128 and do not take an arbitrary additive mask; + anything else falls through to the `math` fallback, which materializes the score matrix. CUDA's + memory-efficient kernel accepts both -- verified on torch 2.7.1+cu128, where + `can_use_efficient_attention` is true for the VAE's 512-wide head and for a masked 128-wide + transformer head, and measured peak stays linear in both cases -- which is why the estimates + were linear to begin with. + """ + return head_dim <= 128 and not has_attn_mask + + +def _materializing(): + return patch("invokeai.backend.util.attention._sdpa_has_fused_kernel", side_effect=_rocm_like_probe) + + +# Any CUDA device object works here: the probe is patched out, so nothing is allocated on it. +MATERIALIZING = torch.device("cuda") + + +class TestMaterializedScoreMatrixIsBudgeted: + """The linear estimates above assume SDPA never builds the O(S^2) score matrix. That is a + property of the *build*, not of FLUX.2: ROCm's fused kernels reject both the VAE's 512-wide + attention head and the dense additive mask regional prompting attaches, and fall back to + `math`. Where that happens the score matrix is the dominant term, so the estimate has to + include it -- otherwise the fix works on CUDA and still OOMs on ROCm. + """ + + def _mock_bf16_vae(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return vae + + def _vae_estimate(self, operation, px, device, tile_size=None): + tensor = torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) + return estimate_vae_working_memory_flux2( + operation=operation, + image_tensor=tensor, + vae=self._mock_bf16_vae(), + tile_size=tile_size, + device=device, + ) + + # (operation, pixel size, mid-block tokens). The VAE attends on the 8x-downsampled grid. + @pytest.mark.parametrize( + "operation, px, tokens", + [ + ("decode", 1024, 128 * 128), + ("decode", 1536, 192 * 192), + ("encode", 1024, 128 * 128), + ("encode", 1328, 166 * 166), + ], + ) + def test_vae_estimate_gains_exactly_the_score_matrix(self, operation, px, tokens): + fused = self._vae_estimate(operation, px, device=FUSED) + with _materializing(): + materializing = self._vae_estimate(operation, px, device=MATERIALIZING) + assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + def test_vae_score_matrix_dominates_at_high_resolution(self): + """The reviewer's case: a 1536px decode is ~9.6GB of linear activations on CUDA, and more + than that again in scores where SDPA has to materialize them. An estimate that omits the + term is not merely tight, it is wrong by a factor of three.""" + with _materializing(): + materializing = self._vae_estimate("decode", 1536, device=MATERIALIZING) + assert materializing > 25 * GB + assert materializing > 2 * self._vae_estimate("decode", 1536, device=FUSED) + + def test_tiled_vae_score_matrix_is_bounded_by_the_tile(self): + """Tiling already bounds the linear term; it must bound the quadratic one too, or the + reference-image encode would reserve as if it ran untiled.""" + with _materializing(): + estimates = [ + self._vae_estimate("encode", px, device=MATERIALIZING, tile_size=512) for px in (1024, 1328, 2024) + ] + untiled = self._vae_estimate("encode", 2024, device=MATERIALIZING) + assert len(set(estimates)) == 1 + tile_tokens = (512 // 8) ** 2 + assert estimates[0] - self._vae_estimate("encode", 1024, device=FUSED, tile_size=512) == ( + tile_tokens * tile_tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + # Tiling turns a 62GB reservation at the 2024px reference cap into under 1GB. + assert estimates[0] < untiled / 50 + + def test_regional_prompting_adds_the_score_matrix(self): + """The dense `S x S` additive bias is what pushes SDPA off its fused kernel. Budgeting only + the bias tensor -- as this PR first did -- under-reserves by the score matrix, which is two + orders of magnitude larger.""" + seq_len = 4096 + 512 + bias_bytes = seq_len * seq_len * 2 + fused = _estimate( + image_seq_len=4096, text_seq_len=512, regional_bias=bias_bytes, has_regional_mask=True, device=FUSED + ) + with _materializing(): + materializing = _estimate( + image_seq_len=4096, + text_seq_len=512, + regional_bias=bias_bytes, + has_regional_mask=True, + device=MATERIALIZING, + ) + assert materializing - fused == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + assert materializing - fused > 50 * bias_bytes + + def test_denoise_without_a_regional_mask_is_unaffected(self): + """FLUX.2's 128-wide attention head is inside every backend's fused limit, so an ordinary + generation -- reference images included -- keeps the plain linear estimate. This term exists + for the masked case; it must not tax the common one.""" + with _materializing(): + assert _estimate(image_seq_len=4096, ref_image_seq_len=12288, device=MATERIALIZING) == _estimate( + image_seq_len=4096, ref_image_seq_len=12288, device=FUSED + ) + + +class TestSdpaBackendProbe: + """`sdpa_score_matrix_bytes` decides the term above, so its defaults are load-bearing.""" + + def test_non_cuda_devices_keep_the_fused_assumption(self): + """torch exposes no eligibility query outside CUDA/ROCm. Guessing `math` there would add + double-digit GB to every estimate on MPS and CPU on no evidence at all.""" + assert ( + sdpa_score_matrix_bytes( + device=torch.device("cpu"), + dtype=torch.bfloat16, + num_heads=1, + head_dim=512, + seq_len=16384, + has_attn_mask=True, + ) + == 0 + ) + + def test_empty_sequences_cost_nothing(self): + with _materializing(): + assert ( + sdpa_score_matrix_bytes( + device=MATERIALIZING, dtype=torch.bfloat16, num_heads=48, head_dim=128, seq_len=0 + ) + == 0 + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="asks the real CUDA/ROCm dispatcher") + def test_this_build_reports_its_own_dispatch(self): + """On CUDA both shapes are fused and this whole term is zero -- the fix is a no-op for the + hardware the constants were measured on. On ROCm the same call reports the fallback and the + term appears. Either answer is correct; the point is that it comes from torch.""" + vae_bytes = sdpa_score_matrix_bytes( + device=torch.device("cuda"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + masked_bytes = sdpa_score_matrix_bytes( + device=torch.device("cuda"), + dtype=torch.bfloat16, + num_heads=48, + head_dim=128, + seq_len=4608, + has_attn_mask=True, + ) + if torch.version.hip is None: + assert vae_bytes == 0 + assert masked_bytes == 0 + else: + assert vae_bytes > 0 + assert masked_bytes > 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="measures real peak reserved memory") + @pytest.mark.parametrize("num_heads, seq_len, head_dim", [(1, 4096, 512), (4, 4096, 128)]) + def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, head_dim): + """Pin the bytes-per-score-element calibration against a real `math` forward, so a future + edit to the constant cannot silently reintroduce the shortfall it exists to cover.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + device = torch.device("cuda") + q = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_reserved() + with sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): + F.scaled_dot_product_attention(q, k, v) + torch.cuda.synchronize() + measured = torch.cuda.max_memory_reserved() - before + + estimate = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimate >= measured + assert estimate <= 2 * measured From 119664a243f993de4db6015884515793da767ef7 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 26 Aug 2026 05:01:00 +0200 Subject: [PATCH 3/5] fix(flux2): ask the real dispatcher which SDPA path a build takes The score-matrix term probed torch's CUDA eligibility helpers and read everything else as fused. That was wrong twice over: MPS has no fused SDPA kernel at all and runs the MPSGraph math transcription, so a 1024px VAE decode was admitted ~3.5GB short; and a failed probe returned "fused" too, turning "we don't know" into the one answer that can OOM. Ask `_fused_sdp_choice` instead -- the same dispatch query `scaled_dot_product_attention` runs to pick its kernel. Torch registers it for CPU, CUDA/ROCm and XPU only, so the call raises on exactly the devices that fall through to `math`, and every other failure lands on the conservative side by the same branch. Diffusers models do not reach torch's SDPA directly, so also consult `dispatch_attention_fn`'s active backend: a user on `_native_math` materializes the score matrix on hardware whose probe reports fused. Only the transformer needs this -- the FLUX.2 VAE's mid-block attention still calls SDPA itself through `AttnProcessor2_0` -- and a test pins that asymmetry. On CUDA with the stock backend every one of these terms remains zero. --- invokeai/app/invocations/flux2_denoise.py | 19 ++- invokeai/backend/util/attention.py | 144 ++++++++++++---- invokeai/backend/util/vae_working_memory.py | 11 +- .../invocations/test_flux2_working_memory.py | 161 ++++++++++++++++-- 4 files changed, 277 insertions(+), 58 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 6dcb412e39e..e6655b5c593 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -643,13 +643,15 @@ def _estimate_working_memory( during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra activation branch per patched layer, so we add a per-LoRA margin. - The linear model holds only while SDPA picks a fused kernel. Regional prompting is where that - stops being a given: it hands the transformer a dense additive ``S x S`` bias, which flash - attention never accepts and which ROCm's memory-efficient kernel rejects as well, leaving the - ``math`` fallback and its materialized ``heads x S x S`` score matrix. We ask torch which path - this build will take for these shapes and add the score matrix only when it is really there -- - on CUDA the memory-efficient kernel takes the bias and the term is zero (verified: peak stays - linear with the bias attached). + The linear model holds only while attention runs on a fused kernel. Regional prompting is + where that stops being a given: it hands the transformer a dense additive ``S x S`` bias, + which flash attention never accepts and which ROCm's memory-efficient kernel rejects as + well, leaving the ``math`` fallback and its materialized ``heads x S x S`` score matrix. The + device decides too -- MPS has no fused SDPA kernel at all -- and so does the diffusers + attention backend this build dispatches through. ``sdpa_score_matrix_bytes`` asks all three + and adds the score matrix only where it is really built: on CUDA with the stock backend the + memory-efficient kernel takes the bias and the term is zero (verified: peak stays linear + with the bias attached). """ GB = 1024**3 MB = 1024**2 @@ -665,6 +667,9 @@ def _estimate_working_memory( head_dim=FLUX2_ATTENTION_HEAD_DIM, seq_len=total_seq_len, has_attn_mask=has_regional_attention_mask, + # The FLUX.2 transformer's attention goes through diffusers' `dispatch_attention_fn`, + # which can route around torch's SDPA entirely -- including to a forced `math` backend. + via_diffusers_dispatch=True, ) if num_loras > 0: estimated += int(0.5 * num_loras * GB) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index dd75a710446..2245f6c750b 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,12 +4,15 @@ for attention mechanism. """ +import warnings from functools import lru_cache import psutil import torch +from torch.nn.attention import SDPBackend from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger def auto_detect_slice_size(latents: torch.Tensor) -> str: @@ -40,12 +43,12 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # SDPA computes attention one of two ways: a fused kernel (flash / memory-efficient / cuDNN) that -# never materializes the O(S^2) score matrix, or the `math` fallback, which does. torch picks per -# call from the dtype, the head dim and whether an attention mask was passed -- and the answer -# differs between builds. CUDA's memory-efficient kernel accepts head dims well past 128 and -# arbitrary additive masks; ROCm's fused kernels reject both and drop to `math`. A working-memory -# estimate that assumes the fused path is therefore only correct on the build it was measured on, -# which is why the helper below asks torch instead of assuming. +# never materializes the O(S^2) score matrix, or the `math` fallback, which does. Which one runs is +# not a property of FLUX.2 -- it is a property of the build, the device, the dtype, the head dim and +# whether an attention mask was passed. CUDA's memory-efficient kernel accepts head dims well past +# 128 and arbitrary additive masks; ROCm's fused kernels reject both; MPS has no fused SDPA kernel +# at all. A working-memory estimate that assumes the fused path is therefore only correct on the +# build it was measured on, which is why the helpers below ask instead of assuming. # Peak *reserved* bytes per element of the materialized score matrix, measured on CUDA with # `SDPBackend.MATH` forced, each point in a fresh process: 12.9 bytes/element at 4k tokens, 10.3 at @@ -53,46 +56,98 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # softmax intermediates are fp32 regardless. So this is an absolute byte count, not a multiple of # the element size. 13 is an upper bound on every measured point from 4k tokens up; below that it # can fall a couple of MB short of the allocator's rounding, which is noise next to the GB-scale -# linear terms this is added to. +# linear terms this is added to. It is a CUDA measurement standing in for every materializing +# backend -- no other was available to calibrate against -- but the intermediates it prices (an +# fp32 score matrix and its softmax) have the same shape wherever the fallback runs. SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 +# `_fused_sdp_choice` reports which kernel `F.scaled_dot_product_attention` would pick. These are +# the answers that mean "a fused kernel"; `MATH` -- and `ERROR`, which torch returns when it cannot +# pick anything at all -- mean the score matrix gets built. +_FUSED_SDP_CHOICES = frozenset( + int(getattr(SDPBackend, name)) + for name in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION", "OVERRIDEABLE") + if hasattr(SDPBackend, name) +) + +_DISPATCH_TORCH = "torch" +_DISPATCH_FUSED = "fused" +_DISPATCH_MATH = "math" + + +@lru_cache(maxsize=None) +def _diffusers_attention_dispatch() -> str: + """Report how the diffusers attention dispatcher will route a diffusers model's attention calls. + + Diffusers models do not call `F.scaled_dot_product_attention` directly -- they go through + `dispatch_attention_fn`, which honours the `DIFFUSERS_ATTN_BACKEND` environment variable and the + `attention_backend()` context manager. Only the default `native` backend hands the call to + torch; the others pin a specific kernel, and `_native_math` pins the materializing one. A + torch-level probe alone would report "fused" for a user who has forced math. + + Returns ``_DISPATCH_TORCH`` when torch decides, ``_DISPATCH_FUSED`` for a backend that never + materializes the score matrix, or ``_DISPATCH_MATH`` when one is built -- including when we + cannot tell, since under-reserving is the failure this whole term exists to prevent. + """ + try: + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + backend, _ = _AttentionBackendRegistry.get_active_backend() + name = str(getattr(backend, "value", backend)) + except Exception: + # A private diffusers attribute that moved, or a selected backend whose kernel failed to + # register. Budget the materializing case, but say so: silently adding several GB to every + # FLUX.2 estimate is not something that should pass unnoticed. + InvokeAILogger.get_logger(__name__).warning( + "Could not determine the active diffusers attention backend; budgeting working memory as if " + "attention materializes its score matrix. Estimates will be conservative." + ) + return _DISPATCH_MATH + + if name == "native": + return _DISPATCH_TORCH + if "math" in name: + return _DISPATCH_MATH + # Every other backend diffusers offers -- flash, sage, xformers, flex, aiter, the pinned + # `_native_*` kernels -- exists precisely to avoid materializing the score matrix. + return _DISPATCH_FUSED + @lru_cache(maxsize=None) -def _sdpa_has_fused_kernel( +def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: - """Ask torch whether any non-materializing SDPA kernel is eligible for these attention shapes. + """Ask torch whether `F.scaled_dot_product_attention` would build the O(S^2) score matrix. + `_fused_sdp_choice` is the same dispatch query torch's own `scaled_dot_product_attention` runs + to pick a kernel, so this is its real answer rather than a reimplementation of its rules. Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence - length, so a tiny probe answers for the real forward. Falls back to ``True`` (the status quo - assumption) whenever torch gives us nothing to go on -- over-reserving many GB on a guess would - push the model out of VRAM and be worse than the shortfall we are trying to avoid. + length, so a tiny probe answers for the real forward. + + Anything that goes wrong reports the materializing path, which is both the conservative answer + and, for the most common cause, the correct one: torch registers `_fused_sdp_choice` for CPU, + CUDA/ROCm and XPU only, so the call raises on MPS -- and MPS is exactly where + `scaled_dot_product_attention` finds no fused kernel either and runs + `_scaled_dot_product_attention_math_for_mps`, an MPSGraph transcription of `Q @ K^T` -> softmax + -> `@ V` that holds the score tensor as a real intermediate. The remaining causes (an allocation + failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and + there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ - if device_type != "cuda": - # `can_use_*_attention` is CUDA/ROCm-only. MPS, XPU and CPU all ship fused SDPA kernels, so - # keep the fused assumption there rather than guessing at their dispatch rules. - return True - try: - from torch.backends.cuda import SDPAParams, can_use_efficient_attention, can_use_flash_attention - device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None - try: - params = SDPAParams(q, q, q, mask, 0.0, False, False) - except TypeError: - # torch < 2.5: no `enable_gqa` field. - params = SDPAParams(q, q, q, mask, 0.0, False) - - checks = [can_use_flash_attention, can_use_efficient_attention] - can_use_cudnn_attention = getattr(torch.backends.cuda, "can_use_cudnn_attention", None) - if can_use_cudnn_attention is not None: - checks.append(can_use_cudnn_attention) - return any(check(params, False) for check in checks) + with warnings.catch_warnings(): + # When no fused kernel is eligible, torch re-runs every check in debug mode to warn why + # each one was rejected. That is the case we are deliberately probing for; we do not + # want a wall of warnings every time an estimate is priced. + warnings.simplefilter("ignore") + choice = int(torch.ops.aten._fused_sdp_choice(q, q, q, mask, 0.0, False)) except Exception: return True + return choice not in _FUSED_SDP_CHOICES + def sdpa_score_matrix_bytes( *, @@ -102,16 +157,35 @@ def sdpa_score_matrix_bytes( head_dim: int, seq_len: int, has_attn_mask: bool = False, + via_diffusers_dispatch: bool = False, ) -> int: """Bytes SDPA spends on a materialized score matrix for one attention call, 0 if fused. Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On - CUDA it is almost always 0; on a build whose fused kernels reject the shapes (notably ROCm, - which caps the head dim at 128 and does not take arbitrary additive masks) it is the dominant - term -- a 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~17GB of them. + CUDA it is almost always 0; where the fused kernels are missing or reject the shapes -- ROCm + caps the head dim at 128 and does not take arbitrary additive masks, MPS ships no fused SDPA + kernel at all -- it is the dominant term: a 1536px FLUX.2 VAE decode materializes 36864^2 + scores, ~17GB of them. + + Set ``via_diffusers_dispatch`` for attention that runs inside a diffusers model (the FLUX.2 + transformer does; the FLUX.2 VAE's mid-block attention does not -- it still reaches + `F.scaled_dot_product_attention` directly through `AttnProcessor2_0`). It consults the + process-wide default backend, which is the one that applies here: estimates are priced before + the model is loaded and outside any `attention_backend()` scope. """ if seq_len <= 0 or num_heads <= 0: return 0 - if _sdpa_has_fused_kernel(device.type, device.index, dtype, head_dim, has_attn_mask): + + score_matrix_bytes = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + if via_diffusers_dispatch: + dispatch = _diffusers_attention_dispatch() + if dispatch == _DISPATCH_FUSED: + return 0 + if dispatch == _DISPATCH_MATH: + return score_matrix_bytes + # _DISPATCH_TORCH: diffusers forwards to `F.scaled_dot_product_attention`, so torch decides. + + if not _torch_sdpa_materializes_score_matrix(device.type, device.index, dtype, head_dim, has_attn_mask): return 0 - return num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + return score_matrix_bytes diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 849efefbad7..960e2d081f9 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -129,10 +129,13 @@ def estimate_vae_working_memory_flux2( That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the - 512-wide head, and measured peak stays linear from 512 to 1536px). A build whose fused kernels - reject the head dim -- ROCm caps it at 128 -- drops to SDPA's ``math`` fallback and materializes - a (pixels/8)^2 score matrix on top of the linear term: ~3.5GB at 1024px and ~17GB at 1536px. We - ask torch which path applies rather than assuming, so the estimate is right on both. + 512-wide head, and measured peak stays linear from 512 to 1536px). A build with no fused kernel + for the shapes -- ROCm caps the head dim at 128, MPS has no fused SDPA kernel at all -- drops to + SDPA's ``math`` fallback and materializes a (pixels/8)^2 score matrix on top of the linear term: + ~3.5GB at 1024px and ~17GB at 1536px. We ask torch which path applies rather than assuming, so + the estimate is right on both. (Unlike the transformer, this attention does not go through + diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls ``F.scaled_dot_product_attention`` + itself -- so torch's own answer is the whole answer here.) When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index aac5409930d..ccc7628c191 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -10,6 +10,7 @@ quantity, including allocator overhead). Every estimate must stay an upper bound on them. """ +from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest @@ -20,19 +21,34 @@ from invokeai.app.invocations.flux2_denoise import FLUX2_MAX_ATTENTION_HEADS, Flux2DenoiseInvocation from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation -from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes +from invokeai.backend.util.attention import ( + SDPA_MATH_BYTES_PER_SCORE_ELEMENT, + _diffusers_attention_dispatch, + _torch_sdpa_materializes_score_matrix, + sdpa_score_matrix_bytes, +) from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 MB = 1024**2 GB = 1024**3 # The measured tables in this module were all taken on CUDA, where SDPA runs a fused kernel and no -# score matrix is materialized. `_sdpa_has_fused_kernel` reports non-CUDA devices as fused, so -# passing a CPU device reproduces that regime without needing a GPU on the test runner. The +# score matrix is materialized. torch reports its CPU flash kernel as eligible for every shape used +# here, so passing a CPU device reproduces that regime without needing a GPU on the test runner. The # materializing regime gets its own class below. FUSED = torch.device("cpu") +@pytest.fixture(autouse=True) +def _clear_dispatch_caches(): + """Both probes are `lru_cache`d for the process; tests that fake one must not leak into the next.""" + _diffusers_attention_dispatch.cache_clear() + _torch_sdpa_materializes_score_matrix.cache_clear() + yield + _diffusers_attention_dispatch.cache_clear() + _torch_sdpa_materializes_score_matrix.cache_clear() + + def _estimate( image_seq_len, ref_image_seq_len=0, @@ -307,20 +323,20 @@ def test_reference_images_raise_the_reservation(self): def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): - """Stand in for `_sdpa_has_fused_kernel` on a build with ROCm's fused-kernel rules. + """Stand in for the torch probe on a build with ROCm's fused-kernel rules. ROCm's fused SDPA kernels cap the head dim at 128 and do not take an arbitrary additive mask; anything else falls through to the `math` fallback, which materializes the score matrix. CUDA's memory-efficient kernel accepts both -- verified on torch 2.7.1+cu128, where - `can_use_efficient_attention` is true for the VAE's 512-wide head and for a masked 128-wide - transformer head, and measured peak stays linear in both cases -- which is why the estimates - were linear to begin with. + `_fused_sdp_choice` reports the efficient kernel for the VAE's 512-wide head and for a masked + 128-wide transformer head, and measured peak stays linear in both cases -- which is why the + estimates were linear to begin with. """ - return head_dim <= 128 and not has_attn_mask + return head_dim > 128 or has_attn_mask def _materializing(): - return patch("invokeai.backend.util.attention._sdpa_has_fused_kernel", side_effect=_rocm_like_probe) + return patch("invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", side_effect=_rocm_like_probe) # Any CUDA device object works here: the probe is patched out, so nothing is allocated on it. @@ -366,6 +382,20 @@ def test_vae_estimate_gains_exactly_the_score_matrix(self, operation, px, tokens materializing = self._vae_estimate(operation, px, device=MATERIALIZING) assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): + """The MPS case end to end, through the real probe rather than a stand-in: on a device torch + cannot answer a dispatch query for, a 1024px decode has to come out ~3.5GB heavier than the + linear term. Reporting those devices as fused -- as this PR first did -- is what let the + decode be admitted to a card that could not run it.""" + with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): + materializing = self._vae_estimate("decode", 1024, device=FUSED) + _torch_sdpa_materializes_score_matrix.cache_clear() + fused = self._vae_estimate("decode", 1024, device=FUSED) + + tokens = 128 * 128 + assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert materializing - fused > 3 * GB + def test_vae_score_matrix_dominates_at_high_resolution(self): """The reviewer's case: a 1536px decode is ~9.6GB of linear activations on CUDA, and more than that again in scores where SDPA has to materialize them. An estimate that omits the @@ -426,9 +456,10 @@ def test_denoise_without_a_regional_mask_is_unaffected(self): class TestSdpaBackendProbe: """`sdpa_score_matrix_bytes` decides the term above, so its defaults are load-bearing.""" - def test_non_cuda_devices_keep_the_fused_assumption(self): - """torch exposes no eligibility query outside CUDA/ROCm. Guessing `math` there would add - double-digit GB to every estimate on MPS and CPU on no evidence at all.""" + def test_cpu_reports_its_fused_flash_kernel(self): + """torch ships a fused flash-attention CPU kernel that takes the VAE's 512-wide head and an + additive mask, so the CPU estimate stays linear -- and the rest of this module can use a CPU + device to stand in for the CUDA regime the constants were measured on.""" assert ( sdpa_score_matrix_bytes( device=torch.device("cpu"), @@ -441,6 +472,40 @@ def test_non_cuda_devices_keep_the_fused_assumption(self): == 0 ) + def test_a_device_torch_cannot_answer_for_is_budgeted_as_math(self): + """MPS is the case that matters: torch registers `_fused_sdp_choice` for CPU, CUDA/ROCm and + XPU only, and it is exactly the devices it cannot answer for that have no fused SDPA kernel + either. A 1024px FLUX.2 VAE decode there materializes 16384^2 scores, ~3.5GB the estimate + used to omit entirely.""" + with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): + estimated = sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimated > 3 * GB + + def test_a_failed_probe_is_budgeted_as_math(self): + """A probe that cannot allocate, or a torch without the op, leaves us knowing nothing. The + old code read that as "fused" and reserved zero; the shortfall it hides is an OOM, so the + unknown answer has to be the expensive one.""" + with patch("torch.empty", side_effect=torch.cuda.OutOfMemoryError("probe could not allocate")): + estimated = sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + def test_the_probe_asks_torch_the_same_question_sdpa_does(self): + """`_fused_sdp_choice` is the dispatch query `F.scaled_dot_product_attention` itself runs, so + a `MATH` answer means the real forward materializes. Reimplementing the eligibility rules + instead would go stale with every torch release.""" + from torch.nn.attention import SDPBackend + + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): + assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) + _torch_sdpa_materializes_score_matrix.cache_clear() + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): + assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) + def test_empty_sequences_cost_nothing(self): with _materializing(): assert ( @@ -495,3 +560,75 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h estimate = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert estimate >= measured assert estimate <= 2 * measured + + +@contextmanager +def _diffusers_backend(name): + """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would. + + The lookup is `lru_cache`d -- it is a process-wide setting read on every estimate -- so the + cache has to be dropped on the way in and on the way out, or the faked backend leaks into the + comparison the test makes against the real one. + """ + from diffusers.models.attention_dispatch import AttentionBackendName + + _diffusers_attention_dispatch.cache_clear() + try: + with patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + return_value=(AttentionBackendName(name), None), + ): + yield + finally: + _diffusers_attention_dispatch.cache_clear() + + +class TestDiffusersAttentionDispatchIsConsulted: + """The FLUX.2 transformer does not call `F.scaled_dot_product_attention` -- it calls diffusers' + `dispatch_attention_fn`, which honours `DIFFUSERS_ATTN_BACKEND` and the `attention_backend()` + context manager. A user on `_native_math` materializes the score matrix on hardware where the + torch probe reports a fused kernel, so asking torch alone is not enough for the transformer. + + The VAE is the other half of the same point: its mid-block attention goes through + `AttnProcessor2_0`, which calls `F.scaled_dot_product_attention` itself, so the diffusers + backend must *not* move its estimate. + """ + + def test_forced_math_backend_reaches_the_denoise_estimate(self): + with _diffusers_backend("_native_math"): + forced_math = _estimate(image_seq_len=4096, device=FUSED) + native = _estimate(image_seq_len=4096, device=FUSED) + seq_len = 4096 + 512 + assert forced_math - native == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + + def test_a_fused_backend_leaves_the_denoise_estimate_linear(self): + """`flash`, `sage`, `xformers` and friends exist precisely to avoid the score matrix; they + must not be taxed for it, on any device.""" + with _diffusers_backend("flash"), _materializing(): + forced_flash = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) + assert forced_flash == _estimate(image_seq_len=4096, has_regional_mask=True, device=FUSED) + + def test_the_vae_estimate_ignores_the_diffusers_backend(self): + """`AttnProcessor2_0` bypasses the dispatcher, so the VAE's answer comes from torch alone.""" + + def estimate(): + v = MagicMock(spec=AutoencoderKLFlux2) + v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return estimate_vae_working_memory_flux2( + operation="decode", image_tensor=torch.zeros(1, 32, 128, 128), vae=v, device=FUSED + ) + + with _diffusers_backend("_native_math"): + forced_math = estimate() + assert forced_math == estimate() + + def test_an_unreadable_dispatcher_is_budgeted_as_math(self): + """`_AttentionBackendRegistry` is private; if diffusers moves it we lose the answer. The + conservative reading is the materializing one, and it is logged rather than silent.""" + with patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + side_effect=AttributeError("moved"), + ): + assert _diffusers_attention_dispatch() == "math" From 12e1c9a5cfd1c1b0bc20caf1d3966a15d1b0878f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 26 Aug 2026 20:16:54 +0200 Subject: [PATCH 4/5] fix(flux2): read the attention backend live instead of caching it once `_diffusers_attention_dispatch()` was `lru_cache`d, so the first estimate in a process pinned the answer forever. A switch to `_native_math` after that kept reserving zero for the S x S score matrix -- the exact case the lookup was added to catch. Read it live; it is a dict lookup against an already-imported module, priced once per invocation. The torch probe had the same defect one level down: its answer depends on the global SDPA kernel toggles, which `sdpa_kernel()` and `enable_flash_sdp()` flip at runtime. That probe allocates and dispatches, so it stays cached -- but keyed on the toggles, so a switch invalidates it. Per-model overrides need no plumbing: `set_attention_backend()` stamps its choice onto the process-wide registry as well as onto the model's processors, deliberately, so the estimate sees it without holding the model it is priced ahead of. A test pins that propagation. --- invokeai/backend/util/attention.py | 58 +++++++++-- .../invocations/test_flux2_working_memory.py | 95 ++++++++++++++----- 2 files changed, 122 insertions(+), 31 deletions(-) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 2245f6c750b..25f0b463cb4 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -75,7 +75,15 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: _DISPATCH_MATH = "math" -@lru_cache(maxsize=None) +@lru_cache(maxsize=1) +def _warn_unknown_diffusers_dispatch() -> None: + """Say once per process that estimates are running blind. Rate-limited, not cached for truth.""" + InvokeAILogger.get_logger(__name__).warning( + "Could not determine the active diffusers attention backend; budgeting working memory as if " + "attention materializes its score matrix. Estimates will be conservative." + ) + + def _diffusers_attention_dispatch() -> str: """Report how the diffusers attention dispatcher will route a diffusers model's attention calls. @@ -85,6 +93,18 @@ def _diffusers_attention_dispatch() -> str: torch; the others pin a specific kernel, and `_native_math` pins the materializing one. A torch-level probe alone would report "fused" for a user who has forced math. + Read live on every estimate, never cached: the active backend is mutable process state, and a + cached answer would keep reserving zero after a switch to `_native_math` -- the one case this + lookup exists to catch. It is a dict lookup against an already-imported module, priced once per + invocation. + + Reading the process-wide backend also covers per-model overrides, which is why the estimate does + not need the model in hand (it is priced before the model is loaded). `set_attention_backend()` + stamps the choice onto the model's attention processors *and* calls + `_AttentionBackendRegistry.set_active_backend()` -- deliberately, "so that it propagates + gracefully throughout". `reset_attention_backend()` clears only the processors, leaving the + registry pinned, which errs towards over-reserving rather than under-reserving. + Returns ``_DISPATCH_TORCH`` when torch decides, ``_DISPATCH_FUSED`` for a backend that never materializes the score matrix, or ``_DISPATCH_MATH`` when one is built -- including when we cannot tell, since under-reserving is the failure this whole term exists to prevent. @@ -98,10 +118,7 @@ def _diffusers_attention_dispatch() -> str: # A private diffusers attribute that moved, or a selected backend whose kernel failed to # register. Budget the materializing case, but say so: silently adding several GB to every # FLUX.2 estimate is not something that should pass unnoticed. - InvokeAILogger.get_logger(__name__).warning( - "Could not determine the active diffusers attention backend; budgeting working memory as if " - "attention materializes its score matrix. Estimates will be conservative." - ) + _warn_unknown_diffusers_dispatch() return _DISPATCH_MATH if name == "native": @@ -113,7 +130,21 @@ def _diffusers_attention_dispatch() -> str: return _DISPATCH_FUSED -@lru_cache(maxsize=None) +def _sdp_kernel_toggles() -> tuple[bool, ...]: + """The global switches that gate each fused SDPA kernel, as `_fused_sdp_choice` sees them. + + `torch.backends.cuda.enable_flash_sdp(False)` and `sdpa_kernel([...])` flip these at runtime and + the dispatch answer flips with them, so they belong in the probe's cache key rather than being + baked into a permanent result. + """ + cuda = torch.backends.cuda + return tuple( + bool(getattr(cuda, name)()) + for name in ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") + if hasattr(cuda, name) + ) + + def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: @@ -133,6 +164,21 @@ def _torch_sdpa_materializes_score_matrix( failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ + return _probe_sdpa_dispatch(device_type, device_index, dtype, head_dim, has_attn_mask, _sdp_kernel_toggles()) + + +@lru_cache(maxsize=None) +def _probe_sdpa_dispatch( + device_type: str, + device_index: int | None, + dtype: torch.dtype, + head_dim: int, + has_attn_mask: bool, + sdp_kernel_toggles: tuple[bool, ...], +) -> bool: + """Cached body of the probe above. Every input torch's answer depends on is part of the key -- + `sdp_kernel_toggles` is not read here, it is carried so a runtime change invalidates the entry. + """ try: device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index ccc7628c191..b876d973c13 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -10,7 +10,6 @@ quantity, including allocator overhead). Every estimate must stay an upper bound on them. """ -from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest @@ -24,6 +23,7 @@ from invokeai.backend.util.attention import ( SDPA_MATH_BYTES_PER_SCORE_ELEMENT, _diffusers_attention_dispatch, + _probe_sdpa_dispatch, _torch_sdpa_materializes_score_matrix, sdpa_score_matrix_bytes, ) @@ -40,13 +40,12 @@ @pytest.fixture(autouse=True) -def _clear_dispatch_caches(): - """Both probes are `lru_cache`d for the process; tests that fake one must not leak into the next.""" - _diffusers_attention_dispatch.cache_clear() - _torch_sdpa_materializes_score_matrix.cache_clear() +def _clear_probe_cache(): + """The torch probe is `lru_cache`d for the process; a faked answer must not leak into the next + test. (The diffusers lookup is deliberately uncached -- see `_diffusers_attention_dispatch`.)""" + _probe_sdpa_dispatch.cache_clear() yield - _diffusers_attention_dispatch.cache_clear() - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() def _estimate( @@ -389,7 +388,7 @@ def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): decode be admitted to a card that could not run it.""" with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): materializing = self._vae_estimate("decode", 1024, device=FUSED) - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() fused = self._vae_estimate("decode", 1024, device=FUSED) tokens = 128 * 128 @@ -494,6 +493,24 @@ def test_a_failed_probe_is_budgeted_as_math(self): ) assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + def test_disabling_the_fused_kernels_at_runtime_invalidates_the_probe(self): + """The torch probe *is* cached -- it allocates and runs a dispatch query -- but its answer + depends on switches callers can flip at runtime (`sdpa_kernel()`, + `torch.backends.cuda.enable_flash_sdp`). Those toggles are part of the cache key, so an + estimate priced after a switch does not inherit the answer from before it. No cache is + cleared between these calls on purpose.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + def estimate(): + return sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + ) + + assert estimate() == 0 + with sdpa_kernel([SDPBackend.MATH]): + assert estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimate() == 0 + def test_the_probe_asks_torch_the_same_question_sdpa_does(self): """`_fused_sdp_choice` is the dispatch query `F.scaled_dot_product_attention` itself runs, so a `MATH` answer means the real forward materializes. Reimplementing the eligibility rules @@ -502,7 +519,7 @@ def test_the_probe_asks_torch_the_same_question_sdpa_does(self): with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) @@ -562,25 +579,14 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h assert estimate <= 2 * measured -@contextmanager def _diffusers_backend(name): - """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would. - - The lookup is `lru_cache`d -- it is a process-wide setting read on every estimate -- so the - cache has to be dropped on the way in and on the way out, or the faked backend leaks into the - comparison the test makes against the real one. - """ + """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would.""" from diffusers.models.attention_dispatch import AttentionBackendName - _diffusers_attention_dispatch.cache_clear() - try: - with patch( - "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", - return_value=(AttentionBackendName(name), None), - ): - yield - finally: - _diffusers_attention_dispatch.cache_clear() + return patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + return_value=(AttentionBackendName(name), None), + ) class TestDiffusersAttentionDispatchIsConsulted: @@ -624,6 +630,45 @@ def estimate(): forced_math = estimate() assert forced_math == estimate() + def test_a_backend_switch_is_not_masked_by_an_earlier_estimate(self): + """The active backend is mutable process state. Caching the first answer would keep + reserving zero for every later estimate in a long-lived process that has since switched to + `_native_math` -- the exact case this lookup exists to catch. Deliberately no cache is + cleared between the two calls here; the production code must not be holding one.""" + native = _estimate(image_seq_len=4096, device=FUSED) + with _diffusers_backend("_native_math"): + after_switch = _estimate(image_seq_len=4096, device=FUSED) + back_to_native = _estimate(image_seq_len=4096, device=FUSED) + + seq_len = 4096 + 512 + assert after_switch - native == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + assert back_to_native == native + + def test_a_model_level_override_reaches_the_registry(self): + """Why the estimator does not need the model in hand: it is priced before the transformer is + loaded, and `set_attention_backend()` stamps its choice onto the process-wide registry as + well as onto the model's attention processors -- deliberately, "so that it propagates + gracefully throughout". If diffusers ever stops doing that, a per-model override could + disagree with the estimate, and this test is where that shows up.""" + from diffusers.configuration_utils import ConfigMixin, register_to_config + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + from diffusers.models.modeling_utils import ModelMixin + + class _Tiny(ModelMixin, ConfigMixin): + @register_to_config + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(2, 2) + + previous = _AttentionBackendRegistry._active_backend + try: + _Tiny().set_attention_backend("_native_math") + assert _diffusers_attention_dispatch() == "math" + finally: + _AttentionBackendRegistry._active_backend = previous + def test_an_unreadable_dispatcher_is_budgeted_as_math(self): """`_AttentionBackendRegistry` is private; if diffusers moves it we lose the answer. The conservative reading is the materializing one, and it is logged rather than silent.""" From df6b3d88df251409fd751159b654c8c07fac4d34 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 02:51:02 +0200 Subject: [PATCH 5/5] fix(flux2): stop caching the SDPA probe and scale the VAE estimate by batch The probe's cache key held the four per-backend enable flags, but torch takes the *first eligible* backend in a priority order that `sdpa_kernel(..., set_priority=True)` reorders while leaving every flag untouched -- measured: same flags, EFFICIENT outside and MATH inside. A fused answer cached before the switch would suppress the score-matrix reservation after it. Rather than adding the priority order to the key -- the next thing to forget is always one more -- drop the cache. The probe costs ~6us against a multi-second forward, so there is nothing to protect. `vae.decode` is also handed whatever batch the latents carry, and a LatentsField is not pinned to one, so an estimate built from H and W alone gave a two-sample decode a single sample's reservation. Measured at 1024px: 4.23GB at batch 1, 7.96GB at 2, 11.89GB at 3 -- linear, slightly sub-linear per sample, so the scaled single-sample estimate stays an upper bound. The score matrix is (batch, heads, S, S) and scales with it. --- invokeai/backend/util/attention.py | 36 +---- invokeai/backend/util/vae_working_memory.py | 12 +- .../invocations/test_flux2_working_memory.py | 134 +++++++++++++++--- 3 files changed, 131 insertions(+), 51 deletions(-) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 25f0b463cb4..9bbda9c290a 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -130,21 +130,6 @@ def _diffusers_attention_dispatch() -> str: return _DISPATCH_FUSED -def _sdp_kernel_toggles() -> tuple[bool, ...]: - """The global switches that gate each fused SDPA kernel, as `_fused_sdp_choice` sees them. - - `torch.backends.cuda.enable_flash_sdp(False)` and `sdpa_kernel([...])` flip these at runtime and - the dispatch answer flips with them, so they belong in the probe's cache key rather than being - baked into a permanent result. - """ - cuda = torch.backends.cuda - return tuple( - bool(getattr(cuda, name)()) - for name in ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") - if hasattr(cuda, name) - ) - - def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: @@ -155,6 +140,12 @@ def _torch_sdpa_materializes_score_matrix( Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence length, so a tiny probe answers for the real forward. + Not cached. The answer turns on global torch state a cache key cannot honestly enumerate: the + per-backend enable flags, but also the *priority order*, which `sdpa_kernel(..., set_priority= + True)` reorders while leaving every flag untouched -- measured, same flags, `EFFICIENT` before + and `MATH` inside. Each item added to such a key is one more thing to get wrong later, and the + probe costs ~6us against a multi-second forward, so it just runs every time. + Anything that goes wrong reports the materializing path, which is both the conservative answer and, for the most common cause, the correct one: torch registers `_fused_sdp_choice` for CPU, CUDA/ROCm and XPU only, so the call raises on MPS -- and MPS is exactly where @@ -164,21 +155,6 @@ def _torch_sdpa_materializes_score_matrix( failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ - return _probe_sdpa_dispatch(device_type, device_index, dtype, head_dim, has_attn_mask, _sdp_kernel_toggles()) - - -@lru_cache(maxsize=None) -def _probe_sdpa_dispatch( - device_type: str, - device_index: int | None, - dtype: torch.dtype, - head_dim: int, - has_attn_mask: bool, - sdp_kernel_toggles: tuple[bool, ...], -) -> bool: - """Cached body of the probe above. Every input torch's answer depends on is part of the key -- - `sdp_kernel_toggles` is not read here, it is carried so a runtime change invalidates the entry. - """ try: device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 960e2d081f9..0c2cf3cfa92 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -140,12 +140,20 @@ def estimate_vae_working_memory_flux2( When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score matrix, if one is materialized at all, is bounded by the tile too. + + Both terms are per sample. `vae.decode` takes whatever batch the latents carry, and a + ``LatentsField`` is not pinned to one, so the batch has to multiply through: measured at 1024px + decode, peak reserved is 4.23GB at batch 1, 7.96GB at batch 2 and 11.89GB at batch 3 -- linear, + and slightly sub-linear per sample, so multiplying the single-sample estimate stays an upper + bound. The score matrix is shaped (batch, heads, S, S), so it scales the same way. The encode + call sites all pass batch 1 today; the shared estimator does not assume it. """ param = next(vae.parameters()) element_size = param.element_size() # Encoding uses ~50% the working memory of decoding. scaling_constant = 2200 if operation == "decode" else 1100 + batch_size = image_tensor.shape[0] if image_tensor.dim() >= 4 else 1 if tile_size is not None: # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. @@ -158,10 +166,12 @@ def estimate_vae_working_memory_flux2( working_memory = out_h * out_w * element_size * scaling_constant mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) + working_memory *= batch_size working_memory += sdpa_score_matrix_bytes( device=device if device is not None else TorchDevice.choose_torch_device(), dtype=param.dtype, - num_heads=_FLUX2_VAE_MID_BLOCK_HEADS, + # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, seq_len=mid_block_seq_len, ) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index b876d973c13..7d3e5783889 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -23,7 +23,6 @@ from invokeai.backend.util.attention import ( SDPA_MATH_BYTES_PER_SCORE_ELEMENT, _diffusers_attention_dispatch, - _probe_sdpa_dispatch, _torch_sdpa_materializes_score_matrix, sdpa_score_matrix_bytes, ) @@ -39,15 +38,6 @@ FUSED = torch.device("cpu") -@pytest.fixture(autouse=True) -def _clear_probe_cache(): - """The torch probe is `lru_cache`d for the process; a faked answer must not leak into the next - test. (The diffusers lookup is deliberately uncached -- see `_diffusers_attention_dispatch`.)""" - _probe_sdpa_dispatch.cache_clear() - yield - _probe_sdpa_dispatch.cache_clear() - - def _estimate( image_seq_len, ref_image_seq_len=0, @@ -179,6 +169,67 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): assert estimates[0] < untiled / 4 +class TestFlux2VaeBatchIsBudgeted: + """`vae.decode` is handed whatever batch the latents carry, and a `LatentsField` is not pinned to + one. An estimate built from H and W alone gives a two-sample decode the same reservation as a + single one, so the cache admits it to a card that cannot run it -- the reservation is there, and + the OOM happens anyway. + + Measured at 1024px on CUDA/bf16, peak reserved, each point in a fresh process: 4.23GB at batch 1, + 7.96GB at batch 2, 11.89GB at batch 3. Linear, and slightly sub-linear per sample, so scaling the + single-sample estimate is an upper bound rather than a fit. + """ + + # (batch, measured peak reserved MB) for a 1024px decode. + MEASURED_DECODE_BATCH = [(1, 4229), (2, 7955), (3, 11889)] + + def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return estimate_vae_working_memory_flux2( + operation="decode", + image_tensor=torch.zeros(batch, 32, px // 8, px // 8), + vae=vae, + tile_size=tile_size, + device=device, + ) + + @pytest.mark.parametrize("batch, measured_mb", MEASURED_DECODE_BATCH) + def test_estimate_is_an_upper_bound_on_the_measured_batch_peak(self, batch, measured_mb): + estimate = self._decode_estimate(batch) + assert estimate >= measured_mb * MB + assert estimate <= 2 * measured_mb * MB + + def test_estimate_scales_with_the_batch(self): + """The regression in one assertion: before this, all three of these were equal.""" + single = self._decode_estimate(1) + assert self._decode_estimate(2) == 2 * single + assert self._decode_estimate(3) == 3 * single + + def test_a_three_dimensional_tensor_is_one_sample(self): + """A bare `(C, H, W)` latent has no batch axis; `shape[0]` would read the channel count.""" + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + unbatched = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=torch.zeros(32, 128, 128), vae=vae, device=FUSED + ) + assert unbatched == self._decode_estimate(1) + + def test_tiling_bounds_the_tile_not_the_batch(self): + """Tiling caps the spatial term at one tile, but every sample still runs through it.""" + single = self._decode_estimate(1, px=1024, tile_size=512) + assert self._decode_estimate(3, px=1024, tile_size=512) == 3 * single + + def test_the_score_matrix_scales_with_the_batch(self): + """It is shaped (batch, heads, S, S), so where it is materialized at all it scales too.""" + tokens = 128 * 128 + score = tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + linear = self._decode_estimate(1) # fused: the spatial term on its own + with _materializing(): + assert self._decode_estimate(1, device=MATERIALIZING) == linear + score + assert self._decode_estimate(3, device=MATERIALIZING) == 3 * (linear + score) + + class TestFlux2VaeInvocationsRequestWorkingMemory: """The estimate is worthless unless it reaches `model_on_device()`.""" @@ -388,7 +439,6 @@ def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): decode be admitted to a card that could not run it.""" with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): materializing = self._vae_estimate("decode", 1024, device=FUSED) - _probe_sdpa_dispatch.cache_clear() fused = self._vae_estimate("decode", 1024, device=FUSED) tokens = 128 * 128 @@ -493,21 +543,66 @@ def test_a_failed_probe_is_budgeted_as_math(self): ) assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - def test_disabling_the_fused_kernels_at_runtime_invalidates_the_probe(self): - """The torch probe *is* cached -- it allocates and runs a dispatch query -- but its answer - depends on switches callers can flip at runtime (`sdpa_kernel()`, - `torch.backends.cuda.enable_flash_sdp`). Those toggles are part of the cache key, so an - estimate priced after a switch does not inherit the answer from before it. No cache is - cleared between these calls on purpose.""" + def _cpu_estimate(self): + return sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + ) + + def test_disabling_the_fused_kernels_at_runtime_changes_the_answer(self): + """The probe is not cached, so an estimate priced after a runtime switch does not inherit the + answer from before it. Nothing is cleared between these calls on purpose.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + assert self._cpu_estimate() == 0 + with sdpa_kernel([SDPBackend.MATH]): + assert self._cpu_estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert self._cpu_estimate() == 0 + + def test_a_priority_reorder_changes_the_answer_with_every_flag_unchanged(self): + """`sdpa_kernel(..., set_priority=True)` puts `MATH` first while leaving all four enable + flags True, and torch takes the first eligible backend in that order. A cache keyed on the + flags -- which is what this probe used to have -- could not see the switch and would keep + reserving zero. Not caching at all is what makes that unrepresentable.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + math_first = [ + SDPBackend.MATH, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + ] + flags = ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") + + assert self._cpu_estimate() == 0 + with sdpa_kernel(math_first, set_priority=True): + # The finding in one line: every flag a key could hold is still True in here. + assert all(getattr(torch.backends.cuda, name)() for name in flags if hasattr(torch.backends.cuda, name)) + # CPU's chooser ignores the priority order, so stand in for the answer CUDA gives. + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): + assert self._cpu_estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert self._cpu_estimate() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="only CUDA's chooser honours the priority order") + def test_this_build_reports_a_real_priority_reorder(self): + """The same case against the real dispatcher rather than a stand-in. Verified on torch + 2.7.1+cu128: `_fused_sdp_choice` answers EFFICIENT outside and MATH inside.""" from torch.nn.attention import SDPBackend, sdpa_kernel def estimate(): return sdpa_score_matrix_bytes( - device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + device=torch.device("cuda"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 ) assert estimate() == 0 - with sdpa_kernel([SDPBackend.MATH]): + with sdpa_kernel( + [ + SDPBackend.MATH, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + ], + set_priority=True, + ): assert estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert estimate() == 0 @@ -519,7 +614,6 @@ def test_the_probe_asks_torch_the_same_question_sdpa_does(self): with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) - _probe_sdpa_dispatch.cache_clear() with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False)