fix(working-memory): reserve VAE working memory per dtype and allocator

docs/src/content/docs/development/Architecture/model-manager.mdx

@@ -987,10 +987,19 @@ An estimator that declares an `EstimatorSpec` returns a `WorkingMemoryEstimate`
 To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has five):
 
 - Register an `EstimatorSpec` at module import, and list the module in `ESTIMATOR_MODULES` (`working_memory/spec.py`); the data test fails on a module that declares a spec without being listed.
-- Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, element size, score-matrix bytes) as a feature. The pure `terms(features, values)` turns features and coefficient values into terms, building fitted terms with `WorkingMemoryTerm.fitted(...)`. The public function returns `SPEC.estimate(features, device)`.
+- Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, the dtype the operation runs in, score-matrix bytes) as a feature, and refuses a dtype the spec has no line for with a `ValueError` naming the estimator, dtype and platform. The pure `terms(features, values)` turns features and coefficient values into terms, building fitted terms with `WorkingMemoryTerm.fitted(...)`. The public function returns `SPEC.estimate(features, device)`.
+- Record the dtype the operation will run in, not the one the cached model holds. The SD1/SDXL nodes cast the shared VAE with `sd_vae_run_dtype`: fp16 for every GPU run that does not ask for fp32 (bf16 never runs), fp32 when asked and for a CPU decode. Their estimate prices the same function's answer, plus what the in-place cast adds to the resident weights. Uncast VAEs compute in their first floating-point parameter wider than a byte, so fp8 or packed SDNQ storage does not count.
 - Bump `version` whenever `terms` changes the formula.
-- Ship one coefficient entry per measured platform value in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry. Unmeasured platform values (MPS, XPU, CPU) are priced as `UNMEASURED_PRICED_AS` in `working_memory/platform.py` says.
-- Ship the measured points next to them: peak reserved bytes (`reserved_peak`), one fresh process per point. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides it, to at most 1.3 times the point, or the point's tighter `max_ratio`. The `device_working_mem_gb` floor counts only where an analytic term decides. Every value an entry carries must be decided by at least one point.
+- Ship one coefficient entry per measured platform in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry, which may carry more values than the borrower reads.
+  - A platform is the convolution backend (`cudnn`, `miopen`, `mps`, `xpu`, `cpu`), torch's caching allocator and the OS. The allocator is torch's backend first (`cudaMallocAsync` ignores expandable segments), then `expandable` if the allocator environment variables enable expandable segments on the native allocator, else `native`; anything but a CUDA/HIP GPU is `native`.
+  - `MEASURED` in `working_memory/platform.py` lists the measured backend, allocator and OS combinations (an OS of `None` for every OS). `priced_as` prices every other pair as one of them, keeping the OS: a backend nothing is measured on as cuDNN with the native allocator, an allocator not measured with its backend as that backend's native allocator. There the entry that carries the values raises them to its `unmeasured_floor`, the pre-dtype slopes.
+  - An entry applies to the OS only when its platform names one, and then takes precedence there over the one that does not. Only MIOpen names one. With the native allocator, Windows demand is concave and up to 1.6x Linux's, so Windows has its own fit and every other OS keeps the Linux one. Expandable segments are measured on Windows only, where they are Invoke's ROCm default, so `MEASURED` lists them for Windows and every other OS prices them as the native allocator there. cuDNN is priced the same on every OS.
+- Ship the measured points next to them, one fresh process per point: `reserved_peak` (the allocator's peak reserved bytes), `driver_peak` (the driver-level device usage), or `demand_peak` (the largest of driver-level, reserved and allocated peak, over two runs, which catches the memory HIP maps beyond the allocator on a process's first operation). A point names the OS only where it measured an OS-specific entry. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides an in-fit point, to at most 1.3 times the point, or the point's tighter `max_ratio`. Points no line can hold within 1.3x are listed with their own ratio in the test's `UPPER_BOUND_EXCEPTIONS`. The `device_working_mem_gb` floor counts only where an analytic term decides.
+- Fit the values, never hand-tune them. The AutoencoderKL estimators price their convolutions as a line per operation and dtype: `k_<operation>_<dtype>` bytes per output pixel plus a base of `b_<operation>_<dtype>` bytes.
+  - `working_memory/fit.py` fits each line on the untiled single-sample in-fit points of every estimator that resolves to the entry. The line is the edge of the points' upper convex hull (with the origin, so the base is never negative) that over-reserves its worst point least, the steeper on a tie and never a falling one. It needs at least two sizes twice apart in pixels. The shipped values add a 5% margin and round up: the slope to a whole byte per pixel, the base to a whole MiB.
+  - A point measured in a regime its platform does not run sets `in_fit: false`: it fits nothing and gets only the lower bound (CUDA with the math attention kernel forced). The test lists every such point.
+  - A line no in-fit point measures is a declared stand-in, listed in the entry's `stand_ins` and never below the pre-dtype slope. MIOpen with the native allocator off Windows prices fp16 and fp32 at the pre-dtype figure, and CogView4's lines there at its expandable-segments ones.
+  - Every value an entry carries is a stand-in or decided by an in-fit point. The data test fails when a fitted value is not the fit of its points, naming the value to ship, and when an in-fit point measures a stand-in's line.
 
 ### Emitting model loading events
 

invokeai/app/invocations/vae/image_to_latents.py

