diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 7a5ce158c98..33725938a4f 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -39,6 +39,7 @@ unpack_flux2, ) from invokeai.backend.flux2.text_conditioning import Flux2TextConditioning +from invokeai.backend.model_manager.configs.flux2_variant import flux2_hidden_size from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import ( @@ -49,8 +50,25 @@ 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 and the head count follows +# the hidden size (Klein 4B: 3072/24, Klein 9B: 4096/32, [dev] 6144/48), so the width is the single +# number that describes both. Only the head dim decides which SDPA kernel is eligible; the head +# count scales the `math` fallback's score matrix. +FLUX2_ATTENTION_HEAD_DIM = 128 +# The width the per-token constant below was measured on. Estimates scale off this. +FLUX2_REFERENCE_HIDDEN_SIZE = 4096 +# Peak reserved activation bytes per attended token at the reference width. Measured slope between +# 4608 and 9216 tokens: 0.3859 MB/tok on CUDA (RTX 4090, torch 2.7.1) and 0.4067 on ROCm (RX 9070 +# XT, gfx1201, torch 2.10). 0.42 is an upper bound on both -- it was 0.4, which the ROCm point +# exceeds. The margin is deliberately small because this is the dominant term at every resolution. +FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH = int(0.42 * 1024**2) +# The widest variant, used when the config does not tell us which one this is -- over-reserving on +# an unknown model beats under-reserving on the largest one. +FLUX2_MAX_HIDDEN_SIZE = 6144 + @invocation( "flux2_denoise", @@ -458,10 +476,47 @@ 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], + # Activation cost per token scales with the transformer's width, and [dev] is 1.5x + # Klein 9B. Fall back to the widest variant when the config cannot tell us. + hidden_size=flux2_hidden_size(getattr(transformer_config, "variant", None)) or FLUX2_MAX_HIDDEN_SIZE, + # 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 +633,94 @@ 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, + hidden_size: int = FLUX2_REFERENCE_HIDDEN_SIZE, + 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 holds from 1.5k to 28k tokens and is independent of the block + count (a no-grad forward frees each block's intermediates). It is *not* independent of the + build: 0.3859 MB/token on CUDA, 0.4067 on ROCm/gfx1201, so the constant is 0.42. + + It is *not* independent of the transformer's width, which is why ``hidden_size`` is a + parameter rather than a constant. Measured slope between 4608 and 9216 tokens, block count + and everything else held fixed: + + CUDA 3072 (Klein 4B) 0.291 MB/tok 4096 (Klein 9B) 0.386 6144 ([dev]) 0.555 + ROCm/gfx1201 3072 0.315 4096 0.407 6144 0.589 + + which is 0.755 / 1.000 / 1.438 on CUDA and 0.774 / 1.000 / 1.448 on ROCm, against width + ratios of 0.75 / 1.00 / 1.50 -- linear in width on both, + and slightly sub-linear at the top so scaling by width stays an upper bound. Calibrating on + Klein 9B alone would have under-reserved [dev] by a third: 1024x1024 with three 1024x1024 + references is 16896 tokens, ~7.6GB reserved against ~10GB needed on both platforms. The head + count follows the same width, so the score-matrix term gets the real one instead of the + widest. + + 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, leaving whatever else the build has -- and if that is + the ``math`` fallback, a materialized ``heads x S x S`` score matrix. Whether a mask forces + that is build-specific and not worth predicting: CUDA's memory-efficient kernel takes it, and + so does ROCm's on gfx1100. 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 term is zero (verified: peak stays linear with the + bias attached). + """ + GB = 1024**3 + per_token_bytes = int(FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH * hidden_size / FLUX2_REFERENCE_HIDDEN_SIZE) + 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=(hidden_size // FLUX2_ATTENTION_HEAD_DIM) * 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..926f3a7c043 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -18,7 +18,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext 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,9 +46,22 @@ 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. + # Use the VAE's intended compute device for both the probe and the encode, matching the + # decode node. `choose_torch_device()` disagrees for a cpu_only VAE, which would both ask the + # dispatch probe about the wrong device and push the tensors onto a device the weights are + # not on (see #9373, and the same note in the decode node). + device = vae_info.compute_device + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + device=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) # Encode using diffusers API 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/model_manager/configs/flux2_variant.py b/invokeai/backend/model_manager/configs/flux2_variant.py index 0ae9244d96d..a7bac8651e2 100644 --- a/invokeai/backend/model_manager/configs/flux2_variant.py +++ b/invokeai/backend/model_manager/configs/flux2_variant.py @@ -60,3 +60,23 @@ def flux2_variant_from_vec_dim(dim: int) -> Flux2VariantType | None: def flux2_variant_from_hidden_size(dim: int) -> Flux2VariantType | None: """Return the distilled FLUX.2 variant for a transformer hidden_size, or ``None`` if unrecognized.""" return _HIDDEN_SIZE_TO_VARIANT.get(dim) + + +# Base variants share their distilled twin's architecture (see the module docstring); only the +# weights differ, so they resolve to the same geometry. +_BASE_TO_DISTILLED: dict[Flux2VariantType, Flux2VariantType] = { + Flux2VariantType.Klein4BBase: Flux2VariantType.Klein4B, + Flux2VariantType.Klein9BBase: Flux2VariantType.Klein9B, +} + + +def flux2_hidden_size(variant: Flux2VariantType | None) -> int | None: + """Return the transformer hidden size for a variant, or ``None`` if unrecognized. + + The forward direction of :func:`flux2_variant_from_hidden_size`. Callers that size activations + need this: the per-token activation cost scales with the transformer's width, so Klein 4B, Klein + 9B and [dev] (3072 / 4096 / 6144) do not cost the same per token. + """ + if variant is None: + return None + return _HIDDEN_SIZE.get(_BASE_TO_DISTILLED.get(variant, variant)) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 1df0f99280b..1ffa8abef72 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,10 +4,16 @@ for attention mechanism. """ +import threading +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 +41,211 @@ 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, and the rules differ per build in ways that are not worth +# hard-coding. Measured, for the FLUX.2 VAE's 512-wide mid-block head: CUDA takes it on the +# memory-efficient kernel, ROCm/gfx1100 drops to `math`, and ROCm/gfx1201 takes it on flash. Two +# cards of the same vendor on the same torch disagree -- so "ROCm materializes" is not a fact to +# hard-code either. Masks are not a discriminator anywhere measured. A working-memory estimate that +# assumes the fused path is 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, with `SDPBackend.MATH` forced, +# each point in a fresh process. The same figures hold 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. +# +# heads seq head_dim CUDA gfx1100 gfx1201 +# 1 4096 512 12.88 13.62 16.38 +# 1 8192 512 10.28 10.78 13.47 +# 1 16384 512 9.71 9.76 11.99 +# 4 4096 128 9.97 10.16 12.84 +# 48 4608 128 9.58 9.59 11.69 +# +# 17 is an upper bound on every measured point on all three. Note how far apart the two ROCm cards +# are: this is not a "CUDA number and a ROCm number", it is per-build, and the cost of guessing low +# is an OOM the estimate exists to prevent. `scripts/calibrate_flux2_working_memory.py` reproduces +# this table on any build. +SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 17 + +# `_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) +) + +# Guards the process-global warning filter list that `catch_warnings` swaps; see the probe below. +_WARNING_FILTER_LOCK = threading.Lock() + +_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 _WARNING_FILTER_LOCK, 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. + # + # `catch_warnings` swaps the process-global filter list, and this probe is deliberately + # uncached, so with concurrent session workers two estimates could interleave their + # enter/exit and leave a stale filter set behind. The lock makes the swap atomic against + # other probes; it costs nothing at ~6us of hold time. + 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 build has no fused kernel for the shapes -- ROCm/gfx1100 + for a 512-wide head, MPS for anything at all -- it is the dominant term: a 1536px FLUX.2 VAE + decode materializes 36864^2 scores, ~21GB of them. + + That is a large term to add on a probe, so it is logged when it fires: a user who suddenly sees + their model pushed out of VRAM should be able to find out why from the log rather than by + reading this file. + + 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 _log_and_return(score_matrix_bytes, device, head_dim, "the diffusers backend") + # _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 _log_and_return(score_matrix_bytes, device, head_dim, "this torch build") + + +@lru_cache(maxsize=None) +def _log_score_matrix_reservation(device_type: str, head_dim: int, reason: str, gib: str) -> None: + """Announce the materializing path once per (device, head dim, reason). It is the difference + between a model that stays resident and one that does not, and nothing else in the log says so. + """ + InvokeAILogger.get_logger(__name__).info( + f"SDPA materializes its attention score matrix on {device_type} for head_dim={head_dim} " + f"({reason}), so working-memory estimates reserve an extra ~{gib} GiB for it." + ) + + +def _log_and_return(score_matrix_bytes: int, device: torch.device, head_dim: int, reason: str) -> int: + _log_score_matrix_reservation(device.type, head_dim, reason, f"{score_matrix_bytes / 1024**3:.1f}") + return score_matrix_bytes diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index bd780d4c0b3..541ed6efbf2 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,151 @@ 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. Either way it is far past the +# 128 head dim some builds cap their fused SDPA kernels at, so whether the score matrix is +# materialized is a per-build question -- `sdpa_score_matrix_bytes` asks rather than assumes. +_FLUX2_VAE_MID_BLOCK_HEADS = 1 +_FLUX2_VAE_MID_BLOCK_HEAD_DIM = 512 +_FLUX2_VAE_SPATIAL_COMPRESSION = 8 + +# Peak reserved bytes per output pixel per element byte, per conv backend and operation. Fitted with +# `scripts/calibrate_flux2_working_memory.py --only vae`, which reproduces this table on any build: +# +# decode encode encode/decode +# cuDNN RTX 4090, torch 2.7.1 2185 1072 0.49 +# MIOpen RX 9070 XT, torch 2.10 3453 2688 0.78 +# MIOpen PRO W7900, torch 2.10 3525 2688 0.76 +# +# Two AMD generations (RDNA3 and RDNA4), the same numbers: the encode column agrees to the byte and +# the decode column to 2%. So this is MIOpen, not a per-card quirk, and the column below is fitted +# to the larger with ~2% headroom. +# +# MIOpen's convolution workspaces are simply larger than cuDNN's. This is not the attention term -- +# it shows up identically on the fused path, and the implied constants are flat across resolution on +# both backends, so the linear model itself holds. And the "encoding costs half of decoding" ratio +# the other estimators in this module use turns out to be a cuDNN property rather than an +# architectural one, which is why the two operations carry their own numbers here instead of a ratio. +# +# Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the +# constant follows the backend. +# +# Caveat for AMD users: `MIOPEN_FIND_MODE=2` selects convolution algorithms heuristically instead of +# by benchmark, and measured a uniform 1.28x more memory at every resolution. It is not the default +# and is not budgeted for here -- raise `device_working_mem_gb` if you set it. +_FLUX2_VAE_SCALING_CONSTANTS: dict[str, dict[str, int]] = { + "cudnn": {"decode": 2200, "encode": 1100}, + "miopen": {"decode": 3600, "encode": 2750}, +} + + +def _flux2_vae_scaling_constant(operation: Literal["encode", "decode"], device: torch.device) -> int: + """Pick the pixel-area constant for the convolution backend this device will actually use. + + A HIP build reports ``device.type == "cuda"``, so the torch build -- not the device string -- is + what separates MIOpen from cuDNN. MPS and CPU are unmeasured and take the cuDNN column; on MPS + that pairs with a score-matrix term the probe always charges there, so the total is not thin. + """ + is_rocm = device.type == "cuda" and torch.version.hip is not None + return _FLUX2_VAE_SCALING_CONSTANTS["miopen" if is_rocm else "cudnn"][operation] + + +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, and + the implied constant is flat across 512-1536px on every backend measured. What is *not* constant + is the constant itself: MIOpen's convolution workspaces cost ~1.6x cuDNN's for a decode and ~2.4x + for an encode, so it is looked up per backend -- see ``_FLUX2_VAE_SCALING_CONSTANTS`` for the + fitted table and the caveats. Peak *reserved* memory is what is measured throughout, the + conservative quantity that includes allocator overhead. + + For reference, decoding 1024x1024 peaks at ~4.3GB on cuDNN and ~6.6GB on MIOpen, 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/gfx1100 reports ``math`` for this 512-wide head, though gfx1201 takes it + on flash, and MPS has no fused kernel at all -- materializes a (pixels/8)^2 score matrix, which + grows quadratically and overtakes the linear term somewhere past 1280px. We ask torch which path + applies rather than assuming, so the estimate is right on all of them. + + The two terms are independent: a build can have a fused kernel and still need the larger + convolution constant, which is exactly what gfx1201 does. They also do not add -- see the + ``max`` at the end of this function for why, and for the measurements behind it. (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() + + device = device if device is not None else TorchDevice.choose_torch_device() + scaling_constant = _flux2_vae_scaling_constant(operation, device) + 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 + score_matrix_bytes = sdpa_score_matrix_bytes( + device=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, + ) + + # max, not sum: the two terms peak in different phases of the same forward. The mid-block sits at + # the 8x-downsampled bottleneck -- first in the decoder, last in the encoder -- so the full- + # resolution convolution feature maps that drive the linear term are not live while the score + # matrix is, and peak *reserved* is a high-water mark, not a running total. Measured (see the + # constants table above for the method): on gfx1100 and gfx1201 forcing `math` moves the measured + # peak by nothing at all up to 1024px, and the totals stay flat-linear in area either way. On + # CUDA the score matrix only pokes above the convolution peak at 1536px, and then by 2.6GB + # against the 21.5GB this term prices standalone, because the attention phase reuses blocks the + # allocator is already holding. Summing them reserved 11.1GB for a 1024px gfx1100 decode that + # measures 6.7GB; taking the max reserves 6.9GB. + # + # Where a max model is weakest is the crossover, where the two terms are near-equal and whatever + # overlap exists is no longer hidden. There is exactly one measured point like that: a 768px + # encode with cuDNN's linear constant and a materializing kernel measures 1.80GB against a + # 1.35GB max. It is not reachable as a shortfall, because the cache floors every reservation at + # `device_working_mem_gb` (3GB by default, see `ModelCache._get_vram_available`) and the whole + # crossover region sits under that floor. On the builds that really do materialize -- gfx1100 and + # gfx1201 -- the MIOpen constant keeps the linear term above the score term across the measured + # range, so the crossover does not arise there at all. + return int(max(working_memory, score_matrix_bytes)) + + def estimate_vae_working_memory_anima( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py new file mode 100644 index 00000000000..f322e228432 --- /dev/null +++ b/scripts/calibrate_flux2_working_memory.py @@ -0,0 +1,671 @@ +"""Calibrate the FLUX.2 working-memory estimates against measured peak CUDA/HIP memory. + +Background +---------- +Four constants decide how much VRAM the model cache keeps free for a FLUX.2 operation, and every one +of them was fitted on CUDA: + +1. ``SDPA_MATH_BYTES_PER_SCORE_ELEMENT`` (``backend/util/attention.py``) -- bytes per element of the + score matrix, charged only where SDPA has no fused kernel and materializes it. +2. ``estimate_vae_working_memory_flux2``'s 2200 / 1100 bytes per pixel per element byte + (``backend/util/vae_working_memory.py``). +3. ``Flux2DenoiseInvocation._estimate_working_memory``'s 0.4 MB per token, defined at the Klein 9B + width (4096) and scaled linearly by the loaded variant's width. +4. The dispatch question underneath all of it: does this build's SDPA fuse these shapes, or does it + build the score matrix? + +The estimate is consumed by the model cache via ``free >= estimate`` to decide what to evict, so it +MUST be an upper bound: this measures peak *reserved* (not merely allocated) memory, the conservative +quantity that includes caching-allocator overhead and kernel scratch. + +Why this script exists +---------------------- +ROCm answers (1), (2) and (4) differently from CUDA -- its fused kernels cap the head dim at 128, so +the FLUX.2 VAE's 512-wide mid-block head falls back to ``math`` and the score matrix becomes a real, +dominant term. The constants shipped today are known to be short on ROCm for VAE decode in the middle +of the resolution range. Recalibrating needs the hardware, so this puts the whole measurement in one +runnable file: run it on an AMD card and paste the output. + +Portability +----------- +Backend-agnostic: only ``torch.cuda.*``, which is the same API on NVIDIA/CUDA and AMD/ROCm (HIP) +builds. Run the SAME script on each backend and compare. The curve *shape* is architectural and +should match; the absolute constants can differ (cuDNN vs MIOpen workspaces, fused-kernel +availability, allocator rounding). Ship the max across backends, plus headroom. + +No checkpoints required. Every measurement uses a randomly initialized model at the real geometry -- +memory depends on shapes, not on weight values, and a stock-config ``AutoencoderKLFlux2`` reproduces +the real-weight 1024px decode measurement to within 2%. Pass ``--vae`` to measure a specific VAE (the +small-decoder variant has different ``block_out_channels``). + +Each point is measured in a FRESH SUBPROCESS so the caching allocator's fragmentation history cannot +contaminate the reserved-delta reading. A point that OOMs is recorded as ``oom`` rather than aborting +the run, so the grid can probe up to the card's ceiling safely. + +Usage +----- + python scripts/calibrate_flux2_working_memory.py + python scripts/calibrate_flux2_working_memory.py --only vae --csv flux2_rocm.csv + python scripts/calibrate_flux2_working_memory.py --only denoise --max-px 1024 +""" + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from invokeai.app.invocations.flux2_denoise import ( + FLUX2_ATTENTION_HEAD_DIM, + FLUX2_REFERENCE_HIDDEN_SIZE, + Flux2DenoiseInvocation, +) +from invokeai.backend.model_manager.configs.flux2_variant import flux2_hidden_size +from invokeai.backend.model_manager.taxonomy import Flux2VariantType +from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes +from invokeai.backend.util.vae_working_memory import ( + _flux2_vae_scaling_constant, + estimate_vae_working_memory_flux2, +) + +GIB = 1024**3 +MIB = 1024**2 + +# The FLUX.2 VAE attends on the 8x-downsampled grid with a single 512-wide head. Mirrors +# `_FLUX2_VAE_*` in vae_working_memory.py. +VAE_SPATIAL_COMPRESSION = 8 +VAE_MID_BLOCK_HEAD_DIM = 512 +LATENT_SCALE_FACTOR = 8 + +# (variant, transformer hidden size, context_in_dim). The source of truth is +# `model_manager/configs/flux2_variant.py`; `_check_variant_table` asserts we have not drifted. +VARIANTS = [ + (Flux2VariantType.Klein4B, 3072, 7680), + (Flux2VariantType.Klein9B, 4096, 12288), + (Flux2VariantType.Dev, 6144, 15360), +] + +DEFAULT_VAE_PX = [512, 768, 1024, 1280, 1536] + +# (heads, seq, head_dim) for the score-matrix constant. Spans the shapes the two estimators actually +# produce: the VAE's single 512-wide head, and the transformer's many 128-wide ones. +DEFAULT_SDPA_SHAPES = [ + (1, 4096, 512), + (1, 8192, 512), + (1, 16384, 512), + (4, 4096, 128), + (48, 4608, 128), +] + +# Sequence lengths the denoise slope is taken between. 4608 is a plain 1024x1024 generation +# (4096 image + 512 text); 9216 is that doubled, which is also what a batch of 2 costs. +DENOISE_SEQ_SHORT = 4608 +DENOISE_SEQ_LONG = 9216 +DENOISE_TEXT_TOKENS = 512 + +DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32} + +# Environment that materially changes a peak-*reserved* reading or which kernel runs. Reported in the +# header so a pasted result cannot be ambiguous about it later. +REPORTED_ENV_PREFIXES = ( + "MIOPEN_", + "PYTORCH_CUDA_ALLOC_CONF", + "PYTORCH_HIP_ALLOC_CONF", + "FLASH_ATTENTION_", + "TORCH_ROCM_", + "MIGRAPHX_", + "HSA_", + "TORCH_BLAS_", +) + +# `SDPBackend` is a pybind11 enum and is not iterable, so name it by hand. +SDP_BACKEND_NAMES = { + int(getattr(SDPBackend, name)): name + for name in ("ERROR", "MATH", "FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION", "OVERRIDEABLE") + if hasattr(SDPBackend, name) +} + + +def _check_variant_table() -> None: + """Fail loudly if the widths hard-coded above ever drift from the model manager's table.""" + for variant, hidden, _ in VARIANTS: + actual = flux2_hidden_size(variant) + if actual != hidden: + raise SystemExit(f"Variant table is stale: {variant.value} is {actual}, not {hidden}.") + + +def _peak_reserved(fn) -> int | None: + """Run ``fn`` and return the growth in peak reserved bytes, or ``None`` if it OOMed. + + Reserved rather than allocated: that is the quantity the estimate has to bound, because it is + what the allocator actually takes off the card. + """ + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_reserved() + try: + fn() + torch.cuda.synchronize() + except (torch.cuda.OutOfMemoryError, RuntimeError) as e: + if "out of memory" not in str(e).lower(): + raise + return None + return torch.cuda.max_memory_reserved() - baseline + + +# -------------------------------------------------------------------------------------------- +# 1. Dispatch: what does this build's SDPA actually do with the shapes the estimators produce? +# -------------------------------------------------------------------------------------------- + + +def measure_dispatch(dtype: torch.dtype) -> list[dict]: + """Report the kernel torch would pick for each shape. Costs nothing; explains everything else. + + This is the table that decides whether the score-matrix term is charged at all, and it is where + build-specific assumptions go to die -- an earlier revision of this feature assumed ROCm rejects + additive masks, which gfx1100 disproves. + """ + device = torch.device("cuda") + rows = [] + for head_dim in (128, 512): + for has_mask in (False, True): + 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_mask else None + try: + choice = int(torch.ops.aten._fused_sdp_choice(q, q, q, mask, 0.0, False)) + name = SDP_BACKEND_NAMES.get(choice, str(choice)) + except Exception as e: # noqa: BLE001 - the failure itself is the datum + name = f"raised ({type(e).__name__})" + rows.append({"head_dim": head_dim, "mask": has_mask, "choice": name}) + return rows + + +# -------------------------------------------------------------------------------------------- +# 2. SDPA_MATH_BYTES_PER_SCORE_ELEMENT +# -------------------------------------------------------------------------------------------- + + +@torch.inference_mode() +def measure_sdpa(num_heads: int, seq_len: int, head_dim: int, dtype: torch.dtype) -> dict: + """Peak reserved bytes per element of a materialized score matrix, with MATH forced.""" + device = torch.device("cuda") + q = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + def run() -> None: + with sdpa_kernel([SDPBackend.MATH]): + F.scaled_dot_product_attention(q, k, v) + + peak = _peak_reserved(run) + row = {"num_heads": num_heads, "seq_len": seq_len, "head_dim": head_dim, "oom": peak is None} + if peak is not None: + elements = num_heads * seq_len * seq_len + row |= {"reserved_delta": peak, "bytes_per_element": peak / elements} + return row + + +# -------------------------------------------------------------------------------------------- +# 3. The VAE's linear constants +# -------------------------------------------------------------------------------------------- + + +def _load_vae(vae_path: str | None, dtype: torch.dtype): + from diffusers import AutoencoderKLFlux2 + + if vae_path: + return AutoencoderKLFlux2.from_pretrained(vae_path, local_files_only=True, torch_dtype=dtype) + # Weight values do not affect activation memory, only shapes do; the stock config is the real + # FLUX.2 VAE geometry (32 latent channels, block_out_channels ending at 512). + return AutoencoderKLFlux2().to(dtype=dtype) + + +@torch.inference_mode() +def measure_vae(operation: str, px: int, dtype: torch.dtype, force_math: bool, vae_path: str | None) -> dict: + """Peak reserved memory for one untiled decode/encode, against what the estimator predicts. + + ``force_math`` runs the mid-block attention through SDPA's ``math`` fallback even where a fused + kernel exists, so a CUDA box can approximate the regime ROCm is in permanently. It is not a + substitute for measuring on ROCm -- the two are not equivalent, which is itself worth showing. + """ + device = torch.device("cuda") + vae = _load_vae(vae_path, dtype).to(device).eval() + vae.disable_tiling() # the decode/encode invocations do not tile; match them. + + param = next(vae.parameters()) + element_size = param.element_size() + if operation == "decode": + latent_channels = int(vae.config.latent_channels) + x = torch.randn(1, latent_channels, px // LATENT_SCALE_FACTOR, px // LATENT_SCALE_FACTOR, **_td(device, dtype)) + else: + x = torch.randn(1, 3, px, px, **_td(device, dtype)) + + def run() -> None: + ctx = sdpa_kernel([SDPBackend.MATH]) if force_math else _null_context() + with ctx: + if operation == "decode": + vae.decode(x, return_dict=False) + else: + vae.encode(x, return_dict=False)[0].mode() + + peak = _peak_reserved(run) + row = { + "operation": operation, + "px": px, + "force_math": force_math, + "element_size": element_size, + "oom": peak is None, + } + if peak is None: + return row + + estimate = estimate_vae_working_memory_flux2(operation=operation, image_tensor=x, vae=vae, device=device) + seq_len = (px // VAE_SPATIAL_COMPRESSION) ** 2 + score_bytes = sdpa_score_matrix_bytes( + device=device, dtype=param.dtype, num_heads=1, head_dim=VAE_MID_BLOCK_HEAD_DIM, seq_len=seq_len + ) + if force_math and score_bytes == 0: + # Fused here, so the estimator charged nothing -- but we forced math, so price it anyway. + score_bytes = seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + estimate = max(estimate, score_bytes) + + return row | { + "reserved_delta": peak, + "estimate": estimate, + "score_term": score_bytes, + # The whole measured peak over the pixel area, which is what the 2200/1100-style literals + # name. The score matrix is deliberately NOT backed out: the two terms peak in different + # phases of the forward, so the estimate takes the larger of them rather than the sum, and + # the convolution phase is what sets the high-water mark until the quadratic term overtakes + # it past ~1280px. Where `score_term` exceeds the measured peak, this column is not the + # quantity being fitted -- the score model is, and it is bounding rather than describing. + "implied_linear_constant": peak / (px * px * element_size), + "covered": peak <= estimate, + } + + +# -------------------------------------------------------------------------------------------- +# 4. The denoise per-token constant, and its width scaling +# -------------------------------------------------------------------------------------------- + + +@torch.inference_mode() +def measure_denoise(hidden: int, context_dim: int, seq_len: int, blocks: int, dtype: torch.dtype) -> dict: + """Peak reserved memory for one transformer forward at a given width and sequence length. + + Uses a reduced block count on purpose: the per-token cost is block-count independent, because a + no-grad forward frees each block's intermediates as it goes. ``--blocks`` measures a second point + so that assumption can be re-checked on this build rather than inherited. + """ + from diffusers import Flux2Transformer2DModel + + device = torch.device("cuda") + model = ( + Flux2Transformer2DModel( + num_layers=blocks, + num_single_layers=blocks, + num_attention_heads=hidden // FLUX2_ATTENTION_HEAD_DIM, + attention_head_dim=FLUX2_ATTENTION_HEAD_DIM, + joint_attention_dim=context_dim, + ) + .to(device=device, dtype=dtype) + .eval() + ) + + img_tokens = seq_len - DENOISE_TEXT_TOKENS + kwargs = { + "hidden_states": torch.randn(1, img_tokens, 128, **_td(device, dtype)), + "encoder_hidden_states": torch.randn(1, DENOISE_TEXT_TOKENS, context_dim, **_td(device, dtype)), + "timestep": torch.full((1,), 0.5, **_td(device, dtype)), + "img_ids": torch.zeros(img_tokens, 4, **_td(device, dtype)), + "txt_ids": torch.zeros(DENOISE_TEXT_TOKENS, 4, **_td(device, dtype)), + "guidance": torch.full((1,), 4.0, **_td(device, dtype)), + "return_dict": False, + } + + # Warm up first, then drop the allocator's cache: the cold call pays one-off weight-cast and + # workspace costs that the per-token slope must not absorb. The slope is a difference between + # two of these, so any constant overhead cancels either way -- warming just reduces the noise. + try: + model(**kwargs) + except (torch.cuda.OutOfMemoryError, RuntimeError) as e: + if "out of memory" not in str(e).lower(): + raise + return {"hidden": hidden, "seq_len": seq_len, "blocks": blocks, "oom": True} + + peak = _peak_reserved(lambda: model(**kwargs)) + row = {"hidden": hidden, "seq_len": seq_len, "blocks": blocks, "oom": peak is None} + if peak is not None: + row |= {"reserved_delta": peak} + return row + + +# -------------------------------------------------------------------------------------------- +# plumbing +# -------------------------------------------------------------------------------------------- + + +def _td(device: torch.device, dtype: torch.dtype) -> dict: + return {"device": device, "dtype": dtype} + + +class _null_context: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +def _report_environment() -> None: + """Print the env vars that can move these numbers, and warn about the one that invalidates them. + + ``garbage_collection_threshold`` makes the allocator release cached blocks once reserved memory + passes a fraction of the card -- which is precisely the quantity being measured here. A run with + it set reports a *lower* peak than the operation actually needs, and the giveaway is a series + that stops rising with resolution. + """ + import os + + interesting = {k: v for k, v in sorted(os.environ.items()) if k.startswith(REPORTED_ENV_PREFIXES)} + if interesting: + print("env: " + ", ".join(f"{k}={v}" for k, v in interesting.items())) + else: + print("env: none of the MIOpen / allocator / flash-attention overrides are set") + for key, value in interesting.items(): + if "garbage_collection_threshold" in value: + print( + f" WARNING: {key} sets garbage_collection_threshold. That releases cached blocks " + "once reserved memory crosses the threshold, so peak-reserved readings near the " + "card's capacity will read LOW. Re-run without it before fitting any constant." + ) + + +def _flag_non_monotonic(rows: list[dict]) -> None: + """Peak cannot fall as the input grows. If it does, the reading is not measuring what it claims.""" + for operation in ("decode", "encode"): + for force_math in (False, True): + series = [ + r for r in rows if r["operation"] == operation and r["force_math"] is force_math and not r.get("oom") + ] + series.sort(key=lambda r: r["px"]) + for earlier, later in zip(series, series[1:], strict=False): + if later["reserved_delta"] < earlier["reserved_delta"]: + print( + f" WARNING: {operation} peak FELL from {earlier['px']}px to {later['px']}px " + f"({earlier['reserved_delta'] / GIB:.3f} -> {later['reserved_delta'] / GIB:.3f} GiB, " + f"force_math={force_math}). Peak cannot decrease as the input grows; the larger " + "points are being clipped (allocator GC threshold, or another process freeing " + "memory). Do not fit a constant to this run." + ) + break + + +def _run_point(args: list[str]) -> dict | None: + """Run one measurement in a fresh subprocess and return its JSON row.""" + proc = subprocess.run([sys.executable, __file__, "--single", *args], capture_output=True, text=True) + line = proc.stdout.strip().splitlines()[-1] if proc.stdout.strip() else "" + try: + return json.loads(line) + except Exception: + tail = proc.stderr.strip().splitlines()[-1:] or ["(no stderr)"] + print(f" FAILED {' '.join(args)}: {tail[0]}") + return None + + +def report_dispatch(dtype_name: str) -> list[dict]: + print("\n=== 1. SDPA dispatch: which kernel would this build pick? ===") + print("The score-matrix term is charged only where the answer is MATH (or the query raises).\n") + rows = measure_dispatch(DTYPES[dtype_name]) or [] + print(f"{'head_dim':>9} {'mask':>6} {'kernel':>22}") + print("-" * 40) + for r in rows: + print(f"{r['head_dim']:>9} {str(r['mask']):>6} {r['choice']:>22}") + print("\n head_dim 512 is the FLUX.2 VAE mid-block; 128 is the transformer, and the masked row") + print(" is regional prompting. MATH on the 512 row means the VAE estimate needs the score term.") + return rows + + +def report_sdpa(shapes: list[tuple[int, int, int]], dtype_name: str) -> list[dict]: + print("\n=== 2. SDPA_MATH_BYTES_PER_SCORE_ELEMENT ===") + print(f"Peak reserved per score element with MATH forced. Shipped constant: {SDPA_MATH_BYTES_PER_SCORE_ELEMENT}.\n") + print(f"{'heads':>6} {'seq':>7} {'head_dim':>9} {'reserved(GiB)':>14} {'bytes/elem':>11}") + print("-" * 52) + rows = [] + for heads, seq, head_dim in shapes: + row = _run_point(["sdpa", str(heads), str(seq), str(head_dim), dtype_name]) + if row is None: + continue + rows.append(row) + if row.get("oom"): + print(f"{heads:>6} {seq:>7} {head_dim:>9} {'OOM':>14}") + continue + print( + f"{heads:>6} {seq:>7} {head_dim:>9} {row['reserved_delta'] / GIB:>14.3f} {row['bytes_per_element']:>11.2f}" + ) + fitted = [r["bytes_per_element"] for r in rows if not r.get("oom")] + if fitted: + worst = max(fitted) + verdict = "OK" if SDPA_MATH_BYTES_PER_SCORE_ELEMENT >= worst else "SHORT" + print(f"\n max = {worst:.2f}; shipped constant is {SDPA_MATH_BYTES_PER_SCORE_ELEMENT} -> {verdict}") + return rows + + +def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[dict]: + decode_k = _flux2_vae_scaling_constant("decode", torch.device("cuda")) + encode_k = _flux2_vae_scaling_constant("encode", torch.device("cuda")) + print(f"\n=== 3. VAE linear constants (this build selects {decode_k} decode / {encode_k} encode) ===") + print("`implied_k` is the measured peak over pixel area, directly comparable to those literals;") + print("fit on the rows whose `math` column matches what this build really does (section 1).") + print("`covered` is the question that matters: is the shipped estimate an upper bound here?") + print("Caveat: forcing math on a build that HAS a fused kernel is not equivalent to a build that") + print("has none -- on cuDNN the forced-math decode measures *below* the fused one until 1536px,") + print("because the memory-efficient kernel's workspace is the larger term. Only a real run on the") + print("materializing build calibrates it.\n") + print( + f"{'op':7} {'px':>5} {'math':>5} {'measured(GiB)':>14} {'estimate(GiB)':>14} {'implied_k':>10} {'covered':>8}" + ) + print("-" * 70) + rows = [] + for operation in ("decode", "encode"): + for px in pxs: + for force_math in (False, True): + args = ["vae", operation, str(px), dtype_name, "1" if force_math else "0"] + if vae_path: + args.append(vae_path) + row = _run_point(args) + if row is None: + continue + rows.append(row) + if row.get("oom"): + print(f"{operation:7} {px:>5} {str(force_math):>5} {'OOM':>14}") + continue + k = row["implied_linear_constant"] + print( + f"{operation:7} {px:>5} {str(force_math):>5} {row['reserved_delta'] / GIB:>14.3f} " + f"{row['estimate'] / GIB:>14.3f} {(f'{k:.0f}' if k else 'n/a'):>10} " + f"{('yes' if row['covered'] else 'NO'):>8}" + ) + print("") + _flag_non_monotonic(rows) + for operation in ("decode", "encode"): + # Compare against the column this build actually selects, not against a hard-coded number. + shipped = _flux2_vae_scaling_constant(operation, torch.device("cuda")) + for force_math in (False, True): + ks = [ + r["implied_linear_constant"] + for r in rows + if r["operation"] == operation + and r["force_math"] is force_math + and not r.get("oom") + # Skip rows where the score matrix, not the convolution phase, is what the estimate + # is bounding: there the linear constant is not the thing under test. + and r["score_term"] < r["reserved_delta"] + ] + if not ks: + continue + mode = "math" if force_math else "fused" + verdict = "OK" if shipped >= max(ks) else "SHORT" + print(f" {operation} ({mode}): implied_k max = {max(ks):.0f}, shipped = {shipped} -> {verdict}") + short = [r for r in rows if not r.get("oom") and not r["covered"]] + if short: + print("\n Points the shipped estimate does NOT cover:") + for r in short: + gap = (r["reserved_delta"] - r["estimate"]) / GIB + # The cache never reserves less than `device_working_mem_gb`, so a point under that floor + # is not actually short in production. + floored = " (absorbed by the 3GB device_working_mem_gb floor)" if r["reserved_delta"] < 3 * GIB else "" + print(f" {r['operation']} {r['px']}px force_math={r['force_math']}: short by {gap:.2f} GiB{floored}") + return rows + + +def report_denoise(dtype_name: str, blocks: int, extra_blocks: int | None) -> list[dict]: + print("\n=== 4. Denoise per-token constant and its width scaling ===") + print(f"Shipped: {0.4:.1f} MB/token at hidden={FLUX2_REFERENCE_HIDDEN_SIZE}, scaled linearly by width.\n") + print( + f"{'variant':10} {'hidden':>7} {'blocks':>7} {'short(GiB)':>11} {'long(GiB)':>10} {'MB/token':>9} {'ratio':>7}" + ) + print("-" * 68) + rows = [] + slopes: dict[int, float] = {} + block_counts = [blocks] + ([extra_blocks] if extra_blocks else []) + for variant, hidden, context_dim in VARIANTS: + for nb in block_counts: + pair = [] + for seq in (DENOISE_SEQ_SHORT, DENOISE_SEQ_LONG): + row = _run_point(["denoise", str(hidden), str(context_dim), str(seq), str(nb), dtype_name]) + if row is None or row.get("oom"): + pair = [] + break + rows.append(row) + pair.append(row["reserved_delta"]) + if not pair: + print(f"{variant.value:10} {hidden:>7} {nb:>7} {'OOM':>11}") + continue + slope = (pair[1] - pair[0]) / (DENOISE_SEQ_LONG - DENOISE_SEQ_SHORT) + if nb == blocks: + slopes[hidden] = slope + print( + f"{variant.value:10} {hidden:>7} {nb:>7} {pair[0] / GIB:>11.3f} {pair[1] / GIB:>10.3f} " + f"{slope / MIB:>9.4f} {slope / slopes.get(FLUX2_REFERENCE_HIDDEN_SIZE, slope):>7.3f}" + ) + + if extra_blocks: + print(f"\n The two block counts should agree per width; if they do not, the '{blocks} blocks stand in") + print(" for the real model' assumption does not hold on this build and the rest is suspect.") + + reference = slopes.get(FLUX2_REFERENCE_HIDDEN_SIZE) + if reference: + shipped = 0.4 * MIB + verdict = "OK" if shipped >= reference else "SHORT" + print( + f"\n hidden={FLUX2_REFERENCE_HIDDEN_SIZE}: measured {reference / MIB:.4f} MB/token, shipped 0.4 -> {verdict}" + ) + for hidden, slope in sorted(slopes.items()): + print( + f" hidden={hidden}: slope ratio {slope / reference:.3f} against width ratio " + f"{hidden / FLUX2_REFERENCE_HIDDEN_SIZE:.3f}" + ) + print("\n Those two columns should match: that is the claim that width scales the constant.") + # What the full estimator would reserve for the case #9500 describes. + for variant, hidden, _ in VARIANTS: + est = Flux2DenoiseInvocation._estimate_working_memory( + None, image_seq_len=4096, ref_image_seq_len=12288, text_seq_len=512, num_loras=0, hidden_size=hidden + ) + need = slopes.get(hidden, reference) * 16896 + print( + f" {variant.value:10} 1024px + 3 refs: estimate {est / GIB:5.2f} GiB, " + f"measured slope implies {need / GIB:5.2f} GiB" + ) + return rows + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--only", + choices=["dispatch", "sdpa", "vae", "denoise"], + action="append", + help="Run only these sections (repeatable). Default: all four.", + ) + parser.add_argument("--dtype", choices=list(DTYPES), default="bfloat16", help="Compute dtype. Default bfloat16.") + parser.add_argument("--max-px", type=int, default=None, help="Skip VAE resolutions above this.") + parser.add_argument("--vae", type=str, default=None, help="Optional AutoencoderKLFlux2 diffusers dir.") + parser.add_argument("--blocks", type=int, default=2, help="Transformer blocks per stream. Default 2.") + parser.add_argument( + "--extra-blocks", + type=int, + default=None, + help="Measure a second block count too, to re-check block-count independence on this build.", + ) + parser.add_argument("--csv", type=str, default=None, help="Write the raw rows to CSV.") + # Internal: measure one point in this process and print one JSON line. + parser.add_argument("--single", nargs="*", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("No CUDA/HIP device available.") + _check_variant_table() + + if args.single: + kind, rest = args.single[0], args.single[1:] + if kind == "sdpa": + heads, seq, head_dim, dtype_name = int(rest[0]), int(rest[1]), int(rest[2]), rest[3] + print(json.dumps(measure_sdpa(heads, seq, head_dim, DTYPES[dtype_name]))) + elif kind == "vae": + operation, px, dtype_name, force_math = rest[0], int(rest[1]), rest[2], rest[3] == "1" + vae_path = rest[4] if len(rest) > 4 else None + print(json.dumps(measure_vae(operation, px, DTYPES[dtype_name], force_math, vae_path))) + elif kind == "denoise": + hidden, context_dim, seq, blocks, dtype_name = ( + int(rest[0]), + int(rest[1]), + int(rest[2]), + int(rest[3]), + rest[4], + ) + print(json.dumps(measure_denoise(hidden, context_dim, seq, blocks, DTYPES[dtype_name]))) + else: + raise SystemExit(f"unknown point kind {kind}") + return + + sections = args.only or ["dispatch", "sdpa", "vae", "denoise"] + print( + f"torch {torch.__version__} | device {torch.cuda.get_device_name(0)} | " + f"hip={torch.version.hip} | dtype={args.dtype}" + ) + _report_environment() + pxs = [p for p in DEFAULT_VAE_PX if args.max_px is None or p <= args.max_px] + + rows: list[dict] = [] + if "dispatch" in sections: + rows += [{"section": "dispatch", **r} for r in report_dispatch(args.dtype)] + if "sdpa" in sections: + rows += [{"section": "sdpa", **r} for r in report_sdpa(DEFAULT_SDPA_SHAPES, args.dtype)] + if "vae" in sections: + rows += [{"section": "vae", **r} for r in report_vae(pxs, args.dtype, args.vae)] + if "denoise" in sections: + rows += [{"section": "denoise", **r} for r in report_denoise(args.dtype, args.blocks, args.extra_blocks)] + + if args.csv: + import csv + + fieldnames: list[str] = [] + for r in rows: + for key in r: + if key not in fieldnames: + fieldnames.append(key) + with Path(args.csv).open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + print(f"\nWrote {args.csv}") + + +if __name__ == "__main__": + main() 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..9cdbc695c24 --- /dev/null +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -0,0 +1,1224 @@ +"""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_ATTENTION_HEAD_DIM, + FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH, + FLUX2_MAX_HIDDEN_SIZE, + FLUX2_REFERENCE_HIDDEN_SIZE, + 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 ( + _FLUX2_VAE_SCALING_CONSTANTS, + estimate_vae_working_memory_flux2, +) + +MB = 1024**2 +GB = 1024**3 + +# The estimator's default width, and the head count that follows from it. The suite's pinned +# `MEASURED_DENOISE` table was taken on Klein 9B, which is that default, so those rows are unmoved +# by the width scaling. +KLEIN_9B_HEADS = FLUX2_REFERENCE_HIDDEN_SIZE // FLUX2_ATTENTION_HEAD_DIM +PER_TOKEN = FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH + + +def _per_token(hidden): + """The estimator's per-token cost at a given width, as it computes it.""" + return int(FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH * hidden / FLUX2_REFERENCE_HIDDEN_SIZE) + + +# 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, + hidden_size=FLUX2_REFERENCE_HIDDEN_SIZE, + 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, + hidden_size=hidden_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 * PER_TOKEN + + 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 * PER_TOKEN + ) + + 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 TestTransformerWidthIsBudgeted: + """Per-token activation cost is linear in the transformer's hidden width, not only in the token + count. Measured slope between 4608 and 9216 tokens on CUDA/bf16, block count and everything else + held fixed, each point in a fresh process: + + 3072 (Klein 4B) 0.2912 MB/tok 4096 (Klein 9B) 0.3859 6144 ([dev]) 0.5547 + + which is 0.755 / 1.000 / 1.438 of the Klein 9B slope against width ratios of 0.75 / 1.00 / 1.50. + Linear, and sub-linear at the top, so scaling the constant by width stays an upper bound. + + The same holds on ROCm/gfx1201, measured independently: 0.3147 / 0.4067 / 0.5890, i.e. ratios of + 0.774 / 1.000 / 1.448. The absolute constant is per-build (which is why it is 0.42, not CUDA's + 0.386); the width ratio is architectural and is the same on both. + + Calibrating on Klein 9B alone -- which this did -- silently under-reserved FLUX.2 [dev] by a + third, and [dev] reaches the denoise node as a first-class path with its own loader and starter + models. The head count follows the same width, so the score-matrix term gets the variant's real + count rather than the widest. + """ + + # (platform, hidden size, measured MB per token). Reproduce with + # `scripts/calibrate_flux2_working_memory.py --only denoise`. + MEASURED_WIDTH_SLOPE = [ + ("cuda-4090", 3072, 0.2912), + ("cuda-4090", 4096, 0.3859), + ("cuda-4090", 6144, 0.5547), + ("rocm-gfx1201", 3072, 0.3147), + ("rocm-gfx1201", 4096, 0.4067), + ("rocm-gfx1201", 6144, 0.5890), + ] + + def _slope_per_token(self, hidden): + short = _estimate(image_seq_len=4096, text_seq_len=512, hidden_size=hidden) + long = _estimate(image_seq_len=8704, text_seq_len=512, hidden_size=hidden) + return (long - short) / (9216 - 4608) + + @pytest.mark.parametrize("platform, hidden, measured_mb", MEASURED_WIDTH_SLOPE) + def test_the_estimate_upper_bounds_the_measured_slope(self, platform, hidden, measured_mb): + """The constant has to bound the worst build measured, not the one it was written on. It was + 0.4, which is under the 0.4067 gfx1201 costs at the reference width; it is 0.42. The upper + guard keeps that from drifting into over-reservation, since this is the dominant term.""" + slope = self._slope_per_token(hidden) + assert slope >= measured_mb * MB + assert slope <= 1.25 * measured_mb * MB + + def test_the_slope_is_linear_in_width(self): + klein_9b = self._slope_per_token(4096) + assert self._slope_per_token(3072) == pytest.approx(0.75 * klein_9b, rel=0.01) + assert self._slope_per_token(6144) == pytest.approx(1.50 * klein_9b, rel=0.01) + + @pytest.mark.parametrize("hidden, heads", [(3072, 24), (4096, 32), (6144, 48)]) + def test_the_score_matrix_uses_the_variants_real_head_count(self, hidden, heads): + """It used to charge 48 heads for every variant. On a materializing backend that is a third + too much on Klein 9B and double on Klein 4B -- and on MPS, where the probe always takes the + materializing branch, that term lands on every single estimate.""" + seq_len = 4096 + 512 + fused = _estimate(image_seq_len=4096, has_regional_mask=True, hidden_size=hidden) + with _materializing(): + materializing = _estimate( + image_seq_len=4096, has_regional_mask=True, hidden_size=hidden, device=MATERIALIZING + ) + assert materializing - fused == heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + +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) * PER_TOKEN + + 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) * PER_TOKEN + ) + + 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) * PER_TOKEN} + + 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) * PER_TOKEN + + 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 * PER_TOKEN + KLEIN_9B_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 TestVaeConstantsFollowTheConvBackend: + """The pixel-area constant is not one number. MIOpen's convolution workspaces cost far more than + cuDNN's for the same decode, and the gap is not the attention term -- it is identical on the + fused path. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, peak reserved, + each point in a fresh process: + + decode encode + cuDNN RTX 4090 2180 2185 2165 1072 1063 1061 + MIOpen RX 9070 XT 3453 3369 3368 2687 2688 2688 + + Flat across 512/768/1024px on both, so the linear model holds; only the coefficient moves. Note + the encode column: cuDNN's is half its decode, MIOpen's is four fifths, so "encoding costs half + of decoding" -- the ratio the other estimators in the module use -- is a cuDNN property, not an + architectural one. + """ + + # (backend, operation, px, measured GiB) + MEASURED = [ + ("cudnn", "decode", 512, 1.064), + ("cudnn", "decode", 768, 2.400), + ("cudnn", "decode", 1024, 4.229), + ("cudnn", "encode", 512, 0.523), + ("cudnn", "encode", 768, 1.168), + ("cudnn", "encode", 1024, 2.072), + ("miopen", "decode", 512, 1.686), + ("miopen", "decode", 768, 3.701), + ("miopen", "decode", 1024, 6.578), + # PRO W7900 (gfx1100). Costs slightly more per pixel than gfx1201 and sets the column. + ("miopen", "decode", 512, 1.721), + ("miopen", "decode", 768, 3.781), + ("miopen", "decode", 1024, 6.703), + ("miopen", "encode", 512, 1.312), + ("miopen", "encode", 768, 2.953), + ("miopen", "encode", 1024, 5.250), + ] + + @pytest.mark.parametrize("backend, operation, px, measured_gib", MEASURED) + def test_the_constant_upper_bounds_its_own_backend(self, backend, operation, px, measured_gib): + """Each column has to bound the hardware it was fitted on. Before this the cuDNN column was + used everywhere, leaving a 2.3 GiB shortfall on a 1024px MIOpen decode.""" + constant = _FLUX2_VAE_SCALING_CONSTANTS[backend][operation] + estimate = px * px * 2 * constant # bf16 element size + assert estimate >= measured_gib * GB + assert estimate <= 1.3 * measured_gib * GB + + def test_the_cudnn_constant_would_not_cover_miopen(self): + """The regression this guards: a 1024px MIOpen decode needs 6.58 GiB and the cuDNN constant + reserves 4.30. That is the same class of shortfall #9500 reports, one backend over.""" + cudnn = _FLUX2_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + assert 1024 * 1024 * 2 * cudnn < 6.578 * GB + + def test_the_encode_ratio_is_not_architectural(self): + """cuDNN's encode is half its decode; MIOpen's is four fifths. A shared ratio cannot express + both, which is why the table carries the two operations separately.""" + ratios = { + backend: consts["encode"] / consts["decode"] for backend, consts in _FLUX2_VAE_SCALING_CONSTANTS.items() + } + assert ratios["cudnn"] == pytest.approx(0.50, abs=0.02) + assert ratios["miopen"] == pytest.approx(0.76, abs=0.03) + + def test_a_rocm_build_selects_the_miopen_column(self): + """A HIP build reports `device.type == "cuda"`, so the torch build is the discriminator.""" + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + latents = torch.zeros(1, 32, 128, 128) + + 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=latents, vae=v, device=torch.device("cuda") + ) + + with ( + patch("torch.version.hip", "7.1.25424"), + patch("invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", return_value=False), + ): + rocm = estimate() + with ( + patch("torch.version.hip", None), + patch("invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", return_value=False), + ): + cuda = estimate() + + assert rocm > cuda + assert rocm / cuda == pytest.approx(3600 / 2200, rel=1e-3) + + +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 with the + batch just as the linear term does -- and since both scale together, the larger of the two + stays the larger at every batch size.""" + 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) == max(linear, score) + assert self._decode_estimate(3, device=MATERIALIZING) == 3 * max(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, variant="klein_9b"): + """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 + from invokeai.backend.model_manager.taxonomy import Flux2VariantType + + context.models.get_config.return_value = MagicMock( + base=BaseModelType.Flux2, + type=ModelType.Main, + format=ModelFormat.Checkpoint, + variant=Flux2VariantType(variant) if variant is not None else None, + ) + 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 * PER_TOKEN + + 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) * PER_TOKEN + + @pytest.mark.parametrize( + "variant, hidden", + [("klein_4b", 3072), ("klein_4b_base", 3072), ("klein_9b", 4096), ("klein_9b_base", 4096), ("dev", 6144)], + ) + def test_the_variants_width_reaches_the_reservation(self, variant, hidden): + """Per-token cost is linear in the transformer's width, so the reservation has to know which + variant it is loading. Calibrating on Klein 9B alone under-reserved [dev] by a third. Base + variants share their distilled twin's geometry.""" + assert self._run(num_ref_tokens=0, variant=variant) == _estimate( + image_seq_len=64 * 64, text_seq_len=512, hidden_size=hidden + ) + + def test_dev_reserves_half_again_what_klein_9b_does(self): + """The blocker in one assertion: 1024x1024 with three 1024x1024 references is 16896 tokens, + where the difference is ~3GB.""" + tokens = 64 * 64 + 12288 + 512 + klein = self._run(num_ref_tokens=12288, variant="klein_9b") + dev = self._run(num_ref_tokens=12288, variant="dev") + assert dev - klein == tokens * (_per_token(6144) - _per_token(4096)) + assert dev - klein > 2 * GB + + def test_an_unreadable_variant_falls_back_to_the_widest(self): + """Over-reserving on a model we cannot identify beats under-reserving on the largest one.""" + assert self._run(num_ref_tokens=0, variant=None) == _estimate( + image_seq_len=64 * 64, text_seq_len=512, hidden_size=FLUX2_MAX_HIDDEN_SIZE + ) + + 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) * PER_TOKEN + + 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) * PER_TOKEN + + +def _materializing_probe(device_type, device_index, dtype, head_dim, has_attn_mask): + """A synthetic stand-in for the torch probe, so both regimes can be exercised without hardware. + + It reports `math` for a head dim above 128 -- which is what ROCm really does, gfx1100 answers + MATH for the VAE's 512-wide head -- and also for an additive mask, which is *not* a hardware + claim: gfx1100 reports the memory-efficient kernel for a masked 128-wide head, same as CUDA. + The mask leg is here only to drive the masked branch of the estimator; which builds take it is + the probe's business, not this module's. + """ + return head_dim > 128 or has_attn_mask + + +class _null: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +def _materializing(): + return patch( + "invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", side_effect=_materializing_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: some builds have no fused kernel for the VAE's 512-wide + head, or for the dense additive mask regional prompting attaches, and fall through to `math`. + Where that happens the score matrix eventually dominates, so the estimate has to account for it + -- otherwise the fix works on CUDA and still OOMs elsewhere. + + How it accounts for it differs between the two estimators, and the difference is measured rather + than assumed. In the transformer the score matrix is live alongside the block activations, so it + adds. In the VAE the mid-block sits alone at the 8x-downsampled bottleneck, so the two peak in + different phases and the estimate takes the larger -- see `TestVaeTermsDoNotAdd`. + """ + + 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", 1536, 192 * 192), + ("encode", 1024, 128 * 128), + ("encode", 1328, 166 * 166), + ], + ) + def test_the_score_matrix_takes_over_where_it_is_the_larger_term(self, operation, px, tokens): + """Quadratic beats linear eventually. At these sizes it already has, so the estimate is the + score matrix exactly -- not the linear term, and not the two added together.""" + with _materializing(): + materializing = self._vae_estimate(operation, px, device=MATERIALIZING) + assert materializing == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert materializing > self._vae_estimate(operation, px, device=FUSED) + + 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, the score matrix has to be priced. 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", 1536, device=FUSED) + fused = self._vae_estimate("decode", 1536, device=FUSED) + + tokens = 192 * 192 + assert materializing == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert materializing > fused + 10 * GB + + def test_vae_score_matrix_dominates_at_high_resolution(self): + """A 1536px decode is ~9.7GB of convolution activations on cuDNN, and more than twice 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 two.""" + with _materializing(): + materializing = self._vae_estimate("decode", 1536, device=MATERIALIZING) + assert materializing > 20 * GB + assert materializing > 2 * self._vae_estimate("decode", 1536, device=FUSED) + + def test_tiling_makes_the_score_matrix_irrelevant(self): + """Tiling caps the mid-block sequence at (tile/8)^2, which drops the quadratic term far below + the tile's own linear cost -- so a tiled estimate is the linear term whether the build + materializes or not, and is flat across input sizes.""" + 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 + assert estimates[0] == self._vae_estimate("encode", 1024, device=FUSED, tile_size=512) + # Tiling turns a 100GB+ reservation at the 2024px reference cap into under 2GB. + assert estimates[0] < untiled / 50 + + +class TestVaeTermsDoNotAdd: + """The two VAE terms peak in different phases of the same forward, so the estimate takes the + larger rather than the sum. The mid-block sits at the 8x-downsampled bottleneck -- first in the + decoder, last in the encoder -- so the full-resolution convolution feature maps that drive the + linear term are not live while the score matrix is, and peak *reserved* is a high-water mark. + + Measured, decode, peak reserved, fresh process per point: + + px cuDNN fused cuDNN forced-math linear(2200) score(17) + 1024 4.229 3.527 4.30 4.25 + 1280 6.590 5.924 6.71 10.38 + 1536 9.353 11.965 9.67 21.52 + + Forcing math measures *below* the fused path until 1536px, and even there it exceeds it by + 2.6GB against the 21.5GB the score term prices standalone -- the attention phase reuses blocks + the allocator is already holding. On gfx1100 and gfx1201, forcing math moves the measured peak + by nothing at all up to 1024px, and the totals stay flat-linear in area either way. + + Summing the terms reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7GB. Taking the + max reserves 6.8GB. + """ + + def _estimate(self, px, device, materializing): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + with _materializing() if materializing else _null(): + return estimate_vae_working_memory_flux2( + operation="decode", image_tensor=torch.zeros(1, 32, px // 8, px // 8), vae=vae, device=device + ) + + # (pixel size, measured GiB) for a decode on the PRO W7900. Its 512px point is what sets the + # MIOpen decode constant: implied 3525, where gfx1201 asks for only 3453. + MEASURED_W7900_DECODE = [(1024, 6.703), (768, 3.781), (512, 1.721)] + + @pytest.mark.parametrize("px, measured_gib", MEASURED_W7900_DECODE) + def test_the_max_model_bounds_the_measured_miopen_peak(self, px, measured_gib): + linear = px * px * 2 * _FLUX2_VAE_SCALING_CONSTANTS["miopen"]["decode"] + score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert max(linear, score) >= measured_gib * GB + assert max(linear, score) <= 1.15 * measured_gib * GB + + def test_the_sum_model_would_have_over_reserved_by_two_thirds(self): + """At 1024px the sum reserves 11.1GiB for a decode that measures 6.7 -- on a 16GB card that + is the difference between the transformer staying resident and being evicted.""" + linear = 1024 * 1024 * 2 * _FLUX2_VAE_SCALING_CONSTANTS["miopen"]["decode"] + score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert (linear + score) / (6.703 * GB) > 1.6 + assert max(linear, score) / (6.703 * GB) < 1.15 + + def test_the_estimate_is_the_larger_term_not_the_sum(self): + """At 1024px the two are within 2% of each other on the cuDNN column, which makes this the + sharpest place to tell the models apart.""" + linear = 1024 * 1024 * 2 * _FLUX2_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert self._estimate(1024, MATERIALIZING, materializing=True) == max(linear, score) + assert self._estimate(1024, MATERIALIZING, materializing=True) < linear + score + + # (operation, px, measured GiB) with cuDNN's linear constant and a materializing kernel -- the + # combination MPS lands in, and the only place the max model is not an upper bound on its own. + MEASURED_CUDNN_FORCED_MATH = [ + ("encode", 512, 0.523), + ("encode", 768, 1.801), + ("encode", 1024, 4.072), + ("decode", 1024, 3.527), + ("decode", 1280, 5.924), + ("decode", 1536, 11.965), + ] + + @pytest.mark.parametrize("operation, px, measured_gib", MEASURED_CUDNN_FORCED_MATH) + def test_the_working_memory_floor_covers_the_crossover(self, operation, px, measured_gib): + """A max model is weakest where the two terms are near-equal, and one measured point shows + it: a 768px encode wants 1.80GiB against a 1.35GiB max. It is not reachable as a shortfall, + because the cache floors every reservation at `device_working_mem_gb` and the whole crossover + region sits below that floor. Reproducible to three decimals across runs, so this is the + model's real shape, not noise -- which is why it is pinned rather than rounded away.""" + constant = _FLUX2_VAE_SCALING_CONSTANTS["cudnn"][operation] + area = px * px if operation == "decode" else px * px + linear = area * 2 * constant + score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + floored = max(max(linear, score), 3 * GB) # what ModelCache._get_vram_available reserves + assert floored >= measured_gib * GB + + def test_a_fused_build_is_unaffected(self): + """The max only ever removes reservation, never adds it: with no score matrix the estimate is + the linear term, exactly as before.""" + linear = 1024 * 1024 * 2 * _FLUX2_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + assert self._estimate(1024, FUSED, materializing=False) == linear + + 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 == (KLEIN_9B_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): + """The head dim is the discriminator that actually holds: CUDA's memory-efficient kernel + takes the VAE's 512-wide head, ROCm caps at 128 and reports `math` for it. Masks are not a + discriminator -- gfx1100 reports the memory-efficient kernel for a masked 128-wide head just + as CUDA does -- so the masked case only has to agree with whatever torch says, which is the + whole point of asking it.""" + 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, + ) + assert vae_bytes == (0 if torch.version.hip is None else 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT) + materializes = _torch_sdpa_materializes_score_matrix( + "cuda", torch.device("cuda").index, torch.bfloat16, 128, True + ) + assert masked_bytes == (48 * 4608 * 4608 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT if materializes else 0) + + # (platform, heads, seq, head_dim, bytes per score element), SDPBackend.MATH forced, each point + # in a fresh process. Reproduce with `scripts/calibrate_flux2_working_memory.py --only sdpa`. + MEASURED_BYTES_PER_ELEMENT = [ + ("cuda-4090", 1, 4096, 512, 12.88), + ("cuda-4090", 1, 8192, 512, 10.28), + ("cuda-4090", 1, 16384, 512, 9.71), + ("cuda-4090", 4, 4096, 128, 9.97), + ("cuda-4090", 48, 4608, 128, 9.58), + ("rocm-gfx1100", 1, 4096, 512, 13.62), + ("rocm-gfx1100", 1, 8192, 512, 10.78), + ("rocm-gfx1100", 1, 16384, 512, 9.76), + ("rocm-gfx1100", 4, 4096, 128, 10.16), + ("rocm-gfx1100", 48, 4608, 128, 9.59), + ("rocm-gfx1201", 1, 4096, 512, 16.38), + ("rocm-gfx1201", 1, 8192, 512, 13.47), + ("rocm-gfx1201", 1, 16384, 512, 11.99), + ("rocm-gfx1201", 4, 4096, 128, 12.84), + ("rocm-gfx1201", 48, 4608, 128, 11.69), + ] + + @pytest.mark.parametrize("platform, num_heads, seq_len, head_dim, measured", MEASURED_BYTES_PER_ELEMENT) + def test_the_constant_upper_bounds_every_measured_platform(self, platform, num_heads, seq_len, head_dim, measured): + """The constant has to bound the worst build anyone has measured, not the one it was written + on. It was 13 (CUDA's worst point rounded up), then 14 when gfx1100 came in at 13.62, and is + 17 because gfx1201 costs 16.38 for the same shape. The spread between the two AMD cards is + wider than the gap between CUDA and either of them, so there is no "the ROCm number".""" + assert SDPA_MATH_BYTES_PER_SCORE_ELEMENT >= measured + + @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, ( + f"{estimate / num_heads / seq_len / seq_len:.2f} bytes/element measured here; " + f"SDPA_MATH_BYTES_PER_SCORE_ELEMENT={SDPA_MATH_BYTES_PER_SCORE_ELEMENT} must be an upper bound" + ) + # Loose on purpose. The constant is the worst of three measured platforms (CUDA ~10 + # bytes/element, gfx1100 ~10-14, gfx1201 ~12-16), so on any one of them it over-shoots by + # design; and this runs in-process, where the allocator may satisfy the forward from blocks + # it already holds and report a smaller delta than a fresh process would. The guard is here + # to catch an order-of-magnitude blunder, not to pin the calibration -- that is what + # `test_the_constant_upper_bounds_every_measured_platform` and the calibration script do. + assert estimate <= 3 * 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 == (KLEIN_9B_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 == (KLEIN_9B_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"