diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 7a5ce158c98..b8c8752049c 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", @@ -458,10 +466,44 @@ 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 + # 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), + # Taken from `x`, not from `b`. `b` is the *noise* tensor's batch, which this node builds + # at 1 from width/height/seed even when the init latents carry more; the img2img preblend + # above then broadcasts the two, so `x` is the only thing that knows how many samples + # actually go through the transformer. Reference latents are repeated to match it + # (`ensure_batch_size` below), so they scale with it too. + batch_size=x.shape[0], + # The mask itself is already allocated; only the additive bias built per forward is new. + regional_attention_bias_bytes=( + 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: # 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 +620,78 @@ 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, + batch_size: int = 1, + 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. + + 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. + + Batch multiplies the token count and nothing else. A batch of B is B independent sequences, + so it enters the linear term exactly as extra sequence does -- measured on the Klein geometry + with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 (9216 + tokens) at 5126MB, and 9728 tokens at B=1 at 5584MB. Batch and sequence are interchangeable + to within the noise. Reference latents are repeated per sample (`ensure_batch_size`), so they + scale with it too, and the score matrix is shaped (batch, heads, S, S). The fixed base does + not scale -- it is about weights, not activations -- and neither does the regional bias, which + is built as (1, 1, S, S) and broadcast across the batch. + + 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 + per_token_bytes = int(0.4 * MB) + total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len + estimated = total_seq_len * batch_size * 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 * batch_size, + 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: + # A sidecar branch is an activation, so it scales with the batch like the rest of them. + estimated += int(0.5 * num_loras * batch_size * 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..f0852f1880f 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, device=vae_info.compute_device + ) + + 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..2da6f38b517 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,16 @@ 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, + device=TorchDevice.choose_torch_device(), + ) + + 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..9184b15c1d2 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,18 @@ 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, + device=TorchDevice.choose_torch_device(), + ) + + 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 +233,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/attention.py b/invokeai/backend/util/attention.py index 1df0f99280b..9bbda9c290a 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,10 +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: @@ -35,3 +40,174 @@ 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. 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 +# 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. 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=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. + + 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. + + 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. + """ + 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. + _warn_unknown_diffusers_dispatch() + 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 + + +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 `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. + + 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 + `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. + """ + 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) + mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None + 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( + *, + device: torch.device, + dtype: torch.dtype, + num_heads: int, + 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; 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 + + 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 score_matrix_bytes diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index bd780d4c0b3..0c2cf3cfa92 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -2,12 +2,15 @@ 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 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 @@ -98,6 +101,84 @@ 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. + 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 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 + 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. + 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 *= batch_size + working_memory += sdpa_score_matrix_bytes( + device=device if device is not None else TorchDevice.choose_torch_device(), + dtype=param.dtype, + # 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, + ) + + 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..0769a293a17 --- /dev/null +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -0,0 +1,881 @@ +"""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 +import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 + +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, + _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. 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") + + +def _estimate( + image_seq_len, + ref_image_seq_len=0, + text_seq_len=512, + num_loras=0, + batch_size=1, + regional_bias=0, + has_regional_mask=False, + device=FUSED, +): + 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, + batch_size=batch_size, + regional_attention_bias_bytes=regional_bias, + has_regional_attention_mask=has_regional_mask, + device=device, + ) + + +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 TestFlux2DenoiseBatchIsBudgeted: + """A batch of B is B independent sequences, so it enters the linear term exactly as extra + sequence does. Measured on the Klein geometry (48 heads x 128, mlp 3.0) with a reduced block + count -- the constant is block-count independent -- peak reserved, each point in a fresh process: + + B=1, 4608 tokens -> 2570MB B=2, 4608 each (9216 total) -> 5126MB + B=1, 9728 tokens -> 5584MB B=2, 9728 each (19456 total) -> 11120MB + B=1, 14336 tokens -> 8284MB B=3, 4608 each (13824 total) -> 7656MB + + Per *total* token that is 0.554-0.578MB across every row: batch and sequence are interchangeable. + (The absolute figure is not comparable to the Klein table elsewhere in this module -- a 3-block + stand-in amortizes per-forward overhead differently. Only the equivalence is being tested.) + + Batched latents reach this node through the API and custom graphs, not the stock UI. + """ + + def test_batch_and_sequence_are_interchangeable(self): + """Two samples of 4608 tokens must cost what one sample of 9216 costs -- the measurement + above says 5126MB against 5584MB, equal to within the allocator's noise.""" + assert _estimate(image_seq_len=4096, text_seq_len=512, batch_size=2) == _estimate( + image_seq_len=8704, text_seq_len=512, batch_size=1 + ) + + @pytest.mark.parametrize("batch", [2, 3, 4]) + def test_each_extra_sample_adds_exactly_its_own_tokens(self, batch): + single = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, batch_size=batch) - single == (batch - 1) * (4096 + 512) * int(0.4 * MB) + + def test_reference_tokens_scale_with_the_batch(self): + """`ensure_batch_size` repeats the reference latents across the batch.""" + single = _estimate(image_seq_len=4096, ref_image_seq_len=12288) + assert _estimate(image_seq_len=4096, ref_image_seq_len=12288, batch_size=2) - single == ( + (4096 + 12288 + 512) * int(0.4 * MB) + ) + + def test_the_fixed_base_does_not_scale_with_the_batch(self): + """It covers transient weight casts and allocator slack -- properties of the weights, not of + how many samples run through them. Scaling it would add a GB per sample for nothing.""" + deltas = { + _estimate(image_seq_len=4096, batch_size=b + 1) - _estimate(image_seq_len=4096, batch_size=b) + for b in (1, 2, 3) + } + assert deltas == {(4096 + 512) * int(0.4 * MB)} + + def test_the_regional_bias_does_not_scale_with_the_batch(self): + """`get_joint_attention_kwargs` builds it as (1, 1, S, S) and lets SDPA broadcast it, so + there is exactly one of them however many samples are in flight.""" + bias = (4096 + 512) ** 2 * 2 + single = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True) + double = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True, batch_size=2) + assert double - single == (4096 + 512) * int(0.4 * MB) + + def test_the_score_matrix_scales_with_the_batch(self): + """Where it is materialized at all it is shaped (batch, heads, S, S).""" + seq_len = 4096 + 512 + with _materializing(): + single = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) + double = _estimate(image_seq_len=4096, has_regional_mask=True, batch_size=2, device=MATERIALIZING) + assert double - single == ( + seq_len * int(0.4 * MB) + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + + def test_the_lora_margin_scales_with_the_batch(self): + """A sidecar patch adds an activation branch, and activations are per sample.""" + assert _estimate(image_seq_len=4096, num_loras=2, batch_size=3) - _estimate( + image_seq_len=4096, num_loras=0, batch_size=3 + ) == 3 * int(1.0 * GB) + + +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(), 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(), device=FUSED + ) + 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, + device=FUSED, + ) + 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(), device=FUSED + ) + 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()`.""" + + 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, batch: int = 1, init_batch: int | None = None): + """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) + + if init_batch is not None: + # img2img: the node loads these, then preblends them with its own batch-1 noise. + context.tensors.load.return_value = torch.zeros(init_batch, 32, 128, 128) + + invocation = Flux2DenoiseInvocation.model_construct( + latents=MagicMock(latents_name="init") if init_batch is not None else 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), + patch.object( + Flux2DenoiseInvocation, "_prepare_noise_tensor", return_value=torch.zeros(batch, 32, 128, 128) + ), + 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) + + def test_the_real_batch_reaches_the_reservation(self): + """A batched latent tensor is reachable through the API and custom graphs. The node has `b` + in hand at the estimate; before this it simply did not pass it, so a two-sample run reserved + one sample's worth and the cache admitted it to a card that could not run it.""" + assert self._run(num_ref_tokens=0, batch=2) == _estimate(image_seq_len=64 * 64, text_seq_len=512, batch_size=2) + + def test_a_batched_run_reserves_more_than_a_single_one(self): + single = self._run(num_ref_tokens=0, batch=1) + assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * int(0.4 * MB) + + def test_a_batched_init_latent_beats_the_batch_1_noise_it_is_blended_with(self): + """img2img takes `x = t_0 * noise + (1 - t_0) * init_latents`, and this node builds its noise + at batch 1 from width/height/seed. Two batched init latents therefore broadcast up to a + two-sample `x` while the noise tensor -- the thing the batch used to be read from -- still + says 1. The reservation has to follow `x`.""" + assert self._run(num_ref_tokens=0, init_batch=2) == _estimate( + image_seq_len=64 * 64, text_seq_len=512, batch_size=2 + ) + + def test_the_blended_batch_is_read_after_the_broadcast(self): + single = self._run(num_ref_tokens=0, init_batch=1) + assert self._run(num_ref_tokens=0, init_batch=3) - single == 2 * (64 * 64 + 512) * int(0.4 * MB) + + def test_repeated_reference_latents_are_counted_per_sample(self): + """`ensure_batch_size` repeats the reference latents across the batch, so their tokens scale + with it as well -- the worst case in #9500, doubled.""" + single = self._run(num_ref_tokens=12288, batch=1) + double = self._run(num_ref_tokens=12288, batch=2) + assert double - single == (64 * 64 + 12288 + 512) * int(0.4 * MB) + + +def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): + """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 + `_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 or has_attn_mask + + +def _materializing(): + 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. +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_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) + 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 + 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_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"), + dtype=torch.bfloat16, + num_heads=1, + head_dim=512, + seq_len=16384, + has_attn_mask=True, + ) + == 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 _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("cuda"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + ) + + assert estimate() == 0 + 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 + + 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) + 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 ( + 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 + + +def _diffusers_backend(name): + """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would.""" + from diffusers.models.attention_dispatch import AttentionBackendName + + return patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + return_value=(AttentionBackendName(name), None), + ) + + +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_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.""" + with patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + side_effect=AttributeError("moved"), + ): + assert _diffusers_attention_dispatch() == "math"