@@ -27,7 +27,7 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
 from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl, sd_vae_run_dtype
 
 """
 SDXL VAE color compensation values determined experimentally to reduce color drift.
@@ -86,9 +86,10 @@ def vae_encode(
         with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL"
             orig_dtype = vae.dtype
-            if upcast:
-                vae.to(dtype=torch.float32)
-
+            # The dtype the estimate above priced: fp32 when upcasting, fp16 otherwise, on every device.
+            run_dtype = sd_vae_run_dtype("encode", upcast, vae_info.compute_device)
+            vae.to(dtype=run_dtype)
+            if run_dtype == torch.float32:
                 use_torch_2_0_or_xformers = hasattr(vae.decoder, "mid_block") and isinstance(
                     vae.decoder.mid_block.attentions[0].processor,
                     (
@@ -107,10 +108,6 @@ def vae_encode(
                 # else:
                 #    latents = latents.float()
 
-            else:
-                vae.to(dtype=torch.float16)
-                # latents = latents.half()
-
             if tiled:
                 vae.enable_tiling()
             else:

invokeai/app/invocations/vae/latents_to_image.py

@@ -20,7 +20,7 @@
 from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
 from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
 from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl, sd_vae_run_dtype
 
 
 @invocation(
@@ -74,16 +74,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
             device = vae_info.compute_device
             latents = latents.to(device)
-            # Force fp32 when running on CPU (e.g. when the VAE is configured cpu_only). fp16 conv does run on CPU with
+            # fp32 when asked, and on the CPU (e.g. when the VAE is configured cpu_only): fp16 conv does run on CPU with
             # the pinned torch, but it's much slower than fp32 there, and SD/SDXL VAE has known fp16 overflow issues
-            # that produce black images — fp32 avoids both.
-            if self.fp32 or device.type == "cpu":
-                # FP32 mode: convert everything to float32 for maximum precision
-                vae.to(dtype=torch.float32)
-                latents = latents.float()
-            else:
-                vae.to(dtype=torch.float16)
-                latents = latents.half()
+            # that produce black images — fp32 avoids both. fp16 otherwise. The estimate above priced the same rule.
+            run_dtype = sd_vae_run_dtype("decode", self.fp32, device)
+            vae.to(dtype=run_dtype)
+            latents = latents.to(dtype=run_dtype)
 
             if use_tiling:
                 vae.enable_tiling()

invokeai/app/util/torch_cuda_allocator.py

@@ -6,13 +6,26 @@
 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.
+# Every variable torch reads its caching-allocator configuration from.
 _ALLOCATOR_CONF_ENV_VARS = ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_HIP_ALLOC_CONF")
 
 ROCM_WINDOWS_ALLOC_CONF = "expandable_segments:True"
 
 
+def expandable_segments_enabled() -> bool:
+    """Whether the torch caching allocator runs in expandable-segments mode.
+
+    The mode is configured through the allocator env vars before torch import (InvokeAI's own
+    `pytorch_cuda_alloc_conf` setting is plumbed into PYTORCH_CUDA_ALLOC_CONF the same way), so
+    parsing them is authoritative for the life of the process. It lives here, beside the variables, because this
+    module imports neither torch nor its callers: the model cache and the working-memory platform.
+    """
+    for var in _ALLOCATOR_CONF_ENV_VARS:
+        if "expandable_segments:true" in os.environ.get(var, "").replace(" ", "").lower():
+            return True
+    return False
+
+
 def _installed_torch_version() -> str | None:
     """The installed torch distribution's version, read from its metadata so torch itself is not imported."""
     try:

invokeai/backend/model_manager/load/model_cache/model_cache.py

@@ -1,6 +1,5 @@
 import gc
 import logging
-import os
 import queue
 import sys
 import threading
@@ -15,6 +14,7 @@
 import psutil
 import torch
 
+from invokeai.app.util.torch_cuda_allocator import expandable_segments_enabled
 from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot
 from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
 from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
@@ -39,20 +39,6 @@
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter
 
-
-def _expandable_segments_enabled() -> bool:
-    """Whether the torch caching allocator runs in expandable-segments mode.
-
-    The mode is configured through the allocator env vars before torch import (InvokeAI's own
-    `pytorch_cuda_alloc_conf` setting is plumbed into PYTORCH_CUDA_ALLOC_CONF the same way), so
-    parsing them is authoritative for the life of the process.
-    """
-    for var in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_HIP_ALLOC_CONF"):
-        if "expandable_segments:true" in os.environ.get(var, "").replace(" ", "").lower():
-            return True
-    return False
-
-
 # Size of a GB in bytes.
 GB = 2**30
 
@@ -2231,7 +2217,7 @@ def _release_allocator_blocks_for_reserve(self, working_mem_bytes: Optional[int]
         driver_free, _ = TorchDevice.cuda_mem_get_info(self._execution_device)
         if driver_free >= reserve:
             return
-        if _expandable_segments_enabled():
+        if expandable_segments_enabled():
             # Freed pages inside a segment are not counted as reclaimable there, but empty_cache() unmaps them.
             releasable = torch.cuda.memory_reserved(self._execution_device) - torch.cuda.memory_allocated(
                 self._execution_device
@@ -2403,7 +2389,7 @@ def _get_reclaimable_allocator_bytes(self) -> int:
         `_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():
+        if expandable_segments_enabled():
             return 0
         try:
             if self._execution_device.type == "cuda":
@@ -2618,7 +2604,7 @@ def _offload_unlocked_models(
         # 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()
+        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())