diff --git a/docs/src/content/docs/configuration/Optimization/low-vram-mode.mdx b/docs/src/content/docs/configuration/Optimization/low-vram-mode.mdx index a0ce7b48f0e..dd4345b2979 100644 --- a/docs/src/content/docs/configuration/Optimization/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/Optimization/low-vram-mode.mdx @@ -39,6 +39,7 @@ Low-VRAM mode and related workload-specific optimizations include: - Dynamic RAM and VRAM cache sizes (`max_cache_ram_gb`, `max_cache_vram_gb`) - Working memory (`device_working_mem_gb`) - Keeping a RAM weight copy (`keep_ram_copy_of_weights`) +- Tiled VAE decode and encode (`auto_tiled_decode`, `force_tiled_decode`) - Wan video memory optimization (`wan_memory_optimization`) - PiD decode activation chunking (`pid_memory_optimization`) @@ -84,6 +85,8 @@ The PyTorch CUDA allocator's behavior can be configured using the `pytorch_cuda_ pytorch_cuda_alloc_conf: "backend:cudaMallocAsync" ``` +When the setting is left unset, PyTorch's defaults apply, with one exception: with a ROCm build of PyTorch on Windows, Invoke uses `expandable_segments:True`. There, Windows does not fail an allocation it cannot fit into contiguous VRAM. It moves the allocation into much slower system memory instead, and fragmented VRAM makes that happen even when enough VRAM is free. Setting `pytorch_cuda_alloc_conf` yourself, or any of the `PYTORCH_ALLOC_CONF` / `PYTORCH_CUDA_ALLOC_CONF` / `PYTORCH_HIP_ALLOC_CONF` environment variables, replaces this default. + A more complete explanation of the available configuration options is [here](https://pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf). ### Dynamic RAM and VRAM cache sizes @@ -112,7 +115,8 @@ But, if your GPU has enough VRAM to hold models fully, you might get a perf boos max_cache_ram_gb: 28 # The default max cache VRAM size is adjusted dynamically based on the amount of available VRAM (taking into -# consideration the VRAM used by other processes). +# consideration the VRAM used by other processes; with a ROCm build on Windows, the video-memory budget Windows grants +# Invoke, beyond which it would move allocations into system memory). # You can override the default value by setting `max_cache_vram_gb`. # CAUTION: Most users should not manually set this value. See warning below. max_cache_vram_gb: 16 @@ -168,6 +172,14 @@ keep_ram_copy_of_weights: false Without tiling, the memory a VAE needs grows with the square of the image size. For FLUX.1, an untiled 3072² encode peaks at about 18.8 GiB on an RTX 4090, and a tiled one at 0.8 GiB, while tiled memory stays roughly flat as the image grows. Tiling costs some speed on smaller images. +Invoke tiles large decodes by itself: when an untiled decode would claim most of the memory the GPU keeps resident, it decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Krea-2, Anima and Wan decodes. Because a tiled decode is not pixel-identical to an untiled one, this only kicks in for images too large for the card: for Z-Image on a 16GB GPU, for example, from about 1850x1850 on Nvidia and from about 1500x1500 on AMD, whose convolution library needs more memory. On a ROCm build under Windows the comparison uses the video-memory budget Windows grants Invoke, which is lower than the card's total and lower again while another program uses the GPU. Anima tiles a little earlier, because its untiled decode pushes the Anima model itself out of VRAM on 8GB cards. + +To always decode untiled unless you ask for tiling, turn this off. A decode too large for the card then runs out of memory: the FLUX.1 and Z-Image decodes retry in tiles, the others fail, and on Windows the driver may move the decode into system memory and run it very slowly instead. Anima is not affected by this setting: there, tiling a large decode is what keeps the Anima model in VRAM, and turning it off would only make those decodes slower. + +```yaml +auto_tiled_decode: false +``` + - **Always tile:** set `force_tiled_decode: true` in `invokeai.yaml`. It covers the VAE decode and encode for SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image, the Anima encode, and the Krea-2 style reference. Other model families (for example SD3, FLUX.2, Wan and LTX-2) ignore it. - **Per node:** the SD, FLUX.1, Z-Image and Qwen-Image **Latents to Image** and **Image to Latents** nodes, and **Image to Latents - Anima**, have **Tiled** and **Tile Size** fields for workflows. The [Upscale panel](/users-guide/image-generation/upscale/) turns tiling on by itself. - **Automatic retry:** if a FLUX.1, Z-Image or Anima decode, or a FLUX.1 encode, runs out of memory, InvokeAI retries it once with tiling and logs that it did. @@ -211,6 +223,15 @@ It is strongly suggested to disable this feature: Unfortunately, the Nvidia implementation is not optimized for applications like Invoke and does more harm than good. ::: +### AMD GPUs on Windows (ROCm) + +Invoke does not officially support ROCm on Windows (see the [System Requirements](/start-here/system-requirements/)). If you run a ROCm build of PyTorch there anyway, Windows behaves as if the Nvidia sysmem fallback were always on: an allocation that does not fit is moved into much slower system memory instead of failing, and it stays there. Task Manager shows it as shared GPU memory. Invoke works around this where it can: + +- The model cache plans against the video-memory budget Windows grants Invoke, not against the card's total VRAM. +- The PyTorch allocator defaults to `expandable_segments:True` (see [PyTorch CUDA allocator config](#pytorch-cuda-allocator-config)), so fragmented VRAM does not force allocations out. +- Decodes too large for the card are tiled before they start (see [Tiled VAE decode and encode](#tiled-vae-decode-and-encode)); waiting for an out-of-memory error does not work there, because none comes. +- When part of Invoke's GPU memory stays in system memory across two generations, the log says so. Other programs using the same GPU, like games or browsers, can push Invoke's memory out as well; close them. If it keeps happening, lower the image size, set `max_cache_vram_gb` lower or `device_working_mem_gb` higher, and restart Invoke to get the memory back into VRAM. + ## Troubleshooting ### Windows page file diff --git a/docs/src/content/docs/start-here/system-requirements.mdx b/docs/src/content/docs/start-here/system-requirements.mdx index 903d213b495..f92cebbb254 100644 --- a/docs/src/content/docs/start-here/system-requirements.mdx +++ b/docs/src/content/docs/start-here/system-requirements.mdx @@ -161,10 +161,18 @@ If, after restarting the app, this doesn't improve your performance, either rest On ROCm builds Invoke routes attention whose head dimension exceeds 256 to PyTorch's math kernel. The fused kernels return wrong output there on torch 2.13 / ROCm 7.2 (measured on a W7900), which showed up as black or noisy images from the image VAEs - their mid-block attention is one 512-wide - head. The math kernel materializes the attention matrix, so large untiled VAE decodes reserve - more VRAM (about 4 GB at 1024px, 21 GB at 1536px) and may push other models out of the cache; - tiled decoding keeps that bounded. Export `INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM` to raise the - threshold on a build whose fused kernels are known to be correct. + head. The math kernel builds the whole attention matrix (about 4 GB for a 1024px VAE decode, 21 GB + at 1536px), so Invoke computes it in pieces of at most 1 GB. The result can differ from a single + math call in the last bit of a bf16 value, and each such call is roughly 10% slower. Export + `INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM` to raise the threshold on a build whose fused kernels are + known to be correct. +::: + +:::note[Attention without a fused kernel] + Where PyTorch has no fused attention kernel for a call, it runs on the math kernel as well, and + Invoke computes that in pieces of at most 1 GB too. That is what makes large models and + resolutions usable on such builds: Z-Image at 1024px would otherwise build 8 GB of attention + scores per call, Krea-2 almost 13 GB. ::: :::note[Idle process pegging one CPU core] diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index e7412b3d059..8d7f930311f 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -584,7 +584,7 @@ { "category": "CACHE", "default": null, - "description": "Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.", + "description": "Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.", "env_var": "INVOKEAI_PYTORCH_CUDA_ALLOC_CONF", "literal_values": [], "name": "pytorch_cuda_alloc_conf", @@ -725,6 +725,17 @@ "type": "", "validation": {} }, + { + "category": "GENERATION", + "default": true, + "description": "Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.", + "env_var": "INVOKEAI_AUTO_TILED_DECODE", + "literal_values": [], + "name": "auto_tiled_decode", + "required": false, + "type": "", + "validation": {} + }, { "category": "GENERATION", "default": 1, diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py index 3b6c3c094f3..d9aca4e851b 100644 --- a/invokeai/app/invocations/krea2/krea2_denoise.py +++ b/invokeai/app/invocations/krea2/krea2_denoise.py @@ -60,6 +60,7 @@ from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Krea2ConditioningInfo from invokeai.backend.util import sdpa_scope +from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger @@ -553,19 +554,30 @@ def _run_diffusion(self, context: InvocationContext): # headroom. Without this hint the cache reserves only the small default working memory and places # the model before LoRA patches are applied, so a model+LoRA combination that just fits the base # forward OOMs once the LoRA's extra activations are added. + regional_attention_mask_bytes = self._regional_attention_mask_bytes( + pos_extension, neg_extension, inference_dtype + ) estimated_working_memory = self._estimate_working_memory( image_seq_len=image_seq_len, positive_text_seq_len=pos_prompt_embeds.shape[1], negative_text_seq_len=neg_prompt_embeds.shape[1] if neg_prompt_embeds is not None else None, do_cfg=do_cfg, num_loras=len(self.transformer.loras), - regional_attention_mask_bytes=self._regional_attention_mask_bytes( - pos_extension, neg_extension, inference_dtype - ), + regional_attention_mask_bytes=regional_attention_mask_bytes, style_reference_kv_bytes=( style_extension.kv_cache_bytes(inference_dtype) if style_extension is not None else 0 ), ) + # The activation estimate assumes a fused kernel. Where the build has none for this attention -- ROCm on + # Windows, MPS -- the score matrix is built as well (chunked on ROCm, see `install_rocm_sdpa_guard`). + estimated_working_memory += self._attention_score_bytes( + transformer_info.model, + seq_len=image_seq_len + + max(pos_prompt_embeds.shape[1], neg_prompt_embeds.shape[1] if neg_prompt_embeds is not None else 0), + has_attn_mask=regional_attention_mask_bytes > 0, + device=device, + dtype=inference_dtype, + ) # An int8_tensorwise build stores its linears quantized and materializes the dequantized, # derotated weight per forward call. That transient is alive alongside the activations above, # so it is added rather than compared -- and it is invisible to an activation estimate, which @@ -801,6 +813,32 @@ def _regional_attention_mask_bytes( peak_attention_bias_bytes *= attention_dtype_bytes return final_mask_bytes + peak_build_scratch_bytes + peak_attention_bias_bytes + @staticmethod + def _attention_score_bytes( + transformer: object, seq_len: int, has_attn_mask: bool, device: torch.device, dtype: torch.dtype + ) -> int: + """The score matrix the main blocks' joint attention materializes, where it does (0 on a fused kernel). + + One call is alive at a time: the blocks run in sequence, and so do the conditional and unconditional passes. + The K/V heads are expanded to the query heads before the call on any build without a fused grouped-query + kernel (`backend/krea2/attention.py`), so the query head count is the one that counts. Heads come from the + loaded model's config. The probe answers for torch's default policy, so forcing the math kernel through the + `INVOKE_KREA2_SDPA_BACKEND` debugging switch is not priced. + """ + config = getattr(transformer, "config", None) + num_heads = getattr(config, "num_attention_heads", None) + head_dim = getattr(config, "attention_head_dim", None) + if not isinstance(num_heads, int) or not isinstance(head_dim, int): + return 0 + return sdpa_score_matrix_bytes( + device=device, + dtype=dtype, + num_heads=num_heads, + head_dim=head_dim, + seq_len=seq_len, + has_attn_mask=has_attn_mask, + ) + def _estimate_working_memory( self, image_seq_len: int, diff --git a/invokeai/app/invocations/vae/anima_latents_to_image.py b/invokeai/app/invocations/vae/anima_latents_to_image.py index 9cc3d4735e7..d6243d888a5 100644 --- a/invokeai/app/invocations/vae/anima_latents_to_image.py +++ b/invokeai/app/invocations/vae/anima_latents_to_image.py @@ -44,6 +44,7 @@ ) from invokeai.backend.util.vae_working_memory import ( estimate_vae_working_memory_anima, + should_pretile_vae_decode, ) # Tile geometry for tiled Wan VAE decode. 512px tiles with a 384px stride (128px blended @@ -54,6 +55,13 @@ ANIMA_VAE_TILE_SIZE = 512 ANIMA_VAE_TILE_STRIDE = 384 +# Lower than the other decodes' `VAE_PRETILE_VRAM_FRACTION`, on purpose. A full 1024x1024 Wan VAE decode reserves ~6GB +# of working memory; on small-VRAM GPUs that evicts the (~4GB) Anima transformer from the model cache and thrashes the +# allocator near the VRAM ceiling (decode times of 7s+ observed on 8GB, vs ~1s tiled with the transformer left +# resident). Above this share of the device a tiled decode wins; below it a single pass is faster (~0.65s vs ~1.05s at +# 1024x1024) and exact. +ANIMA_PRETILE_VRAM_FRACTION = 0.7 + @invocation( "anima_l2i", @@ -69,25 +77,6 @@ class AnimaLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard): latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection) vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) - @staticmethod - def _use_tiled_decode(device: torch.device, full_decode_working_memory: int) -> bool: - """Decide whether to decode in tiles. - - A full 1024x1024 Wan VAE decode reserves ~6GB of working memory. On small-VRAM - GPUs this evicts the (~4GB) Anima transformer from the model cache and thrashes - the allocator near the VRAM ceiling (decode times of 7s+ observed on 8GB, vs - ~1s tiled with the transformer left resident). Tile when the full-decode working - memory would consume most of the device, otherwise a single-pass decode is - faster (~0.65s vs ~1.05s at 1024x1024) and exact. - """ - if device.type == "cuda": - total_vram = torch.cuda.get_device_properties(device).total_memory - elif device.type == "xpu": - total_vram = torch.xpu.get_device_properties(device).total_memory - else: - return False - return full_decode_working_memory > 0.7 * total_vram - @torch.no_grad() def invoke(self, context: InvocationContext) -> ImageOutput: latents = context.tensors.load(self.latents.latents_name) @@ -117,7 +106,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae=vae_info.model, tile_size=None, ) - use_tiling = self._use_tiled_decode(TorchDevice.choose_torch_device(), full_decode_working_memory) + # Not gated on `auto_tiled_decode`: this rule predates that setting and is a speed optimization, not a way + # around an out-of-memory error. Turning the setting off here would hand an 8GB card the 7s decode above + # instead of restoring anything. + use_tiling = should_pretile_vae_decode( + vae_info.compute_device, full_decode_working_memory, ANIMA_PRETILE_VRAM_FRACTION + ) estimated_working_memory = estimate_vae_working_memory_anima( operation="decode", image_tensor=latents, diff --git a/invokeai/app/invocations/vae/cogview4_image_to_latents.py b/invokeai/app/invocations/vae/cogview4_image_to_latents.py index facbc38dd42..09cde601d32 100644 --- a/invokeai/app/invocations/vae/cogview4_image_to_latents.py +++ b/invokeai/app/invocations/vae/cogview4_image_to_latents.py @@ -41,7 +41,7 @@ class CogView4ImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard): def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor: assert isinstance(vae_info.model, AutoencoderKL) estimated_working_memory = estimate_vae_working_memory_cogview4( - operation="encode", image_tensor=image_tensor, vae=vae_info.model + operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) diff --git a/invokeai/app/invocations/vae/cogview4_latents_to_image.py b/invokeai/app/invocations/vae/cogview4_latents_to_image.py index 5d5f7fbb8cc..99fb2f21261 100644 --- a/invokeai/app/invocations/vae/cogview4_latents_to_image.py +++ b/invokeai/app/invocations/vae/cogview4_latents_to_image.py @@ -46,7 +46,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae_info = context.models.load(self.vae.vae) assert isinstance(vae_info.model, (AutoencoderKL)) estimated_working_memory = estimate_vae_working_memory_cogview4( - operation="decode", image_tensor=latents, vae=vae_info.model + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), diff --git a/invokeai/app/invocations/vae/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index d365f88180d..170162e8077 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -19,7 +19,10 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import MIN_TILE_SAMPLE_SIZE, scoped_vae_tiling -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_flux, + should_pretile_vae_decode, +) @invocation( @@ -61,12 +64,33 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents # exercises on purpose. The encode node *is* guarded, because it used to reject every # `AutoencoderKL` and the swap is what widened it. assert isinstance(vae_info.model, AutoencoderKL) - use_tiling = self.tiled or context.config.get().force_tiled_decode + config = context.config.get() + use_tiling = self.tiled or config.force_tiled_decode tile_size = self.tile_size if use_tiling else None estimated_working_memory = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=vae_info.model, tile_size=tile_size + operation="decode", + image_tensor=latents, + vae=vae_info.model, + tile_size=tile_size, + device=vae_info.compute_device, ) + # Where the untiled decode would claim most of what the device keeps resident, tile before it runs: on Windows + # an allocation that does not fit is paged rather than refused, so the retry below never fires. + if ( + not use_tiling + and config.auto_tiled_decode + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + ): + use_tiling = True + tile_size = self.tile_size + estimated_working_memory = estimate_vae_working_memory_flux( + operation="decode", + image_tensor=latents, + vae=vae_info.model, + tile_size=tile_size, + device=vae_info.compute_device, + ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) diff --git a/invokeai/app/invocations/vae/flux_vae_encode.py b/invokeai/app/invocations/vae/flux_vae_encode.py index 7b77f3fc9d1..5984fa737ed 100644 --- a/invokeai/app/invocations/vae/flux_vae_encode.py +++ b/invokeai/app/invocations/vae/flux_vae_encode.py @@ -80,16 +80,22 @@ def vae_encode( ) effective_tile_size = tile_size if tiled else None estimated_working_memory = estimate_vae_working_memory_flux( - operation="encode", image_tensor=image_tensor, vae=vae_info.model, tile_size=effective_tile_size + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + tile_size=effective_tile_size, + device=vae_info.compute_device, ) # Seeded here rather than by the caller on purpose: an OOM retry calls this again, and a # generator held outside would already have been drawn from by the attempt that failed, so # the retry would produce a different latent than a directly tiled run. - generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0) + generator = torch.Generator(device=vae_info.compute_device).manual_seed(0) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype - image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + # The VAE's own device, as in the decode nodes: a cpu_only VAE (or another worker's GPU) + # is not the session device, and encoding there would fail or run on the wrong card (#9373). + image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae_dtype) with scoped_vae_tiling(vae, effective_tile_size): latents = vae.encode(image_tensor).latent_dist.sample(generator) # `AutoencoderKL` leaves the scaling to the caller; `shift_factor` is None on a diff --git a/invokeai/app/invocations/vae/image_to_latents.py b/invokeai/app/invocations/vae/image_to_latents.py index bc8d1403606..71d3951d79c 100644 --- a/invokeai/app/invocations/vae/image_to_latents.py +++ b/invokeai/app/invocations/vae/image_to_latents.py @@ -81,6 +81,7 @@ def vae_encode( vae=vae_info.model, tile_size=tile_size if tiled else None, fp32=upcast, + device=vae_info.compute_device, ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" diff --git a/invokeai/app/invocations/vae/latents_to_image.py b/invokeai/app/invocations/vae/latents_to_image.py index 482c808b350..4fbc87b8659 100644 --- a/invokeai/app/invocations/vae/latents_to_image.py +++ b/invokeai/app/invocations/vae/latents_to_image.py @@ -61,6 +61,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae=vae_info.model, tile_size=self.tile_size if use_tiling else None, fp32=self.fp32, + device=vae_info.compute_device, ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), diff --git a/invokeai/app/invocations/vae/qwen_image_image_to_latents.py b/invokeai/app/invocations/vae/qwen_image_image_to_latents.py index ba8c9353a75..6331494c0f3 100644 --- a/invokeai/app/invocations/vae/qwen_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/qwen_image_image_to_latents.py @@ -75,6 +75,7 @@ def vae_encode( image_tensor=image_tensor, vae=vae_info.model, tile_size=effective_tile_size, + device=vae_info.compute_device, ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): # Reinterpret an Anima-classified Wan VAE as AutoencoderKLQwenImage (identical weights). diff --git a/invokeai/app/invocations/vae/qwen_image_latents_to_image.py b/invokeai/app/invocations/vae/qwen_image_latents_to_image.py index 2e1c648bfa2..88f709dc0cc 100644 --- a/invokeai/app/invocations/vae/qwen_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/qwen_image_latents_to_image.py @@ -22,7 +22,11 @@ patch_qwen_image_vae_tiling, resolve_qwen_image_vae_tile_size, ) -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_qwen_image, + qwen_image_untiled_decode_peak_bytes, + should_pretile_vae_decode, +) @invocation( @@ -53,7 +57,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latents = context.tensors.load(self.latents.latents_name) vae_info = context.models.load(self.vae.vae) - tiled = self.tiled or context.config.get().force_tiled_decode + config = context.config.get() + tiled = self.tiled or config.force_tiled_decode # Resolve tile_size=0 ("model default") before estimating, so the memory the cache reserves # matches the tiles the VAE will actually use. Without this the estimate stays at the # full-frame figure (~21 GB at 2560x1440 on CUDA) and tiling frees nothing: the VAE is @@ -71,7 +76,27 @@ def invoke(self, context: InvocationContext) -> ImageOutput: image_tensor=latents, vae=vae_info.model, tile_size=effective_tile_size, + device=vae_info.compute_device, ) + # The tiling decision is priced from the measured peak, not from the reservation above: that one carries the + # headroom and cross-card conservatism a reservation needs, and tiling is not pixel-identical. + if ( + not tiled + and config.auto_tiled_decode + and should_pretile_vae_decode( + vae_info.compute_device, + qwen_image_untiled_decode_peak_bytes(latents, vae_info.model, vae_info.compute_device), + ) + ): + tiled = True + effective_tile_size = resolve_qwen_image_vae_tile_size(self.tile_size) + estimated_working_memory = estimate_vae_working_memory_qwen_image( + operation="decode", + image_tensor=latents, + vae=vae_info.model, + tile_size=effective_tile_size, + device=vae_info.compute_device, + ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): context.util.signal_progress("Running VAE") # A native-layout qwen_image_vae single file is classified with the Anima base and loaded diff --git a/invokeai/app/invocations/vae/sd3_image_to_latents.py b/invokeai/app/invocations/vae/sd3_image_to_latents.py index 9af641d8bcf..5d558eeed52 100644 --- a/invokeai/app/invocations/vae/sd3_image_to_latents.py +++ b/invokeai/app/invocations/vae/sd3_image_to_latents.py @@ -37,7 +37,7 @@ class SD3ImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard): def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor: assert isinstance(vae_info.model, AutoencoderKL) estimated_working_memory = estimate_vae_working_memory_sd3( - operation="encode", image_tensor=image_tensor, vae=vae_info.model + operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) diff --git a/invokeai/app/invocations/vae/sd3_latents_to_image.py b/invokeai/app/invocations/vae/sd3_latents_to_image.py index 9d7eab5d9dc..442ea0a83b6 100644 --- a/invokeai/app/invocations/vae/sd3_latents_to_image.py +++ b/invokeai/app/invocations/vae/sd3_latents_to_image.py @@ -48,7 +48,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae_info = context.models.load(self.vae.vae) assert isinstance(vae_info.model, (AutoencoderKL)) estimated_working_memory = estimate_vae_working_memory_sd3( - operation="decode", image_tensor=latents, vae=vae_info.model + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), diff --git a/invokeai/app/invocations/vae/wan_latents_to_video.py b/invokeai/app/invocations/vae/wan_latents_to_video.py index ae1b0754a2f..b63ad36694e 100644 --- a/invokeai/app/invocations/vae/wan_latents_to_video.py +++ b/invokeai/app/invocations/vae/wan_latents_to_video.py @@ -39,7 +39,10 @@ patch_qwen_image_vae_tiling, resolve_qwen_image_vae_tile_size, ) -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_wan +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_wan, + should_pretile_vae_decode, +) from invokeai.backend.wan.vae_decode import iter_wan_vae_decode_chunks @@ -133,34 +136,28 @@ def invoke(self, context: InvocationContext) -> VideoOutput: streaming=optimize_memory, ) # Long/high-res clips can need a working set no card fits. When the full-frame - # estimate exceeds the execution device's total VRAM, fall back to spatial tiling - # and budget for the tiled working set instead. (A cpu_only VAE runs in system - # RAM, where the working set is not the constraint.) - use_tiling = False + # estimate would claim most of what the VAE's device keeps resident, fall back to spatial + # tiling and budget for the tiled working set instead. (A cpu_only VAE runs in system RAM, + # where the working set is not the constraint; the helper never tiles off-GPU.) + use_tiling = context.config.get().auto_tiled_decode and should_pretile_vae_decode( + vae_info.compute_device, estimated_working_memory + ) tile_size: int | None = None - if not getattr(vae_info.config, "cpu_only", None): - exec_device = TorchDevice.choose_torch_device() - total_vram: int | None = None - if exec_device.type == "cuda": - total_vram = torch.cuda.get_device_properties(exec_device).total_memory - elif exec_device.type == "xpu": - total_vram = torch.xpu.get_device_properties(exec_device).total_memory - if total_vram is not None and estimated_working_memory > 0.9 * total_vram: - use_tiling = True - # Resolved against the module constant, not read off `vae_info.model`: that is the - # cache's shared instance, so its current geometry is whatever the last invocation - # left rather than what this decode is about to ask for. The sentinel resolves to - # 256/192, which is also this VAE class's own default. - tile_size = resolve_qwen_image_vae_tile_size(0) - estimated_working_memory = estimate_vae_working_memory_wan( - operation="decode", - vae=vae_info.model, - pixel_height=h_pixel, - pixel_width=w_pixel, - pixel_frames=t_pixel, - tile_size=tile_size, - streaming=False, - ) + if use_tiling: + # Resolved against the module constant, not read off `vae_info.model`: that is the + # cache's shared instance, so its current geometry is whatever the last invocation + # left rather than what this decode is about to ask for. The sentinel resolves to + # 256/192, which is also this VAE class's own default. + tile_size = resolve_qwen_image_vae_tile_size(0) + estimated_working_memory = estimate_vae_working_memory_wan( + operation="decode", + vae=vae_info.model, + pixel_height=h_pixel, + pixel_width=w_pixel, + pixel_frames=t_pixel, + tile_size=tile_size, + streaming=False, + ) tmp = tempfile.NamedTemporaryFile(prefix="invokeai_wan_video_", suffix=".mp4", delete=False) tmp.close() diff --git a/invokeai/app/invocations/vae/z_image_image_to_latents.py b/invokeai/app/invocations/vae/z_image_image_to_latents.py index 21e7d91e735..d3aaa13513a 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -16,7 +16,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_tiling_scope import MIN_TILE_SAMPLE_SIZE, scoped_vae_tiling from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux @@ -70,6 +69,7 @@ def vae_encode( image_tensor=image_tensor, vae=vae_info.model, tile_size=effective_tile_size, + device=vae_info.compute_device, ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): @@ -79,7 +79,9 @@ def vae_encode( ) vae_dtype = next(iter(vae.parameters())).dtype - image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + # The VAE's own device, as in the decode nodes: a cpu_only VAE (or another worker's GPU) + # is not the session device, and encoding there would fail or run on the wrong card (#9373). + image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae_dtype) # The VAE belongs to the model cache and is shared with the decode node and with every # other node that reaches this class. Tiling is a property of this one encode, so the diff --git a/invokeai/app/invocations/vae/z_image_latents_to_image.py b/invokeai/app/invocations/vae/z_image_latents_to_image.py index eae50c57958..2c5055402f3 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -19,7 +19,10 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import MIN_TILE_SAMPLE_SIZE, scoped_vae_tiling -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_flux, + should_pretile_vae_decode, +) @invocation( @@ -57,7 +60,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: "Ensure you are using a compatible VAE model." ) - use_tiling = self.tiled or context.config.get().force_tiled_decode + config = context.config.get() + use_tiling = self.tiled or config.force_tiled_decode # Estimate working memory needed for VAE decode estimated_working_memory = estimate_vae_working_memory_flux( @@ -65,7 +69,21 @@ def invoke(self, context: InvocationContext) -> ImageOutput: image_tensor=latents, vae=vae_info.model, tile_size=self.tile_size if use_tiling else None, + device=vae_info.compute_device, ) + if ( + not use_tiling + and config.auto_tiled_decode + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + ): + use_tiling = True + estimated_working_memory = estimate_vae_working_memory_flux( + operation="decode", + image_tensor=latents, + vae=vae_info.model, + tile_size=self.tile_size, + device=vae_info.compute_device, + ) # Seamless applies to every VAE this node accepts. It used to be skipped for the FLUX # autoencoder, which had no circular-padding path -- that class is gone, and a standalone diff --git a/invokeai/app/run_app.py b/invokeai/app/run_app.py index 759c246b14c..c39573a352c 100644 --- a/invokeai/app/run_app.py +++ b/invokeai/app/run_app.py @@ -58,7 +58,10 @@ def run_app() -> None: import uvicorn from invokeai.app.services.config.config_default import get_config - from invokeai.app.util.torch_cuda_allocator import configure_torch_cuda_allocator + from invokeai.app.util.torch_cuda_allocator import ( + apply_rocm_windows_allocator_default, + configure_torch_cuda_allocator, + ) from invokeai.backend.util.logging import InvokeAILogger # Load config. @@ -70,6 +73,8 @@ def run_app() -> None: # NOTE: It is important that this happens before torch is imported. if app_config.pytorch_cuda_alloc_conf: configure_torch_cuda_allocator(app_config.pytorch_cuda_alloc_conf, logger) + else: + apply_rocm_windows_allocator_default(logger) # This import must happen after configure_torch_cuda_allocator() is called, because the module imports torch. from invokeai.app.invocations.baseinvocation import InvocationRegistry diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 39483cc31ab..19f08d36480 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -116,7 +116,7 @@ class InvokeAIAppConfig(BaseSettings): ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable. - pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. + pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, "expandable_segments:True" is used. device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number) precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32` sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. @@ -125,6 +125,7 @@ class InvokeAIAppConfig(BaseSettings): attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8` force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it. + auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there. pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting. max_queue_size: Maximum number of items in the session queue. session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin` @@ -234,7 +235,7 @@ class InvokeAIAppConfig(BaseSettings): lazy_offload: bool = Field(default=True, description="DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.") # PyTorch Memory Allocator - pytorch_cuda_alloc_conf: Optional[str] = Field(default=None, description="Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.") + pytorch_cuda_alloc_conf: Optional[str] = Field(default=None, description="Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.") # DEVICE device: str = Field(default="auto", description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)", pattern=r"^(auto|cpu|mps|xpu(:\d+)?|cuda(:\d+)?)$") @@ -249,6 +250,7 @@ class InvokeAIAppConfig(BaseSettings): attention_type: ATTENTION_TYPE = Field(default="auto", description="Attention type.") attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto", description='Slice size, valid when attention_type=="sliced".') force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.") + auto_tiled_decode: bool = Field(default=True, description="Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.") pil_compress_level: int = Field(default=1, description="The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.") max_queue_size: int = Field(default=10000, gt=0, description="Maximum number of items in the session queue.") session_queue_mode: SESSION_QUEUE_MODE = Field(default="round_robin", description="Session queue mode. Use 'FIFO' for strict first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In round-robin mode, priority orders each user's own jobs, but the user rotation takes precedence: one user's high-priority job does not preempt another user's turn. In single-user mode, jobs are served in submission order either way — except that on multi-GPU systems the default 'round_robin' allows same-priority jobs to be reordered slightly so a freed GPU prefers jobs whose models it already has loaded. Set 'FIFO' to disable that reordering and enforce strict submission order.") diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 140238896cf..05a045882b3 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -50,6 +50,11 @@ idle_device_borrowed, ) from invokeai.backend.util.devices import TorchDevice, disable_conv_benchmark_empty_cache +from invokeai.backend.util.wddm import paged_bytes + +# Paged-out VRAM worth a warning. Weight uploads page a few hundred MB for an instant (0.25-0.36 GiB measured on an +# RX 9060 XT), so the threshold sits above that; a session that ends with more than this paged has lost real VRAM. +_VRAM_PAGING_WARNING_BYTES = 512 * 2**20 # A failed owner lookup is retried before the item is refused, so that a transient error # — a busy-timeout on the shared SQLite connection under multi-GPU write contention, say — @@ -659,6 +664,9 @@ def __init__(self, device: Optional[torch.device], runner: SessionRunnerBase) -> self.cancel_event = ThreadEvent() self.queue_item: Optional[SessionQueueItem] = None self.thread: Optional[Thread] = None + # Paged-out VRAM at the previous session's end and at the last warning; only this worker's thread uses them. + self.paging_last_bytes = 0 + self.paging_warned_bytes = 0 @property def label(self) -> str: @@ -948,6 +956,38 @@ def _release_vram_after_session(self, worker: _SessionWorker) -> None: f"Could not release cached VRAM after the session on {worker.label}", exc_info=True ) + def _warn_if_vram_paged(self, worker: _SessionWorker) -> None: + """Say so when Windows keeps part of this worker's GPU memory in shared system memory across sessions. + + On a ROCm build under Windows an allocation that does not fit the video-memory budget, or finds no contiguous + VRAM, is placed in system memory instead of failing -- and another program on the same GPU can push this + process's memory out the same way. Every generation that touches it slows down, with nothing else in the log to + explain why. Only what was paged at the end of two consecutive sessions counts: an overflow at the end of a + decode can return to VRAM within seconds. Warns once per episode: again only after that grew by the threshold, + re-armed once a reading drops below it. `paged_bytes` answers None everywhere else, which keeps this silent. + """ + try: + paged = paged_bytes(worker.device or TorchDevice.choose_torch_device()) + except Exception: + self._invoker.services.logger.debug(f"Could not read paged VRAM on {worker.label}", exc_info=True) + return + if paged is None: + return + sustained = min(paged, worker.paging_last_bytes) + worker.paging_last_bytes = paged + if paged < _VRAM_PAGING_WARNING_BYTES: + worker.paging_warned_bytes = 0 + return + if sustained < worker.paging_warned_bytes + _VRAM_PAGING_WARNING_BYTES: + return + worker.paging_warned_bytes = sustained + self._invoker.services.logger.warning( + f"Windows kept {sustained / 2**30:.1f} GiB of Invoke's GPU memory on {worker.label} in shared system memory " + "after the last two generations, which makes them much slower. Close other programs that use this GPU. If " + "it keeps happening, lower the image size, set max_cache_vram_gb lower or device_working_mem_gb higher, " + "and restart Invoke to get the memory back into VRAM." + ) + def resume(self) -> SessionProcessorStatus: if not self._resume_event.is_set(): self._resume_event.set() @@ -1167,6 +1207,7 @@ def _process( finally: GENERATION_DEVICE_POOL.release_session(worker.device) self._release_vram_after_session(worker) + self._warn_if_vram_paged(worker) except Exception as e: error_type = e.__class__.__name__ diff --git a/invokeai/app/util/startup_utils.py b/invokeai/app/util/startup_utils.py index 8a5c49399c0..6d3c41306a6 100644 --- a/invokeai/app/util/startup_utils.py +++ b/invokeai/app/util/startup_utils.py @@ -162,10 +162,11 @@ def apply_monkeypatches() -> None: install_peer_aware_empty_cache() # ROCm's fused SDPA kernels return wrong output for heads wider than 256 (the VAE mid-block - # attention), turning decodes into noise or black images. Route those calls to the math kernel. - from invokeai.backend.util.attention import install_rocm_sdpa_head_dim_guard + # attention), turning decodes into noise or black images. Route those calls to the math kernel, + # and compute math-kernel attention -- the only kernel ROCm on Windows has -- in bounded chunks. + from invokeai.backend.util.attention import install_rocm_sdpa_guard - install_rocm_sdpa_head_dim_guard() + install_rocm_sdpa_guard() def register_mime_types() -> None: diff --git a/invokeai/app/util/torch_cuda_allocator.py b/invokeai/app/util/torch_cuda_allocator.py index d1c34cd3ceb..7e94be57900 100644 --- a/invokeai/app/util/torch_cuda_allocator.py +++ b/invokeai/app/util/torch_cuda_allocator.py @@ -1,6 +1,82 @@ +import importlib.metadata +import importlib.util import logging import os +import re import sys +from pathlib import Path + +# Every variable torch reads its caching-allocator configuration from. `ModelCache._expandable_segments_enabled` keeps +# its own copy of this list: the model cache imports torch, and this module must not. +_ALLOCATOR_CONF_ENV_VARS = ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_HIP_ALLOC_CONF") + +ROCM_WINDOWS_ALLOC_CONF = "expandable_segments:True" + + +def _installed_torch_version() -> str | None: + """The installed torch distribution's version, read from its metadata so torch itself is not imported.""" + try: + return importlib.metadata.version("torch") + except importlib.metadata.PackageNotFoundError: + return None + + +def _installed_torch_is_rocm() -> bool: + """Whether the installed torch is a ROCm build, without importing torch. + + AMD's published wheels carry `+rocm` in the distribution version, but a wheel built locally is versioned like + `2.8.0a0+gitfc14c65`. Both record a `hip` version in torch's own `version.py`, which is read as a file -- the same + value `torch.version.hip` returns, which is what `invokeai.backend.util.wddm` gates on. Keeping the two answers + together matters: a build that gets the video-memory budget but not this allocator default fragments its way into + system memory with nothing in the log to say why. + """ + version = _installed_torch_version() + if version is not None and "+rocm" in version: + return True + try: + spec = importlib.util.find_spec("torch") # locates the package; does not execute it + locations = list(spec.submodule_search_locations or []) if spec is not None else [] + for location in locations: + source = Path(location, "version.py") + if source.is_file(): + return ( + re.search(r"^hip\s*(?::[^=]+)?=\s*['\"]", source.read_text(encoding="utf-8"), re.MULTILINE) + is not None + ) + except Exception: + return False + return False + + +def apply_rocm_windows_allocator_default(logger: logging.Logger) -> None: + """Default a ROCm build on Windows to the expandable-segments allocator, unless the allocator is configured already. + + Windows does not fail an allocation it cannot place in contiguous VRAM: it puts all of it in shared system memory, + silently and for as long as the caching allocator keeps the segment. Freed-up VRAM does not pull it back, so one + fragmented allocation keeps slowing every later tensor that reuses the segment. Measured on an RX 9060 XT with + torch 2.12+rocm7.14: a 2 GiB tensor landed in system memory with 4.7 GiB of VRAM free in 256 MiB holes, and Z-Image + at 1024px denoised in 144/139/186 s over three runs, against 37/34/34 s with expandable segments, which map a large + block from small physical pieces instead of needing one contiguous range. + + Must run before torch is imported. Any allocator variable already in the environment, like an explicit + `pytorch_cuda_alloc_conf` (whose caller skips this default), wins. + """ + if sys.platform != "win32" or any(os.environ.get(var) for var in _ALLOCATOR_CONF_ENV_VARS): + return + if not _installed_torch_is_rocm(): + return + if "torch" in sys.modules: + # Setting the variable now would change nothing but the model cache's reading of it. + logger.warning( + "ROCm on Windows: torch was imported before the allocator default could apply; leaving it unset." + ) + return + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ROCM_WINDOWS_ALLOC_CONF + logger.info( + f"ROCm on Windows: using the expandable-segments allocator (PYTORCH_CUDA_ALLOC_CONF={ROCM_WINDOWS_ALLOC_CONF}), " + "so fragmented VRAM does not push allocations into shared system memory. Set pytorch_cuda_alloc_conf to " + "override it, e.g. to 'expandable_segments:False'." + ) def configure_torch_cuda_allocator(pytorch_cuda_alloc_conf: str, logger: logging.Logger): diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3c1dce56416..c64bc9fe69e 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2301,7 +2301,7 @@ def _get_vram_available(self, working_mem_bytes: Optional[int], honor_cap: bool if self._execution_device.type == "cuda": vram_allocated = torch.cuda.memory_allocated(self._execution_device) - vram_free, _vram_total = torch.cuda.mem_get_info(self._execution_device) + vram_free, _vram_total = TorchDevice.cuda_mem_get_info(self._execution_device) # Blocks the caching allocator holds but is not using are just as available to this # process as driver-free memory: the allocator reuses them directly, and empty_cache() # returns whole unoccupied segments to the driver. mem_get_info() alone counts them as @@ -2342,9 +2342,11 @@ def _get_reclaimable_allocator_bytes(self) -> int: Best-effort: 0 when the backend does not expose allocator stats, and 0 under expandable-segments mode — there, freed blocks inside a segment are NOT counted as - inactive splits, empty_cache() reclaims nothing, and a large allocation cannot use the - holes, so the whole (reserved - allocated) figure is untrustworthy (a measured hard OOM - on an allocation the credited budget claimed would fit). + inactive splits and a large allocation cannot use the holes, so the whole + (reserved - allocated) figure is untrustworthy (a measured hard OOM on an allocation the + credited budget claimed would fit). empty_cache() does unmap the freed pages, which is why + `_offload_unlocked_models` runs one after each offload in that mode; what it cannot do is + make the credit trustworthy before it runs. """ if _expandable_segments_enabled(): return 0 @@ -2552,11 +2554,22 @@ def _offload_unlocked_models( f"Offloading unlocked models with goal of making room for {vram_bytes_required / MB:.2f}MB of VRAM." ) vram_bytes_freed = 0 + # Under expandable segments the measurement credits no allocator-held blocks, so an offload + # only shows up once empty_cache() unmaps its pages; without it every unlocked model would + # be unloaded by the full shortfall. When a peer device is mid-session that release is + # deferred, and the bytes just freed are credited to the measurement instead. That credit is + # the same optimistic figure `_get_reclaimable_allocator_bytes` refuses to grant -- the + # allocator holds those bytes, but a large contiguous allocation may not be servable from + # them -- so it is deliberately loop-local: it stops this loop offloading everything, and + # the load that follows re-measures for itself (`lock`, `make_room_in_vram`) and streams + # from RAM if the room did not materialize. + empty_cache_per_offload = _expandable_segments_enabled() + vram_bytes_freed_uncounted = 0 # TODO(ryand): Give more thought to the offloading policy used here. cache_entries_increasing_size = sorted(self._cached_models.values(), key=lambda x: x.cached_model.total_bytes()) for cache_entry in cache_entries_increasing_size: # We do not fully trust the count of bytes freed, so we check again on each iteration. - vram_available = vram_available_fn() + vram_available = vram_available_fn() + vram_bytes_freed_uncounted vram_bytes_to_free = vram_bytes_required - vram_available if vram_bytes_to_free <= 0: break @@ -2569,13 +2582,18 @@ def _offload_unlocked_models( self._logger.debug( f"Unloaded {cache_entry.key} from VRAM to free {(cache_entry_bytes_freed / MB):.0f} MB." ) + if empty_cache_per_offload: + if TorchDevice.empty_cache(): + vram_bytes_freed_uncounted = 0 + else: + vram_bytes_freed_uncounted += cache_entry_bytes_freed vram_bytes_freed += cache_entry_bytes_freed # Only pay for empty_cache() when something was actually offloaded. Paced VRAM moves run # this method once per pass, and on most passes there is nothing left to offload — # an unconditional empty_cache() would return the allocator's blocks to the driver # (and synchronize the device on ROCm) dozens of times per stream for no benefit. - if vram_bytes_freed > 0: + if vram_bytes_freed > 0 and not empty_cache_per_offload: TorchDevice.empty_cache() return vram_bytes_freed diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 444e8260443..202fbf979e5 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -7,7 +7,7 @@ import threading import warnings from functools import lru_cache, wraps -from typing import Any +from typing import Any, Callable import psutil import torch @@ -99,8 +99,9 @@ def auto_detect_slice_size(latents: torch.Tensor) -> str: # # The guard keys on the torch build, not the card: only gfx1100 was measured, and a wrong image is # the worse failure, so every ROCm device takes the math kernel for wide heads until a build is -# known to be fixed. The cost is a materialized score matrix for those calls (~4.3 GiB for a 1024px -# VAE decode, ~21 GiB at 1536px; the working-memory estimators price it). A user on a build where +# known to be fixed. The cost is the math kernel for those calls: slower, and a materialized score +# matrix (~4.3 GiB for a 1024px VAE decode, ~21 GiB at 1536px), which the guard computes in chunks of +# `SDPA_MATH_CHUNK_BYTES` and the working-memory estimators price accordingly. A user on a build where # the fused kernels are correct can raise the threshold, or switch the guard off with a large value: # INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM=100000 _ROCM_FUSED_SDPA_MAX_HEAD_DIM_DEFAULT = 256 @@ -124,7 +125,19 @@ def _read_rocm_fused_sdpa_max_head_dim() -> int: ROCM_FUSED_SDPA_MAX_HEAD_DIM = _read_rocm_fused_sdpa_max_head_dim() _IS_ROCM = torch.version.hip is not None -_ROCM_SDPA_GUARD_SENTINEL = "_invokeai_rocm_head_dim_guard" +_ROCM_SDPA_GUARD_SENTINEL = "_invokeai_rocm_sdpa_guard" + +# Largest score matrix one math-kernel SDPA call may materialize on ROCm; larger calls run in chunks of at most this +# (`_chunked_sdpa`). Where no fused kernel exists -- ROCm on Windows has none, measured on an RX 9060 XT with torch +# 2.12+rocm7.14 -- an unchunked call costs heads x seq^2 x ~12-17 bytes: 8 GiB for Z-Image at 1024px, 12.9 GiB for +# Krea-2, 40 GiB for Z-Image at 1536px, against a 16 GB card. Measured per call at 1 GiB, against unchunked: Z-Image +# shape +8 % time, FLUX VAE mid-block (one 512-wide head, 1024px) +12 %; both within one bf16 ulp of the unchunked +# result. Read at call time, so tests can lower it. +SDPA_MATH_CHUNK_BYTES = 1 << 30 + + +def _rocm_cuda(device_type: str) -> bool: + return _IS_ROCM and device_type == "cuda" def rocm_sdpa_uses_math_kernel(device_type: str, head_dim: int) -> bool: @@ -133,7 +146,19 @@ def rocm_sdpa_uses_math_kernel(device_type: str, head_dim: int) -> bool: The single rule shared by the guard itself and by the working-memory estimators, so a reserved score matrix always corresponds to a kernel that really runs. """ - return _IS_ROCM and device_type == "cuda" and head_dim > ROCM_FUSED_SDPA_MAX_HEAD_DIM + return _rocm_cuda(device_type) and head_dim > ROCM_FUSED_SDPA_MAX_HEAD_DIM + + +def rocm_sdpa_chunks_math(device_type: str) -> bool: + """True when math-kernel attention on this device runs in chunks of at most `SDPA_MATH_CHUNK_BYTES`. + + That is a ROCm device with the guard installed on the live binding -- not a guard that failed to install, and not + a process that never installed it (scripts and tests that skip `apply_monkeypatches`), where the whole score + matrix is still built and must be budgeted in full. + """ + return _rocm_cuda(device_type) and getattr( + torch.nn.functional.scaled_dot_product_attention, _ROCM_SDPA_GUARD_SENTINEL, False + ) def _math_sdpa( @@ -156,9 +181,111 @@ def _math_sdpa( )[0] -def install_rocm_sdpa_head_dim_guard() -> None: - """Rebind ``torch.nn.functional.scaled_dot_product_attention`` so that, on ROCm, a call whose - head_dim exceeds ``ROCM_FUSED_SDPA_MAX_HEAD_DIM`` runs the math kernel instead of a fused one. +def _math_score_bytes(query: torch.Tensor, key: torch.Tensor) -> int: + """Upper bound on the score matrix a math-kernel call over these (4-D) inputs materializes.""" + batch, heads, query_len, _ = query.shape + return batch * heads * query_len * key.shape[-2] * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + +def _chunkable_kv(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> bool: + """Whether `_chunked_sdpa` slices K/V correctly: 4-D and of the query's batch. Other broadcasts that math SDPA + accepts (3-D K/V, a wider K/V batch) stay whole.""" + batch = query.shape[0] + return key.dim() == value.dim() == 4 and key.shape[0] == value.shape[0] == batch + + +def _materializes( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None, + scale: float | None, + enable_gqa: bool, +) -> bool: + """Whether torch would run this exact call on the math kernel, under whatever backend policy is active now. + + `_fused_sdp_choice` is the dispatcher query `F.scaled_dot_product_attention` itself runs, so an `sdpa_policy` + window or diffusers' pinned `sdpa_kernel` is honoured. Anything but a clear fused answer counts as math: when no + backend is eligible at all the probe raises, and the chunks then raise the same error the whole call would have. + """ + try: + choice = torch.ops.aten._fused_sdp_choice( + query, key, value, attn_mask, 0.0, False, scale=scale, enable_gqa=enable_gqa + ) + except Exception: + return True + return int(choice) not in _FUSED_SDP_CHOICES + + +def _chunked_sdpa( + call: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor], + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None, + enable_gqa: bool, +) -> torch.Tensor: + """Run a 4-D `(batch, heads, seq, dim)` attention as calls whose score matrices stay within `SDPA_MATH_CHUNK_BYTES`. + + Softmax normalizes each query row over all keys, so splitting heads or query rows is exact in exact arithmetic, + and both orders agree to within one bf16 ulp in practice. Head groups come first: every call keeps its full + `query x key` shape, which is bitwise the unchunked result wherever the kernel's arithmetic does not depend on + the batch it runs over (it does on CPU; a rocBLAS batched GEMM re-tiles when 30 head-batches become groups of 2, + measured at 0.0005 on a bf16 Z-Image shape). Query rows are split only when a single head (or GQA group) alone + exceeds the budget -- the one-head VAE mid-block -- or when K/V are broadcast across heads rather than grouped. A + mask is sliced along the axis being split wherever it is not broadcast there. The output is allocated from the + first chunk's result, which carries the dtype autocast chose. + """ + budget = SDPA_MATH_CHUNK_BYTES + batch, heads, query_len, _ = query.shape + key_len, kv_heads = key.shape[-2], key.shape[-3] + per_head = batch * query_len * key_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + + if kv_heads == heads: + ratio: int | None = 1 + elif enable_gqa and kv_heads > 0 and heads % kv_heads == 0: + ratio = heads // kv_heads + else: + ratio = None + group = 0 if ratio is None else (budget // per_head) // ratio * ratio + + out: torch.Tensor | None = None + if group > 0: + mask_per_head = attn_mask is not None and attn_mask.dim() >= 3 and attn_mask.shape[-3] == heads + assert ratio is not None + for start in range(0, heads, group): + end = min(start + group, heads) + mask = attn_mask[..., start:end, :, :] if mask_per_head else attn_mask + kv = slice(start // ratio, end // ratio) + part = call(query[:, start:end], key[:, kv], value[:, kv], mask) + if out is None: + out = part.new_empty((*query.shape[:-1], part.shape[-1])) + out[:, start:end] = part + else: + rows = max(1, budget // (batch * heads * key_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT)) + mask_per_row = attn_mask is not None and attn_mask.dim() >= 2 and attn_mask.shape[-2] == query_len + for start in range(0, query_len, rows): + end = min(start + rows, query_len) + mask = attn_mask[..., start:end, :] if mask_per_row else attn_mask + part = call(query[:, :, start:end], key, value, mask) + if out is None: + out = part.new_empty((*query.shape[:-1], part.shape[-1])) + out[:, :, start:end] = part + assert out is not None + return out + + +def install_rocm_sdpa_guard() -> None: + """Rebind ``torch.nn.functional.scaled_dot_product_attention`` on ROCm so that attention runs on a kernel that is + both correct and bounded in memory. + + Two rules, both only for ROCm devices: + + * A call whose head_dim exceeds ``ROCM_FUSED_SDPA_MAX_HEAD_DIM`` runs the math kernel instead of a fused one, which + returns wrong output there. + * A call that runs on the math kernel -- routed by the first rule, or chosen by torch because no fused kernel takes + it (all of them, on ROCm under Windows) -- and would materialize more than ``SDPA_MATH_CHUNK_BYTES`` of scores + runs in chunks (`_chunked_sdpa`). Causal and dropout calls, and anything but 4-D inputs, are left whole. The math kernel is invoked as the aten op directly rather than through ``sdpa_kernel(MATH)``: that context manager writes torch's process-global backend flags, so on a multi-GPU install it @@ -167,10 +294,11 @@ def install_rocm_sdpa_head_dim_guard() -> None: Callers that resolve the function through the module at call time inherit the guard -- that is diffusers' `AttnProcessor2_0` and its attention dispatcher, and every autoencoder in this repo. - Two differences from the fused path for the routed calls only: torch's argument validation is + Two differences from the fused path for the wide-head calls only: torch's argument validation is skipped (an invalid mask dtype is not rejected up front), and the aten op is not on autocast's cast list, so fp32 inputs under ``torch.autocast`` compute in fp32 instead of the autocast - dtype. No current wide-head caller runs under autocast. + dtype. No current wide-head caller runs under autocast. Chunks of narrower heads go through the + original function, keeping its semantics. Idempotent; a no-op on non-ROCm builds. If the aten op is missing or rejects the call shape this wrapper relies on, the guard is not installed and the log says so, rather than failing @@ -187,8 +315,9 @@ def install_rocm_sdpa_head_dim_guard() -> None: _math_sdpa(probe, probe, probe, torch.ones((2, 2), dtype=torch.bool), 0.0, False, None, False) except Exception as e: logger.warning( - f"ROCm head-dim SDPA guard NOT installed: the math attention op rejected the call shape ({e!r}). " - f"Attention with head_dim > {ROCM_FUSED_SDPA_MAX_HEAD_DIM} may return wrong output on this build." + f"ROCm SDPA guard NOT installed: the math attention op rejected the call shape ({e!r}). " + f"Attention with head_dim > {ROCM_FUSED_SDPA_MAX_HEAD_DIM} may return wrong output on this build, and " + "math-kernel attention builds its whole score matrix." ) return original = functional.scaled_dot_product_attention @@ -206,24 +335,56 @@ def guarded_scaled_dot_product_attention( enable_gqa: bool = False, **kwargs: Any, ) -> torch.Tensor: - if not rocm_sdpa_uses_math_kernel(query.device.type, query.shape[-1]): - return original( - query, key, value, attn_mask, dropout_p, is_causal, *args, scale=scale, enable_gqa=enable_gqa, **kwargs - ) - if args or kwargs: + wide_head = rocm_sdpa_uses_math_kernel(query.device.type, query.shape[-1]) + if wide_head and (args or kwargs): # A torch that grew new SDPA arguments would need the math call updated to match; # refuse rather than silently drop them. raise TypeError( - f"scaled_dot_product_attention received arguments the ROCm head-dim guard does not forward: " + f"scaled_dot_product_attention received arguments the ROCm SDPA guard does not forward: " f"{args!r}, {sorted(kwargs)!r}" ) - return _math_sdpa(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa) + chunk = ( + _rocm_cuda(query.device.type) + and query.dim() == 4 + and not query.is_nested + and not is_causal + and dropout_p == 0.0 + and not args + and not kwargs + and _math_score_bytes(query, key) > SDPA_MATH_CHUNK_BYTES + and _chunkable_kv(query, key, value) + and (wide_head or _materializes(query, key, value, attn_mask, scale, enable_gqa)) + ) + if wide_head: + if not chunk: + return _math_sdpa(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa) + return _chunked_sdpa( + lambda q, k, v, m: _math_sdpa(q, k, v, m, 0.0, False, scale, enable_gqa), + query, + key, + value, + attn_mask, + enable_gqa, + ) + if chunk: + return _chunked_sdpa( + lambda q, k, v, m: original(q, k, v, m, 0.0, False, scale=scale, enable_gqa=enable_gqa), + query, + key, + value, + attn_mask, + enable_gqa, + ) + return original( + query, key, value, attn_mask, dropout_p, is_causal, *args, scale=scale, enable_gqa=enable_gqa, **kwargs + ) setattr(guarded_scaled_dot_product_attention, _ROCM_SDPA_GUARD_SENTINEL, True) functional.scaled_dot_product_attention = guarded_scaled_dot_product_attention logger.info( f"ROCm: attention with head_dim > {ROCM_FUSED_SDPA_MAX_HEAD_DIM} runs on the math SDPA kernel " - f"(the fused kernels return wrong output there; {_ROCM_FUSED_SDPA_MAX_HEAD_DIM_ENV} overrides the threshold)." + f"(the fused kernels return wrong output there; {_ROCM_FUSED_SDPA_MAX_HEAD_DIM_ENV} overrides the threshold), " + f"and math-kernel attention runs in chunks of at most {SDPA_MATH_CHUNK_BYTES / 2**30:g} GiB of scores." ) @@ -359,27 +520,37 @@ def sdpa_score_matrix_bytes( `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. + + Where the ROCm guard computes math attention in chunks (`rocm_sdpa_chunks_math`), one chunk is + what is alive at a time, so the term is capped at ``SDPA_MATH_CHUNK_BYTES``. That cap assumes the + call is one the guard will chunk: a plain 4-D, non-causal, dropout-free call with K/V of the + query's batch, which is what every caller here prices. A causal or dropout call, or one with + broadcast K/V, runs whole and would need the full figure. """ if seq_len <= 0 or num_heads <= 0: return 0 score_matrix_bytes = num_heads * seq_len * seq_len * SDPA_MATH_BYTES_PER_SCORE_ELEMENT + chunked = rocm_sdpa_chunks_math(device.type) + if chunked: + score_matrix_bytes = min(score_matrix_bytes, SDPA_MATH_CHUNK_BYTES) + suffix = ", in chunks" if chunked else "" if via_diffusers_dispatch: dispatch = _diffusers_attention_dispatch() if dispatch == _DISPATCH_FUSED: return 0 if dispatch == _DISPATCH_MATH: - return _log_and_return(score_matrix_bytes, device, head_dim, "the diffusers backend") + return _log_and_return(score_matrix_bytes, device, head_dim, f"the diffusers backend{suffix}") # _DISPATCH_TORCH: diffusers forwards to `F.scaled_dot_product_attention`, so torch decides. if rocm_sdpa_uses_math_kernel(device.type, head_dim): - # `install_rocm_sdpa_head_dim_guard` sends these calls to the math kernel regardless of - # what torch's dispatch would pick, so the probe would answer for a kernel that never runs. - return _log_and_return(score_matrix_bytes, device, head_dim, "the ROCm head-dim guard") + # `install_rocm_sdpa_guard` sends these calls to the math kernel regardless of what + # torch's dispatch would pick, so the probe would answer for a kernel that never runs. + return _log_and_return(score_matrix_bytes, device, head_dim, f"the ROCm SDPA guard{suffix}") 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") + return _log_and_return(score_matrix_bytes, device, head_dim, f"this torch build{suffix}") @lru_cache(maxsize=None) diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index ff2e9da8c9b..5b94633ef2b 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -9,6 +9,7 @@ from invokeai.app.services.config.config_default import get_config from invokeai.backend.util.level_zero import xpu_device_is_integrated, xpu_memory_info from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.wddm import video_memory_budget # legacy APIs TorchPrecisionNames = Literal["float32", "float16", "bfloat16"] @@ -333,8 +334,8 @@ def normalize(cls, device: Union[str, torch.device]) -> torch.device: return device @classmethod - def empty_cache(cls) -> None: - """Clear the GPU device cache — unless another generation device is mid-session. + def empty_cache(cls) -> bool: + """Clear the GPU device cache — unless another generation device is mid-session. Says whether it ran. ``torch.cuda.empty_cache()`` is process-global: it takes EVERY device's caching-allocator mutex and cudaFree/hipFrees their cached blocks, and a free on a @@ -366,7 +367,7 @@ def empty_cache(cls) -> None: InvokeAILogger.get_logger(cls.__name__).debug( "Deferring empty_cache: another generation device is mid-session." ) - return + return False # Clear before running: a skip that races in after this point re-sets the flag, so a # request is never lost, only (harmlessly) repeated. cls._empty_cache_deferred.clear() @@ -376,6 +377,7 @@ def empty_cache(cls) -> None: torch.cuda.empty_cache() if _xpu_is_available(): torch.xpu.empty_cache() + return True @classmethod def flush_deferred_empty_cache(cls) -> None: @@ -458,6 +460,21 @@ def xpu_mem_get_info(cls, device: torch.device) -> tuple[int, int]: return (max(total_bytes - reserved_bytes, 0), total_bytes) + @classmethod + def cuda_mem_get_info(cls, device: torch.device) -> tuple[int, int]: + """Return ``(free, total)`` VRAM in bytes for a CUDA or ROCm device. + + ``torch.cuda.mem_get_info``, except that on a ROCm build under Windows the free figure is capped by what the + Windows video-memory budget still allows this process (`wddm.video_memory_budget`). There, torch's figure is + the device total minus this process's own live allocations: it ignores other processes, and Windows pages + allocations into shared system memory well before it is exhausted instead of failing them. + """ + free, total = torch.cuda.mem_get_info(device) + budget = video_memory_budget(device) + if budget is not None: + free = min(free, max(budget - (total - free), 0)) + return free, total + @classmethod def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype: return NAME_TO_PRECISION[precision_name] diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 4aff57b3b74..c4c16bd8766 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -8,10 +8,71 @@ from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR +from invokeai.backend.util import wddm from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.vae_tiling_scope import resolve_tile_size +# Share of the device's memory an untiled decode may claim before the decode nodes tile it up front (see +# `should_pretile_vae_decode`). High enough that common resolutions on common cards stay untiled -- tiled output is +# not pixel-identical -- and low enough that a decode no card of that size can hold is tiled before it starts. +VAE_PRETILE_VRAM_FRACTION = 0.9 + + +def should_pretile_vae_decode( + device: torch.device, full_decode_bytes: int, vram_fraction: float = VAE_PRETILE_VRAM_FRACTION +) -> bool: + """Whether a decode should be tiled up front because its untiled working memory would claim more than + ``vram_fraction`` of the memory ``device`` will keep resident. + + Waiting for an out-of-memory error does not work everywhere: on Windows, drivers page an allocation that does not + fit into system memory instead of failing it (always for ROCm, by default for NVIDIA's sysmem fallback), and the + decode just runs very slowly. ``device`` is where the VAE runs: a ``cpu_only`` VAE (and MPS, sharing system + memory) is never tiled on these grounds. + + On Windows the ceiling is the video-memory budget rather than the card's nameplate total -- the budget is the point + past which Windows pages this process out. With one correction: a budget *below* what this process keeps resident is + Windows asking it to trim itself, and the cache does exactly that to honour the reservation, so the ceiling is then + that residency. Measured on an RX 9060 XT (16 GB): with Invoke holding 12 GiB and nothing else on the card the + budget read 7.62 GiB, and comparing against that tiled a 7.0 GiB decode that fits, changing output that had been + pixel-identical. + + Residency, not ``CurrentUsage``: that also counts what Windows has already paged out, which a foreign process forces + exactly when the budget carries its pressure (next to an 8 GiB holder the budget stepped down to 7.58 GiB while our + ``CurrentUsage`` kept growing past our 7.70 GiB in VRAM). So the override is ``CurrentUsage`` minus this process's + paged bytes, and where either is unknown -- ``CurrentUsage`` beyond the card, or no counter reading -- the budget + alone decides (so a transient counter failure can tile a decode that the next run leaves untiled). + + Known limit: the budget reacts to a foreign process only once this process grows into the memory it holds -- it + stays at 15.09 GiB next to an 8 GiB holder while Invoke holds little -- torch's free figure ignores other processes, + and the adapter-wide performance counter contradicts itself between processes (15.66 GiB read from the holding + process against 0.18 GiB from another, same instance, same instant). A decode sized before that point can still + page, and the end-of-session paging warning is what surfaces it. + """ + if device.type == "cuda": + total_bytes = torch.cuda.get_device_properties(device).total_memory + local = wddm.local_video_memory(device) + if local is not None: + budget_bytes, usage_bytes = local + ceiling_bytes = budget_bytes + paged = wddm.paged_bytes(device) if usage_bytes is not None else None + if paged is not None: + ceiling_bytes = max(budget_bytes, usage_bytes - paged) + total_bytes = min(total_bytes, ceiling_bytes) + elif device.type == "xpu": + total_bytes = torch.xpu.get_device_properties(device).total_memory + else: + return False + if full_decode_bytes <= vram_fraction * total_bytes: + return False + InvokeAILogger.get_logger(__name__).debug( + f"Decoding in tiles: an untiled decode would need ~{full_decode_bytes / 2**30:.1f} GiB of working memory, more " + f"than {vram_fraction:.0%} of the {total_bytes / 2**30:.1f} GiB the GPU keeps resident." + ) + return True + + # The diffusers AutoencoderKL (SD1/SDXL, SD3, CogView4) and the FLUX.1 AutoEncoder run the same # mid-block self-attention as the FLUX.2 VAE: one head over the 512-channel width, on the # 8x-downsampled grid. Where that attention runs on the math kernel (see `sdpa_score_matrix_bytes`) @@ -55,15 +116,17 @@ def estimate_vae_working_memory_sd15_sdxl( vae: AutoencoderKL | AutoencoderTiny, tile_size: int | None, fp32: bool, + device: torch.device | None = None, ) -> int: """Estimate the working memory required to encode or decode the given tensor.""" # It was found experimentally that the peak working memory scales linearly with the number of pixels and the # element size (precision). This estimate is accurate for both SD1 and SDXL. element_size = 4 if fp32 else 2 - # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414 - # Encoding uses ~45% the working memory as decoding. - scaling_constant = 2200 if operation == "decode" else 1100 + # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per + # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where + # the VAE will run; it defaults to the session device. + scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 @@ -88,7 +151,7 @@ def estimate_vae_working_memory_sd15_sdxl( # max, not sum: the mid-block attention and the full-resolution convolutions peak in # different phases of the forward (see the FLUX.2 estimator for the measurements). score_dtype = torch.float32 if fp32 else next(vae.parameters()).dtype - working_memory = max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, score_dtype)) + working_memory = max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, score_dtype, device=device)) if fp32: # If we are running in FP32, then we should account for the likely increase in model size (~250MB). @@ -98,7 +161,10 @@ def estimate_vae_working_memory_sd15_sdxl( def estimate_vae_working_memory_cogview4( - operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL + operation: Literal["encode", "decode"], + image_tensor: torch.Tensor, + vae: AutoencoderKL, + device: torch.device | None = None, ) -> int: """Estimate the working memory required by the invocation in bytes.""" latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 @@ -107,14 +173,17 @@ def estimate_vae_working_memory_cogview4( w = latent_scale_factor_for_operation * image_tensor.shape[-1] element_size = next(vae.parameters()).element_size() - # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414 - # Encoding uses ~45% the working memory as decoding. - scaling_constant = 2200 if operation == "decode" else 1100 + # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per + # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where + # the VAE will run; it defaults to the session device. + scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) working_memory = h * w * element_size * scaling_constant # max, not sum: the mid-block attention and the full-resolution convolutions peak in different # phases of the forward (see the FLUX.2 estimator for the measurements). - return int(max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype))) + return int( + max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) + ) # What a tiled decode does *not* bound: the assembled image, several times over, per output pixel. @@ -153,6 +222,7 @@ def estimate_vae_working_memory_flux( image_tensor: torch.Tensor, vae: AutoencoderKL, tile_size: int | None = None, + device: torch.device | None = None, ) -> int: """Estimate the working memory required by the invocation in bytes. @@ -165,14 +235,18 @@ def estimate_vae_working_memory_flux( Both operations tile. The tile-bounded term is the same for either; what tiling does *not* bound differs, so the residual is chosen by `operation` -- see `_FLUX_VAE_TILED_IMAGE_COPIES` for the decode side. + + `device` is where the VAE will run (a `cpu_only` VAE runs on the CPU); it picks the convolution + backend's constant. Defaults to the session device. """ latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 element_size = next(vae.parameters()).element_size() - # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414 - # Encoding uses ~45% the working memory as decoding. - scaling_constant = 2200 if operation == "decode" else 1100 + # The FLUX.1 autoencoder shares the FLUX.2 VAE's convolution stack and measures the same constants on both + # backends; see `_FLUX_VAE_SCALING_CONSTANTS`. + device = device if device is not None else TorchDevice.choose_torch_device() + scaling_constant = _flux_vae_scaling_constant(operation, device) out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] @@ -231,7 +305,7 @@ def estimate_vae_working_memory_flux( # max, not sum: the mid-block attention and the full-resolution convolutions peak in different # phases of the forward (see the FLUX.2 estimator for the measurements). dtype = next(vae.parameters()).dtype - return int(max(working_memory, _vae_mid_block_score_matrix_bytes(score_h, score_w, dtype))) + return int(max(working_memory, _vae_mid_block_score_matrix_bytes(score_h, score_w, dtype, device=device))) # The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the @@ -243,17 +317,21 @@ def estimate_vae_working_memory_flux( _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: +# Peak reserved bytes per output pixel per element byte, per conv backend and operation, for the FLUX.2 VAE and the +# FLUX.1 autoencoder alike. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, which reproduces +# this table on any build: # -# decode encode encode/decode -# cuDNN RTX 4090, torch 2.7.1 2185 1072 0.49 -# MIOpen RX 9070 XT, torch 2.10 3453 2688 0.78 -# MIOpen PRO W7900, torch 2.10 3525 2688 0.76 +# decode encode encode/decode +# cuDNN RTX 4090, torch 2.7.1 2185 1072 0.49 +# MIOpen RX 9070 XT, torch 2.10 3453 2688 0.78 +# MIOpen PRO W7900, torch 2.10 3525 2688 0.76 +# MIOpen RX 9060 XT, torch 2.12, Windows, FLUX.1 VAE 3451 2688 0.78 # # 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. +# to the larger with ~2% headroom. The FLUX.1 row (bf16, 768-1536px, Invoke's ROCm SDPA guard installed, +# `MIOPEN_FIND_MODE=FAST`) lands on the same numbers: the two VAEs share their convolution stack. The +# cuDNN column is the FLUX.1 constant this estimator shipped with before. # # 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 @@ -264,16 +342,18 @@ def estimate_vae_working_memory_flux( # Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the # constant follows the backend. # -# Caveat for AMD users: `MIOPEN_FIND_MODE=2` selects convolution algorithms heuristically instead of -# by benchmark, and measured a uniform 1.28x more memory at every resolution. It is not the default -# and is not budgeted for here -- raise `device_working_mem_gb` if you set it. -_FLUX2_VAE_SCALING_CONSTANTS: dict[str, dict[str, int]] = { +# Caveat for AMD users: `MIOPEN_FIND_MODE=2` (FAST) selects convolution algorithms heuristically instead of +# by benchmark, and measured a uniform 1.28x more memory at every resolution when these constants were +# first fitted. Invoke has since made FAST its default (`run_app.py`); the RX 9060 XT row above was taken +# under it and did not show the 1.28x. If you see decodes exceed their reservation on another card, raise +# `device_working_mem_gb`. +_FLUX_VAE_SCALING_CONSTANTS: dict[str, dict[str, int]] = { "cudnn": {"decode": 2200, "encode": 1100}, "miopen": {"decode": 3600, "encode": 2750}, } -def _flux2_vae_scaling_constant(operation: Literal["encode", "decode"], device: torch.device) -> int: +def _flux_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 @@ -281,7 +361,7 @@ def _flux2_vae_scaling_constant(operation: Literal["encode", "decode"], device: 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] + return _FLUX_VAE_SCALING_CONSTANTS["miopen" if is_rocm else "cudnn"][operation] def estimate_vae_working_memory_flux2( @@ -296,7 +376,7 @@ def estimate_vae_working_memory_flux2( 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 + for an encode, so it is looked up per backend -- see ``_FLUX_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. @@ -334,7 +414,7 @@ def estimate_vae_working_memory_flux2( element_size = param.element_size() device = device if device is not None else TorchDevice.choose_torch_device() - scaling_constant = _flux2_vae_scaling_constant(operation, device) + scaling_constant = _flux_vae_scaling_constant(operation, device) batch_size = image_tensor.shape[0] if image_tensor.dim() >= 4 else 1 if tile_size is not None: @@ -660,11 +740,89 @@ def estimate_vae_working_memory_wan( return int(per_frame + clip_bytes) +# What a full-frame Qwen-Image VAE decode was actually measured to peak at, per output pixel-byte, against the output +# area it was measured at -- the per-point table in `estimate_vae_working_memory_qwen_image`, taking the larger of the +# measured ROCm cards at each point (the W7900 column; the RX 9060 XT measures 2650-2772 throughout). +# +# This exists because the shipped constants are one flat figure fitted to the worst point of a curve that is not +# flat, across cards that differ by up to 1.9x. That is the right figure to *reserve* -- a reservation that is too +# small is an out-of-memory error or, on Windows, a decode paged into system memory. It is the wrong figure to decide +# tiling with: a tiled decode is not pixel-identical, so a 1536^2 decode whose real peak is 14.4 GiB must not be +# tiled on a 24 GiB card because a flat 5500 prices it at 24.2 GiB. +# +# Neither column is extrapolated. Past its last measured area the decision falls back to the reservation constant, +# which is what "we have not measured this, assume the conservative figure" has to mean here: CUDA's grid ends at +# 1792^2 because 2048^2 ran out of memory on the card it was measured on, and carrying a flat maximum past that point +# priced a 2048^2 decode 0.6 GiB under a 24 GiB card's line while its own reservation exceeded the whole card. +# What the estimator reserves per output pixel-byte: the worst measured point of each backend's column plus ~8% +# headroom (the per-point grid is in `estimate_vae_working_memory_qwen_image`). +_QWEN_VAE_RESERVATION_CONSTANTS: dict[str, dict[str, int]] = { + "rocm": {"decode": 5500, "encode": 6300}, + "cuda": {"decode": 2900, "encode": 1600}, +} + +_QWEN_VAE_MEASURED_DECODE_PEAKS: dict[str, tuple[tuple[int, int], ...]] = { + "rocm": ( + (512 * 512, 5132), + (768 * 768, 4596), + (1024 * 1024, 4570), + (1536 * 1536, 3273), + (1792 * 1792, 3735), + (2048 * 2048, 4813), + ), + "cuda": ( + (512 * 512, 2660), + (768 * 768, 2519), + (1024 * 1024, 2690), + (1536 * 1536, 2671), + (1792 * 1792, 2281), + ), +} + + +def qwen_image_untiled_decode_peak_bytes( + image_tensor: torch.Tensor, vae: AutoencoderKLQwenImage, device: torch.device +) -> int: + """What an untiled decode of this latent is expected to peak at, for the up-front tiling decision only. + + Reservations keep using `estimate_vae_working_memory_qwen_image`, which is deliberately conservative. Between two + measured points this is the chord between them in bytes -- so an interior area can price below either neighbour's + per-pixel constant, which is the point: the constant dips while the bytes rise, and taking the larger neighbour + priced a smaller image above a larger one. Past the last measured area it takes the reservation figure, since a + decode larger than anything measured is not one to guess low about. + """ + h = LATENT_SCALE_FACTOR * image_tensor.shape[-2] + w = LATENT_SCALE_FACTOR * image_tensor.shape[-1] + element_size = next(vae.parameters()).element_size() + is_rocm = device.type == "cuda" and torch.version.hip is not None + points = _QWEN_VAE_MEASURED_DECODE_PEAKS["rocm" if is_rocm else "cuda"] + + area = h * w + if area > points[-1][0]: + constant = _QWEN_VAE_RESERVATION_CONSTANTS["rocm" if is_rocm else "cuda"]["decode"] + return h * w * element_size * constant + # Interpolate the bytes, not the constant: the per-pixel figure dips (ROCm 3273 at 1536^2, CUDA 2281 at 1792^2) + # while the bytes themselves rise with area, so taking the larger neighbouring constant priced a smaller image + # above a larger one -- a 1728x1856 CUDA decode at 15.96 GiB against 1792^2 at 13.64 GiB. Measured points keep + # their measured value. + previous_area, previous_bytes = 0, 0 + for measured_area, peak in points: + measured_bytes = measured_area * peak + if measured_area >= area: + if measured_area == area or previous_area == 0: + return int(area * element_size * peak) + span = (measured_bytes - previous_bytes) / (measured_area - previous_area) + return int(element_size * (previous_bytes + span * (area - previous_area))) + previous_area, previous_bytes = measured_area, measured_bytes + raise AssertionError("areas past the measured range are handled above") + + def estimate_vae_working_memory_qwen_image( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLQwenImage, tile_size: int | None = None, + device: torch.device | None = None, ) -> int: """Estimate the working memory required by the invocation in bytes. @@ -679,7 +837,8 @@ def estimate_vae_working_memory_qwen_image( plus the pixel-space buffers, which stay resident on the execution device either way. ``tile_size`` is the resolved tile size (the nodes' 0 sentinel already substituted), and assumes - the 4:3 tile-to-stride ratio applied by ``patch_qwen_image_vae_tiling``. + the 4:3 tile-to-stride ratio applied by ``patch_qwen_image_vae_tiling``. ``device`` is where the + VAE will run; it selects the constants below and defaults to the session device. """ latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 @@ -717,11 +876,17 @@ def estimate_vae_working_memory_qwen_image( # materialised seq^2 score matrix), i.e. XPU gets an efficient kernel and is in the same # O(area) regime as CUDA. If a future driver regresses to math attention, this branch -- # not the constants -- is what needs to change. - is_rocm = torch.version.hip is not None - if operation == "decode": - scaling_constant = 5500 if is_rocm else 2900 - else: # encode - scaling_constant = 6300 if is_rocm else 1600 + # - A second ROCm card measures far below the shipped pair: on an RX 9060 XT (gfx1200, torch 2.12+rocm7.14, + # fp16, 2026-09-23, the same script with the attention guard installed) the implied constant is flat at + # 2650-2772 for decode and 1541-1552 for encode, all the way to 2048^2. The shipped figures stay at the + # W7900's, which is the card that needs them; what the spread means is that the reservation carries + # cross-card conservatism, so it must not double as the up-front tiling decision -- see + # `qwen_image_untiled_decode_peak_bytes`, which prices that decision from the measured curve instead. + # A named device answers for that device -- a `cpu_only` VAE is not on MIOpen even on a ROCm build. Without one + # the answer stays the build-wide one these constants were always selected by, so callers that pass no device + # (and tests that pin none) keep the figure they had. + is_rocm = torch.version.hip is not None and (device is None or device.type == "cuda") + scaling_constant = _QWEN_VAE_RESERVATION_CONSTANTS["rocm" if is_rocm else "cuda"][operation] if tile_size is not None and tile_size > 0: # Bounded by one tile (plus overlap) rather than the full frame. @@ -744,7 +909,10 @@ def estimate_vae_working_memory_qwen_image( def estimate_vae_working_memory_sd3( - operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL + operation: Literal["encode", "decode"], + image_tensor: torch.Tensor, + vae: AutoencoderKL, + device: torch.device | None = None, ) -> int: """Estimate the working memory required by the invocation in bytes.""" # Encode operations use approximately 50% of the memory required for decode operations @@ -755,12 +923,15 @@ def estimate_vae_working_memory_sd3( w = latent_scale_factor_for_operation * image_tensor.shape[-1] element_size = next(vae.parameters()).element_size() - # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414 - # Encoding uses ~45% the working memory as decoding. - scaling_constant = 2200 if operation == "decode" else 1100 + # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per + # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where + # the VAE will run; it defaults to the session device. + scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) working_memory = h * w * element_size * scaling_constant # max, not sum: the mid-block attention and the full-resolution convolutions peak in different # phases of the forward (see the FLUX.2 estimator for the measurements). - return int(max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype))) + return int( + max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) + ) diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py new file mode 100644 index 00000000000..ec39f7b7389 --- /dev/null +++ b/invokeai/backend/util/wddm.py @@ -0,0 +1,409 @@ +"""The Windows video-memory budget of a ROCm device, read from the WDDM kernel thunks in gdi32. + +On a ROCm build under Windows, ``torch.cuda.mem_get_info`` reports the device total minus this process's own usage. It +does not see other processes, and it keeps reporting free memory after Windows has started to page this process's +allocations out -- Windows does not fail an allocation that exceeds what it will keep resident, it places it in shared +system memory, silently and for as long as the allocation lives. What it will keep resident is the per-process budget +``D3DKMTQueryVideoMemoryInfo`` reports (the figure ``IDXGIAdapter3::QueryVideoMemoryInfo`` returns, without COM). + +Measured on an RX 9060 XT (16 GB, torch 2.12+rocm7.14): alone, the budget is 15.09 of 15.92 GiB and paging starts +once the committed local usage passes budget + 0.38 GiB; once the process itself holds about 12 GiB the budget drops to +7.62 GiB without anything being paged -- Windows asking it to trim, not a limit it has hit. Next to another GPU process +the budget is not "total minus the other process": it stays at 15.09 GiB while this process holds little and steps down +once the two together pass about 12 GiB -- next to an 8 GiB holder, from 15.09 at 3.25 GiB of our own down to 7.58 at +7.5 GiB, where this process then starts to page. So by the time this process could page, the budget is its share of the +card. + +What this process can still allocate before paging is the budget minus its live allocations -- the device total minus +torch's free figure. Not minus ``CurrentUsage``: that also counts memory HIP keeps after a free and hands back to the +next allocation (1-2 GiB measured, still held after ``empty_cache()``), so it would hide what an offload just freed. + +Whether Windows has actually paged this process out is not visible through D3DKMT (``CurrentUsage`` counts paged bytes +as local); it is visible in the ``GPU Process Memory`` performance counters Task Manager shows, read here through PDH. + +Everything here is best-effort: any failure yields ``None`` ("unknown") rather than raising, and callers must keep their +existing behaviour when the answer is unknown. +""" + +import ctypes +import functools +import os +import sys +import threading +from dataclasses import dataclass +from typing import Optional + +import torch + +from invokeai.backend.util.logging import InvokeAILogger + +_KMTQAITYPE_ADAPTERADDRESS = 6 +_MEMORY_SEGMENT_GROUP_LOCAL = 0 +_PDH_FMT_LARGE = 0x00000400 +_PDH_MORE_DATA = 0x800007D2 +# The English counter path; PdhAddEnglishCounterW resolves it on a localized Windows too. +_PDH_SHARED_USAGE_COUNTER = r"\GPU Process Memory(*)\Shared Usage" +# PDH_CSTATUS_VALID_DATA and PDH_CSTATUS_NEW_DATA; any other status leaves `largeValue` undefined. +_PDH_VALID_DATA = (0x00000000, 0x00000001) + + +# Win32 types spelled as fixed-width ctypes, so this module imports on every platform. +_UINT = ctypes.c_uint32 +_ULONG = ctypes.c_uint32 +_LONG = ctypes.c_int32 +_BOOL = ctypes.c_int32 + + +class _Luid(ctypes.Structure): + _fields_ = [("LowPart", _UINT), ("HighPart", _LONG)] + + +class _AdapterInfo(ctypes.Structure): + """``D3DKMT_ADAPTERINFO``.""" + + _fields_ = [ + ("hAdapter", _UINT), + ("AdapterLuid", _Luid), + ("NumOfSources", _ULONG), + ("bPrecisePresentRegionsPreferred", _BOOL), + ] + + +class _EnumAdapters2(ctypes.Structure): + """``D3DKMT_ENUMADAPTERS2``.""" + + _fields_ = [("NumAdapters", _ULONG), ("pAdapters", ctypes.POINTER(_AdapterInfo))] + + +class _QueryAdapterInfo(ctypes.Structure): + """``D3DKMT_QUERYADAPTERINFO``.""" + + _fields_ = [ + ("hAdapter", _UINT), + ("Type", ctypes.c_int), + ("pPrivateDriverData", ctypes.c_void_p), + ("PrivateDriverDataSize", _UINT), + ] + + +class _AdapterAddress(ctypes.Structure): + """``D3DKMT_ADAPTERADDRESS``: the adapter's PCI location.""" + + _fields_ = [("BusNumber", _UINT), ("DeviceNumber", _UINT), ("FunctionNumber", _UINT)] + + +class _CloseAdapter(ctypes.Structure): + """``D3DKMT_CLOSEADAPTER``.""" + + _fields_ = [("hAdapter", _UINT)] + + +class _QueryVideoMemoryInfo(ctypes.Structure): + """``D3DKMT_QUERYVIDEOMEMORYINFO``. A null ``hProcess`` means the calling process.""" + + _fields_ = [ + ("hProcess", ctypes.c_void_p), + ("hAdapter", _UINT), + ("MemorySegmentGroup", ctypes.c_int), + ("Budget", ctypes.c_uint64), + ("CurrentUsage", ctypes.c_uint64), + ("CurrentReservation", ctypes.c_uint64), + ("AvailableForReservation", ctypes.c_uint64), + ("PhysicalAdapterIndex", _UINT), + ] + + +class _PdhCounterValue(ctypes.Structure): + """``PDH_FMT_COUNTERVALUE`` read as ``PDH_FMT_LARGE``.""" + + _fields_ = [("CStatus", ctypes.c_uint32), ("largeValue", ctypes.c_int64)] + + +class _PdhCounterValueItem(ctypes.Structure): + """``PDH_FMT_COUNTERVALUE_ITEM_W``.""" + + _fields_ = [("szName", ctypes.c_wchar_p), ("FmtValue", _PdhCounterValue)] + + +_GDI32_FUNCTIONS = ( + "D3DKMTEnumAdapters2", + "D3DKMTQueryAdapterInfo", + "D3DKMTQueryVideoMemoryInfo", + "D3DKMTCloseAdapter", +) + + +@dataclass(frozen=True) +class _Adapter: + handle: int + luid: int + total_bytes: int + + +_lock = threading.Lock() +# Per torch device index: the matching WDDM adapter, or None once resolution has failed. +_adapters: dict[int, Optional[_Adapter]] = {} + +_pdh_lock = threading.Lock() +# One PDH query per counter path, opened on first use: path -> (query, counter) or None when that counter is absent. +_pdh_queries: dict[str, Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]] = {} + + +def _supported(device: torch.device) -> bool: + return sys.platform == "win32" and torch.version.hip is not None and device.type == "cuda" + + +@functools.lru_cache(maxsize=1) +def _load_gdi32() -> Optional[ctypes.CDLL]: + """Open gdi32 once and declare every thunk as taking a struct pointer and returning an NTSTATUS.""" + try: + lib = ctypes.WinDLL("gdi32") + for name in _GDI32_FUNCTIONS: + fn = getattr(lib, name) + fn.argtypes = [ctypes.c_void_p] + fn.restype = ctypes.c_long + except (AttributeError, OSError) as exc: + InvokeAILogger.get_logger(__name__).debug(f"WDDM thunks unavailable: {exc}") + return None + return lib + + +@functools.lru_cache(maxsize=1) +def _load_pdh() -> Optional[ctypes.CDLL]: + """Open pdh once and declare the few entry points used; every one returns a PDH_STATUS.""" + try: + lib = ctypes.WinDLL("pdh") + prototypes = { + "PdhOpenQueryW": [ctypes.c_wchar_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)], + "PdhAddEnglishCounterW": [ + ctypes.c_void_p, + ctypes.c_wchar_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ], + "PdhCollectQueryData": [ctypes.c_void_p], + "PdhGetFormattedCounterArrayW": [ + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.POINTER(ctypes.c_uint32), + ctypes.POINTER(ctypes.c_uint32), + ctypes.c_void_p, + ], + "PdhCloseQuery": [ctypes.c_void_p], + } + for name, argtypes in prototypes.items(): + fn = getattr(lib, name) + fn.argtypes = argtypes + fn.restype = ctypes.c_uint32 + except (AttributeError, OSError) as exc: + InvokeAILogger.get_logger(__name__).debug(f"PDH unavailable: {exc}") + return None + return lib + + +def _close(lib: ctypes.CDLL, handle: int) -> None: + try: + lib.D3DKMTCloseAdapter(ctypes.byref(_CloseAdapter(handle))) + except Exception: + pass + + +def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: + """Open the WDDM adapter at the torch device's PCI location, closing every other handle the enumeration opened. + + ROCm's device properties carry no LUID, so the PCI bus and device numbers are the link. Anything but exactly one + match is treated as unknown rather than guessed. + """ + props = torch.cuda.get_device_properties(index) + location = (int(props.pci_bus_id), int(props.pci_device_id)) + + enum = _EnumAdapters2(0, None) + if lib.D3DKMTEnumAdapters2(ctypes.byref(enum)) != 0: + return None + infos = (_AdapterInfo * enum.NumAdapters)() + enum.pAdapters = infos + if lib.D3DKMTEnumAdapters2(ctypes.byref(enum)) != 0: + return None + + matches: list[_AdapterInfo] = [] + pending = list(infos[: enum.NumAdapters]) + current: list[_AdapterInfo] = [] + try: + while pending: + # Held in `current` while it is processed: popped from `pending` and not yet in `matches`, a handle whose + # query raises would otherwise be the one handle the cleanup below misses. + current = [pending.pop()] + info = current[0] + address = _AdapterAddress() + query = _QueryAdapterInfo( + info.hAdapter, + _KMTQAITYPE_ADAPTERADDRESS, + ctypes.cast(ctypes.byref(address), ctypes.c_void_p), + ctypes.sizeof(address), + ) + if ( + lib.D3DKMTQueryAdapterInfo(ctypes.byref(query)) == 0 + and (address.BusNumber, address.DeviceNumber) == location + ): + matches.append(info) + else: + _close(lib, info.hAdapter) + current = [] + except Exception: + # Every handle the enumeration opened belongs to this process until it is closed: the ones matched so far, the + # one being processed, and the ones this loop never reached. + for info in matches + current + pending: + _close(lib, info.hAdapter) + raise + + if len(matches) != 1: + for info in matches: + _close(lib, info.hAdapter) + return None + luid = (matches[0].AdapterLuid.HighPart << 32) | matches[0].AdapterLuid.LowPart + return _Adapter(handle=matches[0].hAdapter, luid=luid, total_bytes=int(props.total_memory)) + + +def _adapter_for(device: torch.device) -> Optional[tuple[ctypes.CDLL, _Adapter]]: + lib = _load_gdi32() + if lib is None: + return None + try: + index = device.index if device.index is not None else torch.cuda.current_device() + except Exception: + return None + with _lock: + if index not in _adapters: + try: + _adapters[index] = _resolve_adapter(lib, index) + except Exception as exc: + InvokeAILogger.get_logger(__name__).debug(f"WDDM adapter lookup for {device} failed: {exc}") + _adapters[index] = None + adapter = _adapters[index] + return None if adapter is None else (lib, adapter) + + +def local_video_memory(device: torch.device) -> Optional[tuple[int, Optional[int]]]: + """``(budget, current_usage)`` for this process on the device's adapter, both driver-side figures in bytes. + + ``None`` off Windows ROCm, and whenever the driver cannot answer or answers a budget that is implausible. The usage + alone is ``None`` when it exceeds the adapter (see below). Called per generation and per decode, so every failure + answers "unknown" here rather than reaching those callers. + """ + if not _supported(device): + return None + try: + resolved = _adapter_for(device) + if resolved is None: + return None + lib, adapter = resolved + info = _QueryVideoMemoryInfo(None, adapter.handle, _MEMORY_SEGMENT_GROUP_LOCAL) + if lib.D3DKMTQueryVideoMemoryInfo(ctypes.byref(info)) != 0: + return None + if not 0 < info.Budget <= adapter.total_bytes: + return None + # `CurrentUsage` counts bytes Windows has already paged out, so it exceeds the adapter's own size exactly when + # the process over-commits. It says nothing about residency then, so it is reported as unknown; the budget + # stands, since the model cache's free-VRAM cap is the caller that needs it most in that state. + usage = int(info.CurrentUsage) + return int(info.Budget), usage if 0 <= usage <= adapter.total_bytes else None + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"WDDM memory lookup for {device} failed", exc_info=True) + return None + + +def video_memory_budget(device: torch.device) -> Optional[int]: + """This process's WDDM budget on the device's adapter, in bytes: what Windows keeps resident before paging.""" + local = local_video_memory(device) + return None if local is None else local[0] + + +def _open_counter_query(pdh: ctypes.CDLL, path: str) -> Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]: + query = ctypes.c_void_p() + if pdh.PdhOpenQueryW(None, None, ctypes.byref(query)) != 0: + return None + counter = ctypes.c_void_p() + if pdh.PdhAddEnglishCounterW(query, path, None, ctypes.byref(counter)) != 0: + pdh.PdhCloseQuery(query) + return None + return query, counter + + +def _counter_values(path: str) -> dict[str, int]: + """That counter's value per instance, collected once. ``{}`` when the counter or the loader is unavailable.""" + pdh = _load_pdh() + if pdh is None: + return {} + with _pdh_lock: + try: + if path not in _pdh_queries: + _pdh_queries[path] = _open_counter_query(pdh, path) + opened = _pdh_queries[path] + if opened is None: + return {} + return _usage_by_instance(pdh, *opened) + except Exception: + return {} + + +def _usage_by_instance(pdh: ctypes.CDLL, query: ctypes.c_void_p, counter: ctypes.c_void_p) -> dict[str, int]: + """Collect the query once; that counter's value per instance, keyed by its lower-case instance name.""" + if pdh.PdhCollectQueryData(query) != 0: + return {} + size, count = ctypes.c_uint32(0), ctypes.c_uint32(0) + status = pdh.PdhGetFormattedCounterArrayW(counter, _PDH_FMT_LARGE, ctypes.byref(size), ctypes.byref(count), None) + if status != _PDH_MORE_DATA: + return {} + buffer = (ctypes.c_byte * size.value)() + status = pdh.PdhGetFormattedCounterArrayW(counter, _PDH_FMT_LARGE, ctypes.byref(size), ctypes.byref(count), buffer) + if status != 0: + return {} + items = ctypes.cast(buffer, ctypes.POINTER(_PdhCounterValueItem)) + return { + items[i].szName.lower(): int(items[i].FmtValue.largeValue) + for i in range(count.value) + if items[i].szName and items[i].FmtValue.CStatus in _PDH_VALID_DATA + } + + +def paged_bytes(device: torch.device) -> Optional[int]: + """Bytes of this process's memory on the device's adapter that Windows currently keeps in shared system memory. + + Allocations land there when they did not fit the budget or found no contiguous VRAM, and they stay there; a few + hundred MB also appear briefly while weights are copied in. ``None`` off Windows ROCm and whenever the counters + cannot answer, including before this process has used the adapter at all. + + The first call opens this counter's PDH query (~0.3 s); each later call collects it once (~0.3 ms). + """ + if not _supported(device): + return None + resolved = _adapter_for(device) + if resolved is None: + return None + instance = f"pid_{os.getpid()}_{_adapter_instance(resolved[1].luid)}" + return _counter_values(_PDH_SHARED_USAGE_COUNTER).get(instance) + + +def _adapter_instance(luid: int) -> str: + return f"luid_0x{(luid >> 32) & 0xFFFFFFFF:08x}_0x{luid & 0xFFFFFFFF:08x}_phys_0" + + +def reset_cache() -> None: + """Close the cached adapter handles and the PDH query, and forget the loaders (tests only).""" + with _lock: + opened = [adapter for adapter in _adapters.values() if adapter is not None] + _adapters.clear() + lib = _load_gdi32() if opened else None + if lib is not None: + for adapter in opened: + _close(lib, adapter.handle) + _load_gdi32.cache_clear() + + with _pdh_lock: + opened_queries = [opened for opened in _pdh_queries.values() if opened is not None] + _pdh_queries.clear() + pdh = _load_pdh() if opened_queries else None + if pdh is not None: + for query, _counter in opened_queries: + pdh.PdhCloseQuery(query) + _load_pdh.cache_clear() diff --git a/invokeai/backend/wan/rocm_causal_conv3d.py b/invokeai/backend/wan/rocm_causal_conv3d.py index 7d9499431f5..8280c825787 100644 --- a/invokeai/backend/wan/rocm_causal_conv3d.py +++ b/invokeai/backend/wan/rocm_causal_conv3d.py @@ -140,7 +140,7 @@ def patch_wan_causal_conv3d_for_rocm() -> None: implicit-GEMM conv3d kernels, so CUDA builds keep the stock path. There is deliberately no HIP-version gate (see the module docstring). The row tearing once blamed on this decomposition on HIP 7.2 has two known causes, both since fixed: the - ROCm fused-SDPA head-dim defect (``install_rocm_sdpa_head_dim_guard`` in + ROCm fused-SDPA head-dim defect (``install_rocm_sdpa_guard`` in ``invokeai.backend.util.attention``; the Wan VAE mid-block attention is wider than its threshold) and MIOpen faulting on the strided views the decomposition used to hand it (see ``_decomposed_conv3d``). With both in place the decomposition matches native diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 6cd24a3e9c7..1f5b5465c83 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -58128,7 +58128,7 @@ } ], "title": "Pytorch Cuda Alloc Conf", - "description": "Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally." + "description": "Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used." }, "device": { "type": "string", @@ -58204,6 +58204,12 @@ "description": "Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.", "default": false }, + "auto_tiled_decode": { + "type": "boolean", + "title": "Auto Tiled Decode", + "description": "Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.", + "default": true + }, "pil_compress_level": { "type": "integer", "title": "Pil Compress Level", @@ -58498,7 +58504,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 1af709d1487..70275616943 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -22201,7 +22201,7 @@ export type components = { * ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. * vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. * lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable. - * pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. + * pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, "expandable_segments:True" is used. * device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number) * precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32` * sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. @@ -22210,6 +22210,7 @@ export type components = { * attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` * attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8` * force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it. + * auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there. * pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting. * max_queue_size: Maximum number of items in the session queue. * session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin` @@ -22572,7 +22573,7 @@ export type components = { lazy_offload?: boolean; /** * Pytorch Cuda Alloc Conf - * @description Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. + * @description Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, "expandable_segments:True" is used. */ pytorch_cuda_alloc_conf?: string | null; /** @@ -22638,6 +22639,12 @@ export type components = { * @default false */ force_tiled_decode?: boolean; + /** + * Auto Tiled Decode + * @description Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there. + * @default true + */ + auto_tiled_decode?: boolean; /** * Pil Compress Level * @description The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting. diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index f4d532d60f5..d7d6ca34349 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -68,7 +68,7 @@ from invokeai.backend.model_manager.taxonomy import Flux2VariantType from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes from invokeai.backend.util.vae_working_memory import ( - _flux2_vae_scaling_constant, + _flux_vae_scaling_constant, estimate_vae_working_memory_flux2, ) @@ -457,8 +457,8 @@ 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]: - decode_k = _flux2_vae_scaling_constant("decode", torch.device("cuda")) - encode_k = _flux2_vae_scaling_constant("encode", torch.device("cuda")) + decode_k = _flux_vae_scaling_constant("decode", torch.device("cuda")) + encode_k = _flux_vae_scaling_constant("encode", torch.device("cuda")) print(f"\n=== 3. VAE linear constants (this build selects {decode_k} decode / {encode_k} encode) ===") print("`implied_k` is the measured peak over pixel area, directly comparable to those literals;") print("fit on the rows whose `math` column matches what this build really does (section 1).") @@ -495,7 +495,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di _flag_non_monotonic(rows) for operation in ("decode", "encode"): # Compare against the column this build actually selects, not against a hard-coded number. - shipped = _flux2_vae_scaling_constant(operation, torch.device("cuda")) + shipped = _flux_vae_scaling_constant(operation, torch.device("cuda")) for force_math in (False, True): ks = [ r["implied_linear_constant"] diff --git a/scripts/calibrate_qwen_vae_working_memory.py b/scripts/calibrate_qwen_vae_working_memory.py index 2f81752aece..552c32a16a1 100644 --- a/scripts/calibrate_qwen_vae_working_memory.py +++ b/scripts/calibrate_qwen_vae_working_memory.py @@ -47,6 +47,8 @@ import torch from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage +from invokeai.backend.util.attention import install_rocm_sdpa_guard + LATENT_SCALE_FACTOR = 8 # (height, width) pixel-space resolutions. Squares to test linearity in area, plus non-square @@ -142,6 +144,11 @@ def _build_input(operation: str, h: int, w: int, z_dim: int, dtype: torch.dtype) @torch.inference_mode() def measure_one(vae_path: str, operation: str, h: int, w: int, dtype: torch.dtype) -> dict: """Measure peak reserved-memory growth for a single decode/encode. Runs in a child process.""" + # The production path on ROCm runs attention through this guard, which bounds the score matrix a + # math-kernel call materializes. Measuring without it fits the constants to a peak the app never + # reaches. A no-op on every other build. + install_rocm_sdpa_guard() + vae = _load_vae(vae_path, dtype) vae.to("cuda") vae.disable_tiling() # Qwen invocations never tile; match that. diff --git a/scripts/generate_docs_json.py b/scripts/generate_docs_json.py index d1ece136151..1a09cea5a84 100644 --- a/scripts/generate_docs_json.py +++ b/scripts/generate_docs_json.py @@ -225,7 +225,7 @@ def write_docs_bundle(bundle: dict[str, Any], output_dir: Path = OUTPUT_DIR) -> def _normalize_value(value: Any) -> str | int | float | bool | list[Any] | dict[str, Any] | None: if isinstance(value, Path): - return str(value) + return value.as_posix() if isinstance(value, (str, int, float, bool)) or value is None: return value if isinstance(value, list): diff --git a/tests/app/invocations/test_anima_vae.py b/tests/app/invocations/test_anima_vae.py index a44d37b4960..76c946b764a 100644 --- a/tests/app/invocations/test_anima_vae.py +++ b/tests/app/invocations/test_anima_vae.py @@ -1,7 +1,6 @@ """Tests for the Anima VAE invocations: which VAEs they accept, working-memory estimation, the tiled-decode decision, and the tiled retry on out-of-memory.""" -import math from unittest.mock import MagicMock, patch import accelerate @@ -13,6 +12,7 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.app.invocations.vae.anima_image_to_latents import AnimaImageToLatentsInvocation from invokeai.app.invocations.vae.anima_latents_to_image import ( + ANIMA_PRETILE_VRAM_FRACTION, ANIMA_VAE_TILE_SIZE, ANIMA_VAE_TILE_STRIDE, AnimaLatentsToImageInvocation, @@ -25,6 +25,8 @@ # diffusers-layout Qwen-Image export as AutoencoderKLQwenImage. WAN21_VAE_LAYOUTS = [AutoencoderKLWan, AutoencoderKLQwenImage] +_L2I = "invokeai.app.invocations.vae.anima_latents_to_image" + def _mock_vae(vae_class: type = AutoencoderKLWan, dtype: torch.dtype = torch.float16) -> MagicMock: vae = MagicMock(spec=vae_class) @@ -126,21 +128,6 @@ def test_the_tiled_decode_estimate_is_left_flat(self): assert estimates[0] == estimates[1] -class TestUseTiledDecode: - @pytest.mark.parametrize("device_type", ["cpu", "mps"]) - def test_non_cuda_never_tiles(self, device_type): - assert AnimaLatentsToImageInvocation._use_tiled_decode(torch.device(device_type), 10**12) is False - - def test_cuda_flips_at_70_percent_of_total_vram(self): - total_vram = 8 * 2**30 - boundary = 0.7 * total_vram - device = torch.device("cuda") - with patch("torch.cuda.get_device_properties", return_value=MagicMock(total_memory=total_vram)) as mock_props: - assert AnimaLatentsToImageInvocation._use_tiled_decode(device, math.floor(boundary)) is False - assert AnimaLatentsToImageInvocation._use_tiled_decode(device, math.ceil(boundary) + 1) is True - mock_props.assert_called_with(device) - - def _build_decode_mocks(latents: torch.Tensor, decoded: torch.Tensor, vae_class: type = AutoencoderKLWan): """Mock the Anima VAE decode path: a spec'd VAE, its LoadedModel wrapper, and the invocation context, wired so `AnimaLatentsToImageInvocation.invoke` runs end-to-end on CPU.""" @@ -348,7 +335,7 @@ def test_oom_while_already_tiled_reraises(self): with ( patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cpu")), - patch.object(AnimaLatentsToImageInvocation, "_use_tiled_decode", return_value=True), + patch("invokeai.app.invocations.vae.anima_latents_to_image.should_pretile_vae_decode", return_value=True), ): with pytest.raises(torch.cuda.OutOfMemoryError): _build_l2i_invocation().invoke(context) @@ -374,6 +361,29 @@ def test_decode_requests_estimated_working_memory(self): vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) +class TestAnimaPretiling: + @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_either_way(self, auto): + """Anima asks with its own, lower fraction (measured on 8GB cards) about the VAE's own device, and does so + whatever `auto_tiled_decode` says: tiling here is what makes a decode that size fast (~1s against 7s+ with + the transformer evicted), not a way around an out-of-memory error.""" + decoded = torch.zeros(1, 3, 1, 64, 64) + vae, vae_info, context = _build_decode_mocks(latents=torch.zeros(1, 16, 32, 32), decoded=decoded) + context.config.get.return_value.auto_tiled_decode = auto + full = estimate_vae_working_memory_anima( + operation="decode", image_tensor=torch.zeros(1, 16, 32, 32), vae=vae, tile_size=None + ) + + with ( + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cpu")), + patch(f"{_L2I}.should_pretile_vae_decode", return_value=True) as pretile, + ): + _build_l2i_invocation().invoke(context) + + pretile.assert_called_once_with(vae_info.compute_device, full, ANIMA_PRETILE_VRAM_FRACTION) + vae.enable_tiling.assert_called_once() + + class TestAnimaImageToLatentsEncode: def test_encode_disables_tiling_and_requests_working_memory(self): vae = _mock_vae(dtype=torch.float32) diff --git a/tests/app/invocations/test_anima_vae_tiling.py b/tests/app/invocations/test_anima_vae_tiling.py index a893b3c2d51..0f6db50c53b 100644 --- a/tests/app/invocations/test_anima_vae_tiling.py +++ b/tests/app/invocations/test_anima_vae_tiling.py @@ -122,7 +122,7 @@ def spy_decode(*args, **kwargs): vae.decode = spy_decode with ( patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cpu")), - patch.object(AnimaLatentsToImageInvocation, "_use_tiled_decode", return_value=True), + patch("invokeai.app.invocations.vae.anima_latents_to_image.should_pretile_vae_decode", return_value=True), ): _build_invocation().invoke(context) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 6378b3b8768..286b6b44511 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -36,7 +36,7 @@ sdpa_score_matrix_bytes, ) from invokeai.backend.util.vae_working_memory import ( - _FLUX2_VAE_SCALING_CONSTANTS, + _FLUX_VAE_SCALING_CONSTANTS, estimate_vae_working_memory_flux2, ) @@ -370,7 +370,7 @@ class TestVaeConstantsFollowTheConvBackend: 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] + constant = _FLUX_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 @@ -378,14 +378,14 @@ def test_the_constant_upper_bounds_its_own_backend(self, backend, operation, px, 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"] + cudnn = _FLUX_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() + backend: consts["encode"] / consts["decode"] for backend, consts in _FLUX_VAE_SCALING_CONSTANTS.items() } assert ratios["cudnn"] == pytest.approx(0.50, abs=0.02) assert ratios["miopen"] == pytest.approx(0.76, abs=0.03) @@ -456,7 +456,7 @@ def estimate(px): at_1024 = estimate(1024) at_1536 = estimate(1536) - miopen = _FLUX2_VAE_SCALING_CONSTANTS["miopen"]["decode"] + miopen = _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] assert at_1024 == 1024 * 1024 * 2 * miopen, "at 1024px the MIOpen linear term is the larger" assert at_1024 >= 6.703 * GB, "and it still covers the measured W7900 peak it was fitted to" assert at_1536 == 36864 * 36864 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT, "at 1536px the score matrix is" @@ -937,7 +937,7 @@ def _estimate(self, px, device, materializing): @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"] + linear = px * px * 2 * _FLUX_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 @@ -945,7 +945,7 @@ def test_the_max_model_bounds_the_measured_miopen_peak(self, px, measured_gib): 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"] + linear = 1024 * 1024 * 2 * _FLUX_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 @@ -953,7 +953,7 @@ def test_the_sum_model_would_have_over_reserved_by_two_thirds(self): 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"] + linear = 1024 * 1024 * 2 * _FLUX_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 @@ -976,7 +976,7 @@ def test_the_working_memory_floor_covers_the_crossover(self, operation, px, meas 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] + constant = _FLUX_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 @@ -986,7 +986,7 @@ def test_the_working_memory_floor_covers_the_crossover(self, operation, px, meas 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"] + linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] assert self._estimate(1024, FUSED, materializing=False) == linear def test_regional_prompting_adds_the_score_matrix(self): diff --git a/tests/app/invocations/test_flux_vae_decode_tiling.py b/tests/app/invocations/test_flux_vae_decode_tiling.py index 1f28fbe1bf8..2b6f998067e 100644 --- a/tests/app/invocations/test_flux_vae_decode_tiling.py +++ b/tests/app/invocations/test_flux_vae_decode_tiling.py @@ -120,6 +120,31 @@ def test_the_reservation_matches_the_decode_that_will_run( _build_invocation(tiled=tiled, tile_size=256 if tiled else 0).invoke(context) assert estimate.call_args.kwargs["tile_size"] == expected_tile_size + @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + """On Windows an allocation that does not fit is paged rather than refused, so the OOM retry never fires and + the decision has to be made from the estimate alone.""" + module = "invokeai.app.invocations.vae.flux_vae_decode" + vae_info, context, tiling_state = _build_decode_mocks( + flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) + ) + context.config.get.return_value.auto_tiled_decode = auto + with ( + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, + ): + _build_invocation().invoke(context) + + if auto: + pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) + assert estimate.call_args.kwargs["tile_size"] == 0 + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + assert tiling_state[0]["use_tiling"] is True + assert tiling_state[0]["tile_sample_min_size"] == DEFAULT_TILE_SAMPLE_MIN_SIZE + else: + pretile.assert_not_called() + assert tiling_state[0]["use_tiling"] is False + class TestOomFallback: def test_an_untiled_oom_retries_once_tiled_and_says_so(self, flux_shaped_vae): diff --git a/tests/app/invocations/test_flux_vae_encode_tiling.py b/tests/app/invocations/test_flux_vae_encode_tiling.py index 81cdc375795..19413922d93 100644 --- a/tests/app/invocations/test_flux_vae_encode_tiling.py +++ b/tests/app/invocations/test_flux_vae_encode_tiling.py @@ -48,6 +48,8 @@ def record_and_encode(x, return_dict=True): vae_info = MagicMock() vae_info.model = vae + # A real device: the node places the image and the generator on the VAE's own device. + vae_info.compute_device = torch.device("cpu") cm = MagicMock() cm.__enter__ = MagicMock(return_value=(None, vae)) cm.__exit__ = MagicMock(return_value=None) diff --git a/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py index 793c4c84e08..a2ca338534c 100644 --- a/tests/app/invocations/test_krea2_denoise.py +++ b/tests/app/invocations/test_krea2_denoise.py @@ -1289,3 +1289,57 @@ def test_an_int8_model_reserves_the_transient_on_top_of_the_activations(self, mo transient = peak_dequant_transient_bytes(quantized, torch.float32) assert transient > 0 assert quantized_bytes == dense_bytes + transient + + +class TestTheAttentionScoreMatrixReachesTheReservation: + """The activation estimate assumes a fused attention kernel. Where the build has none -- ROCm on Windows, MPS -- + the main blocks also build a score matrix (12.9 GiB per call for 1024px unchunked), and the reservation must + price it: for the heads of the loaded model and the sequence the transformer really attends over.""" + + def _reservation_for( + self, monkeypatch, tmp_path, score_bytes: int, invocation=None, **context_kwargs + ) -> tuple[int, list[dict], _Transformer]: + import invokeai.app.invocations.krea2.krea2_denoise as denoise_module + + transformer = _Transformer() + transformer.config = SimpleNamespace(num_attention_heads=48, attention_head_dim=128) + info = _TransformerInfo(transformer) + context = _runtime_context(tmp_path, transformer, **context_kwargs) + context.models.load = lambda _identifier: info + _patch_runtime(monkeypatch) + asked: list[dict] = [] + + def score_matrix_bytes(**kwargs) -> int: + asked.append(kwargs) + return score_bytes + + monkeypatch.setattr(denoise_module, "sdpa_score_matrix_bytes", score_matrix_bytes) + (invocation or _runtime_invocation(cfg_scale=1.0))._run_diffusion(context) + return info.working_mem_bytes, asked, transformer + + def test_the_score_matrix_is_reserved_for_the_attended_sequence(self, monkeypatch, tmp_path) -> None: + fused_bytes, _, _ = self._reservation_for(monkeypatch, tmp_path, score_bytes=0) + materializing_bytes, asked, transformer = self._reservation_for(monkeypatch, tmp_path, score_bytes=12345) + + assert materializing_bytes == fused_bytes + 12345 + assert len(asked) == 1 + assert (asked[0]["num_heads"], asked[0]["head_dim"]) == (48, 128) + assert asked[0]["seq_len"] == transformer.combined_sequence_lengths[0] + assert asked[0]["has_attn_mask"] is False + + def test_the_longer_of_the_cfg_passes_and_the_regional_mask_are_priced(self, monkeypatch, tmp_path) -> None: + invocation = _runtime_invocation(cfg_scale=2.0) + invocation.positive_conditioning = Krea2ConditioningField( + conditioning_name="positive", mask=TensorField(tensor_name="positive-region") + ) + invocation.negative_conditioning = Krea2ConditioningField( + conditioning_name="negative", mask=TensorField(tensor_name="negative-region") + ) + + _, asked, transformer = self._reservation_for( + monkeypatch, tmp_path, score_bytes=1, invocation=invocation, negative_text_seq_len=5 + ) + + assert len(set(transformer.combined_sequence_lengths)) == 2, "the passes must attend over different lengths" + assert asked[0]["seq_len"] == max(transformer.combined_sequence_lengths) + assert asked[0]["has_attn_mask"] is True diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index 8be2c0a3208..246477618fd 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -14,7 +14,10 @@ QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE, patch_qwen_image_vae_tiling, ) -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_qwen_image, + qwen_image_untiled_decode_peak_bytes, +) class TestQwenImageWorkingMemoryEstimate: @@ -51,7 +54,7 @@ def test_constant_selected_per_backend(self, operation, is_rocm, expected_consta hip_value = "7.1.0" if is_rocm else None with patch("torch.version.hip", hip_value): result = estimate_vae_working_memory_qwen_image( - operation=operation, image_tensor=image_tensor, vae=mock_vae + operation=operation, image_tensor=image_tensor, vae=mock_vae, device=torch.device("cuda") ) assert result == h * w * 2 * expected_constant @@ -492,6 +495,190 @@ def test_latents_to_image_honours_the_global_force_tiled_decode(self): assert mock_estimate.call_args.kwargs["tile_size"] == QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE mock_vae.enable_tiling.assert_called_once() + @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto): + """Krea-2 decodes through this node, and it has no OOM retry at all: the estimate is the only trigger.""" + module = "invokeai.app.invocations.vae.qwen_image_latents_to_image" + mock_vae, mock_vae_info = self._mock_vae_info() + mock_vae.decode.side_effect = RuntimeError("stop at decode") + mock_vae.config.z_dim = 16 + mock_vae.config.latents_mean = [0.0] * 16 + mock_vae.config.latents_std = [1.0] * 16 + + mock_context = MagicMock() + mock_context.models.load.return_value = mock_vae_info + mock_context.tensors.load.return_value = torch.zeros(1, 16, 1, 64, 64) + mock_context.config.get.return_value.force_tiled_decode = False + mock_context.config.get.return_value.auto_tiled_decode = auto + + with ( + patch(f"{module}.estimate_vae_working_memory_qwen_image", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.qwen_image_untiled_decode_peak_bytes", return_value=14 * 2**30) as peak, + patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, + patch(f"{module}.SeamlessExt.static_patch_model", return_value=nullcontext()), + ): + invocation = QwenImageLatentsToImageInvocation.model_construct( + latents=MagicMock(latents_name="test_latents"), + vae=MagicMock(vae=MagicMock(), seamless_axes=[]), + tiled=False, + tile_size=0, + ) + with pytest.raises(RuntimeError, match="stop at decode"): + invocation.invoke(mock_context) + + if auto: + # The decision is priced from the measured peak, not from the reservation, which is deliberately padded. + peak.assert_called_once() + pretile.assert_called_once_with(mock_vae_info.compute_device, 14 * 2**30) + assert estimate.call_args.kwargs["tile_size"] == QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE + mock_vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + mock_vae.enable_tiling.assert_called_once() + else: + peak.assert_not_called() + pretile.assert_not_called() + assert estimate.call_count == 1 + mock_vae.enable_tiling.assert_not_called() + + +# Peak reserved bytes per output pixel per element byte actually measured for the Qwen-Image VAE +# decode on ROCm, from the per-point table in `estimate_vae_working_memory_qwen_image`. The shipped +# constant is a flat 5500 -- "max observed + ~8% headroom" -- which is the right figure to *reserve*, +# but the real curve is not flat, so it over-predicts the peak by up to 68%. +MEASURED_ROCM_DECODE_CONSTANT = {512: 5132, 768: 4596, 1024: 4570, 1536: 3273, 1792: 3735, 2048: 4813} + + +class TestMeasuredDecodePeak: + """`qwen_image_untiled_decode_peak_bytes` prices the tiling decision from the measured curve, so a decode whose + real peak fits is not tiled by a reservation's headroom.""" + + def _peak_hw(self, h: int, w: int, device: str = "cuda", hip: str | None = "7.2.0") -> int: + mock_vae = MagicMock(spec=AutoencoderKLQwenImage) + mock_vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) + with patch("torch.version.hip", hip): + return qwen_image_untiled_decode_peak_bytes( + torch.zeros(1, 16, 1, h // 8, w // 8), mock_vae, torch.device(device) + ) + + def _peak(self, px: int, device: str = "cuda", hip: str | None = "7.2.0") -> int: + mock_vae = MagicMock(spec=AutoencoderKLQwenImage) + mock_vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) + with patch("torch.version.hip", hip): + return qwen_image_untiled_decode_peak_bytes( + torch.zeros(1, 16, 1, px // 8, px // 8), mock_vae, torch.device(device) + ) + + @pytest.mark.parametrize("px, constant", [(512, 5132), (1024, 4570), (1536, 3273), (2048, 4813)]) + def test_a_measured_point_is_its_measured_peak(self, px, constant): + assert self._peak(px) == px * px * 2 * constant + + def test_between_two_points_interpolates_the_bytes(self): + """The per-pixel figure dips at 1536^2 while the bytes keep rising, so the bytes are what gets interpolated: + taking the larger neighbouring constant priced a smaller image above a larger one.""" + at_1024 = self._peak(1024) + at_1536 = self._peak(1536) + at_1280 = self._peak(1280) + + assert at_1024 < at_1280 < at_1536, "monotonic in area" + # 1280^2 sits between the 1024^2 and 1536^2 measurements; the chord between them decides it. + span = (at_1536 - at_1024) / (1536**2 - 1024**2) + assert at_1280 == pytest.approx(at_1024 + span * (1280**2 - 1024**2), rel=1e-6) + + def test_a_smaller_image_is_never_priced_above_a_larger_one(self): + """1728x1856 against 1792^2 on CUDA: 0.13% less area, and bracket-max interpolation made it 2.3 GiB dearer. + 1600x1472 against 1024x2304 is the ROCm counterpart: 20.05 GiB against 14.38 GiB before.""" + assert self._peak_hw(1728, 1856, hip=None) <= self._peak_hw(1792, 1792, hip=None) + assert self._peak_hw(1600, 1472) <= self._peak_hw(1024, 2304) + + @pytest.mark.parametrize("hip", ["7.2.0", None], ids=["rocm", "cuda"]) + def test_the_price_never_falls_as_the_area_grows(self, hip): + """Across the whole grid of sizes the decision sees, a larger output area is never priced lower -- the + property the single pairs above are instances of, for both backends' curves.""" + sizes = range(512, 2561, 64) + prices = sorted((h * w, self._peak_hw(h, w, hip=hip)) for h in sizes for w in sizes) + inversions = [(a, b) for a, b in zip(prices, prices[1:], strict=False) if b[1] < a[1]] + assert inversions == [] + + def test_past_the_measured_range_takes_the_reservation_constant(self): + """Nothing was measured there, so the decision uses the figure the reservation uses -- guessing low would + leave a decode larger than anything measured untiled.""" + assert self._peak(3072) == 3072 * 3072 * 2 * 5500 + + @pytest.mark.parametrize("px, constant", [(1024, 2690), (1536, 2671), (1792, 2281), (2048, 2900)]) + def test_a_cuda_card_is_priced_from_the_cuda_measurements(self, px, constant): + """CUDA has its own curve, and its grid stops at 1792px because 2048px ran out of memory when it was + measured; past that point the reservation constant takes over rather than the last value seen.""" + assert self._peak(px, hip=None) == px * px * 2 * constant + + def test_the_peak_stays_under_the_reservation(self): + """The reservation must remain the conservative figure of the two, or it would stop being an upper bound.""" + mock_vae = MagicMock(spec=AutoencoderKLQwenImage) + mock_vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) + latents = torch.zeros(1, 16, 1, 192, 192) + with patch("torch.version.hip", "7.2.0"): + reserved = estimate_vae_working_memory_qwen_image( + operation="decode", image_tensor=latents, vae=mock_vae, device=torch.device("cuda") + ) + peak = qwen_image_untiled_decode_peak_bytes(latents, mock_vae, torch.device("cuda")) + + assert peak < reserved + + +class TestPretilingDoesNotFireOnReservationHeadroom: + """A decode whose real peak fits the card must not be tiled. + + Tiling is not pixel-identical, so it must be triggered by a decode that genuinely does not fit -- + not by the headroom baked into a reservation constant. Over-reserving used to cost only cache + eviction; comparing that padded figure against a share of the card turns it into a silent output + change for images that would have decoded in a single pass. + + Krea-2 decodes through this node too, and it has no out-of-memory retry, so the estimate is the + only trigger either way. + """ + + @pytest.mark.parametrize("px, total_gib", [(1536, 24), (1792, 32)]) + def test_a_decode_whose_measured_peak_fits_the_card_is_not_tiled(self, px, total_gib, monkeypatch): + total_bytes = total_gib * 2**30 + latents = torch.zeros(1, 16, 1, px // 8, px // 8) + + # The constants are measured per backend, so pin it rather than inheriting the host's build. + monkeypatch.setattr(torch.version, "hip", "7.2.0") + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + + # The measured table is fp16, and the estimator runs once here and once inside the node, + # so `parameters()` has to hand out a fresh iterator on every call. + mock_vae = MagicMock(spec=AutoencoderKLQwenImage) + mock_vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) + mock_vae_info = MagicMock() + mock_vae_info.model = mock_vae + mock_vae_info.compute_device = torch.device("cuda", 0) + # Stop the moment the decision has been made: everything after this would need a real GPU. + mock_vae_info.model_on_device.side_effect = RuntimeError("stop after the tiling decision") + + untiled = estimate_vae_working_memory_qwen_image( + operation="decode", image_tensor=latents, vae=mock_vae, tile_size=None + ) + measured_peak = px * px * 2 * MEASURED_ROCM_DECODE_CONSTANT[px] + assert measured_peak < 0.9 * total_bytes, "test shape must actually fit the card" + + mock_context = MagicMock() + mock_context.models.load.return_value = mock_vae_info + mock_context.tensors.load.return_value = latents + mock_context.config.get.return_value.force_tiled_decode = False + mock_context.config.get.return_value.auto_tiled_decode = True + + invocation = QwenImageLatentsToImageInvocation.model_construct( + latents=MagicMock(latents_name="test_latents"), + vae=MagicMock(vae=MagicMock(), seamless_axes=[]), + tiled=False, + tile_size=0, + ) + with pytest.raises(RuntimeError, match="stop after the tiling decision"): + invocation.invoke(mock_context) + + # A pre-tiled decode re-estimates against one tile, so the reservation shrinks. Reserving the + # full-frame figure is what "decoded in a single pass" looks like from here. + assert mock_vae_info.model_on_device.call_args.kwargs["working_mem_bytes"] == untiled + class TestQwenImageVaeTiling: """Exercise the tiling parameters against a real (tiny, randomly initialised) Qwen-Image VAE. diff --git a/tests/app/invocations/test_vae_working_memory_backend.py b/tests/app/invocations/test_vae_working_memory_backend.py new file mode 100644 index 00000000000..5829d061afe --- /dev/null +++ b/tests/app/invocations/test_vae_working_memory_backend.py @@ -0,0 +1,63 @@ +"""Which convolution-backend constant the diffusers-AutoencoderKL estimators use. + +SD1/SDXL, SD3 and CogView4 run the same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so they take the +same measured per-backend constants: MIOpen needs ~1.64x what cuDNN does. The change is large enough to matter -- an +SD1.5 1024px fp32 decode reserves 9.49 GB on cuDNN and 15.36 GB on MIOpen -- and the branch is invisible on a CUDA +machine, so it is pinned here rather than left to whichever card the suite runs on. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL + +from invokeai.backend.util.vae_working_memory import ( + estimate_vae_working_memory_cogview4, + estimate_vae_working_memory_sd3, + estimate_vae_working_memory_sd15_sdxl, +) + + +def _vae(dtype: torch.dtype = torch.float16) -> MagicMock: + vae = MagicMock(spec=AutoencoderKL) + vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=dtype)]) + return vae + + +def _estimate(estimator, device: torch.device, hip: str | None): + latents = torch.zeros(1, 4, 64, 64) # 512px output + with ( + patch("torch.version.hip", hip), + # The mid-block score matrix is the other term in these estimates and has its own tests; hold it at zero so + # the assertions below are about the convolution constant alone. + patch("invokeai.backend.util.vae_working_memory._vae_mid_block_score_matrix_bytes", return_value=0), + ): + if estimator is estimate_vae_working_memory_sd15_sdxl: + return estimator( + operation="decode", image_tensor=latents, vae=_vae(), tile_size=None, fp32=False, device=device + ) + return estimator(operation="decode", image_tensor=latents, vae=_vae(), device=device) + + +ESTIMATORS = [ + pytest.param(estimate_vae_working_memory_sd15_sdxl, id="sd15-sdxl"), + pytest.param(estimate_vae_working_memory_sd3, id="sd3"), + pytest.param(estimate_vae_working_memory_cogview4, id="cogview4"), +] + + +@pytest.mark.parametrize("estimator", ESTIMATORS) +def test_a_rocm_vae_reserves_the_miopen_constant(estimator): + assert _estimate(estimator, torch.device("cuda", 0), hip="7.2.0") == 512 * 512 * 2 * 3600 + + +@pytest.mark.parametrize("estimator", ESTIMATORS) +def test_a_cuda_vae_reserves_the_cudnn_constant(estimator): + assert _estimate(estimator, torch.device("cuda", 0), hip=None) == 512 * 512 * 2 * 2200 + + +@pytest.mark.parametrize("estimator", ESTIMATORS) +def test_a_cpu_only_vae_reserves_the_cudnn_constant_even_on_a_rocm_build(estimator): + """A `cpu_only` VAE does not run on MIOpen, so naming its device has to change the answer.""" + assert _estimate(estimator, torch.device("cpu"), hip="7.2.0") == 512 * 512 * 2 * 2200 diff --git a/tests/app/invocations/test_wan_working_memory.py b/tests/app/invocations/test_wan_working_memory.py index 8732f2aadad..108e2263f39 100644 --- a/tests/app/invocations/test_wan_working_memory.py +++ b/tests/app/invocations/test_wan_working_memory.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import PIL.Image +import pytest import torch from diffusers.models.autoencoders import AutoencoderKLWan @@ -302,13 +303,16 @@ def test_latents_to_video_streams_decode_chunks_directly_to_mp4(self): vae.disable_tiling.assert_called_once() vae.decode.assert_not_called() - def test_latents_to_video_falls_back_to_tiling_when_estimate_exceeds_vram(self): + @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) + def test_latents_to_video_falls_back_to_tiling_when_estimate_exceeds_vram(self, auto): vae = _mock_wan_vae() # Not 256: the point of the change is that the tile comes from a module constant rather than # from the shared cache instance, and with 256 on the instance both implementations agree. vae.tile_sample_min_height = 999 vae_info = _mock_vae_info(vae) + vae_info.compute_device = torch.device("cuda") mock_context = self._video_context(vae_info) + mock_context.config.get.return_value.auto_tiled_decode = auto with ( patch( @@ -327,23 +331,29 @@ def test_latents_to_video_falls_back_to_tiling_when_estimate_exceeds_vram(self): except Exception: pass - assert mock_estimate.call_count == 2 - assert mock_estimate.call_args_list[1].kwargs["tile_size"] == 256 - # All four geometry values are set explicitly rather than inherited from the shared module: - # `enable_tiling` falls back to whatever the instance carries for any argument left out, and - # a min below the inherited stride drops whole bands of the frame. - vae.enable_tiling.assert_called_once_with( - tile_sample_min_height=256, - tile_sample_min_width=256, - tile_sample_stride_height=192, - tile_sample_stride_width=192, - ) - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=4 * 2**30) + if auto: + assert mock_estimate.call_count == 2 + assert mock_estimate.call_args_list[1].kwargs["tile_size"] == 256 + # All four geometry values are set explicitly rather than inherited from the shared module: + # `enable_tiling` falls back to whatever the instance carries for any argument left out, and + # a min below the inherited stride drops whole bands of the frame. + vae.enable_tiling.assert_called_once_with( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_sample_stride_height=192, + tile_sample_stride_width=192, + ) + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=4 * 2**30) + else: + assert mock_estimate.call_count == 1 + vae.enable_tiling.assert_not_called() + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=100 * 2**30) def test_latents_to_video_skips_tiling_for_cpu_only_vae(self): - """A cpu_only VAE runs in system RAM; VRAM-based tiling must not kick in.""" + """A cpu_only VAE runs in system RAM; VRAM-based tiling must not kick in, even with a small GPU present.""" vae = _mock_wan_vae() vae_info = _mock_vae_info(vae, cpu_only=True) + vae_info.compute_device = torch.device("cpu") mock_context = self._video_context(vae_info) with ( @@ -353,6 +363,7 @@ def test_latents_to_video_skips_tiling_for_cpu_only_vae(self): ) as mock_estimate, patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cuda")), patch.object(TorchDevice, "empty_cache"), + patch("torch.cuda.get_device_properties", return_value=MagicMock(total_memory=8 * 2**30)), ): invocation = WanLatentsToVideoInvocation.model_construct( latents=MagicMock(latents_name="l"), vae=MagicMock(vae=MagicMock()), fps=16 diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 37f73fb8db0..2f8a46a894d 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -32,11 +32,11 @@ def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: class TestFluxWorkingMemoryEstimate: @pytest.fixture(autouse=True) def _fused_non_rocm_build(self, monkeypatch): - """`estimate_vae_working_memory_flux` takes no device: it resolves one itself and prices the - mid-block score matrix for it. On a ROCm rig that resolves to `cuda`, where the head-dim - guard charges 4.25GiB for the 1024px case below -- which the expected value here clears by - 1%, so these exact-value assertions turn build-dependent the moment either constant moves. - Pin the fused, non-HIP regime they are written for; `TestClassicVaeEstimators` in + """`estimate_vae_working_memory_flux` resolves a device when the caller passes none, and prices the mid-block + score matrix for it. On a ROCm rig that resolves to `cuda`, where the head-dim guard charges 4.25GiB for the + 1024px case below -- which the expected value here clears by 1%, so these exact-value assertions turn + build-dependent the moment either constant moves. Pin the fused, non-HIP regime they are written for, which + is also the cuDNN column of `_FLUX_VAE_SCALING_CONSTANTS`; `TestClassicVaeEstimators` in `tests/backend/util/test_rocm_sdpa_head_dim_guard.py` owns the ROCm side of this estimator. """ import invokeai.backend.util.attention as attention @@ -45,6 +45,30 @@ def _fused_non_rocm_build(self, monkeypatch): monkeypatch.setattr(attention, "_IS_ROCM", False) monkeypatch.setattr(vwm.TorchDevice, "choose_torch_device", classmethod(lambda cls: torch.device("cpu"))) + @pytest.mark.parametrize( + ("hip", "operation", "constant"), + [(None, "decode", 2200), (None, "encode", 1100), ("7.14.0", "decode", 3600), ("7.14.0", "encode", 2750)], + ids=["cudnn-decode", "cudnn-encode", "miopen-decode", "miopen-encode"], + ) + def test_the_convolution_backend_of_the_vaes_device_picks_the_constant(self, monkeypatch, hip, operation, constant): + """MIOpen's workspaces are larger than cuDNN's: measured 3451 decode / 2688 encode on an RX 9060 XT, the same + as the FLUX.2 VAE. A HIP build reports its GPU as "cuda", so the torch build decides, for the VAE's device.""" + import invokeai.backend.util.vae_working_memory as vwm + + monkeypatch.setattr(torch.version, "hip", hip) + monkeypatch.setattr(vwm, "_vae_mid_block_score_matrix_bytes", lambda *args, **kwargs: 0) + tensor = torch.zeros(1, 16, 64, 64) if operation == "decode" else torch.zeros(1, 3, 512, 512) + + estimate = estimate_vae_working_memory_flux( + operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cuda", 0) + ) + on_cpu = estimate_vae_working_memory_flux( + operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cpu") + ) + + assert estimate == 512 * 512 * 2 * constant + assert on_cpu == 512 * 512 * 2 * (2200 if operation == "decode" else 1100), "a cpu_only VAE runs on cuDNN terms" + def test_the_default_reproduces_the_untiled_estimate(self): """Regression guard for the six call sites that pass no tile_size at all.""" latents = torch.zeros(1, 16, 128, 128) @@ -339,6 +363,41 @@ def test_the_tiling_state_is_restored_afterwards(self, flux_shaped_vae): _build_invocation(tiled=True, tile_size=384).invoke(context) assert (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size) == before + @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + """Where the driver pages instead of raising an OOM, the retry below never fires: the estimate decides.""" + module = "invokeai.app.invocations.vae.z_image_latents_to_image" + vae_info, context, state = _build_decode_mocks( + flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) + ) + context.config.get.return_value.auto_tiled_decode = auto + with ( + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, + ): + _build_invocation().invoke(context) + + if auto: + pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) + assert estimate.call_args.kwargs["tile_size"] == 0 + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + assert state[0]["use_tiling"] is True + else: + pretile.assert_not_called() + assert estimate.call_count == 1 + assert state[0]["use_tiling"] is False + + def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae): + module = "invokeai.app.invocations.vae.z_image_latents_to_image" + _, context, _ = _build_decode_mocks(flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) + with ( + patch(f"{module}.estimate_vae_working_memory_flux", return_value=2 * 2**30) as estimate, + patch(f"{module}.should_pretile_vae_decode") as pretile, + ): + _build_invocation(tiled=True).invoke(context) + assert estimate.call_count == 1 + pretile.assert_not_called() + @pytest.mark.parametrize("tiled,expected_tile_size", [(False, None), (True, 0)]) def test_the_estimate_is_tile_bounded_only_when_tiling(self, flux_shaped_vae, tiled, expected_tile_size): path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" @@ -436,6 +495,8 @@ def record(x, return_dict=True): vae_info = MagicMock() vae_info.model = vae + # A real device: the node places the image on the VAE's own device rather than the session's. + vae_info.compute_device = torch.device("cpu") cm = MagicMock() cm.__enter__ = MagicMock(return_value=(None, vae)) cm.__exit__ = MagicMock(return_value=None) diff --git a/tests/app/services/session_processor/conftest.py b/tests/app/services/session_processor/conftest.py new file mode 100644 index 00000000000..2481f2dad54 --- /dev/null +++ b/tests/app/services/session_processor/conftest.py @@ -0,0 +1,10 @@ +import pytest + + +@pytest.fixture(autouse=True) +def _no_real_paged_vram_reading(monkeypatch: pytest.MonkeyPatch) -> None: + """Every completed session asks Windows how much of the GPU memory is paged out. On a Windows ROCm machine that + would be a real PDH query per session; the tests that need an answer script their own.""" + monkeypatch.setattr( + "invokeai.app.services.session_processor.session_processor_default.paged_bytes", lambda device: None + ) diff --git a/tests/app/services/session_processor/test_vram_paging_warning.py b/tests/app/services/session_processor/test_vram_paging_warning.py new file mode 100644 index 00000000000..1746e2e1b4e --- /dev/null +++ b/tests/app/services/session_processor/test_vram_paging_warning.py @@ -0,0 +1,118 @@ +"""The end-of-session warning when Windows keeps part of a worker's VRAM in shared system memory. + +It must count only what stayed paged across two session ends, speak once per episode (again only after that grew by +the threshold, re-armed once a reading fell below it), say nothing where the answer is unknown, and never fail the +worker. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import invokeai.app.services.session_processor.session_processor_default as processor_module +from invokeai.app.services.session_processor.session_processor_default import ( + _VRAM_PAGING_WARNING_BYTES, + DefaultSessionProcessor, + _SessionWorker, +) + +MIB = 2**20 +DEVICE = torch.device("cuda", 0) + + +class _Logger: + def __init__(self) -> None: + self.warnings: list[str] = [] + self.debugs: list[str] = [] + + def warning(self, message: str, **kwargs) -> None: + self.warnings.append(message) + + def debug(self, message: str, **kwargs) -> None: + self.debugs.append(message) + + +@pytest.fixture +def processor() -> tuple[DefaultSessionProcessor, _Logger]: + logger = _Logger() + processor = DefaultSessionProcessor() + processor._invoker = SimpleNamespace(services=SimpleNamespace(logger=logger)) # type: ignore[attr-defined] + return processor, logger + + +def _warned_after_each_session(monkeypatch, processor, readings) -> list[bool]: + """End one session per scripted `paged_bytes` reading on one worker; whether each one warned.""" + processor, logger = processor + worker = _SessionWorker(DEVICE, runner=SimpleNamespace()) # type: ignore[arg-type] + values = iter(readings) + + def paged_bytes(device): + assert device == DEVICE, "the worker's own device is the one asked" + return next(values) + + monkeypatch.setattr(processor_module, "paged_bytes", paged_bytes) + warned = [] + for _ in readings: + before = len(logger.warnings) + processor._warn_if_vram_paged(worker) + warned.append(len(logger.warnings) > before) + return warned + + +def test_warns_once_per_episode_and_again_only_after_growth(monkeypatch, processor): + threshold = _VRAM_PAGING_WARNING_BYTES + readings = [ + 6 * MIB, # an idle process + threshold + 100 * MIB, # paged, once so far + threshold + 300 * MIB, # still paged: warn + 2 * threshold + 400 * MIB, # what stayed paged grew by less than the threshold since the warning + 2 * threshold + 500 * MIB, # grew by the threshold: warn again + 100 * MIB, # recovered below the threshold: re-armed + threshold, + threshold, # a new episode: warn + ] + + warned = _warned_after_each_session(monkeypatch, processor, readings) + + assert warned == [False, False, True, False, True, False, False, True] + assert "cuda:0" in processor[1].warnings[0] + assert "0.6 GiB" in processor[1].warnings[0] # what stayed paged (0.6), not the latest reading (0.8) + + +def test_paging_that_clears_between_sessions_is_not_reported(monkeypatch, processor): + """A decode can overflow for a few seconds at the end of every session and then return to VRAM.""" + paged, cleared = _VRAM_PAGING_WARNING_BYTES + 400 * MIB, 6 * MIB + readings = [paged, cleared, paged, cleared, paged] + + assert _warned_after_each_session(monkeypatch, processor, readings) == [False] * 5 + + +def test_silent_where_the_answer_is_unknown(monkeypatch, processor): + assert _warned_after_each_session(monkeypatch, processor, [None, None]) == [False, False] + + +def test_a_failing_counter_is_logged_at_debug_and_does_not_fail_the_worker(monkeypatch, processor): + processor_instance, logger = processor + + def broken(device): + raise OSError("PDH query failed") + + monkeypatch.setattr(processor_module, "paged_bytes", broken) + + processor_instance._warn_if_vram_paged(_SessionWorker(DEVICE, runner=SimpleNamespace())) # type: ignore[arg-type] + + assert logger.warnings == [] + assert len(logger.debugs) == 1 + + +def test_the_worker_checks_after_each_session(monkeypatch): + from tests.app.services.session_processor.test_session_processor_cancel_guard import _run_guard_scenario + + checked = [] + monkeypatch.setattr(DefaultSessionProcessor, "_warn_if_vram_paged", lambda self, worker: checked.append(worker)) + + _canceled, run_items, worker = _run_guard_scenario(statuses=["in_progress"], set_cancel_event=False) + + assert run_items == [42] + assert checked == [worker] diff --git a/tests/app/util/test_torch_cuda_allocator.py b/tests/app/util/test_torch_cuda_allocator.py index 234002dd6fd..d3461983ec4 100644 --- a/tests/app/util/test_torch_cuda_allocator.py +++ b/tests/app/util/test_torch_cuda_allocator.py @@ -1,8 +1,124 @@ +import os +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock + import pytest import torch +import invokeai.app.util.torch_cuda_allocator as torch_cuda_allocator +from invokeai.app.util.torch_cuda_allocator import ROCM_WINDOWS_ALLOC_CONF, apply_rocm_windows_allocator_default from tests.dangerously_run_function_in_subprocess import dangerously_run_function_in_subprocess +_ALLOCATOR_ENV_VARS = ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_HIP_ALLOC_CONF") + + +@pytest.fixture +def rocm_windows(monkeypatch: pytest.MonkeyPatch, tmp_path): + """A Windows ROCm install with no allocator configuration in the environment. + + `_installed_torch_is_rocm` has two sources -- the distribution version and `hip` in torch's own `version.py` -- + so both are stubbed. Leaving the file lookup real made every case that expects "not ROCm" read the host's torch + and fail on any ROCm developer machine, which is this change's own test bench. + """ + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: "2.12.0+rocm7.14.1") + cuda_build = tmp_path / "torch" + cuda_build.mkdir() + (cuda_build / "version.py").write_text("hip: Optional[str] = None\n", encoding="utf-8") + monkeypatch.setattr( + torch_cuda_allocator.importlib.util, + "find_spec", + lambda name: SimpleNamespace(submodule_search_locations=[str(cuda_build)]), + ) + monkeypatch.delitem(sys.modules, "torch") # as at startup, before anything imports it + for var in _ALLOCATOR_ENV_VARS: + # setenv first so teardown also removes what the code under test sets: delenv on an absent + # variable records nothing to undo, and the leaked value would reach later tests' allocator. + monkeypatch.setenv(var, "") + monkeypatch.delenv(var) + return monkeypatch + + +class TestRocmWindowsAllocatorDefault: + def test_defaults_to_expandable_segments(self, rocm_windows): + logger = MagicMock() + + apply_rocm_windows_allocator_default(logger) + + assert os.environ["PYTORCH_CUDA_ALLOC_CONF"] == ROCM_WINDOWS_ALLOC_CONF + logger.info.assert_called_once() + + @pytest.mark.parametrize("var", _ALLOCATOR_ENV_VARS) + def test_any_existing_allocator_configuration_wins(self, rocm_windows, var): + rocm_windows.setenv(var, "expandable_segments:False") + + apply_rocm_windows_allocator_default(MagicMock()) + + assert os.environ[var] == "expandable_segments:False" + if var != "PYTORCH_CUDA_ALLOC_CONF": + assert "PYTORCH_CUDA_ALLOC_CONF" not in os.environ + + @pytest.mark.parametrize( + ("platform", "version"), + [("linux", "2.13.0+rocm7.2"), ("win32", "2.7.1+cu128"), ("win32", "2.7.1"), ("win32", None)], + ids=["linux-rocm", "windows-cuda", "windows-cpu", "no-torch-metadata"], + ) + def test_other_installs_are_left_alone(self, rocm_windows, platform, version): + rocm_windows.setattr(sys, "platform", platform) + rocm_windows.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: version) + logger = MagicMock() + + apply_rocm_windows_allocator_default(logger) + + assert "PYTORCH_CUDA_ALLOC_CONF" not in os.environ + logger.info.assert_not_called() + + def test_too_late_once_torch_is_imported(self, rocm_windows): + """The allocator reads the variable at import; setting it later would only mislead the model cache, which + parses it to decide whether allocator-held blocks count as free.""" + rocm_windows.setitem(sys.modules, "torch", torch) + logger = MagicMock() + + apply_rocm_windows_allocator_default(logger) + + assert "PYTORCH_CUDA_ALLOC_CONF" not in os.environ + logger.warning.assert_called_once() + + +class TestRocmBuildDetection: + """A locally built Windows ROCm wheel is versioned like `2.8.0a0+gitfc14c65`, so the distribution version alone + would leave it on the default allocator while `wddm` (which reads `torch.version.hip`) still caps its budget.""" + + @pytest.mark.parametrize( + ("hip_line", "expected"), + [("hip: Optional[str] = '7.14.60850'", True), ("hip: Optional[str] = None", False)], + ids=["rocm-build", "cuda-build"], + ) + def test_a_build_without_rocm_in_its_version_is_read_from_torchs_version_file( + self, monkeypatch, tmp_path, hip_line, expected + ): + package = tmp_path / "torch" + package.mkdir() + (package / "version.py").write_text(f"__version__ = '2.8.0a0+gitfc14c65'\n{hip_line}\n", encoding="utf-8") + monkeypatch.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: "2.8.0a0+gitfc14c65") + monkeypatch.setattr( + torch_cuda_allocator.importlib.util, + "find_spec", + lambda name: SimpleNamespace(submodule_search_locations=[str(package)]), + ) + + assert torch_cuda_allocator._installed_torch_is_rocm() is expected + + def test_an_unreadable_installation_is_not_rocm(self, monkeypatch): + monkeypatch.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: "2.8.0a0+gitfc14c65") + monkeypatch.setattr( + torch_cuda_allocator.importlib.util, "find_spec", MagicMock(side_effect=ValueError("no spec")) + ) + + assert torch_cuda_allocator._installed_torch_is_rocm() is False + + # These tests are a bit fiddly, because the depend on the import behaviour of torch. They use subprocesses to isolate # the import behaviour of torch, and then check that the function behaves as expected. We have to hack in some logging # to check that the tested function is behaving as expected. diff --git a/tests/backend/ltx2/test_ltx2_real_files.py b/tests/backend/ltx2/test_ltx2_real_files.py index 56557d790cc..4adb08e4aa2 100644 --- a/tests/backend/ltx2/test_ltx2_real_files.py +++ b/tests/backend/ltx2/test_ltx2_real_files.py @@ -194,9 +194,9 @@ def _accelerator() -> torch.device: device = TorchDevice.choose_torch_device() if device.type == "cpu": pytest.skip("needs an accelerator") - from invokeai.backend.util.attention import install_rocm_sdpa_head_dim_guard + from invokeai.backend.util.attention import install_rocm_sdpa_guard - install_rocm_sdpa_head_dim_guard() + install_rocm_sdpa_guard() return device diff --git a/tests/backend/model_manager/load/model_cache/conftest.py b/tests/backend/model_manager/load/model_cache/conftest.py new file mode 100644 index 00000000000..c4240fb9223 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/conftest.py @@ -0,0 +1,12 @@ +import pytest + + +@pytest.fixture(autouse=True) +def _no_windows_video_memory_budget(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep the cache's VRAM arithmetic on whatever each test patches into ``torch.cuda``. + + Many of these tests fake a ``cuda`` execution device and patch ``torch.cuda.mem_get_info``. On a ROCm build under + Windows, ``TorchDevice.cuda_mem_get_info`` would additionally cap that fake with the real adapter's WDDM budget, + making the expected figures depend on the machine running the suite. Tests of the cap patch it back in. + """ + monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: None) diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py b/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py index 2274fb01090..b7f3023b8cc 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py @@ -15,6 +15,7 @@ import torch from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.util.devices import TorchDevice GB = 1024**3 MB = 1024**2 @@ -187,3 +188,73 @@ def test_physical_availability_shares_the_reclaimable_credit_policy(monkeypatch: monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") assert cache._get_physical_vram_available() == 3 * GB + + +def test_the_windows_video_memory_budget_caps_the_measured_free_vram(monkeypatch: pytest.MonkeyPatch): + """On Windows ROCm, torch reports the device total minus this process's usage. Windows pages allocations into + shared system memory once the process passes its video-memory budget, so the budget's headroom is what the cache + may plan with -- not the larger figure torch reports.""" + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + execution_device="cpu", + storage_device="cpu", + logger=MagicMock(), + shared_cpu_weights=None, + ) + cache._execution_device = torch.device("cuda") # policy only; every VRAM query below is patched out + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (12 * GB, 16 * GB)) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda device: 2 * GB) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda device: 2 * GB) + monkeypatch.setattr(torch.cuda, "memory_stats", lambda device: {"inactive_split_bytes.all.current": 0}) + monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: 11 * GB) + + # 11 GB budget - 4 GB live = 7 GB headroom, + 2 GB allocated - 1 GB working - 2 GB in use; not torch's 12 GB. + assert cache._get_vram_available(None) == 6 * GB + + +@pytest.mark.parametrize("peer_busy", [False, True], ids=["alone", "peer-device-busy"]) +def test_offloading_under_expandable_segments_stops_once_enough_is_free(monkeypatch: pytest.MonkeyPatch, peer_busy): + """Under expandable segments the driver sees an offloaded model's pages only after empty_cache(), and no + allocator credit stands in for them. Re-measuring without it saw no progress and unloaded every unlocked model. + + `empty_cache` is peer-aware: while another generation device is mid-session it defers instead of releasing, so + the measurement cannot see the offload at all and the freed bytes are credited to it directly.""" + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + execution_device="cpu", + storage_device="cpu", + logger=MagicMock(), + shared_cpu_weights=None, + ) + cache._execution_device = torch.device("cuda") # policy only; every VRAM query below is patched out + device = {"free": 0, "allocated": 12 * GB, "unmappable": 0} + + def move_to_ram(entry, bytes_to_free): + device["allocated"] -= 4 * GB + device["unmappable"] += 4 * GB + return 4 * GB + + def empty_cache() -> bool: + if peer_busy: # deferred: nothing is released, and the driver keeps reporting the old figure + return False + device["free"] += device["unmappable"] + device["unmappable"] = 0 + return True + + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda d: (device["free"], 16 * GB)) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda d: device["allocated"]) + monkeypatch.setattr(TorchDevice, "empty_cache", empty_cache) + monkeypatch.setattr(cache, "_move_model_to_ram", MagicMock(side_effect=move_to_ram)) + for key in "ABC": + entry = MagicMock(key=key, is_locked=False) + entry.cached_model.total_bytes.return_value = 4 * GB + cache._cached_models[key] = entry + + # 2 GB needed + 1 GB working memory: one 4 GB model is enough. + assert cache._offload_unlocked_models(2 * GB) == 4 * GB + assert cache._move_model_to_ram.call_count == 1 diff --git a/tests/backend/util/test_devices.py b/tests/backend/util/test_devices.py index cb42672138f..1f01eb2df6a 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -615,6 +615,25 @@ def test_xpu_mem_get_info_estimates_when_native_and_sysman_both_fail(native_erro assert TorchDevice.xpu_mem_get_info(torch.device("xpu")) == (30 * gib, 32 * gib) +@pytest.mark.parametrize( + ("budget_gib", "expected_free_gib"), + [(None, 10), (13, 7), (16, 10), (5, 0)], + ids=["no-budget", "budget-is-tighter", "torch-is-tighter", "over-budget"], +) +def test_cuda_mem_get_info_is_capped_by_the_windows_video_memory_budget(budget_gib, expected_free_gib): + """On Windows ROCm, torch's free figure ignores other processes and the point where Windows starts paging. The + headroom is the budget minus this process's live allocations (16 - 10 = 6 GiB here).""" + gib = 1024**3 + device = torch.device("cuda", 0) + budget = None if budget_gib is None else budget_gib * gib + with ( + patch.object(torch.cuda, "mem_get_info", return_value=(10 * gib, 16 * gib)), + patch("invokeai.backend.util.devices.video_memory_budget", return_value=budget) as mock_budget, + ): + assert TorchDevice.cuda_mem_get_info(device) == (expected_free_gib * gib, 16 * gib) + mock_budget.assert_called_once_with(device) + + def test_get_generation_devices_auto_expands_to_all_xpu(): """With no CUDA, `auto` enumerates every visible XPU device.""" config = get_config() diff --git a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py index 8de8430dd3b..e82d1ccead1 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -1,25 +1,31 @@ -"""The ROCm head-dim guard for `F.scaled_dot_product_attention`. +"""The ROCm SDPA guard for `F.scaled_dot_product_attention`. The fused ROCm kernels are wrong above head_dim 256 (see `ROCM_FUSED_SDPA_MAX_HEAD_DIM`). The guard must send exactly those calls to the math kernel, leave every other call on the original function with its arguments intact, install once, and make the working-memory estimator budget the score -matrix the math kernel materializes. +matrix the math kernel materializes. Math-kernel calls whose score matrix exceeds +`SDPA_MATH_CHUNK_BYTES` must run in chunks that each stay within it and together give the unchunked +result. """ import pytest import torch import torch.nn.functional as F +from torch.nn.attention import SDPBackend import invokeai.backend.util.attention as attention from invokeai.backend.util.attention import ( ROCM_FUSED_SDPA_MAX_HEAD_DIM, - install_rocm_sdpa_head_dim_guard, + install_rocm_sdpa_guard, rocm_sdpa_uses_math_kernel, sdpa_score_matrix_bytes, ) WIDE = ROCM_FUSED_SDPA_MAX_HEAD_DIM + 1 NARROW = ROCM_FUSED_SDPA_MAX_HEAD_DIM +BYTES = attention.SDPA_MATH_BYTES_PER_SCORE_ELEMENT +# torch's own function, captured before any test installs the guard: the unchunked oracle. +TORCH_SDPA = F.scaled_dot_product_attention def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: @@ -60,10 +66,8 @@ def spy(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, sc monkeypatch.setattr(F, "scaled_dot_product_attention", spy) monkeypatch.setattr(attention, "_IS_ROCM", True) # CPU tensors stand in for the ROCm device; the real predicate is covered by TestPredicate. - monkeypatch.setattr( - attention, "rocm_sdpa_uses_math_kernel", lambda device_type, head_dim: head_dim > ROCM_FUSED_SDPA_MAX_HEAD_DIM - ) - install_rocm_sdpa_head_dim_guard() + monkeypatch.setattr(attention, "_rocm_cuda", lambda device_type: True) + install_rocm_sdpa_guard() return calls @@ -123,12 +127,12 @@ def test_wide_heads_refuse_arguments_the_math_call_cannot_forward(self, monkeypa def test_install_is_idempotent_and_inert_off_rocm(self, monkeypatch): _install_as_if_rocm(monkeypatch) guarded = F.scaled_dot_product_attention - install_rocm_sdpa_head_dim_guard() + install_rocm_sdpa_guard() assert F.scaled_dot_product_attention is guarded, "a second install must not wrap the wrapper" monkeypatch.setattr(F, "scaled_dot_product_attention", guarded.__wrapped__) monkeypatch.setattr(attention, "_IS_ROCM", False) - install_rocm_sdpa_head_dim_guard() + install_rocm_sdpa_guard() assert F.scaled_dot_product_attention is guarded.__wrapped__, "no guard on a non-ROCm build" def test_install_backs_off_when_the_math_op_rejects_the_probe(self, monkeypatch, caplog): @@ -138,7 +142,7 @@ def test_install_backs_off_when_the_math_op_rejects_the_probe(self, monkeypatch, monkeypatch.setattr(attention, "_math_sdpa", lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("no such op"))) caplog.set_level("WARNING", logger=attention.__name__) - install_rocm_sdpa_head_dim_guard() + install_rocm_sdpa_guard() assert F.scaled_dot_product_attention is original assert any("NOT installed" in r.message for r in caplog.records) @@ -157,6 +161,180 @@ def test_apply_monkeypatches_installs_the_guard(self, monkeypatch): assert getattr(F.scaled_dot_product_attention, attention._ROCM_SDPA_GUARD_SENTINEL, False) +def _install_chunking( + monkeypatch, budget: int, materializes: bool | None = True, torch_error: str | None = None +) -> list[tuple[torch.Size, torch.Size]]: + """Install the guard as if on ROCm with a small chunk budget; record (query, key) shapes of every call that reaches + torch's function or the math op. ``materializes`` pins the dispatch probe (None leaves it real); ``torch_error`` + makes torch's function raise it, as a build without an eligible kernel does.""" + shapes: list[tuple[torch.Size, torch.Size]] = [] + + def spy(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=None, enable_gqa=False): + shapes.append((query.shape, key.shape)) + if torch_error is not None: + raise RuntimeError(torch_error) + return TORCH_SDPA(query, key, value, attn_mask, dropout_p, is_causal, scale=scale, enable_gqa=enable_gqa) + + real_math = attention._math_sdpa + + def math_spy(query, key, value, *args): + shapes.append((query.shape, key.shape)) + return real_math(query, key, value, *args) + + monkeypatch.setattr(F, "scaled_dot_product_attention", spy) + monkeypatch.setattr(attention, "_math_sdpa", math_spy) + monkeypatch.setattr(attention, "_IS_ROCM", True) + monkeypatch.setattr(attention, "_rocm_cuda", lambda device_type: True) + monkeypatch.setattr(attention, "SDPA_MATH_CHUNK_BYTES", budget) + if materializes is not None: + monkeypatch.setattr(attention, "_materializes", lambda *args: materializes) + install_rocm_sdpa_guard() + shapes.clear() # the install probe's own math call + return shapes + + +def _scores(shape_pair: tuple[torch.Size, torch.Size]) -> int: + query_shape, key_shape = shape_pair + return query_shape[0] * query_shape[1] * query_shape[2] * key_shape[-2] * BYTES + + +class TestChunking: + """Math-kernel attention over the budget runs in chunks; every chunk stays within it, together they are exact.""" + + def test_heads_are_split_first_and_the_result_matches_the_unchunked_call(self, monkeypatch): + """Bitwise here because the CPU kernel's arithmetic does not depend on the batch it runs over; on a ROCm + card the batched GEMM re-tiles with the group count and lands within a bf16 ulp instead.""" + per_head = 64 * 48 * BYTES + shapes = _install_chunking(monkeypatch, budget=2 * per_head) + torch.manual_seed(0) + q = torch.randn(1, 6, 64, 32) + k, v = torch.randn(1, 6, 48, 32), torch.randn(1, 6, 48, 32) + + out = F.scaled_dot_product_attention(q, k, v) + + assert [s[0][1] for s in shapes] == [2, 2, 2], "three head groups, each with every query row" + assert torch.equal(out, TORCH_SDPA(q, k, v)) + + def test_a_single_oversized_head_is_split_by_query_rows(self, monkeypatch): + """The one-head VAE mid-block: only query rows are left to split, on the wide-head math path.""" + per_row = 96 * BYTES + shapes = _install_chunking(monkeypatch, budget=20 * per_row) + torch.manual_seed(1) + q, k, v = (torch.randn(1, 1, 96, WIDE) for _ in range(3)) + + out = F.scaled_dot_product_attention(q, k, v) + + assert [s[0][2] for s in shapes] == [20, 20, 20, 20, 16] + torch.testing.assert_close(out, _reference(q, k, v), atol=1e-5, rtol=1e-4) + + @pytest.mark.parametrize( + ("heads", "kv_heads", "budget_heads", "kwargs"), + [ + (4, 4, 2, {"attn_mask": torch.rand(64, 48) > 0.3}), + (4, 4, 2, {"attn_mask": torch.randn(2, 1, 64, 48)}), + (4, 4, 2, {"attn_mask": torch.randn(2, 4, 64, 48)}), + (4, 4, 2, {"attn_mask": torch.rand(1, 1, 1, 48) > 0.2}), + (4, 4, 2, {"attn_mask": torch.randn(4, 64, 48)}), + (4, 4, 2, {"scale": 0.3}), + (6, 2, 3, {"enable_gqa": True}), + (6, 2, 2, {"enable_gqa": True}), + (4, 1, 2, {}), + (4, 4, 0.5, {"attn_mask": torch.rand(64, 48) > 0.3}), + (4, 4, 0.5, {"attn_mask": torch.randn(2, 1, 64, 48)}), + (4, 4, 0.5, {"attn_mask": torch.randn(2, 4, 64, 48)}), + (4, 4, 0.5, {"attn_mask": torch.rand(2, 1, 1, 48) > 0.2}), + ], + ids=[ + "bool-2d", + "additive-broadcast-heads", + "additive-per-head", + "bool-broadcast-rows", + "3d-per-head", + "scale", + "gqa-whole-groups", + "gqa-group-over-budget", + "kv-broadcast", + "rows-bool-2d", + "rows-additive-broadcast-heads", + "rows-additive-per-head", + "rows-bool-broadcast-rows", + ], + ) + def test_masks_scale_and_grouped_heads_keep_torchs_result(self, monkeypatch, heads, kv_heads, budget_heads, kwargs): + budget = int(budget_heads * 2 * 64 * 48 * BYTES) # under one head, the query rows are split + shapes = _install_chunking(monkeypatch, budget=budget) + torch.manual_seed(2) + q = torch.randn(2, heads, 64, 32) + k, v = torch.randn(2, kv_heads, 48, 32), torch.randn(2, kv_heads, 48, 32) + + out = F.scaled_dot_product_attention(q, k, v, **kwargs) + + assert len(shapes) > 1, "the call must actually have been split" + assert all(_scores(pair) <= budget for pair in shapes), "a chunk exceeded the budget" + torch.testing.assert_close(out, TORCH_SDPA(q, k, v, **kwargs), atol=1e-5, rtol=1e-4) + + @pytest.mark.parametrize( + ("q_shape", "kv_shape"), + [((2, 4, 64, 32), (4, 48, 32)), ((1, 4, 64, 32), (2, 4, 48, 32))], + ids=["3d-kv", "kv-wider-batch"], + ) + def test_kv_broadcasts_the_chunks_cannot_slice_stay_whole(self, monkeypatch, q_shape, kv_shape): + shapes = _install_chunking(monkeypatch, budget=BYTES) + torch.manual_seed(3) + q, k, v = torch.randn(q_shape), torch.randn(kv_shape), torch.randn(kv_shape) + + out = F.scaled_dot_product_attention(q, k, v) + + assert len(shapes) == 1 + torch.testing.assert_close(out, TORCH_SDPA(q, k, v)) + + @pytest.mark.parametrize("kwargs", [{"is_causal": True}, {"dropout_p": 0.1}], ids=["causal", "dropout"]) + def test_causal_and_dropout_calls_stay_whole(self, monkeypatch, kwargs): + shapes = _install_chunking(monkeypatch, budget=BYTES) + q = torch.randn(1, 2, 16, 8) + + F.scaled_dot_product_attention(q, q, q, **kwargs) + + assert shapes == [(q.shape, q.shape)] + + def test_a_call_torch_serves_with_a_fused_kernel_stays_whole(self, monkeypatch): + shapes = _install_chunking(monkeypatch, budget=BYTES, materializes=None) + monkeypatch.setattr(torch.ops.aten, "_fused_sdp_choice", lambda *a, **k: int(SDPBackend.FLASH_ATTENTION)) + q = torch.randn(1, 2, 16, 8) + + F.scaled_dot_product_attention(q, q, q) + + assert shapes == [(q.shape, q.shape)] + + def test_no_eligible_kernel_raises_the_same_error_as_the_whole_call(self, monkeypatch): + """The probe raising means torch has no backend for the call (an exclusive policy, say): chunking must neither + hide that nor route around it onto a kernel the policy forbids.""" + message = "No available kernel. Aborting execution." + shapes = _install_chunking(monkeypatch, budget=16 * 16 * BYTES, materializes=None, torch_error=message) + + def no_backend(*args, **kwargs): + raise RuntimeError("No viable backend for scaled_dot_product_attention was found.") + + monkeypatch.setattr(torch.ops.aten, "_fused_sdp_choice", no_backend) + q = torch.randn(1, 2, 16, 8) + + with pytest.raises(RuntimeError, match=message): + F.scaled_dot_product_attention(q, q, q) + assert shapes[0][0][1] == 1, "the first chunk, run through torch's own function, raised" + + def test_the_output_dtype_follows_autocast(self, monkeypatch): + per_head = 32 * 32 * BYTES + _install_chunking(monkeypatch, budget=per_head) + q = torch.randn(1, 4, 32, 16) + + with torch.autocast("cpu", dtype=torch.bfloat16): + chunked = F.scaled_dot_product_attention(q, q, q) + whole = TORCH_SDPA(q, q, q) + + assert chunked.dtype == whole.dtype + torch.testing.assert_close(chunked.float(), whole.float(), atol=1e-2, rtol=1e-2) + + class TestEstimator: def test_wide_heads_on_rocm_reserve_the_score_matrix_without_asking_torch(self, monkeypatch): monkeypatch.setattr(attention, "_IS_ROCM", True) @@ -193,6 +371,28 @@ def test_wide_heads_off_rocm_still_ask_torch(self, monkeypatch): == 0 ) + def test_the_estimate_is_capped_at_one_chunk_only_while_the_guard_chunks(self, monkeypatch): + """With the guard installed, one chunk is alive at a time; without it -- install failed, or a script that + never installs it -- the whole matrix is built and must be budgeted in full.""" + monkeypatch.setattr(attention, "_IS_ROCM", True) + monkeypatch.setattr(F, "scaled_dot_product_attention", TORCH_SDPA) + device = torch.device("cuda", 0) + full = 16384 * 16384 * BYTES + + def estimate() -> int: + return sdpa_score_matrix_bytes( + device=device, dtype=torch.bfloat16, num_heads=1, head_dim=WIDE, seq_len=16384 + ) + + assert estimate() == full + + install_rocm_sdpa_guard() + assert estimate() == attention.SDPA_MATH_CHUNK_BYTES + assert ( + sdpa_score_matrix_bytes(device=device, dtype=torch.bfloat16, num_heads=1, head_dim=WIDE, seq_len=64) + == 64 * 64 * BYTES + ), "a matrix within one chunk is priced as it is" + needs_rocm = pytest.mark.skipif(torch.version.hip is None or not torch.cuda.is_available(), reason="needs a ROCm GPU") @@ -218,23 +418,33 @@ def test_a_fused_kernel_is_still_wrong_for_the_wide_head(self): reference = attention._math_sdpa(q, k, v, None, 0.0, False, None, False).float() # The unguarded entry point, whether or not the guard is installed in this process. unguarded = getattr(F.scaled_dot_product_attention, "__wrapped__", F.scaled_dot_product_attention) - wrong = 0 + wrong = ran = 0 for backend in (SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION): with sdpa_policy([backend]): try: out = unguarded(q, k, v).float() except RuntimeError: continue # the backend refused the shape: not wrong, just absent + ran += 1 err = (out - reference).abs().max().item() / reference.abs().max().item() wrong += int(not torch.isfinite(out).all() or err > 0.05) + if not ran: + pytest.skip("no fused kernel runs this shape on this build (e.g. gfx1200)") assert wrong > 0, "every fused kernel now agrees with math: the ROCm head-dim guard may be unnecessary" - def test_the_guard_is_exact_under_allocation_churn(self): + def test_the_guard_is_exact_under_allocation_churn(self, monkeypatch): + """The fused kernels' failure depended on allocator state; the guarded path must not. The shape is over one + chunk, so the reference is the guarded result itself (query chunks differ from one unchunked call by up to a + bf16 ulp) -- what must hold is that it never changes.""" import random - install_rocm_sdpa_head_dim_guard() + monkeypatch.setattr(F, "scaled_dot_product_attention", F.scaled_dot_product_attention) + install_rocm_sdpa_guard() q, k, v = self._inputs() - reference = attention._math_sdpa(q, k, v, None, 0.0, False, None, False) + reference = F.scaled_dot_product_attention(q, k, v) + torch.testing.assert_close( + reference.float(), attention._math_sdpa(q, k, v, None, 0.0, False, None, False).float(), atol=2e-3, rtol=0 + ) rnd = random.Random(0) for _ in range(20): junk = [ @@ -247,15 +457,42 @@ def test_the_guard_is_exact_under_allocation_churn(self): torch.cuda.empty_cache() assert torch.equal(out, reference) + @pytest.mark.parametrize( + "shape", + [(1, 30, 4128, 128), (1, 1, 36864, 512)], + ids=["z-image-1024", "flux-vae-mid-block-1536"], + ) + def test_math_attention_peaks_within_one_chunk(self, monkeypatch, shape): + """Z-Image at 1024px would build 8 GiB of scores, a 1536px VAE decode 21 GiB. Where torch has a fused kernel + for the narrow head, the call is not chunked and peaks far lower anyway.""" + monkeypatch.setattr(F, "scaled_dot_product_attention", F.scaled_dot_product_attention) + install_rocm_sdpa_guard() + torch.manual_seed(0) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + F.scaled_dot_product_attention(q, k, v) # warm-up: kernel selection and allocator growth + torch.cuda.synchronize() + base = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + + out = F.scaled_dot_product_attention(q, k, v) + torch.cuda.synchronize() + + peak = torch.cuda.max_memory_allocated() - base + assert peak <= attention.SDPA_MATH_CHUNK_BYTES + out.numel() * out.element_size() + class TestClassicVaeEstimators: """The SD/SDXL and FLUX.1 VAE estimators price the mid-block score matrix the guard forces.""" @pytest.fixture(autouse=True) def _rocm_cuda(self, monkeypatch): + """The guard's ROCm switch and the convolution backend are separate questions, and this class is about the + first: the assertions below carry the cuDNN linear term, so pin `torch.version.hip` (which selects MIOpen's + larger constants) away from it while the attention side stays on ROCm.""" import invokeai.backend.util.vae_working_memory as vwm monkeypatch.setattr(attention, "_IS_ROCM", True) + monkeypatch.setattr(torch.version, "hip", None) monkeypatch.setattr(vwm.TorchDevice, "choose_torch_device", classmethod(lambda cls: torch.device("cuda", 0))) @staticmethod diff --git a/tests/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py new file mode 100644 index 00000000000..4aba3c32e7b --- /dev/null +++ b/tests/backend/util/test_vae_pretile.py @@ -0,0 +1,147 @@ +"""The decode nodes' up-front tiling rule: tile when an untiled decode would claim more than a share of the VAE's GPU. + +Waiting for an out-of-memory error does not work where the driver pages into system memory instead of failing +(Windows ROCm always, NVIDIA's sysmem fallback by default), so the rule has to hold on the estimate alone. +""" + +import math +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION, should_pretile_vae_decode + +TOTAL = 16 * 2**30 + + +@pytest.mark.parametrize("fraction", [0.7, VAE_PRETILE_VRAM_FRACTION], ids=["anima", "image-and-video-vaes"]) +def test_a_gpu_decode_flips_at_the_share_of_its_devices_memory(fraction): + device = torch.device("cuda", 1) + boundary = fraction * TOTAL + with patch("torch.cuda.get_device_properties", return_value=MagicMock(total_memory=TOTAL)) as props: + assert should_pretile_vae_decode(device, math.floor(boundary), fraction) is False + assert should_pretile_vae_decode(device, math.ceil(boundary) + 1, fraction) is True + props.assert_called_with(device) # the VAE's own device, not whichever is current + + +@pytest.mark.parametrize("device_type", ["cpu", "mps"]) +def test_a_decode_off_the_gpu_is_never_tiled_on_these_grounds(device_type): + """A cpu_only VAE runs in system RAM, and MPS shares it: device totals are not the constraint there.""" + assert should_pretile_vae_decode(torch.device(device_type), 10**15, VAE_PRETILE_VRAM_FRACTION) is False + + +def test_1024px_stays_untiled_on_an_8gb_card_and_1536px_tiles_on_16gb(): + """The share is chosen so common sizes keep the exact single-pass decode: a Qwen-Image VAE decode at 1024px + (~5.7 GiB) on an 8 GiB card stays untiled; a FLUX VAE decode at 1536px on ROCm (~15.8 GiB) tiles on 16 GiB.""" + with patch("torch.cuda.get_device_properties", return_value=MagicMock(total_memory=8 * 2**30)): + assert not should_pretile_vae_decode(torch.device("cuda"), 1024 * 1024 * 2 * 2900, VAE_PRETILE_VRAM_FRACTION) + with patch("torch.cuda.get_device_properties", return_value=MagicMock(total_memory=TOTAL)): + assert should_pretile_vae_decode(torch.device("cuda"), 1536 * 1536 * 2 * 3600, VAE_PRETILE_VRAM_FRACTION) + + +def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypatch): + """The gate must measure against what the device keeps resident, not the card's nameplate total. + + This is the platform the feature was written for: the docstring justifies up-front tiling with + "on Windows, drivers page an allocation that does not fit into system memory instead of failing + it (always for ROCm)". Windows pages once the process passes its WDDM budget, so the gate has to + compare against that budget rather than `total_memory`. + + A 16 GiB card whose budget is down to 12.5 GiB pages a 13.3 GiB decode (FLUX.1 at 1408px on + MIOpen). No out-of-memory error is raised there, so the nodes' tiled retry never fires and the + generation just crawls. Measured against the card total the decode looks fine (13.3 < 14.4 GiB), + so it is not tiled. + """ + total_bytes = 16 * 2**30 + budget_bytes = int(12.5 * 2**30) + estimate = 1408 * 1408 * 2 * 3600 # 13.3 GiB: under 90% of the card, over the budget + + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + # This process has allocated nothing yet, so its residency cannot lift the ceiling above the budget. + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: (budget_bytes, 0)) + monkeypatch.setattr("invokeai.backend.util.wddm.paged_bytes", lambda device: 0) + + assert estimate < VAE_PRETILE_VRAM_FRACTION * total_bytes, "test shape must look fine against the card total" + assert estimate > budget_bytes, "test shape must exceed what Windows would keep resident" + + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True + + +def test_a_budget_below_our_own_residency_does_not_tile_a_decode_that_fits(monkeypatch): + """Windows trims the budget as the process itself grows -- measured on an RX 9060 XT, 15.09 GiB until it passes + about 12 of 16 GiB, then 7.62, i.e. below what it already holds. That is an instruction to trim itself, which the + cache does for the reservation, so the ceiling is what it holds: comparing against the bare budget tiled a 7.0 GiB + Z-Image decode against a 6.9 GiB line mid-session and changed output that had been pixel-identical.""" + total_bytes = 16 * 2**30 + estimate = 7 * 2**30 + + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: (int(7.6 * 2**30), 12 * 2**30)) + monkeypatch.setattr("invokeai.backend.util.wddm.paged_bytes", lambda device: 0) # alone: nothing paged out + + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is False + + +def test_a_decode_beyond_both_the_budget_and_our_residency_is_tiled(monkeypatch): + """The combined state -- something else holds part of the card and Invoke holds the rest -- is the one where the + budget does carry the foreign pressure: evicting everything of ours reaches our own usage, not the whole card.""" + total_bytes = 16 * 2**30 + estimate = int(8.93 * 2**30) # a Qwen-Image 1024px decode peak + + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + monkeypatch.setattr( + "invokeai.backend.util.wddm.local_video_memory", lambda device: (int(7.55 * 2**30), int(8.2 * 2**30)) + ) + monkeypatch.setattr("invokeai.backend.util.wddm.paged_bytes", lambda device: 0) + + assert estimate < VAE_PRETILE_VRAM_FRACTION * total_bytes, "the card total alone would let this through" + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True + + +QWEN_1024_DECODE = int(8.93 * 2**30) + + +def _windows_rocm(monkeypatch, budget: float, usage: float | None, paged: float | None) -> None: + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=16 * 2**30)) + usage_bytes = None if usage is None else int(usage * 2**30) + monkeypatch.setattr( + "invokeai.backend.util.wddm.local_video_memory", lambda device: (int(budget * 2**30), usage_bytes) + ) + paged_bytes = None if paged is None else int(paged * 2**30) + monkeypatch.setattr("invokeai.backend.util.wddm.paged_bytes", lambda device: paged_bytes) + + +def test_usage_windows_has_paged_out_does_not_lift_the_ceiling(monkeypatch): + """Next to an 8 GiB holder the budget steps down to its share of the card while `CurrentUsage` keeps counting what + Windows paged out (measured: budget 7.58 GiB, usage 9.65 GiB, 7.70 GiB actually in VRAM). Only the resident part + may override the budget; taking the whole usage priced this decode against a 12 GiB ceiling and let it page.""" + _windows_rocm(monkeypatch, budget=7.6, usage=12.0, paged=4.3) + + assert should_pretile_vae_decode(torch.device("cuda", 0), QWEN_1024_DECODE) is True + + +def test_usage_beyond_the_card_leaves_the_budget_in_charge(monkeypatch): + """An over-committed process reports more usage than the card has; that says nothing about residency, and + clamping it to the card size switched the correction off exactly while Windows was paging.""" + _windows_rocm(monkeypatch, budget=7.6, usage=None, paged=4.3) + + assert should_pretile_vae_decode(torch.device("cuda", 0), QWEN_1024_DECODE) is True + + +def test_residency_unknown_leaves_the_budget_in_charge(monkeypatch): + """Without a paged-bytes reading the resident part is unknown, so a budget below our usage is taken as it is.""" + _windows_rocm(monkeypatch, budget=7.6, usage=12.0, paged=None) + + assert should_pretile_vae_decode(torch.device("cuda", 0), 7 * 2**30) is True + + +def test_the_card_total_stands_when_windows_does_not_answer(monkeypatch): + """Off Windows ROCm, and whenever the driver cannot answer, the rule is the card's own size.""" + total_bytes = 16 * 2**30 + + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: None) + + assert should_pretile_vae_decode(torch.device("cuda", 0), 7 * 2**30) is False + assert should_pretile_vae_decode(torch.device("cuda", 0), 15 * 2**30) is True diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py new file mode 100644 index 00000000000..a2c53b65ea3 --- /dev/null +++ b/tests/backend/util/test_wddm.py @@ -0,0 +1,330 @@ +"""The Windows video-memory readers, against scripted stand-ins for the gdi32 kernel thunks and for PDH. + +The stand-ins are driven through the same ctypes structures the module passes, so struct layout, the two-call adapter +enumeration, the PCI match, handle cleanup and the counter-instance naming are exercised; only the OS behind them is fake. +""" + +import ctypes +import functools +import os +import sys +from types import SimpleNamespace + +import pytest +import torch + +from invokeai.backend.util import wddm +from invokeai.backend.util.devices import TorchDevice + +GIB = 1024**3 +DEVICE = torch.device("cuda", 0) +OUR_BUS, OUR_DEVICE = 3, 0 + + +def _luid_low(handle: int) -> int: + return 0x22A00 + handle + + +class _FakePdh: + """A PDH that reports the given ``GPU Process Memory`` instances' shared usage.""" + + def __init__( + self, + instances: dict[str, int], + unavailable: tuple[str, ...] = (), + invalid: tuple[str, ...] = (), + ): + self.instances = instances + self.unavailable = unavailable + self.invalid = invalid + self.counter_paths: list[str] = [] + self.closed = 0 + self._alive: list[object] = [] + + def PdhOpenQueryW(self, source, user_data, query_ref): + query_ref._obj.value = 1 + return 0 + + def PdhAddEnglishCounterW(self, query, path, user_data, counter_ref): + if path in self.unavailable: + return -1073738824 # PDH_CSTATUS_NO_COUNTER + self.counter_paths.append(path) + handle = 2 + len(self.counter_paths) + counter_ref._obj.value = handle + return 0 + + def PdhCollectQueryData(self, query): + return 0 + + def PdhGetFormattedCounterArrayW(self, counter, fmt, size_ref, count_ref, buffer): + instances = self.instances + count = len(instances) + if buffer is None: + size_ref._obj.value = ctypes.sizeof(wddm._PdhCounterValueItem) * count + count_ref._obj.value = count + return wddm._PDH_MORE_DATA + items = (wddm._PdhCounterValueItem * count).from_buffer(buffer) + for item, (name, value) in zip(items, instances.items(), strict=True): + item.szName = name + item.FmtValue.largeValue = value + if name in self.invalid: + item.FmtValue.CStatus = 0xC0000BB8 # PDH_CSTATUS_INVALID_DATA + self._alive.append(items) # the name buffers live as long as the array object + count_ref._obj.value = count + return 0 + + def PdhCloseQuery(self, query): + self.closed += 1 + return 0 + + +class _FakeGdi32: + """Adapters as (handle, bus, device); every thunk returns NTSTATUS 0 unless told otherwise.""" + + def __init__(self, adapters, budget=15 * GIB, usage=0, query_status=0): + self.adapters = adapters + self.budget = budget + self.usage = usage + self.query_status = query_status + self.enum_calls = 0 + self.closed: list[int] = [] + + def D3DKMTEnumAdapters2(self, ref): + self.enum_calls += 1 + enum = ref._obj + if not enum.pAdapters: + enum.NumAdapters = len(self.adapters) + return 0 + for i, (handle, _bus, _device) in enumerate(self.adapters): + enum.pAdapters[i].hAdapter = handle + enum.pAdapters[i].AdapterLuid.LowPart = _luid_low(handle) + enum.pAdapters[i].AdapterLuid.HighPart = 0 + return 0 + + def D3DKMTQueryAdapterInfo(self, ref): + query = ref._obj + address = wddm._AdapterAddress.from_address(query.pPrivateDriverData) + _handle, bus, device = next(a for a in self.adapters if a[0] == query.hAdapter) + address.BusNumber, address.DeviceNumber = bus, device + return 0 + + def D3DKMTQueryVideoMemoryInfo(self, ref): + info = ref._obj + info.Budget, info.CurrentUsage = self.budget, self.usage + return self.query_status + + def D3DKMTCloseAdapter(self, ref): + self.closed.append(ref._obj.hAdapter) + return 0 + + +@pytest.fixture +def windows_rocm(monkeypatch: pytest.MonkeyPatch): + """A Windows ROCm build whose device 0 sits at PCI bus 3 with 16 GiB. Returns an installer for a fake gdi32.""" + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(torch.version, "hip", "7.14.0") + props = SimpleNamespace(pci_bus_id=OUR_BUS, pci_device_id=OUR_DEVICE, total_memory=16 * GIB) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: props) + wddm.reset_cache() + + def install(fake: _FakeGdi32) -> _FakeGdi32: + monkeypatch.setattr(wddm, "_load_gdi32", functools.lru_cache(maxsize=1)(lambda: fake)) + return fake + + yield install + wddm.reset_cache() + + +def test_reads_the_budget_of_the_adapter_at_the_devices_pci_location(windows_rocm): + gdi32 = windows_rocm(_FakeGdi32([(10, 7, 0), (20, OUR_BUS, OUR_DEVICE), (30, 0xFFFFFFFF, 0xFFFF)], budget=15 * GIB)) + + assert wddm.video_memory_budget(DEVICE) == 15 * GIB + assert sorted(gdi32.closed) == [10, 30], "the other adapters' handles are closed, ours stays open" + + +def test_resolves_the_adapter_once(windows_rocm): + gdi32 = windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + + wddm.video_memory_budget(DEVICE) + gdi32.budget = 12 * GIB # another GPU process started + assert wddm.video_memory_budget(DEVICE) == 12 * GIB + assert gdi32.enum_calls == 2, "one count call and one fill call, for the first query only" + + +def test_usage_beyond_the_card_is_unknown_and_keeps_the_budget(windows_rocm): + """`CurrentUsage` counts bytes Windows has already paged out, so it can exceed the adapter's own size when the + process over-commits. It says nothing about residency then, but the budget must survive: the model cache's + free-VRAM cap needs it most in exactly that state.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)], budget=15 * GIB, usage=17 * GIB)) + + assert wddm.local_video_memory(DEVICE) == (15 * GIB, None) + assert wddm.video_memory_budget(DEVICE) == 15 * GIB + + +@pytest.mark.parametrize( + ("adapters", "kwargs"), + [ + ([(10, 7, 0)], {}), + ([(20, OUR_BUS, OUR_DEVICE), (21, OUR_BUS, OUR_DEVICE)], {}), + ([(20, OUR_BUS, OUR_DEVICE)], {"query_status": -1073741823}), + ([(20, OUR_BUS, OUR_DEVICE)], {"budget": 0}), + ([(20, OUR_BUS, OUR_DEVICE)], {"budget": 17 * GIB}), + ], + ids=[ + "no-adapter-at-the-location", + "two-adapters-at-the-location", + "query-fails", + "zero-budget", + "budget-over-total", + ], +) +def test_unknown_rather_than_guessed(windows_rocm, adapters, kwargs): + gdi32 = windows_rocm(_FakeGdi32(adapters, **kwargs)) + + assert wddm.video_memory_budget(DEVICE) is None + if len(adapters) == 2: + assert sorted(gdi32.closed) == [20, 21], "an ambiguous match must not leak either handle" + + +@pytest.mark.parametrize( + ("platform", "hip", "device"), + [("linux", "7.2.0", DEVICE), ("win32", None, DEVICE), ("win32", "7.14.0", torch.device("cpu"))], + ids=["linux-rocm", "windows-cuda", "cpu-device"], +) +def test_elsewhere_gdi32_is_never_touched(windows_rocm, monkeypatch, platform, hip, device): + def loader(): + raise AssertionError("gdi32 must not be loaded outside Windows ROCm") + + monkeypatch.setattr(wddm, "_load_gdi32", functools.lru_cache(maxsize=1)(loader)) + monkeypatch.setattr(sys, "platform", platform) + monkeypatch.setattr(torch.version, "hip", hip) + + assert wddm.video_memory_budget(device) is None + + +def test_reset_closes_the_cached_handle(windows_rocm): + gdi32 = windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + wddm.video_memory_budget(DEVICE) + + wddm.reset_cache() + + assert gdi32.closed == [20] + + +def _instance(pid: int, handle: int) -> str: + return f"pid_{pid}_luid_0x00000000_0x{_luid_low(handle):08X}_phys_0" + + +def test_paged_bytes_reads_this_process_on_the_devices_adapter(windows_rocm, monkeypatch): + windows_rocm(_FakeGdi32([(10, 7, 0), (20, OUR_BUS, OUR_DEVICE)])) + me = os.getpid() + pdh = _FakePdh( + { + _instance(me, 10): 5 * GIB, # this process, other adapter + _instance(me + 1, 20): 7 * GIB, # another process, our adapter + _instance(me, 20): 3 * GIB, # this process, our adapter (PDH may report upper-case hex) + } + ) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.paged_bytes(DEVICE) == 3 * GIB + assert wddm.paged_bytes(DEVICE) == 3 * GIB + assert pdh.counter_paths == [r"\GPU Process Memory(*)\Shared Usage"], ( + "only the counter this answer needs, opened once" + ) + + wddm.reset_cache() + assert pdh.closed == 1 + + +def test_a_missing_counter_answers_unknown(windows_rocm, monkeypatch): + """A Windows without the ``GPU Process Memory`` counters makes the paged amount unknown instead of failing.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + pdh = _FakePdh({_instance(os.getpid(), 20): 3 * GIB}, unavailable=(r"\GPU Process Memory(*)\Shared Usage",)) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.paged_bytes(DEVICE) is None + + +def test_an_instance_with_invalid_data_is_skipped(windows_rocm, monkeypatch): + """PDH leaves `largeValue` undefined unless the item's status says the data is good.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + me = _instance(os.getpid(), 20) + pdh = _FakePdh({me: 7 * GIB}, invalid=(me,)) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.paged_bytes(DEVICE) is None + + +def test_paged_bytes_is_unknown_before_this_process_used_the_adapter(windows_rocm, monkeypatch): + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + pdh = _FakePdh({_instance(os.getpid() + 1, 20): GIB}) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.paged_bytes(DEVICE) is None + + +def test_paged_bytes_elsewhere_never_touches_pdh(monkeypatch): + def loader(): + raise AssertionError("pdh must not be loaded outside Windows ROCm") + + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(loader)) + monkeypatch.setattr(sys, "platform", "linux") + + assert wddm.paged_bytes(DEVICE) is None + + +def test_structures_match_the_d3dkmt_layout(): + """The kernel reads these by offset: a shifted field would read a neighbour's bytes, not fail.""" + assert ctypes.sizeof(wddm._AdapterInfo) == 20 + assert wddm._QueryVideoMemoryInfo.Budget.offset == 16 + assert ctypes.sizeof(wddm._QueryVideoMemoryInfo) == 56 + # PDH_FMT_COUNTERVALUE_ITEM_W on 64-bit: a name pointer, then the value whose 8-byte union follows a DWORD status. + assert wddm._PdhCounterValueItem.FmtValue.offset == 8 + assert wddm._PdhCounterValue.largeValue.offset == 8 + assert ctypes.sizeof(wddm._PdhCounterValueItem) == 24 + + +needs_windows_rocm = pytest.mark.skipif( + sys.platform != "win32" or torch.version.hip is None or not torch.cuda.is_available(), + reason="needs a ROCm GPU under Windows", +) + + +@pytest.mark.slow +@needs_windows_rocm +def test_on_hardware_freed_memory_counts_as_headroom_again(): + """HIP keeps memory after a free and hands it to the next allocation; WDDM still counts it as this process's usage. + The capped free figure must recover anyway, or an offload the model cache just made would look like no progress.""" + wddm.reset_cache() + device = torch.device("cuda", 0) + torch.empty(1, device=device) + torch.cuda.empty_cache() + assert wddm.video_memory_budget(device) is not None + before, _ = TorchDevice.cuda_mem_get_info(device) + + block = torch.empty(2 * GIB, dtype=torch.uint8, device=device).fill_(1) + torch.cuda.synchronize() + held, _ = TorchDevice.cuda_mem_get_info(device) + del block + torch.cuda.empty_cache() + after, _ = TorchDevice.cuda_mem_get_info(device) + + assert before - held >= 1.9 * GIB + assert after >= before - 64 * 2**20 + + +@pytest.mark.slow +@needs_windows_rocm +def test_on_hardware_the_paged_amount_of_this_process_is_found(): + """The warning depends on finding this process's instance under the name the code builds; if that is wrong, + `paged_bytes` quietly answers None forever.""" + wddm.reset_cache() + device = torch.device("cuda", 0) + block = torch.empty(64 * 2**20, dtype=torch.uint8, device=device).fill_(1) + torch.cuda.synchronize() + + paged = wddm.paged_bytes(device) + + assert paged is not None and paged >= 0 + del block