From 1904479a4f1576e7769b17e80582f8a98b85c458 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 19 Aug 2026 21:36:24 +0200 Subject: [PATCH 01/11] fix(flux2): estimate working memory for denoise and both VAE directions The FLUX.2 path called model_on_device() with no working_mem_bytes anywhere, so the model cache reserved only the default device_working_mem_gb and filled the rest of the card with the model. Reference images make that fatal rather than merely tight: their latents are concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of a 1024x1024 generation -- 6.5GB of activations against a 3GB reservation. Measured on CUDA in bf16 as peak reserved memory: transformer activations scale linearly at ~0.39 MB/token (no O(seq^2) term, SDPA) and are independent of block count; the FLUX.2 VAE costs ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, so a 1024x1024 decode peaks at ~4.3GB. Add Flux2DenoiseInvocation._estimate_working_memory() and estimate_vae_working_memory_flux2(), and pass them at every load site so the cache evicts enough to make room instead of hitting the shortfall as an OOM. Closes #9500 --- invokeai/app/invocations/flux2_denoise.py | 61 +++- invokeai/app/invocations/flux2_vae_decode.py | 10 +- invokeai/app/invocations/flux2_vae_encode.py | 9 +- invokeai/backend/flux2/ref_image_extension.py | 21 +- invokeai/backend/util/vae_working_memory.py | 37 +++ .../invocations/test_flux2_working_memory.py | 287 ++++++++++++++++++ 6 files changed, 418 insertions(+), 7 deletions(-) create mode 100644 tests/app/invocations/test_flux2_working_memory.py diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 7a5ce158c98..25c8b5d8f30 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -458,10 +458,33 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: bn_std=bn_std, ) + # Estimate the peak activation memory the transformer forward will need and ask the model cache + # to keep that much VRAM free. Without this hint the cache reserves only the small default + # working memory and fills the rest of the card with the model, so anything beyond a plain + # low-resolution generation OOMs. Reference images are the dominant term: their latents are + # concatenated onto the image stream, so three 1024x1024 references quadruple the sequence + # (and with it the activation footprint) of a 1024x1024 generation. + ref_image_seq_len = ref_image_extension.ref_image_latents.shape[1] if ref_image_extension is not None else 0 + estimated_working_memory = self._estimate_working_memory( + image_seq_len=packed_h * packed_w, + ref_image_seq_len=ref_image_seq_len, + text_seq_len=max(txt.shape[1], neg_txt.shape[1] if neg_txt is not None else 0), + num_loras=len(self.transformer.loras), + # The mask itself is already allocated; only the additive bias built per forward is new. + # It is skipped entirely when reference images are present (see below). + regional_attention_bias_bytes=( + regional_extension.restricted_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() + if regional_extension.restricted_attn_mask is not None and ref_image_seq_len == 0 + else 0 + ), + ) + with ExitStack() as exit_stack: # Load the transformer model (cached_weights, transformer) = exit_stack.enter_context( - context.models.load(self.transformer.transformer).model_on_device() + context.models.load(self.transformer.transformer).model_on_device( + working_mem_bytes=estimated_working_memory + ) ) config = transformer_config @@ -578,6 +601,42 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) mask = mask.to(device=latents.device, dtype=latents.dtype) return mask.expand_as(latents) + def _estimate_working_memory( + self, + image_seq_len: int, + ref_image_seq_len: int, + text_seq_len: int, + num_loras: int, + regional_attention_bias_bytes: int = 0, + ) -> int: + """Estimate peak transformer activation memory (bytes) so the model cache reserves enough headroom. + + FLUX.2 attention runs through SDPA without materializing the O(seq^2) score matrix, so the + activation footprint scales *linearly* with the total attended sequence -- text tokens, image + tokens, and reference-image tokens alike. Measured on the Klein 9B geometry in bf16 as peak + reserved memory, that slope is ~0.39 MB per token and holds from 1.5k to 28k tokens; it is + also independent of the block count (a no-grad forward frees each block's intermediates), so + the constant applies to both the 4B and 9B variants. + + The reference-image term is what makes this estimate necessary rather than merely nice to + have: a 1024x1024 generation is 4096 image tokens (~1.7GB), but attaching three 1024x1024 + references adds 12288 more for ~6.5GB, and a 1328px tile with three 1328px references reaches + ~10.9GB -- against a default ``device_working_mem_gb`` of 3. + + A fixed base covers resolution-independent overhead (transient fp8/GGUF -> bf16 weight casts + during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra + activation branch per patched layer, so we add a per-LoRA margin. + """ + GB = 1024**3 + MB = 1024**2 + per_token_bytes = int(0.4 * MB) + estimated = (image_seq_len + ref_image_seq_len + text_seq_len) * per_token_bytes + estimated += int(1.0 * GB) + estimated += regional_attention_bias_bytes + if num_loras > 0: + estimated += int(0.5 * num_loras * GB) + return estimated + def _load_text_conditioning( self, context: InvocationContext, diff --git a/invokeai/app/invocations/flux2_vae_decode.py b/invokeai/app/invocations/flux2_vae_decode.py index 58297ec1ae4..d4cad4ad75d 100644 --- a/invokeai/app/invocations/flux2_vae_decode.py +++ b/invokeai/app/invocations/flux2_vae_decode.py @@ -21,6 +21,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 @invocation( @@ -49,7 +50,14 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima Input latents should already be in the correct space after BN denormalization was applied in the denoiser. The VAE expects (B, 32, H, W) format. """ - with vae_info.model_on_device() as (_, vae): + # Decoding at FLUX.2 resolutions costs multiple GB of activations (~4.3GB at 1024x1024), + # far above the default working memory the cache would otherwise reserve. Tell it up front so + # it offloads enough of the (possibly still resident) transformer to leave room. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=latents, vae=vae_info.model + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would diff --git a/invokeai/app/invocations/flux2_vae_encode.py b/invokeai/app/invocations/flux2_vae_encode.py index 1b43483a408..9d92b3819a4 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -19,6 +19,7 @@ from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 @invocation( @@ -46,7 +47,13 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc The VAE encodes to 32-channel latent space. Output latents shape: (B, 32, H/8, W/8). """ - with vae_info.model_on_device() as (_, vae): + # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to + # free room rather than discovering the shortfall as an OOM. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="encode", image_tensor=image_tensor, vae=vae_info.model + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype device = TorchDevice.choose_torch_device() image_tensor = image_tensor.to(device=device, dtype=vae_dtype) diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index 368f3c4452f..ccdb1390e28 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -21,12 +21,16 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux2.sampling_utils import pack_flux2 from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 # Maximum pixel counts for reference images (matches BFL FLUX.2 sampling.py) # Single reference image: 2024² pixels, Multiple: 1024² pixels MAX_PIXELS_SINGLE_REF = 2024**2 # ~4.1M pixels MAX_PIXELS_MULTI_REF = 1024**2 # ~1M pixels +# Tile size (in pixels) forced on the VAE for reference-image encoding, see _prepare_ref_images(). +REF_ENCODE_TILE_SIZE = 512 + def resize_image_to_max_pixels(image: Image.Image, max_pixels: int) -> Image.Image: """Resize image to fit within max_pixels while preserving aspect ratio. @@ -203,8 +207,17 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: image_tensor = image_tensor * 2.0 - 1.0 image_tensor = image_tensor.unsqueeze(0) # Add batch dimension - # Encode using FLUX.2 VAE - with vae_info.model_on_device() as (_, vae): + # Encode using FLUX.2 VAE. The encode below forces REF_ENCODE_TILE_SIZE tiling, so the + # peak is bounded by one tile; tell the cache that up front so it frees the room instead + # of hitting the shortfall as an OOM. + estimated_working_memory = estimate_vae_working_memory_flux2( + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + tile_size=REF_ENCODE_TILE_SIZE, + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) @@ -219,8 +232,8 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: downsample = 2 ** (len(vae.config.block_out_channels) - 1) prev_tiling = (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size) vae.use_tiling = True - vae.tile_sample_min_size = 512 - vae.tile_latent_min_size = 512 // downsample + vae.tile_sample_min_size = REF_ENCODE_TILE_SIZE + vae.tile_latent_min_size = REF_ENCODE_TILE_SIZE // downsample try: # FLUX.2 VAE uses diffusers API latent_dist = vae.encode(image_tensor, return_dict=False)[0] diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index bd780d4c0b3..8ccddbf99d0 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -2,6 +2,7 @@ import torch from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny @@ -98,6 +99,42 @@ def estimate_vae_working_memory_flux( return int(working_memory) +def estimate_vae_working_memory_flux2( + operation: Literal["encode", "decode"], + image_tensor: torch.Tensor, + vae: AutoencoderKLFlux2, + tile_size: int | None = None, +) -> int: + """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. + + Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE -- + ``AutoencoderKLFlux2``'s mid-block attention runs through SDPA, so no O(area^2) term appears. + Measured on CUDA/bf16 as peak *reserved* memory (the conservative quantity, including allocator + overhead), the implied constants are ~2170 (decode) and ~1070 (encode) bytes per pixel per + element byte, flat across 512-1536px; the constants below round those up and match the FLUX.1 + ones. For reference, decoding 1024x1024 peaks at ~4.3GB and 1536x1536 at ~9.6GB -- far above the + default ``device_working_mem_gb``, which is why this estimate must be passed to the model cache. + + When tiling is enabled the peak is bounded by a single tile instead of the full image (measured + ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap). + """ + element_size = next(vae.parameters()).element_size() + + # Encoding uses ~50% the working memory of decoding. + scaling_constant = 2200 if operation == "decode" else 1100 + + if tile_size is not None: + # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. + working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 + else: + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] + out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] + working_memory = out_h * out_w * element_size * scaling_constant + + return int(working_memory) + + def estimate_vae_working_memory_anima( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py new file mode 100644 index 00000000000..816a7f71f37 --- /dev/null +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -0,0 +1,287 @@ +"""FLUX.2 working-memory estimates: the transformer denoise and both VAE directions. + +The FLUX.2 path originally called `model_on_device()` with no `working_mem_bytes` anywhere, so the +model cache reserved only the small default `device_working_mem_gb` and filled the rest of the card +with the model. Reference images make that fatal rather than merely tight: their latents are +concatenated onto the image stream, so three 1024x1024 references quadruple the attended sequence of +a 1024x1024 generation. See https://github.com/invoke-ai/InvokeAI/issues/9500. + +The `MEASURED_*` tables below are peak *reserved* memory measured on CUDA in bf16 (the conservative +quantity, including allocator overhead). Every estimate must stay an upper bound on them. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 + +from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation +from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation +from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 + +MB = 1024**2 +GB = 1024**3 + + +def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, regional_bias=0): + return Flux2DenoiseInvocation._estimate_working_memory( + MagicMock(spec=Flux2DenoiseInvocation), + image_seq_len=image_seq_len, + ref_image_seq_len=ref_image_seq_len, + text_seq_len=text_seq_len, + num_loras=num_loras, + regional_attention_bias_bytes=regional_bias, + ) + + +class TestFlux2DenoiseWorkingMemoryEstimate: + # (image tokens, reference tokens, measured peak reserved MB) on the Klein 9B geometry. + # Token grids are pixels/16, so a 1024px square is 4096 tokens. + MEASURED_DENOISE = [ + (1024, 0, 448), # 512px + (4096, 0, 1702), # 1024px + (4096, 4096, 3324), # 1024px + one 1024px reference + (4096, 8192, 4840), # + two references + (4096, 12288, 6538), # + three references (the tiled-refiner case from #9500) + (6889, 12288, 7528), # 1328px tile + three 1024px references + (6889, 20667, 10954), # 1328px tile + three 1328px references + (16384, 0, 6538), # 2048px, no references + ] + + @pytest.mark.parametrize("image_seq_len, ref_image_seq_len, measured_mb", MEASURED_DENOISE) + def test_estimate_is_an_upper_bound_on_measured_peak(self, image_seq_len, ref_image_seq_len, measured_mb): + """The cache treats the estimate as the amount it must keep free, so under-estimating OOMs.""" + assert _estimate(image_seq_len, ref_image_seq_len) >= measured_mb * MB + + @pytest.mark.parametrize("image_seq_len, ref_image_seq_len, measured_mb", MEASURED_DENOISE) + def test_estimate_does_not_wildly_over_reserve(self, image_seq_len, ref_image_seq_len, measured_mb): + """Over-estimating is not free: the cache offloads the transformer to RAM to honor the + reservation, and a model running over PCIe is indistinguishable from a hang.""" + assert _estimate(image_seq_len, ref_image_seq_len) <= measured_mb * MB + 2 * GB + + def test_reference_image_tokens_are_counted(self): + """The regression this whole module exists for: reference tokens are attended like image + tokens and cost the same per token, so they must enter the estimate.""" + without_refs = _estimate(image_seq_len=4096) + with_refs = _estimate(image_seq_len=4096, ref_image_seq_len=12288) + assert with_refs - without_refs == 12288 * int(0.4 * MB) + + def test_estimate_is_linear_in_total_sequence(self): + """Attention runs through SDPA, so there is no O(seq^2) term to model -- image, reference and + text tokens are interchangeable at the same per-token cost.""" + assert _estimate(image_seq_len=8192) == _estimate(image_seq_len=4096, ref_image_seq_len=4096) + assert _estimate(image_seq_len=4096, text_seq_len=1024) - _estimate(image_seq_len=4096, text_seq_len=512) == ( + 512 * int(0.4 * MB) + ) + + def test_lora_margin_is_added_per_lora(self): + """Sidecar-patched LoRAs add an activation branch per patched layer.""" + base = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, num_loras=1) - base == int(0.5 * GB) + assert _estimate(image_seq_len=4096, num_loras=3) - base == int(1.5 * GB) + + def test_regional_attention_bias_is_added(self): + base = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, regional_bias=123 * MB) - base == 123 * MB + + +class TestFlux2VaeWorkingMemoryEstimate: + # (operation, pixel size, measured peak reserved MB), bf16, untiled. + MEASURED_VAE = [ + ("decode", 512, 1086), + ("decode", 768, 2414), + ("decode", 1024, 4260), + ("decode", 1328, 7146), + ("decode", 1536, 9578), + ("encode", 512, 536), + ("encode", 1024, 2122), + ("encode", 1328, 3022), + ] + + def _mock_bf16_vae(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) # element_size == 2 + return vae + + def _tensor_for(self, operation, px): + # decode receives 32-channel latents at pixels/8; encode receives a pixel image. + return torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) + + @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) + def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae() + ) + assert estimate >= measured_mb * MB + + @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) + def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae() + ) + assert estimate == 1024 * 1024 * 2 * expected_constant + + def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): + """Reference-image encoding forces 512px tiling precisely so the peak stops following the + reference resolution -- measured flat at ~0.55GB from 1024px up to the 2024px reference cap.""" + estimates = [ + estimate_vae_working_memory_flux2( + operation="encode", + image_tensor=torch.zeros(1, 3, px, px), + vae=self._mock_bf16_vae(), + tile_size=512, + ) + for px in (1024, 1328, 2024) + ] + assert len(set(estimates)) == 1 + assert estimates[0] == int(512 * 512 * 2 * 1100 * 1.25) + assert estimates[0] >= 558 * MB # measured tiled peak at 2024px + # The whole point of tiling: it must shrink the reservation, not just bound the VAE. + untiled = estimate_vae_working_memory_flux2( + operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae() + ) + assert estimates[0] < untiled / 4 + + +class TestFlux2VaeInvocationsRequestWorkingMemory: + """The estimate is worthless unless it reaches `model_on_device()`.""" + + def _mock_vae_info(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + + vae_info = MagicMock() + vae_info.model = vae + vae_info.compute_device = torch.device("cpu") + cm = MagicMock() + cm.__enter__ = MagicMock(return_value=(None, vae)) + cm.__exit__ = MagicMock(return_value=None) + vae_info.model_on_device = MagicMock(return_value=cm) + return vae_info + + def test_decode_requests_working_memory(self): + vae_info = self._mock_vae_info() + context = MagicMock() + context.models.load.return_value = vae_info + context.tensors.load.return_value = torch.zeros(1, 32, 128, 128) + + expected = 10 * GB + with patch( + "invokeai.app.invocations.flux2_vae_decode.estimate_vae_working_memory_flux2", return_value=expected + ) as estimate: + invocation = Flux2VaeDecodeInvocation.model_construct( + latents=MagicMock(latents_name="latents"), vae=MagicMock(vae=MagicMock()) + ) + try: + invocation.invoke(context) + except Exception: + # The mocked decode math fails downstream; we only care that the cache was asked to + # reserve the estimate before the device context was entered. + pass + + estimate.assert_called_once() + assert estimate.call_args.kwargs["operation"] == "decode" + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + + def test_encode_requests_working_memory(self): + vae_info = self._mock_vae_info() + context = MagicMock() + context.models.load.return_value = vae_info + + expected = 4 * GB + with ( + patch( + "invokeai.app.invocations.flux2_vae_encode.estimate_vae_working_memory_flux2", return_value=expected + ) as estimate, + patch( + "invokeai.app.invocations.flux2_vae_encode.image_resized_to_grid_as_tensor", + return_value=torch.zeros(3, 1024, 1024), + ), + ): + invocation = Flux2VaeEncodeInvocation.model_construct( + image=MagicMock(image_name="image"), vae=MagicMock(vae=MagicMock()) + ) + try: + invocation.invoke(context) + except Exception: + pass + + estimate.assert_called_once() + assert estimate.call_args.kwargs["operation"] == "encode" + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + + +class _StopBeforeLoad(Exception): + """Raised in place of entering the transformer's device context, to end _run_diffusion early.""" + + +class TestFlux2DenoiseRequestsWorkingMemory: + """The denoise node must hand its estimate to the cache, and that estimate must grow with the + attached reference images -- the combination that #9500 was missing.""" + + def _run(self, num_ref_tokens: int): + """Drive `_run_diffusion` up to the transformer load and return the requested working memory.""" + from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType + from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + ConditioningFieldData, + FLUXConditioningInfo, + ) + + transformer_info = MagicMock() + transformer_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) + + context = MagicMock() + context.models.load.return_value = transformer_info + context.models.get_config.return_value = MagicMock( + base=BaseModelType.Flux2, type=ModelType.Main, format=ModelFormat.Checkpoint + ) + context.conditioning.load.return_value = ConditioningFieldData( + conditionings=[FLUXConditioningInfo(clip_embeds=torch.zeros(1, 768), t5_embeds=torch.zeros(1, 512, 12288))] + ) + + ref_extension = MagicMock() + ref_extension.ref_image_latents = torch.zeros(1, num_ref_tokens, 128) + + invocation = Flux2DenoiseInvocation.model_construct( + latents=None, + noise=None, + denoise_mask=None, + denoising_start=0.0, + denoising_end=1.0, + add_noise=True, + transformer=MagicMock(transformer=MagicMock(), loras=[]), + positive_text_conditioning=MagicMock(conditioning_name="pos", mask=None), + negative_text_conditioning=None, + guidance=4.0, + cfg_scale=1.0, + width=1024, + height=1024, + num_steps=4, + scheduler="euler", + seed=0, + vae=MagicMock(vae=MagicMock()), + kontext_conditioning=MagicMock() if num_ref_tokens else None, + ) + + with ( + patch.object(Flux2DenoiseInvocation, "_get_bn_stats", return_value=None), + patch("invokeai.backend.util.devices.TorchDevice.choose_torch_device", return_value=torch.device("cpu")), + patch("invokeai.app.invocations.flux2_denoise.Flux2RefImageExtension", return_value=ref_extension), + pytest.raises(_StopBeforeLoad), + ): + invocation._run_diffusion(context) + + transformer_info.model_on_device.assert_called_once() + return transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"] + + def test_estimate_reaches_the_model_cache(self): + """Without this the cache reserves only the default `device_working_mem_gb`.""" + assert self._run(num_ref_tokens=0) == _estimate(image_seq_len=64 * 64, text_seq_len=512) + + def test_reference_images_raise_the_reservation(self): + """Three 1024x1024 references add 12288 tokens to a 1024x1024 generation's 4096.""" + without_refs = self._run(num_ref_tokens=0) + with_refs = self._run(num_ref_tokens=12288) + assert with_refs - without_refs == 12288 * int(0.4 * MB) From b457f60cae4a38527ec6ebe0e37237d62c238d8e Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 24 Aug 2026 22:40:46 +0200 Subject: [PATCH 02/11] fix(flux2): budget SDPA's materialized score matrix where it is real The FLUX.2 working-memory estimates were linear in the sequence length, which holds only while SDPA picks a fused kernel. That is a property of the torch build, not of FLUX.2: ROCm's fused kernels cap the head dim at 128 and reject arbitrary additive masks, so both the VAE's 512-wide mid-block head and the dense S x S bias regional prompting attaches fall through to the math fallback and materialize the score matrix -- ~17GB for a 1536px decode, and heads x S^2 for a masked forward. Rather than assume either way, ask torch: sdpa_score_matrix_bytes() queries can_use_flash/efficient/cudnn_attention for the real head dim, dtype and mask, and adds 13 bytes per score element only when no fused kernel is eligible. Measured on CUDA with SDPBackend.MATH forced: 12.9 bytes/element at 4k tokens, 10.3 at 8k, 9.7 at 16k, identical for bf16, fp16 and fp32 because the fallback's softmax intermediates are always fp32. On CUDA every shape reports fused, so the term is zero and the existing calibration is untouched. Non-CUDA devices keep the fused assumption -- torch exposes no equivalent query there, and guessing would reserve double-digit GB on no evidence. --- invokeai/app/invocations/flux2_denoise.py | 41 +++- invokeai/app/invocations/flux2_vae_decode.py | 2 +- invokeai/app/invocations/flux2_vae_encode.py | 5 +- invokeai/backend/flux2/ref_image_extension.py | 1 + invokeai/backend/util/attention.py | 80 +++++++ invokeai/backend/util/vae_working_memory.py | 39 ++- .../invocations/test_flux2_working_memory.py | 222 +++++++++++++++++- 7 files changed, 374 insertions(+), 16 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 25c8b5d8f30..6dcb412e39e 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -49,8 +49,16 @@ from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo +from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice +# FLUX.2 attention geometry. The head dim is 128 across every variant; the head count follows the +# hidden size (Klein 4B: 24, Klein 9B: 32, FLUX.2 dev: 48). Only the head dim decides which SDPA +# kernel is eligible; the head count scales the `math` fallback's score matrix, and since the +# working-memory estimate is computed before the transformer is loaded, we use the largest. +FLUX2_ATTENTION_HEAD_DIM = 128 +FLUX2_MAX_ATTENTION_HEADS = 48 + @invocation( "flux2_denoise", @@ -465,18 +473,23 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # concatenated onto the image stream, so three 1024x1024 references quadruple the sequence # (and with it the activation footprint) of a 1024x1024 generation. ref_image_seq_len = ref_image_extension.ref_image_latents.shape[1] if ref_image_extension is not None else 0 + # The additive bias is skipped entirely when reference images are present (see below), so the + # mask only costs anything -- storage, and possibly a materialized score matrix -- without them. + regional_attn_mask = regional_extension.restricted_attn_mask if ref_image_seq_len == 0 else None estimated_working_memory = self._estimate_working_memory( image_seq_len=packed_h * packed_w, ref_image_seq_len=ref_image_seq_len, text_seq_len=max(txt.shape[1], neg_txt.shape[1] if neg_txt is not None else 0), num_loras=len(self.transformer.loras), # The mask itself is already allocated; only the additive bias built per forward is new. - # It is skipped entirely when reference images are present (see below). regional_attention_bias_bytes=( - regional_extension.restricted_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() - if regional_extension.restricted_attn_mask is not None and ref_image_seq_len == 0 + regional_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() + if regional_attn_mask is not None else 0 ), + has_regional_attention_mask=regional_attn_mask is not None, + device=device, + dtype=inference_dtype, ) with ExitStack() as exit_stack: @@ -608,6 +621,9 @@ def _estimate_working_memory( text_seq_len: int, num_loras: int, regional_attention_bias_bytes: int = 0, + has_regional_attention_mask: bool = False, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, ) -> int: """Estimate peak transformer activation memory (bytes) so the model cache reserves enough headroom. @@ -626,13 +642,30 @@ def _estimate_working_memory( A fixed base covers resolution-independent overhead (transient fp8/GGUF -> bf16 weight casts during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra activation branch per patched layer, so we add a per-LoRA margin. + + The linear model holds only while SDPA picks a fused kernel. Regional prompting is where that + stops being a given: it hands the transformer a dense additive ``S x S`` bias, which flash + attention never accepts and which ROCm's memory-efficient kernel rejects as well, leaving the + ``math`` fallback and its materialized ``heads x S x S`` score matrix. We ask torch which path + this build will take for these shapes and add the score matrix only when it is really there -- + on CUDA the memory-efficient kernel takes the bias and the term is zero (verified: peak stays + linear with the bias attached). """ GB = 1024**3 MB = 1024**2 per_token_bytes = int(0.4 * MB) - estimated = (image_seq_len + ref_image_seq_len + text_seq_len) * per_token_bytes + total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len + estimated = total_seq_len * per_token_bytes estimated += int(1.0 * GB) estimated += regional_attention_bias_bytes + estimated += sdpa_score_matrix_bytes( + device=device if device is not None else TorchDevice.choose_torch_device(), + dtype=dtype, + num_heads=FLUX2_MAX_ATTENTION_HEADS, + head_dim=FLUX2_ATTENTION_HEAD_DIM, + seq_len=total_seq_len, + has_attn_mask=has_regional_attention_mask, + ) if num_loras > 0: estimated += int(0.5 * num_loras * GB) return estimated diff --git a/invokeai/app/invocations/flux2_vae_decode.py b/invokeai/app/invocations/flux2_vae_decode.py index d4cad4ad75d..f0852f1880f 100644 --- a/invokeai/app/invocations/flux2_vae_decode.py +++ b/invokeai/app/invocations/flux2_vae_decode.py @@ -54,7 +54,7 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima # far above the default working memory the cache would otherwise reserve. Tell it up front so # it offloads enough of the (possibly still resident) transformer to leave room. estimated_working_memory = estimate_vae_working_memory_flux2( - operation="decode", image_tensor=latents, vae=vae_info.model + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/app/invocations/flux2_vae_encode.py b/invokeai/app/invocations/flux2_vae_encode.py index 9d92b3819a4..2da6f38b517 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -50,7 +50,10 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to # free room rather than discovering the shortfall as an OOM. estimated_working_memory = estimate_vae_working_memory_flux2( - operation="encode", image_tensor=image_tensor, vae=vae_info.model + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + device=TorchDevice.choose_torch_device(), ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index ccdb1390e28..9184b15c1d2 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -215,6 +215,7 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: image_tensor=image_tensor, vae=vae_info.model, tile_size=REF_ENCODE_TILE_SIZE, + device=TorchDevice.choose_torch_device(), ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 1df0f99280b..dd75a710446 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,6 +4,8 @@ for attention mechanism. """ +from functools import lru_cache + import psutil import torch @@ -35,3 +37,81 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: return "max" else: return "balanced" + + +# SDPA computes attention one of two ways: a fused kernel (flash / memory-efficient / cuDNN) that +# never materializes the O(S^2) score matrix, or the `math` fallback, which does. torch picks per +# call from the dtype, the head dim and whether an attention mask was passed -- and the answer +# differs between builds. CUDA's memory-efficient kernel accepts head dims well past 128 and +# arbitrary additive masks; ROCm's fused kernels reject both and drop to `math`. A working-memory +# estimate that assumes the fused path is therefore only correct on the build it was measured on, +# which is why the helper below asks torch instead of assuming. + +# Peak *reserved* bytes per element of the materialized score matrix, measured on CUDA with +# `SDPBackend.MATH` forced, each point in a fresh process: 12.9 bytes/element at 4k tokens, 10.3 at +# 8k, 9.7 at 16k -- and the same figures for bf16, fp16 and fp32 inputs, because the fallback's +# softmax intermediates are fp32 regardless. So this is an absolute byte count, not a multiple of +# the element size. 13 is an upper bound on every measured point from 4k tokens up; below that it +# can fall a couple of MB short of the allocator's rounding, which is noise next to the GB-scale +# linear terms this is added to. +SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 + + +@lru_cache(maxsize=None) +def _sdpa_has_fused_kernel( + device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool +) -> bool: + """Ask torch whether any non-materializing SDPA kernel is eligible for these attention shapes. + + Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence + length, so a tiny probe answers for the real forward. Falls back to ``True`` (the status quo + assumption) whenever torch gives us nothing to go on -- over-reserving many GB on a guess would + push the model out of VRAM and be worse than the shortfall we are trying to avoid. + """ + if device_type != "cuda": + # `can_use_*_attention` is CUDA/ROCm-only. MPS, XPU and CPU all ship fused SDPA kernels, so + # keep the fused assumption there rather than guessing at their dispatch rules. + return True + + try: + from torch.backends.cuda import SDPAParams, can_use_efficient_attention, can_use_flash_attention + + device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) + q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) + mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None + try: + params = SDPAParams(q, q, q, mask, 0.0, False, False) + except TypeError: + # torch < 2.5: no `enable_gqa` field. + params = SDPAParams(q, q, q, mask, 0.0, False) + + checks = [can_use_flash_attention, can_use_efficient_attention] + can_use_cudnn_attention = getattr(torch.backends.cuda, "can_use_cudnn_attention", None) + if can_use_cudnn_attention is not None: + checks.append(can_use_cudnn_attention) + return any(check(params, False) for check in checks) + except Exception: + return True + + +def sdpa_score_matrix_bytes( + *, + device: torch.device, + dtype: torch.dtype, + num_heads: int, + head_dim: int, + seq_len: int, + has_attn_mask: bool = False, +) -> int: + """Bytes SDPA spends on a materialized score matrix for one attention call, 0 if fused. + + Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On + CUDA it is almost always 0; on a build whose fused kernels reject the shapes (notably ROCm, + which caps the head dim at 128 and does not take arbitrary additive masks) it is the dominant + term -- a 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~17GB of them. + """ + if seq_len <= 0 or num_heads <= 0: + return 0 + if _sdpa_has_fused_kernel(device.type, device.index, dtype, head_dim, has_attn_mask): + return 0 + return num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 8ccddbf99d0..849efefbad7 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -9,6 +9,8 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.backend.flux.modules.autoencoder import AutoEncoder +from invokeai.backend.util.attention import sdpa_score_matrix_bytes +from invokeai.backend.util.devices import TorchDevice _WAN_VAE_SINGLE_FRAME_DECODE_SCALING_CONSTANT = 2900 _WAN_VAE_VIDEO_DECODE_SCALING_CONSTANT_A14B = 6500 @@ -99,26 +101,45 @@ def estimate_vae_working_memory_flux( return int(working_memory) +# The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the +# decoder, on the 8x-downsampled grid. Both are single-head, with the head dim set to the block +# width: 512 for the stock VAE and 384 for the small-decoder variant. The distinction does not +# matter here -- what matters is that both sit far above the 128 head dim ROCm's fused SDPA kernels +# accept, so only the value's side of that limit is load-bearing. +_FLUX2_VAE_MID_BLOCK_HEADS = 1 +_FLUX2_VAE_MID_BLOCK_HEAD_DIM = 512 +_FLUX2_VAE_SPATIAL_COMPRESSION = 8 + + def estimate_vae_working_memory_flux2( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLFlux2, tile_size: int | None = None, + device: torch.device | None = None, ) -> int: """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. - Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE -- - ``AutoencoderKLFlux2``'s mid-block attention runs through SDPA, so no O(area^2) term appears. + Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE. Measured on CUDA/bf16 as peak *reserved* memory (the conservative quantity, including allocator overhead), the implied constants are ~2170 (decode) and ~1070 (encode) bytes per pixel per element byte, flat across 512-1536px; the constants below round those up and match the FLUX.1 ones. For reference, decoding 1024x1024 peaks at ~4.3GB and 1536x1536 at ~9.6GB -- far above the default ``device_working_mem_gb``, which is why this estimate must be passed to the model cache. + That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a + fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the + 512-wide head, and measured peak stays linear from 512 to 1536px). A build whose fused kernels + reject the head dim -- ROCm caps it at 128 -- drops to SDPA's ``math`` fallback and materializes + a (pixels/8)^2 score matrix on top of the linear term: ~3.5GB at 1024px and ~17GB at 1536px. We + ask torch which path applies rather than assuming, so the estimate is right on both. + When tiling is enabled the peak is bounded by a single tile instead of the full image (measured - ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap). + ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score + matrix, if one is materialized at all, is bounded by the tile too. """ - element_size = next(vae.parameters()).element_size() + param = next(vae.parameters()) + element_size = param.element_size() # Encoding uses ~50% the working memory of decoding. scaling_constant = 2200 if operation == "decode" else 1100 @@ -126,11 +147,21 @@ def estimate_vae_working_memory_flux2( if tile_size is not None: # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 + mid_block_seq_len = (tile_size // _FLUX2_VAE_SPATIAL_COMPRESSION) ** 2 else: latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] working_memory = out_h * out_w * element_size * scaling_constant + mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) + + working_memory += sdpa_score_matrix_bytes( + device=device if device is not None else TorchDevice.choose_torch_device(), + dtype=param.dtype, + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS, + head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, + seq_len=mid_block_seq_len, + ) return int(working_memory) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 816a7f71f37..aac5409930d 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -14,18 +14,34 @@ import pytest import torch +import torch.nn.functional as F from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 -from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation +from invokeai.app.invocations.flux2_denoise import FLUX2_MAX_ATTENTION_HEADS, Flux2DenoiseInvocation from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation +from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 MB = 1024**2 GB = 1024**3 - -def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, regional_bias=0): +# The measured tables in this module were all taken on CUDA, where SDPA runs a fused kernel and no +# score matrix is materialized. `_sdpa_has_fused_kernel` reports non-CUDA devices as fused, so +# passing a CPU device reproduces that regime without needing a GPU on the test runner. The +# materializing regime gets its own class below. +FUSED = torch.device("cpu") + + +def _estimate( + image_seq_len, + ref_image_seq_len=0, + text_seq_len=512, + num_loras=0, + regional_bias=0, + has_regional_mask=False, + device=FUSED, +): return Flux2DenoiseInvocation._estimate_working_memory( MagicMock(spec=Flux2DenoiseInvocation), image_seq_len=image_seq_len, @@ -33,6 +49,8 @@ def _estimate(image_seq_len, ref_image_seq_len=0, text_seq_len=512, num_loras=0, text_seq_len=text_seq_len, num_loras=num_loras, regional_attention_bias_bytes=regional_bias, + has_regional_attention_mask=has_regional_mask, + device=device, ) @@ -112,14 +130,14 @@ def _tensor_for(self, operation, px): @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae() + operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae(), device=FUSED ) assert estimate >= measured_mb * MB @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae() + operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae(), device=FUSED ) assert estimate == 1024 * 1024 * 2 * expected_constant @@ -132,6 +150,7 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): image_tensor=torch.zeros(1, 3, px, px), vae=self._mock_bf16_vae(), tile_size=512, + device=FUSED, ) for px in (1024, 1328, 2024) ] @@ -140,7 +159,7 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): assert estimates[0] >= 558 * MB # measured tiled peak at 2024px # The whole point of tiling: it must shrink the reservation, not just bound the VAE. untiled = estimate_vae_working_memory_flux2( - operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae() + operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae(), device=FUSED ) assert estimates[0] < untiled / 4 @@ -285,3 +304,194 @@ def test_reference_images_raise_the_reservation(self): without_refs = self._run(num_ref_tokens=0) with_refs = self._run(num_ref_tokens=12288) assert with_refs - without_refs == 12288 * int(0.4 * MB) + + +def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): + """Stand in for `_sdpa_has_fused_kernel` on a build with ROCm's fused-kernel rules. + + ROCm's fused SDPA kernels cap the head dim at 128 and do not take an arbitrary additive mask; + anything else falls through to the `math` fallback, which materializes the score matrix. CUDA's + memory-efficient kernel accepts both -- verified on torch 2.7.1+cu128, where + `can_use_efficient_attention` is true for the VAE's 512-wide head and for a masked 128-wide + transformer head, and measured peak stays linear in both cases -- which is why the estimates + were linear to begin with. + """ + return head_dim <= 128 and not has_attn_mask + + +def _materializing(): + return patch("invokeai.backend.util.attention._sdpa_has_fused_kernel", side_effect=_rocm_like_probe) + + +# Any CUDA device object works here: the probe is patched out, so nothing is allocated on it. +MATERIALIZING = torch.device("cuda") + + +class TestMaterializedScoreMatrixIsBudgeted: + """The linear estimates above assume SDPA never builds the O(S^2) score matrix. That is a + property of the *build*, not of FLUX.2: ROCm's fused kernels reject both the VAE's 512-wide + attention head and the dense additive mask regional prompting attaches, and fall back to + `math`. Where that happens the score matrix is the dominant term, so the estimate has to + include it -- otherwise the fix works on CUDA and still OOMs on ROCm. + """ + + def _mock_bf16_vae(self): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return vae + + def _vae_estimate(self, operation, px, device, tile_size=None): + tensor = torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) + return estimate_vae_working_memory_flux2( + operation=operation, + image_tensor=tensor, + vae=self._mock_bf16_vae(), + tile_size=tile_size, + device=device, + ) + + # (operation, pixel size, mid-block tokens). The VAE attends on the 8x-downsampled grid. + @pytest.mark.parametrize( + "operation, px, tokens", + [ + ("decode", 1024, 128 * 128), + ("decode", 1536, 192 * 192), + ("encode", 1024, 128 * 128), + ("encode", 1328, 166 * 166), + ], + ) + def test_vae_estimate_gains_exactly_the_score_matrix(self, operation, px, tokens): + fused = self._vae_estimate(operation, px, device=FUSED) + with _materializing(): + materializing = self._vae_estimate(operation, px, device=MATERIALIZING) + assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + def test_vae_score_matrix_dominates_at_high_resolution(self): + """The reviewer's case: a 1536px decode is ~9.6GB of linear activations on CUDA, and more + than that again in scores where SDPA has to materialize them. An estimate that omits the + term is not merely tight, it is wrong by a factor of three.""" + with _materializing(): + materializing = self._vae_estimate("decode", 1536, device=MATERIALIZING) + assert materializing > 25 * GB + assert materializing > 2 * self._vae_estimate("decode", 1536, device=FUSED) + + def test_tiled_vae_score_matrix_is_bounded_by_the_tile(self): + """Tiling already bounds the linear term; it must bound the quadratic one too, or the + reference-image encode would reserve as if it ran untiled.""" + with _materializing(): + estimates = [ + self._vae_estimate("encode", px, device=MATERIALIZING, tile_size=512) for px in (1024, 1328, 2024) + ] + untiled = self._vae_estimate("encode", 2024, device=MATERIALIZING) + assert len(set(estimates)) == 1 + tile_tokens = (512 // 8) ** 2 + assert estimates[0] - self._vae_estimate("encode", 1024, device=FUSED, tile_size=512) == ( + tile_tokens * tile_tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + # Tiling turns a 62GB reservation at the 2024px reference cap into under 1GB. + assert estimates[0] < untiled / 50 + + def test_regional_prompting_adds_the_score_matrix(self): + """The dense `S x S` additive bias is what pushes SDPA off its fused kernel. Budgeting only + the bias tensor -- as this PR first did -- under-reserves by the score matrix, which is two + orders of magnitude larger.""" + seq_len = 4096 + 512 + bias_bytes = seq_len * seq_len * 2 + fused = _estimate( + image_seq_len=4096, text_seq_len=512, regional_bias=bias_bytes, has_regional_mask=True, device=FUSED + ) + with _materializing(): + materializing = _estimate( + image_seq_len=4096, + text_seq_len=512, + regional_bias=bias_bytes, + has_regional_mask=True, + device=MATERIALIZING, + ) + assert materializing - fused == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + assert materializing - fused > 50 * bias_bytes + + def test_denoise_without_a_regional_mask_is_unaffected(self): + """FLUX.2's 128-wide attention head is inside every backend's fused limit, so an ordinary + generation -- reference images included -- keeps the plain linear estimate. This term exists + for the masked case; it must not tax the common one.""" + with _materializing(): + assert _estimate(image_seq_len=4096, ref_image_seq_len=12288, device=MATERIALIZING) == _estimate( + image_seq_len=4096, ref_image_seq_len=12288, device=FUSED + ) + + +class TestSdpaBackendProbe: + """`sdpa_score_matrix_bytes` decides the term above, so its defaults are load-bearing.""" + + def test_non_cuda_devices_keep_the_fused_assumption(self): + """torch exposes no eligibility query outside CUDA/ROCm. Guessing `math` there would add + double-digit GB to every estimate on MPS and CPU on no evidence at all.""" + assert ( + sdpa_score_matrix_bytes( + device=torch.device("cpu"), + dtype=torch.bfloat16, + num_heads=1, + head_dim=512, + seq_len=16384, + has_attn_mask=True, + ) + == 0 + ) + + def test_empty_sequences_cost_nothing(self): + with _materializing(): + assert ( + sdpa_score_matrix_bytes( + device=MATERIALIZING, dtype=torch.bfloat16, num_heads=48, head_dim=128, seq_len=0 + ) + == 0 + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="asks the real CUDA/ROCm dispatcher") + def test_this_build_reports_its_own_dispatch(self): + """On CUDA both shapes are fused and this whole term is zero -- the fix is a no-op for the + hardware the constants were measured on. On ROCm the same call reports the fallback and the + term appears. Either answer is correct; the point is that it comes from torch.""" + vae_bytes = sdpa_score_matrix_bytes( + device=torch.device("cuda"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + masked_bytes = sdpa_score_matrix_bytes( + device=torch.device("cuda"), + dtype=torch.bfloat16, + num_heads=48, + head_dim=128, + seq_len=4608, + has_attn_mask=True, + ) + if torch.version.hip is None: + assert vae_bytes == 0 + assert masked_bytes == 0 + else: + assert vae_bytes > 0 + assert masked_bytes > 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="measures real peak reserved memory") + @pytest.mark.parametrize("num_heads, seq_len, head_dim", [(1, 4096, 512), (4, 4096, 128)]) + def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, head_dim): + """Pin the bytes-per-score-element calibration against a real `math` forward, so a future + edit to the constant cannot silently reintroduce the shortfall it exists to cover.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + device = torch.device("cuda") + q = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_reserved() + with sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): + F.scaled_dot_product_attention(q, k, v) + torch.cuda.synchronize() + measured = torch.cuda.max_memory_reserved() - before + + estimate = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimate >= measured + assert estimate <= 2 * measured From 119664a243f993de4db6015884515793da767ef7 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 26 Aug 2026 05:01:00 +0200 Subject: [PATCH 03/11] fix(flux2): ask the real dispatcher which SDPA path a build takes The score-matrix term probed torch's CUDA eligibility helpers and read everything else as fused. That was wrong twice over: MPS has no fused SDPA kernel at all and runs the MPSGraph math transcription, so a 1024px VAE decode was admitted ~3.5GB short; and a failed probe returned "fused" too, turning "we don't know" into the one answer that can OOM. Ask `_fused_sdp_choice` instead -- the same dispatch query `scaled_dot_product_attention` runs to pick its kernel. Torch registers it for CPU, CUDA/ROCm and XPU only, so the call raises on exactly the devices that fall through to `math`, and every other failure lands on the conservative side by the same branch. Diffusers models do not reach torch's SDPA directly, so also consult `dispatch_attention_fn`'s active backend: a user on `_native_math` materializes the score matrix on hardware whose probe reports fused. Only the transformer needs this -- the FLUX.2 VAE's mid-block attention still calls SDPA itself through `AttnProcessor2_0` -- and a test pins that asymmetry. On CUDA with the stock backend every one of these terms remains zero. --- invokeai/app/invocations/flux2_denoise.py | 19 ++- invokeai/backend/util/attention.py | 144 ++++++++++++---- invokeai/backend/util/vae_working_memory.py | 11 +- .../invocations/test_flux2_working_memory.py | 161 ++++++++++++++++-- 4 files changed, 277 insertions(+), 58 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 6dcb412e39e..e6655b5c593 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -643,13 +643,15 @@ def _estimate_working_memory( during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra activation branch per patched layer, so we add a per-LoRA margin. - The linear model holds only while SDPA picks a fused kernel. Regional prompting is where that - stops being a given: it hands the transformer a dense additive ``S x S`` bias, which flash - attention never accepts and which ROCm's memory-efficient kernel rejects as well, leaving the - ``math`` fallback and its materialized ``heads x S x S`` score matrix. We ask torch which path - this build will take for these shapes and add the score matrix only when it is really there -- - on CUDA the memory-efficient kernel takes the bias and the term is zero (verified: peak stays - linear with the bias attached). + The linear model holds only while attention runs on a fused kernel. Regional prompting is + where that stops being a given: it hands the transformer a dense additive ``S x S`` bias, + which flash attention never accepts and which ROCm's memory-efficient kernel rejects as + well, leaving the ``math`` fallback and its materialized ``heads x S x S`` score matrix. The + device decides too -- MPS has no fused SDPA kernel at all -- and so does the diffusers + attention backend this build dispatches through. ``sdpa_score_matrix_bytes`` asks all three + and adds the score matrix only where it is really built: on CUDA with the stock backend the + memory-efficient kernel takes the bias and the term is zero (verified: peak stays linear + with the bias attached). """ GB = 1024**3 MB = 1024**2 @@ -665,6 +667,9 @@ def _estimate_working_memory( head_dim=FLUX2_ATTENTION_HEAD_DIM, seq_len=total_seq_len, has_attn_mask=has_regional_attention_mask, + # The FLUX.2 transformer's attention goes through diffusers' `dispatch_attention_fn`, + # which can route around torch's SDPA entirely -- including to a forced `math` backend. + via_diffusers_dispatch=True, ) if num_loras > 0: estimated += int(0.5 * num_loras * GB) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index dd75a710446..2245f6c750b 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,12 +4,15 @@ for attention mechanism. """ +import warnings from functools import lru_cache import psutil import torch +from torch.nn.attention import SDPBackend from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger def auto_detect_slice_size(latents: torch.Tensor) -> str: @@ -40,12 +43,12 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # SDPA computes attention one of two ways: a fused kernel (flash / memory-efficient / cuDNN) that -# never materializes the O(S^2) score matrix, or the `math` fallback, which does. torch picks per -# call from the dtype, the head dim and whether an attention mask was passed -- and the answer -# differs between builds. CUDA's memory-efficient kernel accepts head dims well past 128 and -# arbitrary additive masks; ROCm's fused kernels reject both and drop to `math`. A working-memory -# estimate that assumes the fused path is therefore only correct on the build it was measured on, -# which is why the helper below asks torch instead of assuming. +# never materializes the O(S^2) score matrix, or the `math` fallback, which does. Which one runs is +# not a property of FLUX.2 -- it is a property of the build, the device, the dtype, the head dim and +# whether an attention mask was passed. CUDA's memory-efficient kernel accepts head dims well past +# 128 and arbitrary additive masks; ROCm's fused kernels reject both; MPS has no fused SDPA kernel +# at all. A working-memory estimate that assumes the fused path is therefore only correct on the +# build it was measured on, which is why the helpers below ask instead of assuming. # Peak *reserved* bytes per element of the materialized score matrix, measured on CUDA with # `SDPBackend.MATH` forced, each point in a fresh process: 12.9 bytes/element at 4k tokens, 10.3 at @@ -53,46 +56,98 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # softmax intermediates are fp32 regardless. So this is an absolute byte count, not a multiple of # the element size. 13 is an upper bound on every measured point from 4k tokens up; below that it # can fall a couple of MB short of the allocator's rounding, which is noise next to the GB-scale -# linear terms this is added to. +# linear terms this is added to. It is a CUDA measurement standing in for every materializing +# backend -- no other was available to calibrate against -- but the intermediates it prices (an +# fp32 score matrix and its softmax) have the same shape wherever the fallback runs. SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 +# `_fused_sdp_choice` reports which kernel `F.scaled_dot_product_attention` would pick. These are +# the answers that mean "a fused kernel"; `MATH` -- and `ERROR`, which torch returns when it cannot +# pick anything at all -- mean the score matrix gets built. +_FUSED_SDP_CHOICES = frozenset( + int(getattr(SDPBackend, name)) + for name in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION", "OVERRIDEABLE") + if hasattr(SDPBackend, name) +) + +_DISPATCH_TORCH = "torch" +_DISPATCH_FUSED = "fused" +_DISPATCH_MATH = "math" + + +@lru_cache(maxsize=None) +def _diffusers_attention_dispatch() -> str: + """Report how the diffusers attention dispatcher will route a diffusers model's attention calls. + + Diffusers models do not call `F.scaled_dot_product_attention` directly -- they go through + `dispatch_attention_fn`, which honours the `DIFFUSERS_ATTN_BACKEND` environment variable and the + `attention_backend()` context manager. Only the default `native` backend hands the call to + torch; the others pin a specific kernel, and `_native_math` pins the materializing one. A + torch-level probe alone would report "fused" for a user who has forced math. + + Returns ``_DISPATCH_TORCH`` when torch decides, ``_DISPATCH_FUSED`` for a backend that never + materializes the score matrix, or ``_DISPATCH_MATH`` when one is built -- including when we + cannot tell, since under-reserving is the failure this whole term exists to prevent. + """ + try: + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + backend, _ = _AttentionBackendRegistry.get_active_backend() + name = str(getattr(backend, "value", backend)) + except Exception: + # A private diffusers attribute that moved, or a selected backend whose kernel failed to + # register. Budget the materializing case, but say so: silently adding several GB to every + # FLUX.2 estimate is not something that should pass unnoticed. + InvokeAILogger.get_logger(__name__).warning( + "Could not determine the active diffusers attention backend; budgeting working memory as if " + "attention materializes its score matrix. Estimates will be conservative." + ) + return _DISPATCH_MATH + + if name == "native": + return _DISPATCH_TORCH + if "math" in name: + return _DISPATCH_MATH + # Every other backend diffusers offers -- flash, sage, xformers, flex, aiter, the pinned + # `_native_*` kernels -- exists precisely to avoid materializing the score matrix. + return _DISPATCH_FUSED + @lru_cache(maxsize=None) -def _sdpa_has_fused_kernel( +def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: - """Ask torch whether any non-materializing SDPA kernel is eligible for these attention shapes. + """Ask torch whether `F.scaled_dot_product_attention` would build the O(S^2) score matrix. + `_fused_sdp_choice` is the same dispatch query torch's own `scaled_dot_product_attention` runs + to pick a kernel, so this is its real answer rather than a reimplementation of its rules. Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence - length, so a tiny probe answers for the real forward. Falls back to ``True`` (the status quo - assumption) whenever torch gives us nothing to go on -- over-reserving many GB on a guess would - push the model out of VRAM and be worse than the shortfall we are trying to avoid. + length, so a tiny probe answers for the real forward. + + Anything that goes wrong reports the materializing path, which is both the conservative answer + and, for the most common cause, the correct one: torch registers `_fused_sdp_choice` for CPU, + CUDA/ROCm and XPU only, so the call raises on MPS -- and MPS is exactly where + `scaled_dot_product_attention` finds no fused kernel either and runs + `_scaled_dot_product_attention_math_for_mps`, an MPSGraph transcription of `Q @ K^T` -> softmax + -> `@ V` that holds the score tensor as a real intermediate. The remaining causes (an allocation + failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and + there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ - if device_type != "cuda": - # `can_use_*_attention` is CUDA/ROCm-only. MPS, XPU and CPU all ship fused SDPA kernels, so - # keep the fused assumption there rather than guessing at their dispatch rules. - return True - try: - from torch.backends.cuda import SDPAParams, can_use_efficient_attention, can_use_flash_attention - device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None - try: - params = SDPAParams(q, q, q, mask, 0.0, False, False) - except TypeError: - # torch < 2.5: no `enable_gqa` field. - params = SDPAParams(q, q, q, mask, 0.0, False) - - checks = [can_use_flash_attention, can_use_efficient_attention] - can_use_cudnn_attention = getattr(torch.backends.cuda, "can_use_cudnn_attention", None) - if can_use_cudnn_attention is not None: - checks.append(can_use_cudnn_attention) - return any(check(params, False) for check in checks) + with warnings.catch_warnings(): + # When no fused kernel is eligible, torch re-runs every check in debug mode to warn why + # each one was rejected. That is the case we are deliberately probing for; we do not + # want a wall of warnings every time an estimate is priced. + warnings.simplefilter("ignore") + choice = int(torch.ops.aten._fused_sdp_choice(q, q, q, mask, 0.0, False)) except Exception: return True + return choice not in _FUSED_SDP_CHOICES + def sdpa_score_matrix_bytes( *, @@ -102,16 +157,35 @@ def sdpa_score_matrix_bytes( head_dim: int, seq_len: int, has_attn_mask: bool = False, + via_diffusers_dispatch: bool = False, ) -> int: """Bytes SDPA spends on a materialized score matrix for one attention call, 0 if fused. Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On - CUDA it is almost always 0; on a build whose fused kernels reject the shapes (notably ROCm, - which caps the head dim at 128 and does not take arbitrary additive masks) it is the dominant - term -- a 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~17GB of them. + CUDA it is almost always 0; where the fused kernels are missing or reject the shapes -- ROCm + caps the head dim at 128 and does not take arbitrary additive masks, MPS ships no fused SDPA + kernel at all -- it is the dominant term: a 1536px FLUX.2 VAE decode materializes 36864^2 + scores, ~17GB of them. + + Set ``via_diffusers_dispatch`` for attention that runs inside a diffusers model (the FLUX.2 + transformer does; the FLUX.2 VAE's mid-block attention does not -- it still reaches + `F.scaled_dot_product_attention` directly through `AttnProcessor2_0`). It consults the + process-wide default backend, which is the one that applies here: estimates are priced before + the model is loaded and outside any `attention_backend()` scope. """ if seq_len <= 0 or num_heads <= 0: return 0 - if _sdpa_has_fused_kernel(device.type, device.index, dtype, head_dim, has_attn_mask): + + score_matrix_bytes = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + if via_diffusers_dispatch: + dispatch = _diffusers_attention_dispatch() + if dispatch == _DISPATCH_FUSED: + return 0 + if dispatch == _DISPATCH_MATH: + return score_matrix_bytes + # _DISPATCH_TORCH: diffusers forwards to `F.scaled_dot_product_attention`, so torch decides. + + if not _torch_sdpa_materializes_score_matrix(device.type, device.index, dtype, head_dim, has_attn_mask): return 0 - return num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + return score_matrix_bytes diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 849efefbad7..960e2d081f9 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -129,10 +129,13 @@ def estimate_vae_working_memory_flux2( That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the - 512-wide head, and measured peak stays linear from 512 to 1536px). A build whose fused kernels - reject the head dim -- ROCm caps it at 128 -- drops to SDPA's ``math`` fallback and materializes - a (pixels/8)^2 score matrix on top of the linear term: ~3.5GB at 1024px and ~17GB at 1536px. We - ask torch which path applies rather than assuming, so the estimate is right on both. + 512-wide head, and measured peak stays linear from 512 to 1536px). A build with no fused kernel + for the shapes -- ROCm caps the head dim at 128, MPS has no fused SDPA kernel at all -- drops to + SDPA's ``math`` fallback and materializes a (pixels/8)^2 score matrix on top of the linear term: + ~3.5GB at 1024px and ~17GB at 1536px. We ask torch which path applies rather than assuming, so + the estimate is right on both. (Unlike the transformer, this attention does not go through + diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls ``F.scaled_dot_product_attention`` + itself -- so torch's own answer is the whole answer here.) When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index aac5409930d..ccc7628c191 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -10,6 +10,7 @@ quantity, including allocator overhead). Every estimate must stay an upper bound on them. """ +from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest @@ -20,19 +21,34 @@ from invokeai.app.invocations.flux2_denoise import FLUX2_MAX_ATTENTION_HEADS, Flux2DenoiseInvocation from invokeai.app.invocations.flux2_vae_decode import Flux2VaeDecodeInvocation from invokeai.app.invocations.flux2_vae_encode import Flux2VaeEncodeInvocation -from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes +from invokeai.backend.util.attention import ( + SDPA_MATH_BYTES_PER_SCORE_ELEMENT, + _diffusers_attention_dispatch, + _torch_sdpa_materializes_score_matrix, + sdpa_score_matrix_bytes, +) from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 MB = 1024**2 GB = 1024**3 # The measured tables in this module were all taken on CUDA, where SDPA runs a fused kernel and no -# score matrix is materialized. `_sdpa_has_fused_kernel` reports non-CUDA devices as fused, so -# passing a CPU device reproduces that regime without needing a GPU on the test runner. The +# score matrix is materialized. torch reports its CPU flash kernel as eligible for every shape used +# here, so passing a CPU device reproduces that regime without needing a GPU on the test runner. The # materializing regime gets its own class below. FUSED = torch.device("cpu") +@pytest.fixture(autouse=True) +def _clear_dispatch_caches(): + """Both probes are `lru_cache`d for the process; tests that fake one must not leak into the next.""" + _diffusers_attention_dispatch.cache_clear() + _torch_sdpa_materializes_score_matrix.cache_clear() + yield + _diffusers_attention_dispatch.cache_clear() + _torch_sdpa_materializes_score_matrix.cache_clear() + + def _estimate( image_seq_len, ref_image_seq_len=0, @@ -307,20 +323,20 @@ def test_reference_images_raise_the_reservation(self): def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): - """Stand in for `_sdpa_has_fused_kernel` on a build with ROCm's fused-kernel rules. + """Stand in for the torch probe on a build with ROCm's fused-kernel rules. ROCm's fused SDPA kernels cap the head dim at 128 and do not take an arbitrary additive mask; anything else falls through to the `math` fallback, which materializes the score matrix. CUDA's memory-efficient kernel accepts both -- verified on torch 2.7.1+cu128, where - `can_use_efficient_attention` is true for the VAE's 512-wide head and for a masked 128-wide - transformer head, and measured peak stays linear in both cases -- which is why the estimates - were linear to begin with. + `_fused_sdp_choice` reports the efficient kernel for the VAE's 512-wide head and for a masked + 128-wide transformer head, and measured peak stays linear in both cases -- which is why the + estimates were linear to begin with. """ - return head_dim <= 128 and not has_attn_mask + return head_dim > 128 or has_attn_mask def _materializing(): - return patch("invokeai.backend.util.attention._sdpa_has_fused_kernel", side_effect=_rocm_like_probe) + return patch("invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", side_effect=_rocm_like_probe) # Any CUDA device object works here: the probe is patched out, so nothing is allocated on it. @@ -366,6 +382,20 @@ def test_vae_estimate_gains_exactly_the_score_matrix(self, operation, px, tokens materializing = self._vae_estimate(operation, px, device=MATERIALIZING) assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): + """The MPS case end to end, through the real probe rather than a stand-in: on a device torch + cannot answer a dispatch query for, a 1024px decode has to come out ~3.5GB heavier than the + linear term. Reporting those devices as fused -- as this PR first did -- is what let the + decode be admitted to a card that could not run it.""" + with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): + materializing = self._vae_estimate("decode", 1024, device=FUSED) + _torch_sdpa_materializes_score_matrix.cache_clear() + fused = self._vae_estimate("decode", 1024, device=FUSED) + + tokens = 128 * 128 + assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert materializing - fused > 3 * GB + def test_vae_score_matrix_dominates_at_high_resolution(self): """The reviewer's case: a 1536px decode is ~9.6GB of linear activations on CUDA, and more than that again in scores where SDPA has to materialize them. An estimate that omits the @@ -426,9 +456,10 @@ def test_denoise_without_a_regional_mask_is_unaffected(self): class TestSdpaBackendProbe: """`sdpa_score_matrix_bytes` decides the term above, so its defaults are load-bearing.""" - def test_non_cuda_devices_keep_the_fused_assumption(self): - """torch exposes no eligibility query outside CUDA/ROCm. Guessing `math` there would add - double-digit GB to every estimate on MPS and CPU on no evidence at all.""" + def test_cpu_reports_its_fused_flash_kernel(self): + """torch ships a fused flash-attention CPU kernel that takes the VAE's 512-wide head and an + additive mask, so the CPU estimate stays linear -- and the rest of this module can use a CPU + device to stand in for the CUDA regime the constants were measured on.""" assert ( sdpa_score_matrix_bytes( device=torch.device("cpu"), @@ -441,6 +472,40 @@ def test_non_cuda_devices_keep_the_fused_assumption(self): == 0 ) + def test_a_device_torch_cannot_answer_for_is_budgeted_as_math(self): + """MPS is the case that matters: torch registers `_fused_sdp_choice` for CPU, CUDA/ROCm and + XPU only, and it is exactly the devices it cannot answer for that have no fused SDPA kernel + either. A 1024px FLUX.2 VAE decode there materializes 16384^2 scores, ~3.5GB the estimate + used to omit entirely.""" + with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): + estimated = sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimated > 3 * GB + + def test_a_failed_probe_is_budgeted_as_math(self): + """A probe that cannot allocate, or a torch without the op, leaves us knowing nothing. The + old code read that as "fused" and reserved zero; the shortfall it hides is an OOM, so the + unknown answer has to be the expensive one.""" + with patch("torch.empty", side_effect=torch.cuda.OutOfMemoryError("probe could not allocate")): + estimated = sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=512, seq_len=16384 + ) + assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + def test_the_probe_asks_torch_the_same_question_sdpa_does(self): + """`_fused_sdp_choice` is the dispatch query `F.scaled_dot_product_attention` itself runs, so + a `MATH` answer means the real forward materializes. Reimplementing the eligibility rules + instead would go stale with every torch release.""" + from torch.nn.attention import SDPBackend + + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): + assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) + _torch_sdpa_materializes_score_matrix.cache_clear() + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): + assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) + def test_empty_sequences_cost_nothing(self): with _materializing(): assert ( @@ -495,3 +560,75 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h estimate = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert estimate >= measured assert estimate <= 2 * measured + + +@contextmanager +def _diffusers_backend(name): + """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would. + + The lookup is `lru_cache`d -- it is a process-wide setting read on every estimate -- so the + cache has to be dropped on the way in and on the way out, or the faked backend leaks into the + comparison the test makes against the real one. + """ + from diffusers.models.attention_dispatch import AttentionBackendName + + _diffusers_attention_dispatch.cache_clear() + try: + with patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + return_value=(AttentionBackendName(name), None), + ): + yield + finally: + _diffusers_attention_dispatch.cache_clear() + + +class TestDiffusersAttentionDispatchIsConsulted: + """The FLUX.2 transformer does not call `F.scaled_dot_product_attention` -- it calls diffusers' + `dispatch_attention_fn`, which honours `DIFFUSERS_ATTN_BACKEND` and the `attention_backend()` + context manager. A user on `_native_math` materializes the score matrix on hardware where the + torch probe reports a fused kernel, so asking torch alone is not enough for the transformer. + + The VAE is the other half of the same point: its mid-block attention goes through + `AttnProcessor2_0`, which calls `F.scaled_dot_product_attention` itself, so the diffusers + backend must *not* move its estimate. + """ + + def test_forced_math_backend_reaches_the_denoise_estimate(self): + with _diffusers_backend("_native_math"): + forced_math = _estimate(image_seq_len=4096, device=FUSED) + native = _estimate(image_seq_len=4096, device=FUSED) + seq_len = 4096 + 512 + assert forced_math - native == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + + def test_a_fused_backend_leaves_the_denoise_estimate_linear(self): + """`flash`, `sage`, `xformers` and friends exist precisely to avoid the score matrix; they + must not be taxed for it, on any device.""" + with _diffusers_backend("flash"), _materializing(): + forced_flash = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) + assert forced_flash == _estimate(image_seq_len=4096, has_regional_mask=True, device=FUSED) + + def test_the_vae_estimate_ignores_the_diffusers_backend(self): + """`AttnProcessor2_0` bypasses the dispatcher, so the VAE's answer comes from torch alone.""" + + def estimate(): + v = MagicMock(spec=AutoencoderKLFlux2) + v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return estimate_vae_working_memory_flux2( + operation="decode", image_tensor=torch.zeros(1, 32, 128, 128), vae=v, device=FUSED + ) + + with _diffusers_backend("_native_math"): + forced_math = estimate() + assert forced_math == estimate() + + def test_an_unreadable_dispatcher_is_budgeted_as_math(self): + """`_AttentionBackendRegistry` is private; if diffusers moves it we lose the answer. The + conservative reading is the materializing one, and it is logged rather than silent.""" + with patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + side_effect=AttributeError("moved"), + ): + assert _diffusers_attention_dispatch() == "math" From 12e1c9a5cfd1c1b0bc20caf1d3966a15d1b0878f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 26 Aug 2026 20:16:54 +0200 Subject: [PATCH 04/11] fix(flux2): read the attention backend live instead of caching it once `_diffusers_attention_dispatch()` was `lru_cache`d, so the first estimate in a process pinned the answer forever. A switch to `_native_math` after that kept reserving zero for the S x S score matrix -- the exact case the lookup was added to catch. Read it live; it is a dict lookup against an already-imported module, priced once per invocation. The torch probe had the same defect one level down: its answer depends on the global SDPA kernel toggles, which `sdpa_kernel()` and `enable_flash_sdp()` flip at runtime. That probe allocates and dispatches, so it stays cached -- but keyed on the toggles, so a switch invalidates it. Per-model overrides need no plumbing: `set_attention_backend()` stamps its choice onto the process-wide registry as well as onto the model's processors, deliberately, so the estimate sees it without holding the model it is priced ahead of. A test pins that propagation. --- invokeai/backend/util/attention.py | 58 +++++++++-- .../invocations/test_flux2_working_memory.py | 95 ++++++++++++++----- 2 files changed, 122 insertions(+), 31 deletions(-) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 2245f6c750b..25f0b463cb4 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -75,7 +75,15 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: _DISPATCH_MATH = "math" -@lru_cache(maxsize=None) +@lru_cache(maxsize=1) +def _warn_unknown_diffusers_dispatch() -> None: + """Say once per process that estimates are running blind. Rate-limited, not cached for truth.""" + InvokeAILogger.get_logger(__name__).warning( + "Could not determine the active diffusers attention backend; budgeting working memory as if " + "attention materializes its score matrix. Estimates will be conservative." + ) + + def _diffusers_attention_dispatch() -> str: """Report how the diffusers attention dispatcher will route a diffusers model's attention calls. @@ -85,6 +93,18 @@ def _diffusers_attention_dispatch() -> str: torch; the others pin a specific kernel, and `_native_math` pins the materializing one. A torch-level probe alone would report "fused" for a user who has forced math. + Read live on every estimate, never cached: the active backend is mutable process state, and a + cached answer would keep reserving zero after a switch to `_native_math` -- the one case this + lookup exists to catch. It is a dict lookup against an already-imported module, priced once per + invocation. + + Reading the process-wide backend also covers per-model overrides, which is why the estimate does + not need the model in hand (it is priced before the model is loaded). `set_attention_backend()` + stamps the choice onto the model's attention processors *and* calls + `_AttentionBackendRegistry.set_active_backend()` -- deliberately, "so that it propagates + gracefully throughout". `reset_attention_backend()` clears only the processors, leaving the + registry pinned, which errs towards over-reserving rather than under-reserving. + Returns ``_DISPATCH_TORCH`` when torch decides, ``_DISPATCH_FUSED`` for a backend that never materializes the score matrix, or ``_DISPATCH_MATH`` when one is built -- including when we cannot tell, since under-reserving is the failure this whole term exists to prevent. @@ -98,10 +118,7 @@ def _diffusers_attention_dispatch() -> str: # A private diffusers attribute that moved, or a selected backend whose kernel failed to # register. Budget the materializing case, but say so: silently adding several GB to every # FLUX.2 estimate is not something that should pass unnoticed. - InvokeAILogger.get_logger(__name__).warning( - "Could not determine the active diffusers attention backend; budgeting working memory as if " - "attention materializes its score matrix. Estimates will be conservative." - ) + _warn_unknown_diffusers_dispatch() return _DISPATCH_MATH if name == "native": @@ -113,7 +130,21 @@ def _diffusers_attention_dispatch() -> str: return _DISPATCH_FUSED -@lru_cache(maxsize=None) +def _sdp_kernel_toggles() -> tuple[bool, ...]: + """The global switches that gate each fused SDPA kernel, as `_fused_sdp_choice` sees them. + + `torch.backends.cuda.enable_flash_sdp(False)` and `sdpa_kernel([...])` flip these at runtime and + the dispatch answer flips with them, so they belong in the probe's cache key rather than being + baked into a permanent result. + """ + cuda = torch.backends.cuda + return tuple( + bool(getattr(cuda, name)()) + for name in ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") + if hasattr(cuda, name) + ) + + def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: @@ -133,6 +164,21 @@ def _torch_sdpa_materializes_score_matrix( failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ + return _probe_sdpa_dispatch(device_type, device_index, dtype, head_dim, has_attn_mask, _sdp_kernel_toggles()) + + +@lru_cache(maxsize=None) +def _probe_sdpa_dispatch( + device_type: str, + device_index: int | None, + dtype: torch.dtype, + head_dim: int, + has_attn_mask: bool, + sdp_kernel_toggles: tuple[bool, ...], +) -> bool: + """Cached body of the probe above. Every input torch's answer depends on is part of the key -- + `sdp_kernel_toggles` is not read here, it is carried so a runtime change invalidates the entry. + """ try: device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index ccc7628c191..b876d973c13 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -10,7 +10,6 @@ quantity, including allocator overhead). Every estimate must stay an upper bound on them. """ -from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest @@ -24,6 +23,7 @@ from invokeai.backend.util.attention import ( SDPA_MATH_BYTES_PER_SCORE_ELEMENT, _diffusers_attention_dispatch, + _probe_sdpa_dispatch, _torch_sdpa_materializes_score_matrix, sdpa_score_matrix_bytes, ) @@ -40,13 +40,12 @@ @pytest.fixture(autouse=True) -def _clear_dispatch_caches(): - """Both probes are `lru_cache`d for the process; tests that fake one must not leak into the next.""" - _diffusers_attention_dispatch.cache_clear() - _torch_sdpa_materializes_score_matrix.cache_clear() +def _clear_probe_cache(): + """The torch probe is `lru_cache`d for the process; a faked answer must not leak into the next + test. (The diffusers lookup is deliberately uncached -- see `_diffusers_attention_dispatch`.)""" + _probe_sdpa_dispatch.cache_clear() yield - _diffusers_attention_dispatch.cache_clear() - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() def _estimate( @@ -389,7 +388,7 @@ def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): decode be admitted to a card that could not run it.""" with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): materializing = self._vae_estimate("decode", 1024, device=FUSED) - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() fused = self._vae_estimate("decode", 1024, device=FUSED) tokens = 128 * 128 @@ -494,6 +493,24 @@ def test_a_failed_probe_is_budgeted_as_math(self): ) assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + def test_disabling_the_fused_kernels_at_runtime_invalidates_the_probe(self): + """The torch probe *is* cached -- it allocates and runs a dispatch query -- but its answer + depends on switches callers can flip at runtime (`sdpa_kernel()`, + `torch.backends.cuda.enable_flash_sdp`). Those toggles are part of the cache key, so an + estimate priced after a switch does not inherit the answer from before it. No cache is + cleared between these calls on purpose.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + def estimate(): + return sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + ) + + assert estimate() == 0 + with sdpa_kernel([SDPBackend.MATH]): + assert estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert estimate() == 0 + def test_the_probe_asks_torch_the_same_question_sdpa_does(self): """`_fused_sdp_choice` is the dispatch query `F.scaled_dot_product_attention` itself runs, so a `MATH` answer means the real forward materializes. Reimplementing the eligibility rules @@ -502,7 +519,7 @@ def test_the_probe_asks_torch_the_same_question_sdpa_does(self): with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) - _torch_sdpa_materializes_score_matrix.cache_clear() + _probe_sdpa_dispatch.cache_clear() with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) @@ -562,25 +579,14 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h assert estimate <= 2 * measured -@contextmanager def _diffusers_backend(name): - """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would. - - The lookup is `lru_cache`d -- it is a process-wide setting read on every estimate -- so the - cache has to be dropped on the way in and on the way out, or the faked backend leaks into the - comparison the test makes against the real one. - """ + """Force the process-wide diffusers attention backend, as `DIFFUSERS_ATTN_BACKEND` would.""" from diffusers.models.attention_dispatch import AttentionBackendName - _diffusers_attention_dispatch.cache_clear() - try: - with patch( - "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", - return_value=(AttentionBackendName(name), None), - ): - yield - finally: - _diffusers_attention_dispatch.cache_clear() + return patch( + "diffusers.models.attention_dispatch._AttentionBackendRegistry.get_active_backend", + return_value=(AttentionBackendName(name), None), + ) class TestDiffusersAttentionDispatchIsConsulted: @@ -624,6 +630,45 @@ def estimate(): forced_math = estimate() assert forced_math == estimate() + def test_a_backend_switch_is_not_masked_by_an_earlier_estimate(self): + """The active backend is mutable process state. Caching the first answer would keep + reserving zero for every later estimate in a long-lived process that has since switched to + `_native_math` -- the exact case this lookup exists to catch. Deliberately no cache is + cleared between the two calls here; the production code must not be holding one.""" + native = _estimate(image_seq_len=4096, device=FUSED) + with _diffusers_backend("_native_math"): + after_switch = _estimate(image_seq_len=4096, device=FUSED) + back_to_native = _estimate(image_seq_len=4096, device=FUSED) + + seq_len = 4096 + 512 + assert after_switch - native == ( + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + assert back_to_native == native + + def test_a_model_level_override_reaches_the_registry(self): + """Why the estimator does not need the model in hand: it is priced before the transformer is + loaded, and `set_attention_backend()` stamps its choice onto the process-wide registry as + well as onto the model's attention processors -- deliberately, "so that it propagates + gracefully throughout". If diffusers ever stops doing that, a per-model override could + disagree with the estimate, and this test is where that shows up.""" + from diffusers.configuration_utils import ConfigMixin, register_to_config + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + from diffusers.models.modeling_utils import ModelMixin + + class _Tiny(ModelMixin, ConfigMixin): + @register_to_config + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(2, 2) + + previous = _AttentionBackendRegistry._active_backend + try: + _Tiny().set_attention_backend("_native_math") + assert _diffusers_attention_dispatch() == "math" + finally: + _AttentionBackendRegistry._active_backend = previous + def test_an_unreadable_dispatcher_is_budgeted_as_math(self): """`_AttentionBackendRegistry` is private; if diffusers moves it we lose the answer. The conservative reading is the materializing one, and it is logged rather than silent.""" From df6b3d88df251409fd751159b654c8c07fac4d34 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 02:51:02 +0200 Subject: [PATCH 05/11] fix(flux2): stop caching the SDPA probe and scale the VAE estimate by batch The probe's cache key held the four per-backend enable flags, but torch takes the *first eligible* backend in a priority order that `sdpa_kernel(..., set_priority=True)` reorders while leaving every flag untouched -- measured: same flags, EFFICIENT outside and MATH inside. A fused answer cached before the switch would suppress the score-matrix reservation after it. Rather than adding the priority order to the key -- the next thing to forget is always one more -- drop the cache. The probe costs ~6us against a multi-second forward, so there is nothing to protect. `vae.decode` is also handed whatever batch the latents carry, and a LatentsField is not pinned to one, so an estimate built from H and W alone gave a two-sample decode a single sample's reservation. Measured at 1024px: 4.23GB at batch 1, 7.96GB at 2, 11.89GB at 3 -- linear, slightly sub-linear per sample, so the scaled single-sample estimate stays an upper bound. The score matrix is (batch, heads, S, S) and scales with it. --- invokeai/backend/util/attention.py | 36 +---- invokeai/backend/util/vae_working_memory.py | 12 +- .../invocations/test_flux2_working_memory.py | 134 +++++++++++++++--- 3 files changed, 131 insertions(+), 51 deletions(-) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 25f0b463cb4..9bbda9c290a 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -130,21 +130,6 @@ def _diffusers_attention_dispatch() -> str: return _DISPATCH_FUSED -def _sdp_kernel_toggles() -> tuple[bool, ...]: - """The global switches that gate each fused SDPA kernel, as `_fused_sdp_choice` sees them. - - `torch.backends.cuda.enable_flash_sdp(False)` and `sdpa_kernel([...])` flip these at runtime and - the dispatch answer flips with them, so they belong in the probe's cache key rather than being - baked into a permanent result. - """ - cuda = torch.backends.cuda - return tuple( - bool(getattr(cuda, name)()) - for name in ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") - if hasattr(cuda, name) - ) - - def _torch_sdpa_materializes_score_matrix( device_type: str, device_index: int | None, dtype: torch.dtype, head_dim: int, has_attn_mask: bool ) -> bool: @@ -155,6 +140,12 @@ def _torch_sdpa_materializes_score_matrix( Eligibility depends on the dtype, the head dim and the presence of a mask, not on the sequence length, so a tiny probe answers for the real forward. + Not cached. The answer turns on global torch state a cache key cannot honestly enumerate: the + per-backend enable flags, but also the *priority order*, which `sdpa_kernel(..., set_priority= + True)` reorders while leaving every flag untouched -- measured, same flags, `EFFICIENT` before + and `MATH` inside. Each item added to such a key is one more thing to get wrong later, and the + probe costs ~6us against a multi-second forward, so it just runs every time. + Anything that goes wrong reports the materializing path, which is both the conservative answer and, for the most common cause, the correct one: torch registers `_fused_sdp_choice` for CPU, CUDA/ROCm and XPU only, so the call raises on MPS -- and MPS is exactly where @@ -164,21 +155,6 @@ def _torch_sdpa_materializes_score_matrix( failure inside the probe, a torch that predates the op) leave us knowing nothing at all, and there the asymmetry decides: a shortfall costs an OOM, an over-estimate costs some residency. """ - return _probe_sdpa_dispatch(device_type, device_index, dtype, head_dim, has_attn_mask, _sdp_kernel_toggles()) - - -@lru_cache(maxsize=None) -def _probe_sdpa_dispatch( - device_type: str, - device_index: int | None, - dtype: torch.dtype, - head_dim: int, - has_attn_mask: bool, - sdp_kernel_toggles: tuple[bool, ...], -) -> bool: - """Cached body of the probe above. Every input torch's answer depends on is part of the key -- - `sdp_kernel_toggles` is not read here, it is carried so a runtime change invalidates the entry. - """ try: device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 960e2d081f9..0c2cf3cfa92 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -140,12 +140,20 @@ def estimate_vae_working_memory_flux2( When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score matrix, if one is materialized at all, is bounded by the tile too. + + Both terms are per sample. `vae.decode` takes whatever batch the latents carry, and a + ``LatentsField`` is not pinned to one, so the batch has to multiply through: measured at 1024px + decode, peak reserved is 4.23GB at batch 1, 7.96GB at batch 2 and 11.89GB at batch 3 -- linear, + and slightly sub-linear per sample, so multiplying the single-sample estimate stays an upper + bound. The score matrix is shaped (batch, heads, S, S), so it scales the same way. The encode + call sites all pass batch 1 today; the shared estimator does not assume it. """ param = next(vae.parameters()) element_size = param.element_size() # Encoding uses ~50% the working memory of decoding. scaling_constant = 2200 if operation == "decode" else 1100 + batch_size = image_tensor.shape[0] if image_tensor.dim() >= 4 else 1 if tile_size is not None: # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. @@ -158,10 +166,12 @@ def estimate_vae_working_memory_flux2( working_memory = out_h * out_w * element_size * scaling_constant mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) + working_memory *= batch_size working_memory += sdpa_score_matrix_bytes( device=device if device is not None else TorchDevice.choose_torch_device(), dtype=param.dtype, - num_heads=_FLUX2_VAE_MID_BLOCK_HEADS, + # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, seq_len=mid_block_seq_len, ) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index b876d973c13..7d3e5783889 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -23,7 +23,6 @@ from invokeai.backend.util.attention import ( SDPA_MATH_BYTES_PER_SCORE_ELEMENT, _diffusers_attention_dispatch, - _probe_sdpa_dispatch, _torch_sdpa_materializes_score_matrix, sdpa_score_matrix_bytes, ) @@ -39,15 +38,6 @@ FUSED = torch.device("cpu") -@pytest.fixture(autouse=True) -def _clear_probe_cache(): - """The torch probe is `lru_cache`d for the process; a faked answer must not leak into the next - test. (The diffusers lookup is deliberately uncached -- see `_diffusers_attention_dispatch`.)""" - _probe_sdpa_dispatch.cache_clear() - yield - _probe_sdpa_dispatch.cache_clear() - - def _estimate( image_seq_len, ref_image_seq_len=0, @@ -179,6 +169,67 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): assert estimates[0] < untiled / 4 +class TestFlux2VaeBatchIsBudgeted: + """`vae.decode` is handed whatever batch the latents carry, and a `LatentsField` is not pinned to + one. An estimate built from H and W alone gives a two-sample decode the same reservation as a + single one, so the cache admits it to a card that cannot run it -- the reservation is there, and + the OOM happens anyway. + + Measured at 1024px on CUDA/bf16, peak reserved, each point in a fresh process: 4.23GB at batch 1, + 7.96GB at batch 2, 11.89GB at batch 3. Linear, and slightly sub-linear per sample, so scaling the + single-sample estimate is an upper bound rather than a fit. + """ + + # (batch, measured peak reserved MB) for a 1024px decode. + MEASURED_DECODE_BATCH = [(1, 4229), (2, 7955), (3, 11889)] + + def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + return estimate_vae_working_memory_flux2( + operation="decode", + image_tensor=torch.zeros(batch, 32, px // 8, px // 8), + vae=vae, + tile_size=tile_size, + device=device, + ) + + @pytest.mark.parametrize("batch, measured_mb", MEASURED_DECODE_BATCH) + def test_estimate_is_an_upper_bound_on_the_measured_batch_peak(self, batch, measured_mb): + estimate = self._decode_estimate(batch) + assert estimate >= measured_mb * MB + assert estimate <= 2 * measured_mb * MB + + def test_estimate_scales_with_the_batch(self): + """The regression in one assertion: before this, all three of these were equal.""" + single = self._decode_estimate(1) + assert self._decode_estimate(2) == 2 * single + assert self._decode_estimate(3) == 3 * single + + def test_a_three_dimensional_tensor_is_one_sample(self): + """A bare `(C, H, W)` latent has no batch axis; `shape[0]` would read the channel count.""" + vae = MagicMock(spec=AutoencoderKLFlux2) + vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + unbatched = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=torch.zeros(32, 128, 128), vae=vae, device=FUSED + ) + assert unbatched == self._decode_estimate(1) + + def test_tiling_bounds_the_tile_not_the_batch(self): + """Tiling caps the spatial term at one tile, but every sample still runs through it.""" + single = self._decode_estimate(1, px=1024, tile_size=512) + assert self._decode_estimate(3, px=1024, tile_size=512) == 3 * single + + def test_the_score_matrix_scales_with_the_batch(self): + """It is shaped (batch, heads, S, S), so where it is materialized at all it scales too.""" + tokens = 128 * 128 + score = tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + linear = self._decode_estimate(1) # fused: the spatial term on its own + with _materializing(): + assert self._decode_estimate(1, device=MATERIALIZING) == linear + score + assert self._decode_estimate(3, device=MATERIALIZING) == 3 * (linear + score) + + class TestFlux2VaeInvocationsRequestWorkingMemory: """The estimate is worthless unless it reaches `model_on_device()`.""" @@ -388,7 +439,6 @@ def test_mps_style_dispatch_failure_reserves_the_vae_score_matrix(self): decode be admitted to a card that could not run it.""" with patch("torch.ops.aten._fused_sdp_choice", side_effect=NotImplementedError("no MPS kernel")): materializing = self._vae_estimate("decode", 1024, device=FUSED) - _probe_sdpa_dispatch.cache_clear() fused = self._vae_estimate("decode", 1024, device=FUSED) tokens = 128 * 128 @@ -493,21 +543,66 @@ def test_a_failed_probe_is_budgeted_as_math(self): ) assert estimated == 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - def test_disabling_the_fused_kernels_at_runtime_invalidates_the_probe(self): - """The torch probe *is* cached -- it allocates and runs a dispatch query -- but its answer - depends on switches callers can flip at runtime (`sdpa_kernel()`, - `torch.backends.cuda.enable_flash_sdp`). Those toggles are part of the cache key, so an - estimate priced after a switch does not inherit the answer from before it. No cache is - cleared between these calls on purpose.""" + def _cpu_estimate(self): + return sdpa_score_matrix_bytes( + device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + ) + + def test_disabling_the_fused_kernels_at_runtime_changes_the_answer(self): + """The probe is not cached, so an estimate priced after a runtime switch does not inherit the + answer from before it. Nothing is cleared between these calls on purpose.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + assert self._cpu_estimate() == 0 + with sdpa_kernel([SDPBackend.MATH]): + assert self._cpu_estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert self._cpu_estimate() == 0 + + def test_a_priority_reorder_changes_the_answer_with_every_flag_unchanged(self): + """`sdpa_kernel(..., set_priority=True)` puts `MATH` first while leaving all four enable + flags True, and torch takes the first eligible backend in that order. A cache keyed on the + flags -- which is what this probe used to have -- could not see the switch and would keep + reserving zero. Not caching at all is what makes that unrepresentable.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + math_first = [ + SDPBackend.MATH, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + ] + flags = ("flash_sdp_enabled", "mem_efficient_sdp_enabled", "math_sdp_enabled", "cudnn_sdp_enabled") + + assert self._cpu_estimate() == 0 + with sdpa_kernel(math_first, set_priority=True): + # The finding in one line: every flag a key could hold is still True in here. + assert all(getattr(torch.backends.cuda, name)() for name in flags if hasattr(torch.backends.cuda, name)) + # CPU's chooser ignores the priority order, so stand in for the answer CUDA gives. + with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): + assert self._cpu_estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + assert self._cpu_estimate() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="only CUDA's chooser honours the priority order") + def test_this_build_reports_a_real_priority_reorder(self): + """The same case against the real dispatcher rather than a stand-in. Verified on torch + 2.7.1+cu128: `_fused_sdp_choice` answers EFFICIENT outside and MATH inside.""" from torch.nn.attention import SDPBackend, sdpa_kernel def estimate(): return sdpa_score_matrix_bytes( - device=torch.device("cpu"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 + device=torch.device("cuda"), dtype=torch.bfloat16, num_heads=1, head_dim=128, seq_len=4096 ) assert estimate() == 0 - with sdpa_kernel([SDPBackend.MATH]): + with sdpa_kernel( + [ + SDPBackend.MATH, + SDPBackend.FLASH_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + ], + set_priority=True, + ): assert estimate() == 4096 * 4096 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert estimate() == 0 @@ -519,7 +614,6 @@ def test_the_probe_asks_torch_the_same_question_sdpa_does(self): with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.MATH)): assert _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) - _probe_sdpa_dispatch.cache_clear() with patch("torch.ops.aten._fused_sdp_choice", return_value=int(SDPBackend.EFFICIENT_ATTENTION)): assert not _torch_sdpa_materializes_score_matrix("cpu", None, torch.bfloat16, 128, False) From 596eda9d742be0b7081d2ce3a41937476c3adf6d Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 28 Aug 2026 15:48:00 +0200 Subject: [PATCH 06/11] fix(flux2): scale the denoise reservation by the latent batch The node had `b` in hand from preparing the latents and never passed it, so a two-sample run reserved one sample's activations and the cache admitted it to a card that could not run it. Batched latents do not come from the stock UI, but the API and custom graphs reach this node. Batch multiplies the token count and nothing else. Measured on the Klein geometry with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 at 5126MB, and 9728 tokens at B=1 at 5584MB -- per total token that is 0.554-0.578MB across every combination, so batch and sequence are interchangeable. Reference latents are repeated per sample by `ensure_batch_size`, so they scale too, and the score matrix is (batch, heads, S, S). The fixed base does not scale -- it covers weight casts and allocator slack -- and neither does the regional bias, built as (1, 1, S, S) and broadcast. --- invokeai/app/invocations/flux2_denoise.py | 20 +++- .../invocations/test_flux2_working_memory.py | 93 ++++++++++++++++++- 2 files changed, 109 insertions(+), 4 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index e6655b5c593..eb423d99e26 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -481,6 +481,9 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: 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), + # A batched latent/noise tensor is reachable through the API and custom graphs, and the + # reference latents get repeated to match it (`ensure_batch_size` below). + batch_size=b, # 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() @@ -620,6 +623,7 @@ def _estimate_working_memory( ref_image_seq_len: int, text_seq_len: int, num_loras: int, + batch_size: int = 1, regional_attention_bias_bytes: int = 0, has_regional_attention_mask: bool = False, device: torch.device | None = None, @@ -643,6 +647,15 @@ def _estimate_working_memory( during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra activation branch per patched layer, so we add a per-LoRA margin. + Batch multiplies the token count and nothing else. A batch of B is B independent sequences, + so it enters the linear term exactly as extra sequence does -- measured on the Klein geometry + with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 (9216 + tokens) at 5126MB, and 9728 tokens at B=1 at 5584MB. Batch and sequence are interchangeable + to within the noise. Reference latents are repeated per sample (`ensure_batch_size`), so they + scale with it too, and the score matrix is shaped (batch, heads, S, S). The fixed base does + not scale -- it is about weights, not activations -- and neither does the regional bias, which + is built as (1, 1, S, S) and broadcast across the batch. + The linear model holds only while attention runs on a fused kernel. Regional prompting is where that stops being a given: it hands the transformer a dense additive ``S x S`` bias, which flash attention never accepts and which ROCm's memory-efficient kernel rejects as @@ -657,13 +670,13 @@ def _estimate_working_memory( MB = 1024**2 per_token_bytes = int(0.4 * MB) total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len - estimated = total_seq_len * per_token_bytes + estimated = total_seq_len * batch_size * per_token_bytes estimated += int(1.0 * GB) estimated += regional_attention_bias_bytes estimated += sdpa_score_matrix_bytes( device=device if device is not None else TorchDevice.choose_torch_device(), dtype=dtype, - num_heads=FLUX2_MAX_ATTENTION_HEADS, + num_heads=FLUX2_MAX_ATTENTION_HEADS * batch_size, head_dim=FLUX2_ATTENTION_HEAD_DIM, seq_len=total_seq_len, has_attn_mask=has_regional_attention_mask, @@ -672,7 +685,8 @@ def _estimate_working_memory( via_diffusers_dispatch=True, ) if num_loras > 0: - estimated += int(0.5 * num_loras * GB) + # 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( diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 7d3e5783889..497811d044b 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -43,6 +43,7 @@ def _estimate( ref_image_seq_len=0, text_seq_len=512, num_loras=0, + batch_size=1, regional_bias=0, has_regional_mask=False, device=FUSED, @@ -53,6 +54,7 @@ def _estimate( ref_image_seq_len=ref_image_seq_len, text_seq_len=text_seq_len, num_loras=num_loras, + batch_size=batch_size, regional_attention_bias_bytes=regional_bias, has_regional_attention_mask=has_regional_mask, device=device, @@ -110,6 +112,75 @@ def test_regional_attention_bias_is_added(self): assert _estimate(image_seq_len=4096, regional_bias=123 * MB) - base == 123 * MB +class TestFlux2DenoiseBatchIsBudgeted: + """A batch of B is B independent sequences, so it enters the linear term exactly as extra + sequence does. Measured on the Klein geometry (48 heads x 128, mlp 3.0) with a reduced block + count -- the constant is block-count independent -- peak reserved, each point in a fresh process: + + B=1, 4608 tokens -> 2570MB B=2, 4608 each (9216 total) -> 5126MB + B=1, 9728 tokens -> 5584MB B=2, 9728 each (19456 total) -> 11120MB + B=1, 14336 tokens -> 8284MB B=3, 4608 each (13824 total) -> 7656MB + + Per *total* token that is 0.554-0.578MB across every row: batch and sequence are interchangeable. + (The absolute figure is not comparable to the Klein table elsewhere in this module -- a 3-block + stand-in amortizes per-forward overhead differently. Only the equivalence is being tested.) + + Batched latents reach this node through the API and custom graphs, not the stock UI. + """ + + def test_batch_and_sequence_are_interchangeable(self): + """Two samples of 4608 tokens must cost what one sample of 9216 costs -- the measurement + above says 5126MB against 5584MB, equal to within the allocator's noise.""" + assert _estimate(image_seq_len=4096, text_seq_len=512, batch_size=2) == _estimate( + image_seq_len=8704, text_seq_len=512, batch_size=1 + ) + + @pytest.mark.parametrize("batch", [2, 3, 4]) + def test_each_extra_sample_adds_exactly_its_own_tokens(self, batch): + single = _estimate(image_seq_len=4096) + assert _estimate(image_seq_len=4096, batch_size=batch) - single == (batch - 1) * (4096 + 512) * int(0.4 * MB) + + def test_reference_tokens_scale_with_the_batch(self): + """`ensure_batch_size` repeats the reference latents across the batch.""" + single = _estimate(image_seq_len=4096, ref_image_seq_len=12288) + assert _estimate(image_seq_len=4096, ref_image_seq_len=12288, batch_size=2) - single == ( + (4096 + 12288 + 512) * int(0.4 * MB) + ) + + def test_the_fixed_base_does_not_scale_with_the_batch(self): + """It covers transient weight casts and allocator slack -- properties of the weights, not of + how many samples run through them. Scaling it would add a GB per sample for nothing.""" + deltas = { + _estimate(image_seq_len=4096, batch_size=b + 1) - _estimate(image_seq_len=4096, batch_size=b) + for b in (1, 2, 3) + } + assert deltas == {(4096 + 512) * int(0.4 * MB)} + + def test_the_regional_bias_does_not_scale_with_the_batch(self): + """`get_joint_attention_kwargs` builds it as (1, 1, S, S) and lets SDPA broadcast it, so + there is exactly one of them however many samples are in flight.""" + bias = (4096 + 512) ** 2 * 2 + single = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True) + double = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True, batch_size=2) + assert double - single == (4096 + 512) * int(0.4 * MB) + + def test_the_score_matrix_scales_with_the_batch(self): + """Where it is materialized at all it is shaped (batch, heads, S, S).""" + seq_len = 4096 + 512 + with _materializing(): + single = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) + double = _estimate(image_seq_len=4096, has_regional_mask=True, batch_size=2, device=MATERIALIZING) + assert double - single == ( + seq_len * int(0.4 * MB) + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + ) + + def test_the_lora_margin_scales_with_the_batch(self): + """A sidecar patch adds an activation branch, and activations are per sample.""" + assert _estimate(image_seq_len=4096, num_loras=2, batch_size=3) - _estimate( + image_seq_len=4096, num_loras=0, batch_size=3 + ) == 3 * int(1.0 * GB) + + class TestFlux2VaeWorkingMemoryEstimate: # (operation, pixel size, measured peak reserved MB), bf16, untiled. MEASURED_VAE = [ @@ -306,7 +377,7 @@ 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): + def _run(self, num_ref_tokens: int, batch: int = 1): """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 ( @@ -354,6 +425,9 @@ def _run(self, num_ref_tokens: int): 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) @@ -371,6 +445,23 @@ def test_reference_images_raise_the_reservation(self): with_refs = self._run(num_ref_tokens=12288) assert with_refs - without_refs == 12288 * int(0.4 * MB) + def test_the_real_batch_reaches_the_reservation(self): + """A batched latent tensor is reachable through the API and custom graphs. The node has `b` + in hand at the estimate; before this it simply did not pass it, so a two-sample run reserved + one sample's worth and the cache admitted it to a card that could not run it.""" + assert self._run(num_ref_tokens=0, batch=2) == _estimate(image_seq_len=64 * 64, text_seq_len=512, batch_size=2) + + def test_a_batched_run_reserves_more_than_a_single_one(self): + single = self._run(num_ref_tokens=0, batch=1) + assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * int(0.4 * MB) + + def test_repeated_reference_latents_are_counted_per_sample(self): + """`ensure_batch_size` repeats the reference latents across the batch, so their tokens scale + with it as well -- the worst case in #9500, doubled.""" + single = self._run(num_ref_tokens=12288, batch=1) + double = self._run(num_ref_tokens=12288, batch=2) + assert double - single == (64 * 64 + 12288 + 512) * int(0.4 * MB) + def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): """Stand in for the torch probe on a build with ROCm's fused-kernel rules. From 58ef4ae7fc78c9e437e00297f4ccd2a3a65aa602 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 29 Aug 2026 01:37:47 +0200 Subject: [PATCH 07/11] fix(flux2): take the reservation's batch from the blended latents `b` was read from the noise tensor, which this node builds at batch 1 from width/height/seed whenever `add_noise` is set. Batched init latents then broadcast against it in the img2img preblend, producing a two-sample `x` against a one-sample reservation. Read `x.shape[0]` instead. It is already in scope at the estimate -- the blend, the pack and the BN normalize all run above it, and none of them change the batch -- and it is the only thing that knows how many samples reach the transformer. Expanding the noise to match would have worked too, but that changes the noise, and with it the output. Note that the same `b` still feeds `generate_img_ids_flux2`, which is a correctness question rather than a memory one and is left alone here. --- invokeai/app/invocations/flux2_denoise.py | 9 +++++--- .../invocations/test_flux2_working_memory.py | 21 +++++++++++++++++-- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index eb423d99e26..b8c8752049c 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -481,9 +481,12 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: 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), - # A batched latent/noise tensor is reachable through the API and custom graphs, and the - # reference latents get repeated to match it (`ensure_batch_size` below). - batch_size=b, + # Taken from `x`, not from `b`. `b` is the *noise* tensor's batch, which this node builds + # at 1 from width/height/seed even when the init latents carry more; the img2img preblend + # above then broadcasts the two, so `x` is the only thing that knows how many samples + # actually go through the transformer. Reference latents are repeated to match it + # (`ensure_batch_size` below), so they scale with it too. + batch_size=x.shape[0], # The mask itself is already allocated; only the additive bias built per forward is new. regional_attention_bias_bytes=( regional_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size() diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 497811d044b..0769a293a17 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -377,7 +377,7 @@ 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): + def _run(self, num_ref_tokens: int, batch: int = 1, init_batch: int | None = None): """Drive `_run_diffusion` up to the transformer load and return the requested working memory.""" from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( @@ -400,8 +400,12 @@ def _run(self, num_ref_tokens: int, batch: int = 1): 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=None, + latents=MagicMock(latents_name="init") if init_batch is not None else None, noise=None, denoise_mask=None, denoising_start=0.0, @@ -455,6 +459,19 @@ def test_a_batched_run_reserves_more_than_a_single_one(self): single = self._run(num_ref_tokens=0, batch=1) assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * int(0.4 * MB) + def test_a_batched_init_latent_beats_the_batch_1_noise_it_is_blended_with(self): + """img2img takes `x = t_0 * noise + (1 - t_0) * init_latents`, and this node builds its noise + at batch 1 from width/height/seed. Two batched init latents therefore broadcast up to a + two-sample `x` while the noise tensor -- the thing the batch used to be read from -- still + says 1. The reservation has to follow `x`.""" + assert self._run(num_ref_tokens=0, init_batch=2) == _estimate( + image_seq_len=64 * 64, text_seq_len=512, batch_size=2 + ) + + def test_the_blended_batch_is_read_after_the_broadcast(self): + single = self._run(num_ref_tokens=0, init_batch=1) + assert self._run(num_ref_tokens=0, init_batch=3) - single == 2 * (64 * 64 + 512) * int(0.4 * MB) + def test_repeated_reference_latents_are_counted_per_sample(self): """`ensure_batch_size` repeats the reference latents across the batch, so their tokens scale with it as well -- the worst case in #9500, doubled.""" From b6a70b624aadb76bc9dfc1a98e85f13f42444adc Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 29 Aug 2026 19:39:17 +0200 Subject: [PATCH 08/11] fix(flux2): scale the estimate by transformer width, not just token count Per-token activation cost is linear in the transformer's hidden width, and the constant was calibrated on Klein 9B (4096) but applied to every variant. FLUX.2 [dev] is 6144 and reaches this node as a first-class path, so 1024x1024 with three references reserved 7.6GB against ~10GB needed -- the same shortfall #9500 describes, on the model where partial loading makes the estimate decide residency. Measured slope between 4608 and 9216 tokens, everything else held fixed: 0.291 MB/tok at 3072, 0.386 at 4096, 0.555 at 6144 -- 0.755 / 1.00 / 1.438 against width ratios of 0.75 / 1.00 / 1.50. Scale by width, and take the head count from the same number instead of always charging the widest. Also: raise SDPA_MATH_BYTES_PER_SCORE_ELEMENT to 14, which ROCm's 13.62 at the smallest measured shape needs; drop the claim that ROCm rejects additive masks, which gfx1100 disproves; log at info when the score-matrix term fires, since it decides residency and nothing else said so; guard the warning-filter swap with a lock now that the probe runs on every estimate; and give the VAE encode node the same compute device as the decode node. --- invokeai/app/invocations/flux2_denoise.py | 53 +- invokeai/app/invocations/flux2_vae_encode.py | 9 +- .../model_manager/configs/flux2_variant.py | 20 + invokeai/backend/util/attention.py | 76 ++- invokeai/backend/util/vae_working_memory.py | 13 +- scripts/calibrate_flux2_working_memory.py | 598 ++++++++++++++++++ .../invocations/test_flux2_working_memory.py | 174 ++++- 7 files changed, 868 insertions(+), 75 deletions(-) create mode 100644 scripts/calibrate_flux2_working_memory.py diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index b8c8752049c..9472d94d7f2 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 ( @@ -52,12 +53,16 @@ from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice -# FLUX.2 attention geometry. The head dim is 128 across every variant; the head count follows the -# hidden size (Klein 4B: 24, Klein 9B: 32, FLUX.2 dev: 48). Only the head dim decides which SDPA -# kernel is eligible; the head count scales the `math` fallback's score matrix, and since the -# working-memory estimate is computed before the transformer is loaded, we use the largest. +# 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 -FLUX2_MAX_ATTENTION_HEADS = 48 +# The width the per-token constant below was measured on. Estimates scale off this. +FLUX2_REFERENCE_HIDDEN_SIZE = 4096 +# 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( @@ -487,6 +492,9 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # 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() @@ -627,6 +635,7 @@ def _estimate_working_memory( 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, @@ -638,8 +647,19 @@ def _estimate_working_memory( activation footprint scales *linearly* with the total attended sequence -- text tokens, image tokens, and reference-image tokens alike. Measured on the Klein 9B geometry in bf16 as peak reserved memory, that slope is ~0.39 MB per token and holds from 1.5k to 28k tokens; it is - also independent of the block count (a no-grad forward frees each block's intermediates), so - the constant applies to both the 4B and 9B variants. + also independent of the block count (a no-grad forward frees each block's intermediates). + + 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: + + 3072 (Klein 4B) 0.291 MB/tok 4096 (Klein 9B) 0.386 6144 ([dev]) 0.555 + + which is 0.755 / 1.000 / 1.438 against width ratios of 0.75 / 1.00 / 1.50 -- linear in width, + 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. 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 @@ -661,17 +681,18 @@ def _estimate_working_memory( The linear model holds only while attention runs on a fused kernel. Regional prompting is where that stops being a given: it hands the transformer a dense additive ``S x S`` bias, - which flash attention never accepts and which ROCm's memory-efficient kernel rejects as - well, leaving the ``math`` fallback and its materialized ``heads x S x S`` score matrix. The - device decides too -- MPS has no fused SDPA kernel at all -- and so does the diffusers - attention backend this build dispatches through. ``sdpa_score_matrix_bytes`` asks all three - and adds the score matrix only where it is really built: on CUDA with the stock backend the - memory-efficient kernel takes the bias and the term is zero (verified: peak stays linear - with the bias attached). + 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 MB = 1024**2 - per_token_bytes = int(0.4 * MB) + per_token_bytes = int(0.4 * MB * 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) @@ -679,7 +700,7 @@ def _estimate_working_memory( estimated += sdpa_score_matrix_bytes( device=device if device is not None else TorchDevice.choose_torch_device(), dtype=dtype, - num_heads=FLUX2_MAX_ATTENTION_HEADS * batch_size, + 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, diff --git a/invokeai/app/invocations/flux2_vae_encode.py b/invokeai/app/invocations/flux2_vae_encode.py index 2da6f38b517..926f3a7c043 100644 --- a/invokeai/app/invocations/flux2_vae_encode.py +++ b/invokeai/app/invocations/flux2_vae_encode.py @@ -18,7 +18,6 @@ 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 @@ -49,16 +48,20 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc """ # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to # free room rather than discovering the shortfall as an OOM. + # 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=TorchDevice.choose_torch_device(), + 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/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 9bbda9c290a..907aefde902 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -4,6 +4,7 @@ for attention mechanism. """ +import threading import warnings from functools import lru_cache @@ -45,21 +46,29 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # SDPA computes attention one of two ways: a fused kernel (flash / memory-efficient / cuDNN) that # never materializes the O(S^2) score matrix, or the `math` fallback, which does. Which one runs is # not a property of FLUX.2 -- it is a property of the build, the device, the dtype, the head dim and -# whether an attention mask was passed. CUDA's memory-efficient kernel accepts head dims well past -# 128 and arbitrary additive masks; ROCm's fused kernels reject both; MPS has no fused SDPA kernel -# at all. A working-memory estimate that assumes the fused path is therefore only correct on the -# build it was measured on, which is why the helpers below ask instead of assuming. - -# Peak *reserved* bytes per element of the materialized score matrix, measured on CUDA with -# `SDPBackend.MATH` forced, each point in a fresh process: 12.9 bytes/element at 4k tokens, 10.3 at -# 8k, 9.7 at 16k -- and the same figures for bf16, fp16 and fp32 inputs, because the fallback's -# softmax intermediates are fp32 regardless. So this is an absolute byte count, not a multiple of -# the element size. 13 is an upper bound on every measured point from 4k tokens up; below that it -# can fall a couple of MB short of the allocator's rounding, which is noise next to the GB-scale -# linear terms this is added to. It is a CUDA measurement standing in for every materializing -# backend -- no other was available to calibrate against -- but the intermediates it prices (an -# fp32 score matrix and its softmax) have the same shape wherever the fallback runs. -SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 13 +# whether an attention mask was passed, and the rules differ per build in ways that are not worth +# hard-coding. CUDA takes head dims well past 128 on its memory-efficient kernel; ROCm caps them at +# 128 and drops to `math` above that; MPS has no fused SDPA kernel at all. (Masks are *not* a +# reliable discriminator: a 128-wide masked head reports the memory-efficient kernel on gfx1100 too.) +# A working-memory estimate that assumes the fused path is therefore only correct on the build it was +# measured on, which is why the helpers below ask instead of assuming. + +# Peak *reserved* bytes per element of the materialized score matrix, 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 ROCm (gfx1100, torch 2.10) +# 1 4096 512 12.9 13.62 +# 1 8192 512 10.3 10.78 +# 1 16384 512 9.7 9.76 +# 4 4096 128 - 10.16 +# 48 4608 128 - 9.59 +# +# 14 is an upper bound on every measured point on both. It was 13, which held on CUDA but came in +# 0.62 under ROCm's worst point -- 4.7MB on a GB-scale estimate, so not a memory bug, but it made +# the pinned test red on ROCm and the docstring's "upper bound" claim false. +SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 14 # `_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 @@ -70,6 +79,9 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: 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" @@ -159,10 +171,15 @@ def _torch_sdpa_materializes_score_matrix( device = torch.device(device_type) if device_index is None else torch.device(device_type, device_index) q = torch.empty((1, 1, 8, head_dim), device=device, dtype=dtype) mask = torch.empty((1, 1, 8, 8), device=device, dtype=dtype) if has_attn_mask else None - with warnings.catch_warnings(): + 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: @@ -185,9 +202,12 @@ def sdpa_score_matrix_bytes( Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On CUDA it is almost always 0; where the fused kernels are missing or reject the shapes -- ROCm - caps the head dim at 128 and does not take arbitrary additive masks, MPS ships no fused SDPA - kernel at all -- it is the dominant term: a 1536px FLUX.2 VAE decode materializes 36864^2 - scores, ~17GB of them. + caps the head dim at 128, MPS ships no fused SDPA kernel at all -- it is the dominant term: a + 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~18GB 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 @@ -205,9 +225,25 @@ def sdpa_score_matrix_bytes( if dispatch == _DISPATCH_FUSED: return 0 if dispatch == _DISPATCH_MATH: - return score_matrix_bytes + 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 0c2cf3cfa92..eb733a3aa86 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -130,10 +130,15 @@ def estimate_vae_working_memory_flux2( That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the 512-wide head, and measured peak stays linear from 512 to 1536px). A build with no fused kernel - for the shapes -- ROCm caps the head dim at 128, MPS has no fused SDPA kernel at all -- drops to - SDPA's ``math`` fallback and materializes a (pixels/8)^2 score matrix on top of the linear term: - ~3.5GB at 1024px and ~17GB at 1536px. We ask torch which path applies rather than assuming, so - the estimate is right on both. (Unlike the transformer, this attention does not go through + for the shapes -- ROCm caps the head dim at 128 and reports ``math`` for this 512-wide one, MPS + has no fused SDPA kernel at all -- materializes a (pixels/8)^2 score matrix on top of the linear + term: ~3.7GB at 1024px and ~18GB at 1536px. We ask torch which path applies rather than assuming, + so the estimate is right on both. + + Known gap: the linear constants below were calibrated on CUDA, and on ROCm the same decode costs + more than they predict -- measured there at 768px, 3.78GB against 3.45GB estimated. 512px is + covered by the 3GB ``device_working_mem_gb`` floor and 1024px and up have margin, so the shortfall + is confined to the middle of the range and is ~0.33GB. Recalibrating for it needs ROCm hardware. (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.) diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py new file mode 100644 index 00000000000..736fcf65c1d --- /dev/null +++ b/scripts/calibrate_flux2_working_memory.py @@ -0,0 +1,598 @@ +"""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 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} + +# `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) + # The estimate is linear_term + score_matrix. Back the score term out so the linear constant can + # be fitted on its own -- it is the one the 2200/1100 literals name. + 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 += score_bytes + + return row | { + "reserved_delta": peak, + "estimate": estimate, + "score_term": score_bytes, + # Only meaningful while the peak actually exceeds the score term. Where it does not, the + # additive (linear + score) model does not decompose on this build and the number is noise. + "implied_linear_constant": ((peak - score_bytes) / (px * px * element_size)) if peak > score_bytes else None, + "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 _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]: + print("\n=== 3. VAE linear constants (2200 decode / 1100 encode) ===") + print("`implied_k` backs the score-matrix term out, so it is comparable to those literals; fit the") + print("constant 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 CUDA the forced-math decode measures *below* the fused one, because the") + print("memory-efficient kernel's workspace is the larger term there. 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("") + for operation, shipped in (("decode", 2200), ("encode", 1100)): + 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") + and r["implied_linear_constant"] + ] + 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 + print(f" {r['operation']} {r['px']}px force_math={r['force_math']}: short by {gap:.2f} GiB") + 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}" + ) + 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 index 0769a293a17..42c235333ad 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -17,7 +17,12 @@ import torch.nn.functional as F from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 -from invokeai.app.invocations.flux2_denoise import FLUX2_MAX_ATTENTION_HEADS, Flux2DenoiseInvocation +from invokeai.app.invocations.flux2_denoise import ( + FLUX2_ATTENTION_HEAD_DIM, + 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 ( @@ -31,6 +36,11 @@ 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 + # 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 @@ -44,6 +54,7 @@ def _estimate( 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, @@ -55,6 +66,7 @@ def _estimate( 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, @@ -112,6 +124,55 @@ def test_regional_attention_bias_is_added(self): 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. + + 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. + """ + + # (hidden size, measured MB per token) + MEASURED_WIDTH_SLOPE = [(3072, 0.2912), (4096, 0.3859), (6144, 0.5547)] + + 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("hidden, measured_mb", MEASURED_WIDTH_SLOPE) + def test_the_estimate_upper_bounds_the_measured_slope(self, hidden, measured_mb): + 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 @@ -171,7 +232,7 @@ def test_the_score_matrix_scales_with_the_batch(self): single = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) double = _estimate(image_seq_len=4096, has_regional_mask=True, batch_size=2, device=MATERIALIZING) assert double - single == ( - seq_len * int(0.4 * MB) + FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + seq_len * int(0.4 * MB) + KLEIN_9B_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT ) def test_the_lora_margin_scales_with_the_batch(self): @@ -377,7 +438,7 @@ 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): + 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 ( @@ -390,8 +451,13 @@ def _run(self, num_ref_tokens: int, batch: int = 1, init_batch: int | None = Non 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 + 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))] @@ -459,6 +525,33 @@ def test_a_batched_run_reserves_more_than_a_single_one(self): single = self._run(num_ref_tokens=0, batch=1) assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * int(0.4 * MB) + @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 * (int(0.4 * MB * 6144 / 4096) - int(0.4 * MB)) + 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 @@ -480,21 +573,22 @@ def test_repeated_reference_latents_are_counted_per_sample(self): assert double - single == (64 * 64 + 12288 + 512) * int(0.4 * MB) -def _rocm_like_probe(device_type, device_index, dtype, head_dim, has_attn_mask): - """Stand in for the torch probe on a build with ROCm's fused-kernel rules. +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. - ROCm's fused SDPA kernels cap the head dim at 128 and do not take an arbitrary additive mask; - anything else falls through to the `math` fallback, which materializes the score matrix. CUDA's - memory-efficient kernel accepts both -- verified on torch 2.7.1+cu128, where - `_fused_sdp_choice` reports the efficient kernel for the VAE's 512-wide head and for a masked - 128-wide transformer head, and measured peak stays linear in both cases -- which is why the - estimates were linear to begin with. + 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 def _materializing(): - return patch("invokeai.backend.util.attention._torch_sdpa_materializes_score_matrix", side_effect=_rocm_like_probe) + 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. @@ -595,9 +689,7 @@ def test_regional_prompting_adds_the_score_matrix(self): has_regional_mask=True, device=MATERIALIZING, ) - assert materializing - fused == ( - FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - ) + assert materializing - fused == (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): @@ -736,9 +828,11 @@ def test_empty_sequences_cost_nothing(self): @pytest.mark.skipif(not torch.cuda.is_available(), reason="asks the real CUDA/ROCm dispatcher") def test_this_build_reports_its_own_dispatch(self): - """On CUDA both shapes are fused and this whole term is zero -- the fix is a no-op for the - hardware the constants were measured on. On ROCm the same call reports the fallback and the - term appears. Either answer is correct; the point is that it comes from torch.""" + """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 ) @@ -750,12 +844,29 @@ def test_this_build_reports_its_own_dispatch(self): seq_len=4608, has_attn_mask=True, ) - if torch.version.hip is None: - assert vae_bytes == 0 - assert masked_bytes == 0 - else: - assert vae_bytes > 0 - assert masked_bytes > 0 + 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) + + # (heads, seq, head_dim, bytes per score element) measured with SDPBackend.MATH forced on + # ROCm 7.1 / torch 2.10 / gfx1100. The CUDA points live in the constant's own comment; the worst + # of either platform is ROCm's 13.62 at the smallest shape, which is why the constant is 14. + MEASURED_BYTES_PER_ELEMENT_ROCM = [ + (1, 4096, 512, 13.62), + (1, 8192, 512, 10.78), + (1, 16384, 512, 9.76), + (4, 4096, 128, 10.16), + (48, 4608, 128, 9.59), + ] + + @pytest.mark.parametrize("num_heads, seq_len, head_dim, measured", MEASURED_BYTES_PER_ELEMENT_ROCM) + def test_the_constant_upper_bounds_every_measured_platform(self, num_heads, seq_len, head_dim, measured): + """13 held on CUDA and came in under ROCm's smallest shape. The shortfall was 4.7MB -- noise + against the GB-scale linear terms -- but it made the pinned bound below red on ROCm, and the + constant's "upper bound on every measured point" claim false.""" + 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)]) @@ -777,7 +888,10 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h measured = torch.cuda.max_memory_reserved() - before estimate = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert estimate >= measured + assert estimate >= 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" + ) assert estimate <= 2 * measured @@ -807,9 +921,7 @@ def test_forced_math_backend_reaches_the_denoise_estimate(self): forced_math = _estimate(image_seq_len=4096, device=FUSED) native = _estimate(image_seq_len=4096, device=FUSED) seq_len = 4096 + 512 - assert forced_math - native == ( - FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - ) + 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 @@ -843,9 +955,7 @@ def test_a_backend_switch_is_not_masked_by_an_earlier_estimate(self): back_to_native = _estimate(image_seq_len=4096, device=FUSED) seq_len = 4096 + 512 - assert after_switch - native == ( - FLUX2_MAX_ATTENTION_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - ) + assert 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): From 600ac3e32f9ce41b69c4760264b6ae9cb5cf4662 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 29 Aug 2026 22:31:38 +0200 Subject: [PATCH 09/11] fix(flux2): raise both calibrated constants to bound the AMD measurements Two ROCm runs came in through the new calibration script and both shipped constants were under their worst point. SDPA_MATH_BYTES_PER_SCORE_ELEMENT goes 14 -> 17 (gfx1201 costs 16.38 for the shape where CUDA costs 12.88), and the per-token activation constant 0.40 -> 0.42 MB (gfx1201 measures 0.4067 at the reference width). Both are now pinned against all three platforms' measured points rather than only being self-consistent. The runs also disprove the ROCm framing this feature carried: gfx1100 reports MATH for the VAE's 512-wide head, gfx1201 reports FLASH. Two cards, same vendor, same torch, opposite answers -- which is the case for asking torch rather than hard-coding a rule, but it means the docstrings could not keep saying "ROCm caps the head dim at 128". The VAE linear constants are left alone despite measuring short on gfx1201. That run had MIOPEN_FIND_MODE=2 and a HIP allocator garbage_collection threshold set, and its series is non-monotonic above 1024px -- peak reserved falls as resolution rises, which is what a GC threshold does to this measurement. The two AMD cards are also 1.8x apart. The gap is documented where the constants are defined; the script now reports the environment and flags a non-monotonic series so the next run cannot be ambiguous about it. --- invokeai/app/invocations/flux2_denoise.py | 24 +++-- invokeai/backend/util/attention.py | 38 +++---- invokeai/backend/util/vae_working_memory.py | 27 +++-- scripts/calibrate_flux2_working_memory.py | 59 ++++++++++ .../invocations/test_flux2_working_memory.py | 101 ++++++++++++------ 5 files changed, 183 insertions(+), 66 deletions(-) diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 9472d94d7f2..33725938a4f 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/invokeai/app/invocations/flux2_denoise.py @@ -60,6 +60,11 @@ 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 @@ -646,20 +651,24 @@ def _estimate_working_memory( FLUX.2 attention runs through SDPA without materializing the O(seq^2) score matrix, so the activation footprint scales *linearly* with the total attended sequence -- text tokens, image tokens, and reference-image tokens alike. Measured on the Klein 9B geometry in bf16 as peak - reserved memory, that slope is ~0.39 MB per token and holds from 1.5k to 28k tokens; it is - also independent of the block count (a no-grad forward frees each block's intermediates). + 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: - 3072 (Klein 4B) 0.291 MB/tok 4096 (Klein 9B) 0.386 6144 ([dev]) 0.555 + 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 against width ratios of 0.75 / 1.00 / 1.50 -- linear in width, + 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. The head count follows the - same width, so the score-matrix term gets the real one instead of the widest. + 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 @@ -691,8 +700,7 @@ def _estimate_working_memory( bias attached). """ GB = 1024**3 - MB = 1024**2 - per_token_bytes = int(0.4 * MB * hidden_size / FLUX2_REFERENCE_HIDDEN_SIZE) + 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) diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 907aefde902..1ffa8abef72 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -47,28 +47,30 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # 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. CUDA takes head dims well past 128 on its memory-efficient kernel; ROCm caps them at -# 128 and drops to `math` above that; MPS has no fused SDPA kernel at all. (Masks are *not* a -# reliable discriminator: a 128-wide masked head reports the memory-efficient kernel on gfx1100 too.) -# A working-memory estimate that assumes the fused path is therefore only correct on the build it was -# measured on, which is why the helpers below ask instead of assuming. +# 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 ROCm (gfx1100, torch 2.10) -# 1 4096 512 12.9 13.62 -# 1 8192 512 10.3 10.78 -# 1 16384 512 9.7 9.76 -# 4 4096 128 - 10.16 -# 48 4608 128 - 9.59 +# 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 # -# 14 is an upper bound on every measured point on both. It was 13, which held on CUDA but came in -# 0.62 under ROCm's worst point -- 4.7MB on a GB-scale estimate, so not a memory bug, but it made -# the pinned test red on ROCm and the docstring's "upper bound" claim false. -SDPA_MATH_BYTES_PER_SCORE_ELEMENT = 14 +# 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 @@ -201,9 +203,9 @@ def sdpa_score_matrix_bytes( """Bytes SDPA spends on a materialized score matrix for one attention call, 0 if fused. Add this to a working-memory estimate whose linear term was calibrated on a fused kernel. On - CUDA it is almost always 0; where the fused kernels are missing or reject the shapes -- ROCm - caps the head dim at 128, MPS ships no fused SDPA kernel at all -- it is the dominant term: a - 1536px FLUX.2 VAE decode materializes 36864^2 scores, ~18GB of them. + 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 diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index eb733a3aa86..419323836e3 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -130,15 +130,24 @@ def estimate_vae_working_memory_flux2( That linear term holds only while ``AutoencoderKLFlux2``'s mid-block attention runs through a fused SDPA kernel, which is what CUDA does (verified: the memory-efficient kernel takes the 512-wide head, and measured peak stays linear from 512 to 1536px). A build with no fused kernel - for the shapes -- ROCm caps the head dim at 128 and reports ``math`` for this 512-wide one, MPS - has no fused SDPA kernel at all -- materializes a (pixels/8)^2 score matrix on top of the linear - term: ~3.7GB at 1024px and ~18GB at 1536px. We ask torch which path applies rather than assuming, - so the estimate is right on both. - - Known gap: the linear constants below were calibrated on CUDA, and on ROCm the same decode costs - more than they predict -- measured there at 768px, 3.78GB against 3.45GB estimated. 512px is - covered by the 3GB ``device_working_mem_gb`` floor and 1024px and up have margin, so the shortfall - is confined to the middle of the range and is ~0.33GB. Recalibrating for it needs ROCm hardware. (Unlike the transformer, this attention does not go through + 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 on top + of the linear term: ~4.5GB at 1024px and ~21GB at 1536px. We ask torch which path applies rather + than assuming, so the estimate is right on all of them. + + KNOWN GAP -- the linear constants below are CUDA numbers and are short on AMD, by a lot and by an + amount that is not yet pinned down. Implied constant for decode, fitted by + ``scripts/calibrate_flux2_working_memory.py``: ~2180 on CUDA (which is what 2200 was fitted to), + ~2430 on ROCm/gfx1100, ~4300 on ROCm/gfx1201 at 512-1024px. The last of those means an 8.4GB + decode against a 4.3GB reservation. It is a conv-workspace difference (MIOpen vs cuDNN), not the + attention term -- it shows up on the *fused* path too. + + Raising the constant to cover gfx1201 would double every CUDA reservation for nothing, and the + two AMD cards are 1.8x apart, so there is not yet one number to ship. The gfx1201 run also had + ``MIOPEN_FIND_MODE=2`` and ``PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.7`` set, and + its series is non-monotonic above 1024px (peak *falls* as resolution rises), which is what a GC + threshold does to peak-reserved readings -- so those numbers need a clean re-run before anything + is fitted to them. Tracked rather than guessed at. (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.) diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index 736fcf65c1d..8c339828744 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -106,6 +106,19 @@ 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 @@ -344,6 +357,50 @@ 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) @@ -428,6 +485,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di f"{('yes' if row['covered'] else 'NO'):>8}" ) print("") + _flag_non_monotonic(rows) for operation, shipped in (("decode", 2200), ("encode", 1100)): for force_math in (False, True): ks = [ @@ -567,6 +625,7 @@ def main() -> None: 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] = [] diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 42c235333ad..2ad30600f6b 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -19,6 +19,7 @@ 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, @@ -40,6 +41,13 @@ # `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 @@ -103,14 +111,14 @@ def test_reference_image_tokens_are_counted(self): tokens and cost the same per token, so they must enter the estimate.""" without_refs = _estimate(image_seq_len=4096) with_refs = _estimate(image_seq_len=4096, ref_image_seq_len=12288) - assert with_refs - without_refs == 12288 * int(0.4 * MB) + 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 * int(0.4 * MB) + 512 * PER_TOKEN ) def test_lora_margin_is_added_per_lora(self): @@ -134,22 +142,37 @@ class TestTransformerWidthIsBudgeted: 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. """ - # (hidden size, measured MB per token) - MEASURED_WIDTH_SLOPE = [(3072, 0.2912), (4096, 0.3859), (6144, 0.5547)] + # (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("hidden, measured_mb", MEASURED_WIDTH_SLOPE) - def test_the_estimate_upper_bounds_the_measured_slope(self, hidden, measured_mb): + @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 @@ -199,13 +222,13 @@ def test_batch_and_sequence_are_interchangeable(self): @pytest.mark.parametrize("batch", [2, 3, 4]) def test_each_extra_sample_adds_exactly_its_own_tokens(self, batch): single = _estimate(image_seq_len=4096) - assert _estimate(image_seq_len=4096, batch_size=batch) - single == (batch - 1) * (4096 + 512) * int(0.4 * MB) + 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) * int(0.4 * MB) + (4096 + 12288 + 512) * PER_TOKEN ) def test_the_fixed_base_does_not_scale_with_the_batch(self): @@ -215,7 +238,7 @@ def test_the_fixed_base_does_not_scale_with_the_batch(self): _estimate(image_seq_len=4096, batch_size=b + 1) - _estimate(image_seq_len=4096, batch_size=b) for b in (1, 2, 3) } - assert deltas == {(4096 + 512) * int(0.4 * MB)} + 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 @@ -223,7 +246,7 @@ def test_the_regional_bias_does_not_scale_with_the_batch(self): bias = (4096 + 512) ** 2 * 2 single = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True) double = _estimate(image_seq_len=4096, regional_bias=bias, has_regional_mask=True, batch_size=2) - assert double - single == (4096 + 512) * int(0.4 * MB) + 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).""" @@ -232,7 +255,7 @@ def test_the_score_matrix_scales_with_the_batch(self): single = _estimate(image_seq_len=4096, has_regional_mask=True, device=MATERIALIZING) double = _estimate(image_seq_len=4096, has_regional_mask=True, batch_size=2, device=MATERIALIZING) assert double - single == ( - seq_len * int(0.4 * MB) + KLEIN_9B_HEADS * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + 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): @@ -513,7 +536,7 @@ def test_reference_images_raise_the_reservation(self): """Three 1024x1024 references add 12288 tokens to a 1024x1024 generation's 4096.""" without_refs = self._run(num_ref_tokens=0) with_refs = self._run(num_ref_tokens=12288) - assert with_refs - without_refs == 12288 * int(0.4 * MB) + 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` @@ -523,7 +546,7 @@ def test_the_real_batch_reaches_the_reservation(self): def test_a_batched_run_reserves_more_than_a_single_one(self): single = self._run(num_ref_tokens=0, batch=1) - assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * int(0.4 * MB) + assert self._run(num_ref_tokens=0, batch=2) - single == (64 * 64 + 512) * PER_TOKEN @pytest.mark.parametrize( "variant, hidden", @@ -543,7 +566,7 @@ def test_dev_reserves_half_again_what_klein_9b_does(self): 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 * (int(0.4 * MB * 6144 / 4096) - int(0.4 * MB)) + 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): @@ -563,14 +586,14 @@ def test_a_batched_init_latent_beats_the_batch_1_noise_it_is_blended_with(self): def test_the_blended_batch_is_read_after_the_broadcast(self): single = self._run(num_ref_tokens=0, init_batch=1) - assert self._run(num_ref_tokens=0, init_batch=3) - single == 2 * (64 * 64 + 512) * int(0.4 * MB) + 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) * int(0.4 * MB) + assert double - single == (64 * 64 + 12288 + 512) * PER_TOKEN def _materializing_probe(device_type, device_index, dtype, head_dim, has_attn_mask): @@ -850,22 +873,32 @@ def test_this_build_reports_its_own_dispatch(self): ) assert masked_bytes == (48 * 4608 * 4608 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT if materializes else 0) - # (heads, seq, head_dim, bytes per score element) measured with SDPBackend.MATH forced on - # ROCm 7.1 / torch 2.10 / gfx1100. The CUDA points live in the constant's own comment; the worst - # of either platform is ROCm's 13.62 at the smallest shape, which is why the constant is 14. - MEASURED_BYTES_PER_ELEMENT_ROCM = [ - (1, 4096, 512, 13.62), - (1, 8192, 512, 10.78), - (1, 16384, 512, 9.76), - (4, 4096, 128, 10.16), - (48, 4608, 128, 9.59), + # (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("num_heads, seq_len, head_dim, measured", MEASURED_BYTES_PER_ELEMENT_ROCM) - def test_the_constant_upper_bounds_every_measured_platform(self, num_heads, seq_len, head_dim, measured): - """13 held on CUDA and came in under ROCm's smallest shape. The shortfall was 4.7MB -- noise - against the GB-scale linear terms -- but it made the pinned bound below red on ROCm, and the - constant's "upper bound on every measured point" claim false.""" + @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") @@ -892,7 +925,13 @@ def test_constant_upper_bounds_a_forced_math_forward(self, num_heads, seq_len, h 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" ) - assert estimate <= 2 * measured + # 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): From cd41358620d152177f7232677a7fdd05f00f1c26 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 30 Aug 2026 00:19:36 +0200 Subject: [PATCH 10/11] fix(flux2): fit the VAE constants per convolution backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A clean ROCm run — the earlier one had MIOPEN_FIND_MODE=2, worth a uniform 1.28x, and a HIP allocator GC threshold that clipped the high-resolution points — puts the gfx1201 numbers at 3453/2688 bytes per pixel per element byte against cuDNN's 2185/1072. Flat across 512-1024px on both, so the linear model holds; only the coefficient moves. It is MIOpen's convolution workspaces, not the attention term: identical on the fused path. Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the constant follows the backend, keyed on torch.version.hip rather than the device string (a HIP build reports device.type == "cuda"). The two operations also stop sharing a ratio. "Encoding costs half of decoding" holds on cuDNN (0.49) and not on MIOpen (0.78), so it was a backend property masquerading as an architectural one. MIOPEN_FIND_MODE=2 is deliberately not budgeted for: it is not the default and would tax everyone else. Noted where the constants are defined. --- invokeai/backend/util/vae_working_memory.py | 80 +++++++++++------ scripts/calibrate_flux2_working_memory.py | 13 ++- .../invocations/test_flux2_working_memory.py | 89 ++++++++++++++++++- 3 files changed, 153 insertions(+), 29 deletions(-) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 419323836e3..6d74683f109 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -103,13 +103,50 @@ def estimate_vae_working_memory_flux( # The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the # decoder, on the 8x-downsampled grid. Both are single-head, with the head dim set to the block -# width: 512 for the stock VAE and 384 for the small-decoder variant. The distinction does not -# matter here -- what matters is that both sit far above the 128 head dim ROCm's fused SDPA kernels -# accept, so only the value's side of that limit is load-bearing. +# 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 RX 7900 XTX, torch 2.10 ~2433 - - +# +# 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. The two AMD cards differ by 1.4x, so the MIOpen column is fitted to +# the larger and over-reserves on the smaller; that is the safe direction. +# +# 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": 3500, "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"], @@ -120,12 +157,16 @@ def estimate_vae_working_memory_flux2( ) -> int: """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. - Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE. - Measured on CUDA/bf16 as peak *reserved* memory (the conservative quantity, including allocator - overhead), the implied constants are ~2170 (decode) and ~1070 (encode) bytes per pixel per - element byte, flat across 512-1536px; the constants below round those up and match the FLUX.1 - ones. For reference, decoding 1024x1024 peaks at ~4.3GB and 1536x1536 at ~9.6GB -- far above the - default ``device_working_mem_gb``, which is why this estimate must be passed to the model cache. + 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 @@ -135,19 +176,8 @@ def estimate_vae_working_memory_flux2( of the linear term: ~4.5GB at 1024px and ~21GB at 1536px. We ask torch which path applies rather than assuming, so the estimate is right on all of them. - KNOWN GAP -- the linear constants below are CUDA numbers and are short on AMD, by a lot and by an - amount that is not yet pinned down. Implied constant for decode, fitted by - ``scripts/calibrate_flux2_working_memory.py``: ~2180 on CUDA (which is what 2200 was fitted to), - ~2430 on ROCm/gfx1100, ~4300 on ROCm/gfx1201 at 512-1024px. The last of those means an 8.4GB - decode against a 4.3GB reservation. It is a conv-workspace difference (MIOpen vs cuDNN), not the - attention term -- it shows up on the *fused* path too. - - Raising the constant to cover gfx1201 would double every CUDA reservation for nothing, and the - two AMD cards are 1.8x apart, so there is not yet one number to ship. The gfx1201 run also had - ``MIOPEN_FIND_MODE=2`` and ``PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.7`` set, and - its series is non-monotonic above 1024px (peak *falls* as resolution rises), which is what a GC - threshold does to peak-reserved readings -- so those numbers need a clean re-run before anything - is fitted to them. Tracked rather than guessed at. (Unlike the transformer, this attention does not go through + 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. (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.) @@ -165,8 +195,8 @@ def estimate_vae_working_memory_flux2( param = next(vae.parameters()) element_size = param.element_size() - # Encoding uses ~50% the working memory of decoding. - scaling_constant = 2200 if operation == "decode" else 1100 + 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: @@ -182,7 +212,7 @@ def estimate_vae_working_memory_flux2( working_memory *= batch_size working_memory += sdpa_score_matrix_bytes( - device=device if device is not None else TorchDevice.choose_torch_device(), + 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, diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index 8c339828744..13469bea6fd 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -67,7 +67,10 @@ 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 estimate_vae_working_memory_flux2 +from invokeai.backend.util.vae_working_memory import ( + _flux2_vae_scaling_constant, + estimate_vae_working_memory_flux2, +) GIB = 1024**3 MIB = 1024**2 @@ -452,7 +455,9 @@ def report_sdpa(shapes: list[tuple[int, int, int]], dtype_name: str) -> list[dic def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[dict]: - print("\n=== 3. VAE linear constants (2200 decode / 1100 encode) ===") + 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` backs the score-matrix term out, so it is comparable to those literals; fit the") print("constant 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?") @@ -486,7 +491,9 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di ) print("") _flag_non_monotonic(rows) - for operation, shipped in (("decode", 2200), ("encode", 1100)): + 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"] diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 2ad30600f6b..5c2cc4757b8 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -32,7 +32,10 @@ _torch_sdpa_materializes_score_matrix, sdpa_score_matrix_bytes, ) -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2 +from invokeai.backend.util.vae_working_memory import ( + _FLUX2_VAE_SCALING_CONSTANTS, + estimate_vae_working_memory_flux2, +) MB = 1024**2 GB = 1024**3 @@ -324,6 +327,90 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): 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), + ("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.79, 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(3500 / 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 From 72a628a871d081e6b25aab943c48fcb0dc322a8c Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 30 Aug 2026 10:42:05 +0200 Subject: [PATCH 11/11] fix(flux2): take the larger VAE term, not the sum, and refit MIOpen The W7900 run shows forced math and the fused path measuring identically to three decimals at every resolution, with the total flat-linear in area. The score matrix does not add to the convolution peak: the mid-block sits alone at the 8x-downsampled bottleneck, so the full-resolution feature maps are not live while it runs, and peak reserved is a high-water mark rather than a running total. Measured on cuDNN, forcing math stays *below* the fused path until 1536px and then exceeds it by 2.6GB against the 21.5GB the term prices standalone. Summing reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7. Take the max: 7.0GB. A max model is weakest at the crossover, and one measured point sits there -- a 768px encode with cuDNN's constant and a materializing kernel wants 1.80GB against a 1.35GB max, reproducibly. It is not reachable as a shortfall: the cache floors every reservation at device_working_mem_gb and the whole crossover region is below it. Pinned rather than rounded away. The MIOpen decode constant also goes 3500 -> 3600; the W7900 asks for 3525 at 512px where gfx1201 asks for 3453. Its encode column agrees with gfx1201's to the byte, so this is MIOpen rather than a per-card quirk. --- invokeai/backend/util/vae_working_memory.py | 43 ++++- scripts/calibrate_flux2_working_memory.py | 31 ++-- .../invocations/test_flux2_working_memory.py | 173 ++++++++++++++---- 3 files changed, 192 insertions(+), 55 deletions(-) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 6d74683f109..541ed6efbf2 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -116,7 +116,11 @@ def estimate_vae_working_memory_flux( # 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 RX 7900 XTX, torch 2.10 ~2433 - - +# 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 @@ -125,15 +129,14 @@ def estimate_vae_working_memory_flux( # 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. The two AMD cards differ by 1.4x, so the MIOpen column is fitted to -# the larger and over-reserves on the smaller; that is the safe direction. +# 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": 3500, "encode": 2750}, + "miopen": {"decode": 3600, "encode": 2750}, } @@ -172,12 +175,13 @@ def estimate_vae_working_memory_flux2( 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 on top - of the linear term: ~4.5GB at 1024px and ~21GB at 1536px. We ask torch which path applies rather - than assuming, so the estimate is right on all of them. + 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. (Unlike the transformer, this attention does not go through + 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.) @@ -211,7 +215,7 @@ def estimate_vae_working_memory_flux2( mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) working_memory *= batch_size - working_memory += sdpa_score_matrix_bytes( + 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. @@ -220,7 +224,26 @@ def estimate_vae_working_memory_flux2( seq_len=mid_block_seq_len, ) - return int(working_memory) + # 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( diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index 13469bea6fd..f322e228432 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -265,8 +265,6 @@ def run() -> None: return row estimate = estimate_vae_working_memory_flux2(operation=operation, image_tensor=x, vae=vae, device=device) - # The estimate is linear_term + score_matrix. Back the score term out so the linear constant can - # be fitted on its own -- it is the one the 2200/1100 literals name. 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 @@ -274,15 +272,19 @@ def run() -> None: 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 += score_bytes + estimate = max(estimate, score_bytes) return row | { "reserved_delta": peak, "estimate": estimate, "score_term": score_bytes, - # Only meaningful while the peak actually exceeds the score term. Where it does not, the - # additive (linear + score) model does not decompose on this build and the number is noise. - "implied_linear_constant": ((peak - score_bytes) / (px * px * element_size)) if peak > score_bytes else None, + # 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, } @@ -458,12 +460,12 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di 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` backs the score-matrix term out, so it is comparable to those literals; fit the") - print("constant on the rows whose `math` column matches what this build really does (section 1).") + 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 CUDA the forced-math decode measures *below* the fused one, because the") - print("memory-efficient kernel's workspace is the larger term there. Only a real run on the") + 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}" @@ -501,7 +503,9 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di if r["operation"] == operation and r["force_math"] is force_math and not r.get("oom") - and r["implied_linear_constant"] + # 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 @@ -513,7 +517,10 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di print("\n Points the shipped estimate does NOT cover:") for r in short: gap = (r["reserved_delta"] - r["estimate"]) / GIB - print(f" {r['operation']} {r['px']}px force_math={r['force_math']}: short by {gap:.2f} 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 diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 5c2cc4757b8..9cdbc695c24 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -354,6 +354,10 @@ class TestVaeConstantsFollowTheConvBackend: ("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), @@ -381,7 +385,7 @@ def test_the_encode_ratio_is_not_architectural(self): 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.79, abs=0.03) + 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.""" @@ -408,7 +412,7 @@ def estimate(): cuda = estimate() assert rocm > cuda - assert rocm / cuda == pytest.approx(3500 / 2200, rel=1e-3) + assert rocm / cuda == pytest.approx(3600 / 2200, rel=1e-3) class TestFlux2VaeBatchIsBudgeted: @@ -463,13 +467,15 @@ def test_tiling_bounds_the_tile_not_the_batch(self): assert self._decode_estimate(3, px=1024, tile_size=512) == 3 * single def test_the_score_matrix_scales_with_the_batch(self): - """It is shaped (batch, heads, S, S), so where it is materialized at all it scales too.""" + """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) == linear + score - assert self._decode_estimate(3, device=MATERIALIZING) == 3 * (linear + score) + assert self._decode_estimate(1, device=MATERIALIZING) == max(linear, score) + assert self._decode_estimate(3, device=MATERIALIZING) == 3 * max(linear, score) class TestFlux2VaeInvocationsRequestWorkingMemory: @@ -695,6 +701,14 @@ def _materializing_probe(device_type, device_index, dtype, head_dim, has_attn_ma 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 @@ -707,10 +721,15 @@ def _materializing(): class TestMaterializedScoreMatrixIsBudgeted: """The linear estimates above assume SDPA never builds the O(S^2) score matrix. That is a - property of the *build*, not of FLUX.2: ROCm's fused kernels reject both the VAE's 512-wide - attention head and the dense additive mask regional prompting attaches, and fall back to - `math`. Where that happens the score matrix is the dominant term, so the estimate has to - include it -- otherwise the fix works on CUDA and still OOMs on ROCm. + 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): @@ -732,56 +751,144 @@ def _vae_estimate(self, operation, px, device, tile_size=None): @pytest.mark.parametrize( "operation, px, tokens", [ - ("decode", 1024, 128 * 128), ("decode", 1536, 192 * 192), ("encode", 1024, 128 * 128), ("encode", 1328, 166 * 166), ], ) - def test_vae_estimate_gains_exactly_the_score_matrix(self, operation, px, tokens): - fused = self._vae_estimate(operation, px, device=FUSED) + 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 - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + 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, a 1024px decode has to come out ~3.5GB heavier than the - linear term. Reporting those devices as fused -- as this PR first did -- is what let the - decode be admitted to a card that could not run it.""" + 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", 1024, device=FUSED) - fused = self._vae_estimate("decode", 1024, device=FUSED) + materializing = self._vae_estimate("decode", 1536, device=FUSED) + fused = self._vae_estimate("decode", 1536, device=FUSED) - tokens = 128 * 128 - assert materializing - fused == tokens * tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert materializing - fused > 3 * GB + 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): - """The reviewer's case: a 1536px decode is ~9.6GB of linear activations on CUDA, and more - than that again in scores where SDPA has to materialize them. An estimate that omits the - term is not merely tight, it is wrong by a factor of three.""" + """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 > 25 * GB + assert materializing > 20 * GB assert materializing > 2 * self._vae_estimate("decode", 1536, device=FUSED) - def test_tiled_vae_score_matrix_is_bounded_by_the_tile(self): - """Tiling already bounds the linear term; it must bound the quadratic one too, or the - reference-image encode would reserve as if it ran untiled.""" + 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 - tile_tokens = (512 // 8) ** 2 - assert estimates[0] - self._vae_estimate("encode", 1024, device=FUSED, tile_size=512) == ( - tile_tokens * tile_tokens * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - ) - # Tiling turns a 62GB reservation at the 2024px reference cap into under 1GB. + assert estimates[0] == 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