ROCm on Windows: keep VRAM resident and bound attention and VAE memory

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

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]

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": "<class 'bool'>",
       "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": "<class 'bool'>",
+      "validation": {}
+    },
     {
       "category": "GENERATION",
       "default": 1,

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,

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,

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)

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),

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)