From a32c6aa487f0241b829dce74c5a93ada8bc2344c Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:14:01 +0200 Subject: [PATCH 01/21] feat(rocm): default to expandable segments on Windows ROCm With pytorch_cuda_alloc_conf unset, a ROCm build on Windows now runs with expandable_segments:True, so fragmented VRAM no longer pushes allocations into shared system memory. An explicit allocator setting or allocator env var still wins; docs and generated config artifacts updated. --- .../docs/configuration/low-vram-mode.mdx | 2 + docs/src/generated/settings.json | 6 +-- invokeai/app/run_app.py | 7 ++- .../app/services/config/config_default.py | 4 +- invokeai/app/util/torch_cuda_allocator.py | 41 ++++++++++++++ invokeai/frontend/web/openapi.json | 4 +- .../frontend/web/src/services/api/schema.ts | 4 +- tests/app/util/test_torch_cuda_allocator.py | 54 +++++++++++++++++++ 8 files changed, 112 insertions(+), 10 deletions(-) diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index f7f39435b3c..e4cc56492de 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -84,6 +84,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 diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index f23aaf9fc22..cf71fadd7b3 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -160,7 +160,7 @@ }, { "category": "PATHS", - "default": "models/.convert_cache", + "default": "models\\.convert_cache", "description": "Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).", "env_var": "INVOKEAI_CONVERT_CACHE_DIR", "literal_values": [], @@ -171,7 +171,7 @@ }, { "category": "PATHS", - "default": "models/.download_cache", + "default": "models\\.download_cache", "description": "Path to the directory that contains dynamically downloaded models.", "env_var": "INVOKEAI_DOWNLOAD_CACHE_DIR", "literal_values": [], @@ -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", 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 f4640c36537..6cdd2ffff1e 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. @@ -234,7 +234,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+)?)$") diff --git a/invokeai/app/util/torch_cuda_allocator.py b/invokeai/app/util/torch_cuda_allocator.py index d1c34cd3ceb..b7dd513dc5f 100644 --- a/invokeai/app/util/torch_cuda_allocator.py +++ b/invokeai/app/util/torch_cuda_allocator.py @@ -1,7 +1,48 @@ +import importlib.metadata import logging import os import sys +# 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 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 + torch_version = _installed_torch_version() + if torch_version is None or "+rocm" not in torch_version: + 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): """Configure the PyTorch CUDA memory allocator. See diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 63fa48f4f3a..3a919fe1970 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -55014,7 +55014,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", @@ -55384,7 +55384,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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index af0650ccda1..a6721ac5944 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -21120,7 +21120,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. @@ -21491,7 +21491,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; /** diff --git a/tests/app/util/test_torch_cuda_allocator.py b/tests/app/util/test_torch_cuda_allocator.py index 234002dd6fd..92df4d53274 100644 --- a/tests/app/util/test_torch_cuda_allocator.py +++ b/tests/app/util/test_torch_cuda_allocator.py @@ -1,8 +1,62 @@ +import os +import sys +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): + """A Windows ROCm install with no allocator configuration in the environment.""" + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: "2.12.0+rocm7.14.1") + for var in _ALLOCATOR_ENV_VARS: + monkeypatch.delenv(var, raising=False) + 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() + + # 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. From 9b0a53429983ee3760c28983e6a1b7bf71392df5 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:19:06 +0200 Subject: [PATCH 02/21] feat(rocm): keep the model cache within the Windows video-memory budget On ROCm under Windows, free VRAM is now capped by the WDDM budget (gdi32 D3DKMTQueryVideoMemoryInfo), past which Windows pages allocations into system memory instead of failing them. torch's figure there ignores other processes and runs up to the physical total; other platforms are unchanged. --- .../docs/configuration/low-vram-mode.mdx | 3 +- .../load/model_cache/model_cache.py | 2 +- invokeai/backend/util/devices.py | 17 ++ invokeai/backend/util/wddm.py | 235 ++++++++++++++++++ .../load/model_cache/conftest.py | 12 + .../test_model_cache_vram_accounting.py | 24 ++ tests/backend/util/test_devices.py | 18 ++ tests/backend/util/test_wddm.py | 172 +++++++++++++ 8 files changed, 481 insertions(+), 2 deletions(-) create mode 100644 invokeai/backend/util/wddm.py create mode 100644 tests/backend/model_manager/load/model_cache/conftest.py create mode 100644 tests/backend/util/test_wddm.py diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index e4cc56492de..abe6e0d2b28 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -114,7 +114,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 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..7a72611f461 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 diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index 7a28973344c..aa0c60f99a7 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 local_video_memory # legacy APIs TorchPrecisionNames = Literal["float32", "float16", "bfloat16"] @@ -458,6 +459,22 @@ 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.local_video_memory`). There, torch's figure is + the device total minus this process's own usage: 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 = local_video_memory(device) + if budget is not None: + budget_bytes, usage_bytes = budget + free = min(free, max(budget_bytes - usage_bytes, 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/wddm.py b/invokeai/backend/util/wddm.py new file mode 100644 index 00000000000..93a6a6dbef7 --- /dev/null +++ b/invokeai/backend/util/wddm.py @@ -0,0 +1,235 @@ +"""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; next to another GPU process, Windows lowers the budget within +about a second. ``CurrentUsage`` is the committed usage, including bytes already paged out, so budget minus usage is +what this process can still allocate before paging. + +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 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 + + +# 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), + ] + + +_GDI32_FUNCTIONS = ( + "D3DKMTEnumAdapters2", + "D3DKMTQueryAdapterInfo", + "D3DKMTQueryVideoMemoryInfo", + "D3DKMTCloseAdapter", +) + + +@dataclass(frozen=True) +class _Adapter: + handle: 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]] = {} + + +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 + + +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] = [] + for info in infos[: enum.NumAdapters]: + 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) + + if len(matches) != 1: + for info in matches: + _close(lib, info.hAdapter) + return None + return _Adapter(handle=matches[0].hAdapter, 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, int]]: + """This process's WDDM budget and committed local usage on the device's adapter, as ``(budget, usage)`` in bytes. + + ``None`` off Windows ROCm, and whenever the driver cannot answer or answers something implausible. + """ + if not _supported(device): + return None + resolved = _adapter_for(device) + if resolved is None: + return None + lib, adapter = resolved + info = _QueryVideoMemoryInfo(None, adapter.handle, _MEMORY_SEGMENT_GROUP_LOCAL) + try: + if lib.D3DKMTQueryVideoMemoryInfo(ctypes.byref(info)) != 0: + return None + except Exception: + return None + if not 0 < info.Budget <= adapter.total_bytes: + return None + return int(info.Budget), int(info.CurrentUsage) + + +def reset_cache() -> None: + """Close the cached adapter handles and forget the loader (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() 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..fdfc3de15f0 --- /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.local_video_memory", 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..0a39ddfa0a9 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 @@ -187,3 +187,27 @@ 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.local_video_memory", lambda device: (15 * GB, 7 * GB)) + + # 8 GB of budget headroom + 2 GB allocated - 1 GB working - 2 GB in use, not the 12 GB torch reports. + assert cache._get_vram_available(None) == 7 * GB diff --git a/tests/backend/util/test_devices.py b/tests/backend/util/test_devices.py index b4716084ae5..1a1847f8dcf 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -550,6 +550,24 @@ 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", "expected_free_gib"), + [(None, 10), ((15, 8), 7), ((15, 4), 10), ((15, 16), 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, expected_free_gib): + """On Windows ROCm, torch's free figure ignores other processes and the point where Windows starts paging.""" + gib = 1024**3 + device = torch.device("cuda", 0) + reported = None if budget is None else (budget[0] * gib, budget[1] * gib) + with ( + patch.object(torch.cuda, "mem_get_info", return_value=(10 * gib, 16 * gib)), + patch("invokeai.backend.util.devices.local_video_memory", return_value=reported) 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_wddm.py b/tests/backend/util/test_wddm.py new file mode 100644 index 00000000000..d35427e3191 --- /dev/null +++ b/tests/backend/util/test_wddm.py @@ -0,0 +1,172 @@ +"""The Windows video-memory budget reader, against a scripted stand-in for the gdi32 kernel thunks. + +The thunks are driven through the same ctypes structures the module passes, so struct layout, the two-call adapter +enumeration, the PCI match and handle cleanup are exercised; only the kernel behind them is fake. +""" + +import ctypes +import functools +import sys +from types import SimpleNamespace + +import pytest +import torch + +from invokeai.backend.util import wddm + +GIB = 1024**3 +DEVICE = torch.device("cuda", 0) +OUR_BUS, OUR_DEVICE = 3, 0 + + +class _FakeGdi32: + """Adapters as (handle, bus, device); every thunk returns NTSTATUS 0 unless told otherwise.""" + + def __init__(self, adapters, budget=15 * GIB, usage=2 * GIB, 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 + 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, usage=2 * GIB) + ) + + assert wddm.local_video_memory(DEVICE) == (15 * GIB, 2 * 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.local_video_memory(DEVICE) + gdi32.usage = 5 * GIB + assert wddm.local_video_memory(DEVICE) == (15 * GIB, 5 * GIB) + assert gdi32.enum_calls == 2, "one count call and one fill call, for the first query only" + + +@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.local_video_memory(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.local_video_memory(device) is None + + +def test_reset_closes_the_cached_handle(windows_rocm): + gdi32 = windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + wddm.local_video_memory(DEVICE) + + wddm.reset_cache() + + assert gdi32.closed == [20] + + +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 + + +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_the_usage_follows_an_allocation(): + wddm.reset_cache() + device = torch.device("cuda", 0) + torch.empty(1, device=device) + before = wddm.local_video_memory(device) + assert before is not None and before[0] > 0 + + block = torch.empty(256 * 2**20, dtype=torch.uint8, device=device).fill_(1) + torch.cuda.synchronize() + after = wddm.local_video_memory(device) + + assert after is not None + assert after[1] - before[1] >= 0.9 * block.numel() From 0d0f708636da481710545150e8e8f797a7901bab Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:23:42 +0200 Subject: [PATCH 03/21] feat(rocm): warn when Windows keeps VRAM in system memory After each session, a worker on a Windows ROCm device reads the process's shared GPU memory (PDH) and warns once per episode when more than 512 MiB of it sits in system memory, where every generation that touches it slows down. Adds a low-VRAM docs section on AMD GPUs under Windows. --- .../docs/configuration/low-vram-mode.mdx | 8 ++ .../session_processor_default.py | 37 +++++ invokeai/backend/util/wddm.py | 131 +++++++++++++++++- .../test_vram_paging_warning.py | 92 ++++++++++++ tests/backend/util/test_wddm.py | 95 ++++++++++++- 5 files changed, 358 insertions(+), 5 deletions(-) create mode 100644 tests/app/services/session_processor/test_vram_paging_warning.py diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index abe6e0d2b28..2d31c70e88f 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -204,6 +204,14 @@ 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. +- When a generation ends with part of Invoke's GPU memory in system memory, the log says so. Other programs using the same GPU, like games or browsers, can push Invoke's memory out as well. Close them, or restart Invoke to get the memory back into VRAM. + ## Troubleshooting ### Windows page file diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index affefc84d63..6580e2e19cb 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -43,6 +43,11 @@ from invokeai.app.util.profiler import Profiler from invokeai.backend.util.device_pool import GENERATION_DEVICE_POOL 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 — @@ -554,6 +559,8 @@ 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 last warning; only this worker's thread reads or writes it. + self.paging_warned_bytes = 0 @property def label(self) -> str: @@ -841,6 +848,35 @@ 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 after a session. + + 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. It stays there, and every generation that touches it slows down, with + nothing else in the log to explain why. Warns once per episode: again only after it grew by the threshold, + and re-armed once it drops below it. Measured after the session's VRAM release, so the transient paging of a + weight upload is over. `paged_bytes` answers None everywhere else, which keeps this silent there. + """ + 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 + if paged < _VRAM_PAGING_WARNING_BYTES: + worker.paging_warned_bytes = 0 + return + if paged < worker.paging_warned_bytes + _VRAM_PAGING_WARNING_BYTES: + return + worker.paging_warned_bytes = paged + self._invoker.services.logger.warning( + f"Windows is keeping {paged / 2**30:.1f} GiB of Invoke's GPU memory on {worker.label} in shared system " + "memory, which makes generations much slower. Close other programs that use this GPU; if it persists, " + "restart Invoke to get the memory back into VRAM." + ) + def resume(self) -> SessionProcessorStatus: if not self._resume_event.is_set(): self._resume_event.set() @@ -1054,6 +1090,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/backend/util/wddm.py b/invokeai/backend/util/wddm.py index 93a6a6dbef7..bde1c06989d 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -11,12 +11,16 @@ about a second. ``CurrentUsage`` is the committed usage, including bytes already paged out, so budget minus usage is what this process can still allocate before paging. +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 @@ -28,6 +32,10 @@ _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" # Win32 types spelled as fixed-width ctypes, so this module imports on every platform. @@ -96,6 +104,18 @@ class _QueryVideoMemoryInfo(ctypes.Structure): ] +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", @@ -107,6 +127,7 @@ class _QueryVideoMemoryInfo(ctypes.Structure): @dataclass(frozen=True) class _Adapter: handle: int + luid: int total_bytes: int @@ -114,6 +135,11 @@ class _Adapter: # 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 over every process's shared usage, as (query, counter), opened on first use; None if that failed. +_pdh_query: Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]] = None +_pdh_attempted = False + def _supported(device: torch.device) -> bool: return sys.platform == "win32" and torch.version.hip is not None and device.type == "cuda" @@ -134,6 +160,39 @@ def _load_gdi32() -> Optional[ctypes.CDLL]: 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))) @@ -179,7 +238,8 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: for info in matches: _close(lib, info.hAdapter) return None - return _Adapter(handle=matches[0].hAdapter, total_bytes=int(props.total_memory)) + 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]]: @@ -223,8 +283,68 @@ def local_video_memory(device: torch.device) -> Optional[tuple[int, int]]: return int(info.Budget), int(info.CurrentUsage) +def _open_shared_usage_query(pdh: ctypes.CDLL) -> 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, _PDH_SHARED_USAGE_COUNTER, None, ctypes.byref(counter)) != 0: + pdh.PdhCloseQuery(query) + return None + return query, counter + + +def _shared_usage_by_instance(pdh: ctypes.CDLL, query: ctypes.c_void_p, counter: ctypes.c_void_p) -> dict[str, int]: + """Collect the query once; every ``GPU Process Memory`` instance's shared usage, keyed by its lower-case 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} + + +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 the PDH query (~0.3 s); each later call collects it once (~0.3 ms). + """ + global _pdh_query, _pdh_attempted + + if not _supported(device): + return None + resolved = _adapter_for(device) + pdh = _load_pdh() + if resolved is None or pdh is None: + return None + luid = resolved[1].luid + instance = f"pid_{os.getpid()}_luid_0x{(luid >> 32) & 0xFFFFFFFF:08x}_0x{luid & 0xFFFFFFFF:08x}_phys_0" + with _pdh_lock: + if not _pdh_attempted: + _pdh_attempted = True + _pdh_query = _open_shared_usage_query(pdh) + if _pdh_query is None: + return None + try: + return _shared_usage_by_instance(pdh, *_pdh_query).get(instance) + except Exception: + return None + + def reset_cache() -> None: - """Close the cached adapter handles and forget the loader (tests only).""" + """Close the cached adapter handles and the PDH query, and forget the loaders (tests only).""" + global _pdh_query, _pdh_attempted + with _lock: opened = [adapter for adapter in _adapters.values() if adapter is not None] _adapters.clear() @@ -233,3 +353,10 @@ def reset_cache() -> None: for adapter in opened: _close(lib, adapter.handle) _load_gdi32.cache_clear() + + with _pdh_lock: + query, _pdh_query, _pdh_attempted = _pdh_query, None, False + pdh = _load_pdh() if query is not None else None + if pdh is not None and query is not None: + pdh.PdhCloseQuery(query[0]) + _load_pdh.cache_clear() 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..a12a0386d14 --- /dev/null +++ b/tests/app/services/session_processor/test_vram_paging_warning.py @@ -0,0 +1,92 @@ +"""The end-of-session warning when Windows keeps part of a worker's VRAM in shared system memory. + +It must speak once per episode (again only after the paged amount grew by the threshold, re-armed once it 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: warn + threshold + 300 * MIB, # grew by less than the threshold since the warning + 2 * threshold + 200 * MIB, # grew by the threshold: warn again + 100 * MIB, # recovered below the threshold: re-armed + threshold, # a new episode: warn + ] + + assert _warned_after_each_session(monkeypatch, processor, readings) == [False, True, False, True, False, True] + assert "cuda:0" in processor[1].warnings[0] + + +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 diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py index d35427e3191..7b9035c3317 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -1,11 +1,12 @@ -"""The Windows video-memory budget reader, against a scripted stand-in for the gdi32 kernel thunks. +"""The Windows video-memory readers, against scripted stand-ins for the gdi32 kernel thunks and for PDH. -The thunks are driven through the same ctypes structures the module passes, so struct layout, the two-call adapter -enumeration, the PCI match and handle cleanup are exercised; only the kernel behind them is fake. +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 @@ -19,6 +20,50 @@ 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]): + self.instances = instances + 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): + self.counter_paths.append(path) + counter_ref._obj.value = 2 + return 0 + + def PdhCollectQueryData(self, query): + return 0 + + def PdhGetFormattedCounterArrayW(self, counter, fmt, size_ref, count_ref, buffer): + count = len(self.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, self.instances.items(), strict=True): + item.szName = name + item.FmtValue.largeValue = value + 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.""" @@ -38,6 +83,8 @@ def D3DKMTEnumAdapters2(self, ref): 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): @@ -142,6 +189,48 @@ def test_reset_closes_the_cached_handle(windows_rocm): 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"], "one query, opened on first use" + + wddm.reset_cache() + assert pdh.closed == 1 + + +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 From ace24360ec37b7627e8d06a4fcb71a95a2cd0c81 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:31:27 +0200 Subject: [PATCH 04/21] perf(rocm): compute math-kernel attention in bounded chunks The ROCm SDPA guard now splits any math-kernel call over 1 GiB of scores into head groups (bitwise identical) or query rows (one oversized head), and working-memory estimates price one chunk instead of the whole matrix. Z-Image 1024px peaks at 0.5 GiB instead of 8 GiB per attention on an RX 9060 XT; renamed to install_rocm_sdpa_guard. --- .../docs/start-here/system-requirements.mdx | 16 +- invokeai/app/util/startup_utils.py | 7 +- invokeai/backend/util/attention.py | 208 ++++++++++++++-- .../util/test_rocm_sdpa_head_dim_guard.py | 231 +++++++++++++++++- 4 files changed, 417 insertions(+), 45 deletions(-) diff --git a/docs/src/content/docs/start-here/system-requirements.mdx b/docs/src/content/docs/start-here/system-requirements.mdx index 4ddc9d50fc1..6a62a5efe66 100644 --- a/docs/src/content/docs/start-here/system-requirements.mdx +++ b/docs/src/content/docs/start-here/system-requirements.mdx @@ -159,10 +159,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. ::: Run `rocm-smi` on your system's command line verify that drivers and ROCm are installed. If this command fails, or doesn't report versions, you will need to install them. 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/backend/util/attention.py b/invokeai/backend/util/attention.py index 6be3a663d8d..18e8d8ca488 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -8,7 +8,7 @@ import threading import warnings from functools import lru_cache, wraps -from typing import Any +from typing import Any, Callable import psutil import torch @@ -100,8 +100,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 @@ -125,7 +126,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, bitwise identical output; FLUX VAE mid-block (one 512-wide head, 1024px) +12 %, within one bf16 +# ulp. 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: @@ -134,7 +147,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( @@ -157,9 +182,102 @@ 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 _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. Head groups come + first: every call keeps its full `query x key` shape, so the result is bitwise the unchunked one. 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; there the GEMM tiling changes by up to one bf16 ulp. 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 @@ -168,10 +286,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 @@ -188,8 +307,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 @@ -207,24 +327,55 @@ 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 (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." ) @@ -360,27 +511,34 @@ 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``. """ 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/tests/backend/util/test_rocm_sdpa_head_dim_guard.py b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py index 6dd88058799..0b3419d3730 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,155 @@ 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_is_bitwise_unchunked(self, monkeypatch): + 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, {}), + ], + 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", + ], + ) + def test_masks_scale_and_grouped_heads_keep_torchs_result(self, monkeypatch, heads, kv_heads, budget_heads, kwargs): + budget = budget_heads * 2 * 64 * 48 * BYTES + 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("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 +346,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") @@ -229,12 +404,19 @@ def test_a_fused_kernel_is_still_wrong_for_the_wide_head(self): wrong += int(not torch.isfinite(out).all() or err > 0.05) 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,6 +429,29 @@ 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.""" From b98ee98182f8d0fd4f8ba09ea21244a04d133ac6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:33:39 +0200 Subject: [PATCH 05/21] fix(krea2): budget the attention score matrix The Krea-2 working-memory estimate now adds the score matrix where the build materializes it (no fused kernel), priced from the loaded model's heads and the attended sequence; one 1 GiB chunk on ROCm, nothing on CUDA. Before, Krea-2 fp8 on an RX 9060 XT loaded the whole transformer and paged each 12.9 GiB attention into system RAM. --- .../app/invocations/krea2/krea2_denoise.py | 44 +++++++++++++++++-- tests/app/invocations/test_krea2_denoise.py | 34 ++++++++++++++ 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py index 5e7aee2be0a..c794d514683 100644 --- a/invokeai/app/invocations/krea2/krea2_denoise.py +++ b/invokeai/app/invocations/krea2/krea2_denoise.py @@ -58,6 +58,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 @@ -490,15 +491,26 @@ 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, + ) + # 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, @@ -631,6 +643,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/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py index 13fac5dbd8c..2a5588042b4 100644 --- a/tests/app/invocations/test_krea2_denoise.py +++ b/tests/app/invocations/test_krea2_denoise.py @@ -843,3 +843,37 @@ 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) -> 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.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) + _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] From 67f23477694a50ea273d7300f58327cd33d6212c Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:43:07 +0200 Subject: [PATCH 06/21] fix(vae): MIOpen working-memory constants for the FLUX.1 autoencoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On ROCm the FLUX.1 VAE (also Z-Image's) now budgets 3600 decode / 2750 encode bytes per pixel·byte instead of 2200/1100, measured 3451/2688 on an RX 9060 XT: the same convolution stack and numbers as the FLUX.2 VAE, so both share one table. The estimate follows the VAE's compute device, so a cpu_only VAE keeps the cuDNN column. --- .../app/invocations/vae/flux_vae_decode.py | 6 ++- .../app/invocations/vae/flux_vae_encode.py | 2 +- .../vae/z_image_image_to_latents.py | 1 + .../vae/z_image_latents_to_image.py | 1 + invokeai/backend/util/vae_working_memory.py | 49 ++++++++++++------- scripts/calibrate_flux2_working_memory.py | 8 +-- .../invocations/test_flux2_working_memory.py | 18 +++---- .../invocations/test_z_image_tiled_decode.py | 31 ++++++++++++ 8 files changed, 82 insertions(+), 34 deletions(-) diff --git a/invokeai/app/invocations/vae/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index 69a59564aa0..fdcb507c461 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -54,7 +54,11 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents # Only estimate working memory for BFL AutoEncoder (diffusers VAE handles this internally) if isinstance(vae_info.model, AutoEncoder): 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, ) else: estimated_working_memory = 0 diff --git a/invokeai/app/invocations/vae/flux_vae_encode.py b/invokeai/app/invocations/vae/flux_vae_encode.py index 4ec0365c2cb..cfa471697a9 100644 --- a/invokeai/app/invocations/vae/flux_vae_encode.py +++ b/invokeai/app/invocations/vae/flux_vae_encode.py @@ -44,7 +44,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso # should be used for VAE encode sampling. assert isinstance(vae_info.model, AutoEncoder) estimated_working_memory = estimate_vae_working_memory_flux( - 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 ) generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): 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 263346e2962..48d65d3ee7a 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -53,6 +53,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso 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): 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 91805d0a1d8..74213f2233b 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -68,6 +68,7 @@ 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, ) # FLUX VAE doesn't support seamless, so only apply for AutoencoderKL diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index e2e2dd44426..715799fa11c 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -137,6 +137,7 @@ def estimate_vae_working_memory_flux( image_tensor: torch.Tensor, vae: AutoEncoder, tile_size: int | None = None, + device: torch.device | None = None, ) -> int: """Estimate the working memory required by the invocation in bytes. @@ -148,14 +149,18 @@ def estimate_vae_working_memory_flux( Only `decode` tiles: `AutoEncoder.encode` never consults `use_tiling`. So a `tile_size` with `operation="encode"` would price a tiled encode that cannot happen, and no call site passes one. + + `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] @@ -186,7 +191,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 @@ -198,17 +203,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 @@ -219,16 +228,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 @@ -236,7 +247,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( @@ -251,7 +262,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. @@ -289,7 +300,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: 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/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index dea23d130d3..cd1e3e0ace8 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -34,7 +34,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, ) @@ -368,7 +368,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 @@ -376,14 +376,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) @@ -856,7 +856,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 @@ -864,7 +864,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 @@ -872,7 +872,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 @@ -895,7 +895,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 @@ -905,7 +905,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_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 9bc792a3e14..9fafbb2fad1 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -22,6 +22,37 @@ def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: class TestFluxWorkingMemoryEstimate: + @pytest.fixture(autouse=True) + def _cudnn_column(self, monkeypatch): + """These pin the cuDNN-column arithmetic; on a ROCm machine the session device would select MIOpen's.""" + import invokeai.backend.util.vae_working_memory as vwm + + 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) From fe707b55a834fecd2f2888ffc151faade5f28b2f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 02:57:02 +0200 Subject: [PATCH 07/21] feat(vae): tile large decodes up front, switchable with auto_tiled_decode The FLUX.1, Z-Image, Qwen-Image (Krea-2) and Wan decodes now tile when the untiled estimate exceeds 90% of the VAE's GPU memory (Anima keeps its measured 70%), instead of relying on an OOM that Windows paging never raises. New auto_tiled_decode setting (default on) turns it off; force_tiled_decode and the tiled field still win. --- .../docs/configuration/low-vram-mode.mdx | 18 ++++++++ docs/src/generated/settings.json | 11 +++++ .../invocations/vae/anima_latents_to_image.py | 31 +++++-------- .../app/invocations/vae/flux_vae_decode.py | 25 +++++++++- .../vae/qwen_image_latents_to_image.py | 21 ++++++++- .../invocations/vae/wan_latents_to_video.py | 46 +++++++++---------- .../vae/z_image_latents_to_image.py | 22 ++++++++- .../app/services/config/config_default.py | 2 + invokeai/backend/util/vae_working_memory.py | 32 +++++++++++++ invokeai/frontend/web/openapi.json | 8 +++- .../frontend/web/src/services/api/schema.ts | 7 +++ tests/app/invocations/test_anima_vae.py | 18 +------- .../app/invocations/test_anima_vae_tiling.py | 2 +- .../test_flux_vae_decode_tiling.py | 21 +++++++++ .../test_qwen_image_working_memory.py | 42 ++++++++++++++++- .../invocations/test_wan_working_memory.py | 1 + .../invocations/test_z_image_tiled_decode.py | 35 +++++++++++++- tests/backend/util/test_vae_pretile.py | 40 ++++++++++++++++ 18 files changed, 311 insertions(+), 71 deletions(-) create mode 100644 tests/backend/util/test_vae_pretile.py diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 2d31c70e88f..24e9d7e42c7 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -38,6 +38,7 @@ Low-VRAM mode and related workload-specific optimizations include: - PyTorch CUDA allocator config (`pytorch_cuda_alloc_conf`) - Dynamic RAM and VRAM cache sizes (`max_cache_ram_gb`, `max_cache_vram_gb`) - Working memory (`device_working_mem_gb`) +- Tiled VAE decoding (`auto_tiled_decode`, `force_tiled_decode`) - Keeping a RAM weight copy (`keep_ram_copy_of_weights`) - Wan video memory optimization (`wan_memory_optimization`) - PiD decode activation chunking (`pid_memory_optimization`) @@ -158,6 +159,22 @@ device_working_mem_gb: 4 Once decoding completes, the model manager "reclaims" the extra VRAM allocated as working memory for future model loading operations. ::: +### Tiled VAE decoding + +Turning latents into the final image (the VAE decode) needs working memory that grows with the image's pixel count. When an untiled decode would claim most of the GPU's memory, Invoke decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Krea-2, Anima and Wan decodes. A tiled decode is not pixel-identical to an untiled one, so 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. 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 it off: + +```yaml +auto_tiled_decode: false +``` + +To tile every decode regardless of size: + +```yaml +force_tiled_decode: true +``` + ### Keeping a RAM weight copy Invoke has the option of keeping a RAM copy of all model weights, even when they are loaded onto the GPU. This optimization is _on_ by default, and enables faster model switching and LoRA patching. Disabling this feature will reduce the average RAM load while running Invoke (peak RAM likely won't change), at the cost of slower model switching and LoRA patching. If you have limited RAM, you can disable this optimization: @@ -210,6 +227,7 @@ Invoke does not officially support ROCm on Windows (see the [System Requirements - 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 decoding](#tiled-vae-decoding)); waiting for an out-of-memory error does not work there, because none comes. - When a generation ends with part of Invoke's GPU memory in system memory, the log says so. Other programs using the same GPU, like games or browsers, can push Invoke's memory out as well. Close them, or restart Invoke to get the memory back into VRAM. ## Troubleshooting diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index cf71fadd7b3..b20e4b44142 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -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, Anima 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.", + "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/vae/anima_latents_to_image.py b/invokeai/app/invocations/vae/anima_latents_to_image.py index 9cc3d4735e7..10433e38830 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,9 @@ 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) + use_tiling = context.config.get().auto_tiled_decode and 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/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index fdcb507c461..ad3d4738616 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -20,7 +20,11 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import ( + VAE_PRETILE_VRAM_FRACTION, + estimate_vae_working_memory_flux, + should_pretile_vae_decode, +) @invocation( @@ -48,7 +52,8 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents # This node has no tiling fields, so the config switch is the only way a user can ask for a # tiled decode here -- the same one the Z-Image, SD and Qwen-Image decode nodes honour. 0 is # the "use the VAE's default tile" sentinel shared by the estimator and `scoped_vae_tiling`. - use_tiling = context.config.get().force_tiled_decode + config = context.config.get() + use_tiling = config.force_tiled_decode tile_size = 0 if use_tiling else None # Only estimate working memory for BFL AutoEncoder (diffusers VAE handles this internally) @@ -60,6 +65,22 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents tile_size=tile_size, 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, VAE_PRETILE_VRAM_FRACTION + ) + ): + use_tiling = True + tile_size = 0 + 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, + ) else: estimated_working_memory = 0 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 8459c31e861..cb91594300e 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 ( + VAE_PRETILE_VRAM_FRACTION, + estimate_vae_working_memory_qwen_image, + 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 @@ -72,6 +77,18 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae=vae_info.model, tile_size=effective_tile_size, ) + if ( + not tiled + and config.auto_tiled_decode + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory, VAE_PRETILE_VRAM_FRACTION) + ): + 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, + ) 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/wan_latents_to_video.py b/invokeai/app/invocations/vae/wan_latents_to_video.py index 25d2e63e92e..07289a95e23 100644 --- a/invokeai/app/invocations/vae/wan_latents_to_video.py +++ b/invokeai/app/invocations/vae/wan_latents_to_video.py @@ -35,7 +35,11 @@ from invokeai.app.util.video_encoding import make_mp4_writer from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.util.devices import TorchDevice -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_wan +from invokeai.backend.util.vae_working_memory import ( + VAE_PRETILE_VRAM_FRACTION, + estimate_vae_working_memory_wan, + should_pretile_vae_decode, +) from invokeai.backend.wan.vae_decode import iter_wan_vae_decode_chunks @@ -129,29 +133,23 @@ 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 - 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 - tile_size = int(getattr(vae_info.model, "tile_sample_min_height", 256)) - 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, - ) + # estimate would claim most of the VAE's device, 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, VAE_PRETILE_VRAM_FRACTION + ) + if use_tiling: + tile_size = int(getattr(vae_info.model, "tile_sample_min_height", 256)) + 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_latents_to_image.py b/invokeai/app/invocations/vae/z_image_latents_to_image.py index 74213f2233b..ff609ea343b 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -23,7 +23,11 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import ( + VAE_PRETILE_VRAM_FRACTION, + estimate_vae_working_memory_flux, + should_pretile_vae_decode, +) # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder ZImageVAE = Union[AutoencoderKL, FluxAutoEncoder] @@ -60,7 +64,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) is_flux_vae = isinstance(vae_info.model, FluxAutoEncoder) - 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( @@ -70,6 +75,19 @@ def invoke(self, context: InvocationContext) -> ImageOutput: 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, VAE_PRETILE_VRAM_FRACTION) + ): + 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, + ) # FLUX VAE doesn't support seamless, so only apply for AutoencoderKL seamless_context = ( diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 6cdd2ffff1e..321b12e9f7f 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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. 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` @@ -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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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.") 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/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 715799fa11c..e9800a276cc 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -11,6 +11,38 @@ from invokeai.backend.flux.modules.autoencoder import AutoEncoder, resolve_tile_size from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger + +# 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) -> bool: + """Whether a decode should be tiled up front because its untiled working memory would claim more than + ``vram_fraction`` of ``device``'s memory. + + 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. + """ + if device.type == "cuda": + total_bytes = torch.cuda.get_device_properties(device).total_memory + 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__).info( + 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 GPU's {total_bytes / 2**30:.1f} GiB. Set auto_tiled_decode to false to " + "decode untiled." + ) + 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 diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 3a919fe1970..0dd671f9551 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -55090,6 +55090,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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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.", + "default": true + }, "pil_compress_level": { "type": "integer", "title": "Pil Compress Level", @@ -55384,7 +55390,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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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.\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/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index a6721ac5944..a38c2d0f543 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -21129,6 +21129,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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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. * 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` @@ -21557,6 +21558,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, Anima 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. + * @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/tests/app/invocations/test_anima_vae.py b/tests/app/invocations/test_anima_vae.py index 855f0638ba9..17d79a89529 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 @@ -97,21 +96,6 @@ def test_estimate_scales_with_element_size(self): assert fp32 == 2 * fp16 -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.""" @@ -316,7 +300,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) diff --git a/tests/app/invocations/test_anima_vae_tiling.py b/tests/app/invocations/test_anima_vae_tiling.py index 5ff52f3088f..3052f08a519 100644 --- a/tests/app/invocations/test_anima_vae_tiling.py +++ b/tests/app/invocations/test_anima_vae_tiling.py @@ -121,7 +121,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_flux_vae_decode_tiling.py b/tests/app/invocations/test_flux_vae_decode_tiling.py index 9dadc9682d2..c87f6fd4929 100644 --- a/tests/app/invocations/test_flux_vae_decode_tiling.py +++ b/tests/app/invocations/test_flux_vae_decode_tiling.py @@ -13,6 +13,7 @@ from invokeai.app.invocations.vae.flux_vae_decode import FluxVaeDecodeInvocation from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION def _build_decode_mocks(latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): @@ -78,6 +79,26 @@ def test_the_reservation_matches_the_decode_that_will_run(self, force_tiled_deco _build_invocation().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): + module = "invokeai.app.invocations.vae.flux_vae_decode" + vae, vae_info, context = _build_decode_mocks(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, VAE_PRETILE_VRAM_FRACTION) + assert estimate.call_args.kwargs["tile_size"] == 0 + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) + else: + pretile.assert_not_called() + vae.disable_tiling.assert_called_once() + class TestOomFallback: def test_an_untiled_oom_retries_once_tiled_and_says_so(self): diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index a9267641d0a..1aaa3c81f3a 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -14,7 +14,7 @@ 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 VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_qwen_image class TestQwenImageWorkingMemoryEstimate: @@ -416,6 +416,46 @@ 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}.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: + pretile.assert_called_once_with(mock_vae_info.compute_device, 20 * 2**30, VAE_PRETILE_VRAM_FRACTION) + 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: + pretile.assert_not_called() + assert estimate.call_count == 1 + mock_vae.enable_tiling.assert_not_called() + class TestQwenImageVaeTiling: """Exercise the tiling parameters against a real (tiny, randomly initialised) Qwen-Image VAE. diff --git a/tests/app/invocations/test_wan_working_memory.py b/tests/app/invocations/test_wan_working_memory.py index 967bbefb788..0239b424f38 100644 --- a/tests/app/invocations/test_wan_working_memory.py +++ b/tests/app/invocations/test_wan_working_memory.py @@ -299,6 +299,7 @@ def test_latents_to_video_streams_decode_chunks_directly_to_mp4(self): def test_latents_to_video_falls_back_to_tiling_when_estimate_exceeds_vram(self): vae = _mock_wan_vae() vae_info = _mock_vae_info(vae) + vae_info.compute_device = torch.device("cuda") mock_context = self._video_context(vae_info) with ( diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 9fafbb2fad1..75dd323ebca 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -9,7 +9,7 @@ from invokeai.app.invocations.vae.z_image_latents_to_image import ZImageLatentsToImageInvocation from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE, MIN_TILE_SAMPLE_SIZE from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_flux def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: @@ -250,6 +250,39 @@ def test_a_requested_tile_size_is_passed_through(self): _build_invocation(tiled=True, tile_size=384).invoke(context) vae.enable_tiling.assert_called_once_with(tile_sample_min_size=384) + @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): + """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, vae_info, context = _build_decode_mocks(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, VAE_PRETILE_VRAM_FRACTION) + assert estimate.call_args.kwargs["tile_size"] == 0 + vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) + else: + pretile.assert_not_called() + assert estimate.call_count == 1 + vae.disable_tiling.assert_called_once() + + def test_requested_tiling_skips_the_untiled_estimate(self): + module = "invokeai.app.invocations.vae.z_image_latents_to_image" + _, _, context = _build_decode_mocks(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, tiled, expected_tile_size): path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" diff --git a/tests/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py new file mode 100644 index 00000000000..6ef2fc10982 --- /dev/null +++ b/tests/backend/util/test_vae_pretile.py @@ -0,0 +1,40 @@ +"""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) From b424e554c13e82cc1826d4752a14abf7f3b8f189 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 03:48:10 +0200 Subject: [PATCH 08/21] fix(rocm): count freed VRAM as headroom on Windows - Cap free VRAM at the WDDM budget minus live allocations: HIP keeps freed memory that CurrentUsage still counts - Empty the allocator cache after each offload under expandable segments so the cache stops once enough is free - Hardware tests for the headroom after a free and for the PDH lookup --- .../load/model_cache/model_cache.py | 8 ++- invokeai/backend/util/devices.py | 13 ++-- invokeai/backend/util/wddm.py | 12 ++-- .../load/model_cache/conftest.py | 2 +- .../test_model_cache_vram_accounting.py | 46 ++++++++++++- tests/backend/util/test_devices.py | 13 ++-- tests/backend/util/test_wddm.py | 64 +++++++++++++------ 7 files changed, 115 insertions(+), 43 deletions(-) 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 7a72611f461..eea4522cfe7 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2552,6 +2552,10 @@ 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. + empty_cache_per_offload = _expandable_segments_enabled() # 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: @@ -2569,13 +2573,15 @@ 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: + TorchDevice.empty_cache() 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/devices.py b/invokeai/backend/util/devices.py index aa0c60f99a7..93de46aa4b2 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -9,7 +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 local_video_memory +from invokeai.backend.util.wddm import video_memory_budget # legacy APIs TorchPrecisionNames = Literal["float32", "float16", "bfloat16"] @@ -464,15 +464,14 @@ 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.local_video_memory`). There, torch's figure is - the device total minus this process's own usage: it ignores other processes, and Windows pages allocations - into shared system memory well before it is exhausted instead of failing them. + 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 = local_video_memory(device) + budget = video_memory_budget(device) if budget is not None: - budget_bytes, usage_bytes = budget - free = min(free, max(budget_bytes - usage_bytes, 0)) + free = min(free, max(budget - (total - free), 0)) return free, total @classmethod diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index bde1c06989d..be9c37859d2 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -8,8 +8,10 @@ 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; next to another GPU process, Windows lowers the budget within -about a second. ``CurrentUsage`` is the committed usage, including bytes already paged out, so budget minus usage is -what this process can still allocate before paging. +about a second. 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. @@ -261,8 +263,8 @@ def _adapter_for(device: torch.device) -> Optional[tuple[ctypes.CDLL, _Adapter]] return None if adapter is None else (lib, adapter) -def local_video_memory(device: torch.device) -> Optional[tuple[int, int]]: - """This process's WDDM budget and committed local usage on the device's adapter, as ``(budget, usage)`` in bytes. +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. ``None`` off Windows ROCm, and whenever the driver cannot answer or answers something implausible. """ @@ -280,7 +282,7 @@ def local_video_memory(device: torch.device) -> Optional[tuple[int, int]]: return None if not 0 < info.Budget <= adapter.total_bytes: return None - return int(info.Budget), int(info.CurrentUsage) + return int(info.Budget) def _open_shared_usage_query(pdh: ctypes.CDLL) -> Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]: diff --git a/tests/backend/model_manager/load/model_cache/conftest.py b/tests/backend/model_manager/load/model_cache/conftest.py index fdfc3de15f0..c4240fb9223 100644 --- a/tests/backend/model_manager/load/model_cache/conftest.py +++ b/tests/backend/model_manager/load/model_cache/conftest.py @@ -9,4 +9,4 @@ def _no_windows_video_memory_budget(monkeypatch: pytest.MonkeyPatch) -> None: 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.local_video_memory", lambda device: None) + 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 0a39ddfa0a9..933d539c9bd 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 @@ -207,7 +208,46 @@ def test_the_windows_video_memory_budget_caps_the_measured_free_vram(monkeypatch 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.local_video_memory", lambda device: (15 * GB, 7 * GB)) + monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: 11 * GB) - # 8 GB of budget headroom + 2 GB allocated - 1 GB working - 2 GB in use, not the 12 GB torch reports. - assert cache._get_vram_available(None) == 7 * 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 + + +def test_offloading_under_expandable_segments_stops_once_enough_is_free(monkeypatch: pytest.MonkeyPatch): + """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.""" + 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(): + device["free"] += device["unmappable"] + device["unmappable"] = 0 + + 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 1a1847f8dcf..1a22cd0b42e 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -551,18 +551,19 @@ def test_xpu_mem_get_info_estimates_when_native_and_sysman_both_fail(native_erro @pytest.mark.parametrize( - ("budget", "expected_free_gib"), - [(None, 10), ((15, 8), 7), ((15, 4), 10), ((15, 16), 0)], + ("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, expected_free_gib): - """On Windows ROCm, torch's free figure ignores other processes and the point where Windows starts paging.""" +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) - reported = None if budget is None else (budget[0] * gib, budget[1] * gib) + 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.local_video_memory", return_value=reported) as mock_budget, + 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) diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py index 7b9035c3317..5db03b5f52c 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -14,6 +14,7 @@ import torch from invokeai.backend.util import wddm +from invokeai.backend.util.devices import TorchDevice GIB = 1024**3 DEVICE = torch.device("cuda", 0) @@ -67,10 +68,9 @@ def PdhCloseQuery(self, query): class _FakeGdi32: """Adapters as (handle, bus, device); every thunk returns NTSTATUS 0 unless told otherwise.""" - def __init__(self, adapters, budget=15 * GIB, usage=2 * GIB, query_status=0): + def __init__(self, adapters, budget=15 * GIB, 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] = [] @@ -96,7 +96,7 @@ def D3DKMTQueryAdapterInfo(self, ref): def D3DKMTQueryVideoMemoryInfo(self, ref): info = ref._obj - info.Budget, info.CurrentUsage = self.budget, self.usage + info.Budget = self.budget return self.query_status def D3DKMTCloseAdapter(self, ref): @@ -122,20 +122,18 @@ def install(fake: _FakeGdi32) -> _FakeGdi32: 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, usage=2 * GIB) - ) + gdi32 = windows_rocm(_FakeGdi32([(10, 7, 0), (20, OUR_BUS, OUR_DEVICE), (30, 0xFFFFFFFF, 0xFFFF)], budget=15 * GIB)) - assert wddm.local_video_memory(DEVICE) == (15 * GIB, 2 * 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.local_video_memory(DEVICE) - gdi32.usage = 5 * GIB - assert wddm.local_video_memory(DEVICE) == (15 * GIB, 5 * GIB) + 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" @@ -159,7 +157,7 @@ def test_resolves_the_adapter_once(windows_rocm): def test_unknown_rather_than_guessed(windows_rocm, adapters, kwargs): gdi32 = windows_rocm(_FakeGdi32(adapters, **kwargs)) - assert wddm.local_video_memory(DEVICE) is None + 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" @@ -177,12 +175,12 @@ def loader(): monkeypatch.setattr(sys, "platform", platform) monkeypatch.setattr(torch.version, "hip", hip) - assert wddm.local_video_memory(device) is None + 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.local_video_memory(DEVICE) + wddm.video_memory_budget(DEVICE) wddm.reset_cache() @@ -236,6 +234,10 @@ def test_structures_match_the_d3dkmt_layout(): 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( @@ -246,16 +248,38 @@ def test_structures_match_the_d3dkmt_layout(): @pytest.mark.slow @needs_windows_rocm -def test_on_hardware_the_usage_follows_an_allocation(): +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) - before = wddm.local_video_memory(device) - assert before is not None and before[0] > 0 + 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 + - block = torch.empty(256 * 2**20, dtype=torch.uint8, device=device).fill_(1) +@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() - after = wddm.local_video_memory(device) - assert after is not None - assert after[1] - before[1] >= 0.9 * block.numel() + paged = wddm.paged_bytes(device) + + assert paged is not None and paged >= 0 + del block From 97668ebd4bc44e94103c1c108282aa3c8452f4fe Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 03:48:11 +0200 Subject: [PATCH 09/21] fix(rocm): harden allocator default, paging warning and SDPA chunking - Skip the allocator default once torch is imported; stop the test fixture leaking it - Warn only about memory paged across two sessions, with actionable advice - Leave K/V broadcasts the chunks cannot slice whole; cover masked row chunks and Krea-2 CFG/regional pricing --- .../session_processor_default.py | 26 +++++++------ invokeai/app/util/torch_cuda_allocator.py | 6 +++ invokeai/backend/util/attention.py | 8 ++++ tests/app/invocations/test_krea2_denoise.py | 26 +++++++++++-- .../services/session_processor/conftest.py | 10 +++++ .../test_vram_paging_warning.py | 38 ++++++++++++++++--- tests/app/util/test_torch_cuda_allocator.py | 17 ++++++++- .../util/test_rocm_sdpa_head_dim_guard.py | 30 ++++++++++++++- 8 files changed, 138 insertions(+), 23 deletions(-) create mode 100644 tests/app/services/session_processor/conftest.py diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 6580e2e19cb..e4151b0dbd1 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -559,7 +559,8 @@ 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 last warning; only this worker's thread reads or writes it. + # 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 @@ -849,14 +850,14 @@ def _release_vram_after_session(self, worker: _SessionWorker) -> None: ) 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 after a session. + """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. It stays there, and every generation that touches it slows down, with - nothing else in the log to explain why. Warns once per episode: again only after it grew by the threshold, - and re-armed once it drops below it. Measured after the session's VRAM release, so the transient paging of a - weight upload is over. `paged_bytes` answers None everywhere else, which keeps this silent there. + 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()) @@ -865,16 +866,19 @@ def _warn_if_vram_paged(self, worker: _SessionWorker) -> None: 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 paged < worker.paging_warned_bytes + _VRAM_PAGING_WARNING_BYTES: + if sustained < worker.paging_warned_bytes + _VRAM_PAGING_WARNING_BYTES: return - worker.paging_warned_bytes = paged + worker.paging_warned_bytes = sustained self._invoker.services.logger.warning( - f"Windows is keeping {paged / 2**30:.1f} GiB of Invoke's GPU memory on {worker.label} in shared system " - "memory, which makes generations much slower. Close other programs that use this GPU; if it persists, " - "restart Invoke to get the memory back into VRAM." + 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: diff --git a/invokeai/app/util/torch_cuda_allocator.py b/invokeai/app/util/torch_cuda_allocator.py index b7dd513dc5f..a6f1d30fbfd 100644 --- a/invokeai/app/util/torch_cuda_allocator.py +++ b/invokeai/app/util/torch_cuda_allocator.py @@ -36,6 +36,12 @@ def apply_rocm_windows_allocator_default(logger: logging.Logger) -> None: torch_version = _installed_torch_version() if torch_version is None or "+rocm" not in torch_version: 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}), " diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index 18e8d8ca488..d98f558b900 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -188,6 +188,13 @@ def _math_score_bytes(query: torch.Tensor, key: torch.Tensor) -> int: 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, @@ -344,6 +351,7 @@ def guarded_scaled_dot_product_attention( 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: diff --git a/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py index 2a5588042b4..39ebc1e8e33 100644 --- a/tests/app/invocations/test_krea2_denoise.py +++ b/tests/app/invocations/test_krea2_denoise.py @@ -850,13 +850,15 @@ class TestTheAttentionScoreMatrixReachesTheReservation: 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) -> tuple[int, list[dict], _Transformer]: + 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 = _runtime_context(tmp_path, transformer, **context_kwargs) context.models.load = lambda _identifier: info _patch_runtime(monkeypatch) asked: list[dict] = [] @@ -866,7 +868,7 @@ def score_matrix_bytes(**kwargs) -> int: return score_bytes monkeypatch.setattr(denoise_module, "sdpa_score_matrix_bytes", score_matrix_bytes) - _runtime_invocation(cfg_scale=1.0)._run_diffusion(context) + (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: @@ -877,3 +879,21 @@ def test_the_score_matrix_is_reserved_for_the_attended_sequence(self, monkeypatc 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/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 index a12a0386d14..1746e2e1b4e 100644 --- a/tests/app/services/session_processor/test_vram_paging_warning.py +++ b/tests/app/services/session_processor/test_vram_paging_warning.py @@ -1,7 +1,8 @@ """The end-of-session warning when Windows keeps part of a worker's VRAM in shared system memory. -It must speak once per episode (again only after the paged amount grew by the threshold, re-armed once it fell below -it), say nothing where the answer is unknown, and never fail the worker. +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 @@ -63,15 +64,28 @@ def test_warns_once_per_episode_and_again_only_after_growth(monkeypatch, process threshold = _VRAM_PAGING_WARNING_BYTES readings = [ 6 * MIB, # an idle process - threshold + 100 * MIB, # paged: warn - threshold + 300 * MIB, # grew by less than the threshold since the warning - 2 * threshold + 200 * MIB, # grew by the threshold: warn again + 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 ] - assert _warned_after_each_session(monkeypatch, processor, readings) == [False, True, False, True, False, True] + 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): @@ -90,3 +104,15 @@ def broken(device): 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 92df4d53274..3e2ddf8715e 100644 --- a/tests/app/util/test_torch_cuda_allocator.py +++ b/tests/app/util/test_torch_cuda_allocator.py @@ -17,8 +17,12 @@ def rocm_windows(monkeypatch: pytest.MonkeyPatch): """A Windows ROCm install with no allocator configuration in the environment.""" monkeypatch.setattr(sys, "platform", "win32") monkeypatch.setattr(torch_cuda_allocator, "_installed_torch_version", lambda: "2.12.0+rocm7.14.1") + monkeypatch.delitem(sys.modules, "torch") # as at startup, before anything imports it for var in _ALLOCATOR_ENV_VARS: - monkeypatch.delenv(var, raising=False) + # 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 @@ -56,6 +60,17 @@ def test_other_installs_are_left_alone(self, rocm_windows, platform, version): 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() + # 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 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 0b3419d3730..d02979d0147 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -237,6 +237,10 @@ def test_a_single_oversized_head_is_split_by_query_rows(self, monkeypatch): (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", @@ -248,10 +252,14 @@ def test_a_single_oversized_head_is_split_by_query_rows(self, monkeypatch): "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 = budget_heads * 2 * 64 * 48 * BYTES + 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) @@ -263,6 +271,21 @@ def test_masks_scale_and_grouped_heads_keep_torchs_result(self, monkeypatch, hea 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) @@ -393,15 +416,18 @@ 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, monkeypatch): From 2ab7aaaea25aa6b28c568ddcd298ae1e3fee6e77 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 19 Sep 2026 03:48:11 +0200 Subject: [PATCH 10/21] fix(vae): budget diffusers FLUX.1 decodes and quiet up-front tiling - Price and pre-tile FLUX.1 decodes with a diffusers-layout VAE - Log up-front tiling at debug; document what auto_tiled_decode off means - Emit POSIX path defaults in the generated docs settings --- .../docs/configuration/low-vram-mode.mdx | 4 +- docs/src/generated/settings.json | 6 +-- .../app/invocations/vae/flux_vae_decode.py | 37 ++++++++----------- .../vae/qwen_image_latents_to_image.py | 3 +- .../invocations/vae/wan_latents_to_video.py | 3 +- .../vae/z_image_latents_to_image.py | 3 +- .../app/services/config/config_default.py | 4 +- invokeai/backend/util/vae_working_memory.py | 9 +++-- invokeai/frontend/web/openapi.json | 4 +- .../frontend/web/src/services/api/schema.ts | 4 +- scripts/generate_docs_json.py | 2 +- tests/app/invocations/test_anima_vae.py | 28 ++++++++++++++ .../test_flux_vae_decode_tiling.py | 25 +++++++++++-- .../test_qwen_image_working_memory.py | 4 +- .../invocations/test_wan_working_memory.py | 22 ++++++++--- .../invocations/test_z_image_tiled_decode.py | 4 +- 16 files changed, 105 insertions(+), 57 deletions(-) diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 24e9d7e42c7..77208d4e531 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -163,7 +163,7 @@ device_working_mem_gb: 4 Turning latents into the final image (the VAE decode) needs working memory that grows with the image's pixel count. When an untiled decode would claim most of the GPU's memory, Invoke decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Krea-2, Anima and Wan decodes. A tiled decode is not pixel-identical to an untiled one, so 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. 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 it off: +To always decode untiled unless you ask for tiling, turn it off. A decode too large for the card then runs out of memory: the FLUX.1, Z-Image and Anima 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. ```yaml auto_tiled_decode: false @@ -228,7 +228,7 @@ Invoke does not officially support ROCm on Windows (see the [System Requirements - 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 decoding](#tiled-vae-decoding)); waiting for an out-of-memory error does not work there, because none comes. -- When a generation ends with part of Invoke's GPU memory in system memory, the log says so. Other programs using the same GPU, like games or browsers, can push Invoke's memory out as well. Close them, or restart Invoke to get the memory back into VRAM. +- 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 diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index b20e4b44142..dae131d13da 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -160,7 +160,7 @@ }, { "category": "PATHS", - "default": "models\\.convert_cache", + "default": "models/.convert_cache", "description": "Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).", "env_var": "INVOKEAI_CONVERT_CACHE_DIR", "literal_values": [], @@ -171,7 +171,7 @@ }, { "category": "PATHS", - "default": "models\\.download_cache", + "default": "models/.download_cache", "description": "Path to the directory that contains dynamically downloaded models.", "env_var": "INVOKEAI_DOWNLOAD_CACHE_DIR", "literal_values": [], @@ -728,7 +728,7 @@ { "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, Anima 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.", + "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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.", "env_var": "INVOKEAI_AUTO_TILED_DECODE", "literal_values": [], "name": "auto_tiled_decode", diff --git a/invokeai/app/invocations/vae/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index ad3d4738616..5118429b8c3 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -21,7 +21,6 @@ from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling from invokeai.backend.util.vae_working_memory import ( - VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_flux, should_pretile_vae_decode, ) @@ -56,8 +55,22 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents use_tiling = config.force_tiled_decode tile_size = 0 if use_tiling else None - # Only estimate working memory for BFL AutoEncoder (diffusers VAE handles this internally) - if isinstance(vae_info.model, AutoEncoder): + # Priced for either class, as in the Z-Image decode: a diffusers-layout FLUX VAE runs the same network, and a + # decode with nothing reserved is one Windows pages into system memory rather than failing. + 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, + ) + 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 = 0 estimated_working_memory = estimate_vae_working_memory_flux( operation="decode", image_tensor=latents, @@ -65,24 +78,6 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents tile_size=tile_size, 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, VAE_PRETILE_VRAM_FRACTION - ) - ): - use_tiling = True - tile_size = 0 - 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, - ) - else: - estimated_working_memory = 0 with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, (AutoEncoder, AutoencoderKL)) 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 cb91594300e..2c3034cd5e4 100644 --- a/invokeai/app/invocations/vae/qwen_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/qwen_image_latents_to_image.py @@ -23,7 +23,6 @@ resolve_qwen_image_vae_tile_size, ) from invokeai.backend.util.vae_working_memory import ( - VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_qwen_image, should_pretile_vae_decode, ) @@ -80,7 +79,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: if ( not tiled and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory, VAE_PRETILE_VRAM_FRACTION) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) ): effective_tile_size = resolve_qwen_image_vae_tile_size(self.tile_size) estimated_working_memory = estimate_vae_working_memory_qwen_image( diff --git a/invokeai/app/invocations/vae/wan_latents_to_video.py b/invokeai/app/invocations/vae/wan_latents_to_video.py index 07289a95e23..2658a33ff65 100644 --- a/invokeai/app/invocations/vae/wan_latents_to_video.py +++ b/invokeai/app/invocations/vae/wan_latents_to_video.py @@ -36,7 +36,6 @@ from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import ( - VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_wan, should_pretile_vae_decode, ) @@ -137,7 +136,7 @@ def invoke(self, context: InvocationContext) -> VideoOutput: # 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, VAE_PRETILE_VRAM_FRACTION + vae_info.compute_device, estimated_working_memory ) if use_tiling: tile_size = int(getattr(vae_info.model, "tile_sample_min_height", 256)) 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 ff609ea343b..dadf047a407 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -24,7 +24,6 @@ from invokeai.backend.util.oom import is_oom_error from invokeai.backend.util.vae_tiling_scope import scoped_vae_tiling from invokeai.backend.util.vae_working_memory import ( - VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_flux, should_pretile_vae_decode, ) @@ -78,7 +77,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: if ( not use_tiling and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory, VAE_PRETILE_VRAM_FRACTION) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) ): use_tiling = True estimated_working_memory = estimate_vae_working_memory_flux( diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 321b12e9f7f..d0e9152a341 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -125,7 +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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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. + 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. 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` @@ -250,7 +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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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.") + 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.") 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/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index e9800a276cc..16c7015f858 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -19,7 +19,9 @@ VAE_PRETILE_VRAM_FRACTION = 0.9 -def should_pretile_vae_decode(device: torch.device, full_decode_bytes: int, vram_fraction: float) -> bool: +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 ``device``'s memory. @@ -36,10 +38,9 @@ def should_pretile_vae_decode(device: torch.device, full_decode_bytes: int, vram return False if full_decode_bytes <= vram_fraction * total_bytes: return False - InvokeAILogger.get_logger(__name__).info( + 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 GPU's {total_bytes / 2**30:.1f} GiB. Set auto_tiled_decode to false to " - "decode untiled." + f"than {vram_fraction:.0%} of the GPU's {total_bytes / 2**30:.1f} GiB." ) return True diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 0dd671f9551..a31844effe9 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -55093,7 +55093,7 @@ "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, Anima 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.", + "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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.", "default": true }, "pil_compress_level": { @@ -55390,7 +55390,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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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.\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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.\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/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index a38c2d0f543..8a8f990b7d0 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -21129,7 +21129,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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima 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. + * 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. * 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` @@ -21560,7 +21560,7 @@ export type components = { 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, Anima 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. + * @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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. * @default true */ auto_tiled_decode?: boolean; 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 17d79a89529..40061116c19 100644 --- a/tests/app/invocations/test_anima_vae.py +++ b/tests/app/invocations/test_anima_vae.py @@ -12,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, @@ -24,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) @@ -326,6 +329,31 @@ 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_unless_switched_off(self, auto): + """Anima asks with its own, lower fraction (measured on 8GB cards) about the VAE's own device.""" + 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) + + if auto: + pretile.assert_called_once_with(vae_info.compute_device, full, ANIMA_PRETILE_VRAM_FRACTION) + vae.enable_tiling.assert_called_once() + else: + pretile.assert_not_called() + vae.enable_tiling.assert_not_called() + + 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_flux_vae_decode_tiling.py b/tests/app/invocations/test_flux_vae_decode_tiling.py index c87f6fd4929..c80e1c3a5f7 100644 --- a/tests/app/invocations/test_flux_vae_decode_tiling.py +++ b/tests/app/invocations/test_flux_vae_decode_tiling.py @@ -9,16 +9,18 @@ import pytest import torch +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.app.invocations.vae.flux_vae_decode import FluxVaeDecodeInvocation from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder -from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION -def _build_decode_mocks(latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): +def _build_decode_mocks( + latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False, vae_class: type = FluxAutoEncoder +): """Wire FluxVaeDecodeInvocation.invoke to run end-to-end on CPU against a mocked FLUX VAE.""" - vae = MagicMock(spec=FluxAutoEncoder) + vae = MagicMock(spec=vae_class) # A fresh iterator per call: the decode path reads `parameters()` after the estimator already # has, and a single stored iterator would be exhausted by then. vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) @@ -91,7 +93,7 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se _build_invocation().invoke(context) if auto: - pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30, VAE_PRETILE_VRAM_FRACTION) + 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) vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) @@ -100,6 +102,21 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se vae.disable_tiling.assert_called_once() +class TestDiffusersLayoutVae: + def test_a_diffusers_layout_flux_vae_reserves_its_working_memory(self): + """It runs the same network; a decode with nothing reserved is the one Windows pages instead of failing.""" + vae, vae_info, context = _build_decode_mocks( + torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512), vae_class=AutoencoderKL + ) + vae.config = MagicMock(scaling_factor=0.3611, shift_factor=0.1159) + vae.decode.return_value = (torch.zeros(1, 3, 512, 512),) + + _build_invocation().invoke(context) + + (reservation,) = vae_info.model_on_device.call_args.kwargs.values() + assert reservation > 0 + + class TestOomFallback: def test_an_untiled_oom_retries_once_tiled_and_says_so(self): decoded = torch.zeros(1, 3, 512, 512) diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index 1aaa3c81f3a..dd8e4e8a759 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -14,7 +14,7 @@ QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE, patch_qwen_image_vae_tiling, ) -from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_qwen_image +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image class TestQwenImageWorkingMemoryEstimate: @@ -447,7 +447,7 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se invocation.invoke(mock_context) if auto: - pretile.assert_called_once_with(mock_vae_info.compute_device, 20 * 2**30, VAE_PRETILE_VRAM_FRACTION) + pretile.assert_called_once_with(mock_vae_info.compute_device, 20 * 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() diff --git a/tests/app/invocations/test_wan_working_memory.py b/tests/app/invocations/test_wan_working_memory.py index 0239b424f38..82ed50cbb7f 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 @@ -296,11 +297,13 @@ 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() 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( @@ -319,16 +322,22 @@ 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 - vae.enable_tiling.assert_called_once() + if auto: + assert mock_estimate.call_count == 2 + assert mock_estimate.call_args_list[1].kwargs["tile_size"] == 256 + vae.enable_tiling.assert_called_once() + 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) vae.disable_tiling.assert_called_once() - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=4 * 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 ( @@ -338,6 +347,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 75dd323ebca..de36fd1ffad 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -9,7 +9,7 @@ from invokeai.app.invocations.vae.z_image_latents_to_image import ZImageLatentsToImageInvocation from invokeai.backend.flux.modules.autoencoder import DEFAULT_TILE_SAMPLE_MIN_SIZE, MIN_TILE_SAMPLE_SIZE from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder -from invokeai.backend.util.vae_working_memory import VAE_PRETILE_VRAM_FRACTION, estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: @@ -263,7 +263,7 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se _build_invocation().invoke(context) if auto: - pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30, VAE_PRETILE_VRAM_FRACTION) + 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) vae.enable_tiling.assert_called_once_with(tile_sample_min_size=DEFAULT_TILE_SAMPLE_MIN_SIZE) From 33e16af9891896c3e763d0fd7e8344aba9b5c428 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 21 Sep 2026 19:36:59 -0400 Subject: [PATCH 11/21] test(vae): add failing coverage for two up-front tiling defects These tests fail on purpose. They pin behaviour the current up-front tiling rule does not have, and should go green with the fix. - A Qwen-Image/Krea-2 decode whose measured peak fits the card is tiled anyway, because the gate is fed the padded reservation figure (a flat 5500 B/pixel-byte) instead of the expected peak (3273 measured at 1536px). Tiling is not pixel-identical, so this silently changes output for images that would have decoded in a single pass. - should_pretile_vae_decode compares against the card's nameplate total, not the Windows video-memory budget this PR added TorchDevice.cuda_mem_get_info for, so a decode Windows will page into system memory is left untiled. --- .../test_qwen_image_working_memory.py | 64 +++++++++++++++++++ tests/backend/util/test_vae_pretile.py | 34 ++++++++++ 2 files changed, 98 insertions(+) diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index dd8e4e8a759..6ff4dee1e1b 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -457,6 +457,70 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se 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 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/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py index 6ef2fc10982..a0213b28089 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -38,3 +38,37 @@ def test_1024px_stays_untiled_on_an_8gb_card_and_1536px_tiles_on_16gb(): 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, and another GPU + process lowers that budget within about a second -- which is exactly why this PR added + `TorchDevice.cuda_mem_get_info` to cap the model cache's free figure by `wddm.video_memory_budget`. + The tiling gate consults neither that budget nor free memory, only `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)) + # torch cannot see the shortfall: on Windows ROCm its free figure is the device total minus this + # process's own live allocations, and this process has allocated nothing yet. + monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (total_bytes, total_bytes)) + # Patch both binding sites: `devices` imports the name directly, so a fix routing through + # `TorchDevice.cuda_mem_get_info` and one calling `wddm.video_memory_budget` are both covered. + monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: budget_bytes) + monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: budget_bytes) + + 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 From 40a78452df1d7cf4d3d4172e7dc9612463512e12 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 05:54:48 +0200 Subject: [PATCH 12/21] fix(vae): tile on what the GPU keeps resident, not on reservation headroom - Compare against the Windows video-memory budget where there is one, not the card's total - Qwen-Image ROCm constants now follow the bounded math attention this branch ships: measured 2650-2772 decode and 1541-1552 encode across 512-2048px, against 5500/6300 fitted while a call built its whole score matrix - Calibrate with the attention guard installed, so the script measures the path the app runs --- .../vae/qwen_image_image_to_latents.py | 1 + .../vae/qwen_image_latents_to_image.py | 3 + invokeai/backend/util/vae_working_memory.py | 69 +++++++++++++------ scripts/calibrate_qwen_vae_working_memory.py | 7 ++ .../test_qwen_image_working_memory.py | 36 +++++++--- 5 files changed, 87 insertions(+), 29 deletions(-) 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 1c94243d514..9f834092ec3 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 2c3034cd5e4..f1f83123804 100644 --- a/invokeai/app/invocations/vae/qwen_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/qwen_image_latents_to_image.py @@ -75,18 +75,21 @@ def invoke(self, context: InvocationContext) -> ImageOutput: image_tensor=latents, vae=vae_info.model, tile_size=effective_tile_size, + device=vae_info.compute_device, ) if ( not tiled and config.auto_tiled_decode and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) ): + 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") diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 16c7015f858..1bb90aa16c6 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -9,7 +9,8 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.backend.flux.modules.autoencoder import AutoEncoder, resolve_tile_size -from invokeai.backend.util.attention import sdpa_score_matrix_bytes +from invokeai.backend.util import wddm +from invokeai.backend.util.attention import rocm_sdpa_chunks_math, sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger @@ -23,15 +24,22 @@ 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 ``device``'s memory. + ``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. + + The comparison is against the Windows video-memory budget where there is one, not the card's nameplate total: + that budget is the point past which Windows pages this process out, and another GPU process lowers it within + about a second. A decode measured against the total alone would look like it fits and then crawl. """ if device.type == "cuda": total_bytes = torch.cuda.get_device_properties(device).total_memory + budget_bytes = wddm.video_memory_budget(device) + if budget_bytes is not None: + total_bytes = min(total_bytes, budget_bytes) elif device.type == "xpu": total_bytes = torch.xpu.get_device_properties(device).total_memory else: @@ -40,7 +48,7 @@ def should_pretile_vae_decode( 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 GPU's {total_bytes / 2**30:.1f} GiB." + f"than {vram_fraction:.0%} of the {total_bytes / 2**30:.1f} GiB the GPU keeps resident." ) return True @@ -78,15 +86,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 @@ -121,7 +131,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 @@ -130,9 +143,10 @@ 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 @@ -540,6 +554,7 @@ def estimate_vae_working_memory_qwen_image( 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. @@ -554,7 +569,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 @@ -592,11 +608,20 @@ 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 + # On ROCm the figures above were measured while a math-kernel attention call built its whole score matrix. Where + # the guard bounds that call (`rocm_sdpa_chunks_math`), the super-linear term is gone and the curve is flat and + # linear again -- measured 2026-09-23 on an RX 9060 XT (gfx1200, torch 2.12+rocm7.14, fp16, the same script): + # decode 2650-2772 across 512^2 to 2048^2, encode 1541-1552, i.e. what CUDA measures (2660 decode). Shipping the + # unbounded figures there costs two things at once: eviction of models the decode does not need out of the way, + # and -- since the same figure decides it -- tiling a decode that fits, which is not pixel-identical. + device = device if device is not None else TorchDevice.choose_torch_device() + is_rocm = device.type == "cuda" and torch.version.hip is not None + if is_rocm and rocm_sdpa_chunks_math(device.type): + scaling_constant = 3000 if operation == "decode" else 1700 + elif is_rocm: + scaling_constant = 5500 if operation == "decode" else 6300 + else: + scaling_constant = 2900 if operation == "decode" else 1600 if tile_size is not None and tile_size > 0: # Bounded by one tile (plus overlap) rather than the full frame. @@ -619,7 +644,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 @@ -630,9 +658,10 @@ 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 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/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index 6ff4dee1e1b..b476221672e 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -28,15 +28,27 @@ class TestQwenImageWorkingMemoryEstimate: # (operation, latent_h, latent_w) -> the estimator scales pixel area (latent * 8 for decode, # raw for encode) by element_size and the constant. @pytest.mark.parametrize( - "operation, is_rocm, expected_constant", + "operation, is_rocm, chunked, expected_constant", [ - ("decode", True, 5500), - ("decode", False, 2900), - ("encode", True, 6300), - ("encode", False, 1600), + ("decode", True, True, 3000), + ("decode", True, False, 5500), + ("decode", False, False, 2900), + ("encode", True, True, 1700), + ("encode", True, False, 6300), + ("encode", False, False, 1600), + ], + ids=[ + "rocm-decode-chunked", + "rocm-decode-whole", + "cuda-decode", + "rocm-encode-chunked", + "rocm-encode-whole", + "cuda-encode", ], ) - def test_constant_selected_per_backend(self, operation, is_rocm, expected_constant): + def test_constant_selected_per_backend(self, operation, is_rocm, chunked, expected_constant): + """The ROCm figures split by whether math attention is chunked: the larger pair was measured while a single + call built its whole score matrix, which is the super-linear term the guard removes.""" mock_vae = MagicMock(spec=AutoencoderKLQwenImage) mock_vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.float16)]) # element_size == 2 @@ -49,9 +61,12 @@ def test_constant_selected_per_backend(self, operation, is_rocm, expected_consta h = w = 512 hip_value = "7.1.0" if is_rocm else None - with patch("torch.version.hip", hip_value): + with ( + patch("torch.version.hip", hip_value), + patch("invokeai.backend.util.vae_working_memory.rocm_sdpa_chunks_math", return_value=chunked), + ): 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 @@ -481,8 +496,11 @@ def test_a_decode_whose_measured_peak_fits_the_card_is_not_tiled(self, px, total 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. + # The constants are measured per backend, so pin it rather than inheriting the host's build. A running app + # always has the ROCm attention guard installed (`apply_monkeypatches`), and the ROCm constants differ by + # whether it is: measured 2650-2772 with math attention chunked against 5500 with the whole score matrix. monkeypatch.setattr(torch.version, "hip", "7.2.0") + monkeypatch.setattr("invokeai.backend.util.vae_working_memory.rocm_sdpa_chunks_math", lambda device: True) 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, From 39baf58bfd7500cb92407875b029af8d74e0258f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 05:54:48 +0200 Subject: [PATCH 13/21] fix(rocm): make the offload, the budget lookup and the ROCm checks hold up - Credit freed bytes while a peer device defers empty_cache, so multi-GPU stops over-unloading - Never let the budget or paged-bytes lookups raise, and close adapter handles if enumeration throws - Detect a ROCm build from torch's version.py too, so a locally built wheel gets the allocator default - Correct the bitwise-identical claim for head-group chunking --- invokeai/app/util/torch_cuda_allocator.py | 33 +++++++++- .../load/model_cache/model_cache.py | 20 ++++-- invokeai/backend/util/attention.py | 23 ++++--- invokeai/backend/util/devices.py | 7 +- invokeai/backend/util/wddm.py | 64 +++++++++++-------- invokeai/backend/wan/rocm_causal_conv3d.py | 2 +- tests/app/util/test_torch_cuda_allocator.py | 34 ++++++++++ .../test_model_cache_vram_accounting.py | 13 +++- .../util/test_rocm_sdpa_head_dim_guard.py | 4 +- 9 files changed, 147 insertions(+), 53 deletions(-) diff --git a/invokeai/app/util/torch_cuda_allocator.py b/invokeai/app/util/torch_cuda_allocator.py index a6f1d30fbfd..7e94be57900 100644 --- a/invokeai/app/util/torch_cuda_allocator.py +++ b/invokeai/app/util/torch_cuda_allocator.py @@ -1,7 +1,10 @@ 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. @@ -18,6 +21,33 @@ def _installed_torch_version() -> str | None: 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. @@ -33,8 +63,7 @@ def apply_rocm_windows_allocator_default(logger: logging.Logger) -> None: """ if sys.platform != "win32" or any(os.environ.get(var) for var in _ALLOCATOR_CONF_ENV_VARS): return - torch_version = _installed_torch_version() - if torch_version is None or "+rocm" not in torch_version: + 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. 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 eea4522cfe7..fa049ef6593 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -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 @@ -2554,13 +2556,16 @@ def _offload_unlocked_models( 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. + # 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 -- they sit in + # this process's allocator, which reuses them for the load this is making room for. 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 @@ -2574,7 +2579,10 @@ def _offload_unlocked_models( f"Unloaded {cache_entry.key} from VRAM to free {(cache_entry_bytes_freed / MB):.0f} MB." ) if empty_cache_per_offload: - TorchDevice.empty_cache() + 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 diff --git a/invokeai/backend/util/attention.py b/invokeai/backend/util/attention.py index d98f558b900..35b1f048c2d 100644 --- a/invokeai/backend/util/attention.py +++ b/invokeai/backend/util/attention.py @@ -132,8 +132,8 @@ def _read_rocm_fused_sdpa_max_head_dim() -> int: # (`_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, bitwise identical output; FLUX VAE mid-block (one 512-wide head, 1024px) +12 %, within one bf16 -# ulp. Read at call time, so tests can lower it. +# 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 @@ -228,12 +228,14 @@ def _chunked_sdpa( ) -> 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. Head groups come - first: every call keeps its full `query x key` shape, so the result is bitwise the unchunked one. 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; there the GEMM tiling changes by up to one bf16 ulp. 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. + 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 @@ -521,7 +523,10 @@ def sdpa_score_matrix_bytes( 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``. + 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 diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index 93de46aa4b2..ff68067eb87 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -334,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 @@ -367,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() @@ -377,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: diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index be9c37859d2..fc042211cd8 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -220,21 +220,27 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: return None matches: list[_AdapterInfo] = [] - for info in infos[: enum.NumAdapters]: - 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: + try: + for info in infos[: enum.NumAdapters]: + 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) + except Exception: + # Every handle opened by the enumeration belongs to this process until it is closed. + for info in matches: _close(lib, info.hAdapter) + raise if len(matches) != 1: for info in matches: @@ -266,23 +272,25 @@ def _adapter_for(device: torch.device) -> Optional[tuple[ctypes.CDLL, _Adapter]] 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. - ``None`` off Windows ROCm, and whenever the driver cannot answer or answers something implausible. + ``None`` off Windows ROCm, and whenever the driver cannot answer or answers something implausible. It is called + per generation and per decode, so every failure answers "unknown" here rather than reaching those callers. """ if not _supported(device): return None - resolved = _adapter_for(device) - if resolved is None: - return None - lib, adapter = resolved - info = _QueryVideoMemoryInfo(None, adapter.handle, _MEMORY_SEGMENT_GROUP_LOCAL) 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 + return int(info.Budget) except Exception: + InvokeAILogger.get_logger(__name__).debug(f"WDDM budget lookup for {device} failed", exc_info=True) return None - if not 0 < info.Budget <= adapter.total_bytes: - return None - return int(info.Budget) def _open_shared_usage_query(pdh: ctypes.CDLL) -> Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]: @@ -332,12 +340,12 @@ def paged_bytes(device: torch.device) -> Optional[int]: luid = resolved[1].luid instance = f"pid_{os.getpid()}_luid_0x{(luid >> 32) & 0xFFFFFFFF:08x}_0x{luid & 0xFFFFFFFF:08x}_phys_0" with _pdh_lock: - if not _pdh_attempted: - _pdh_attempted = True - _pdh_query = _open_shared_usage_query(pdh) - if _pdh_query is None: - return None try: + if not _pdh_attempted: + _pdh_attempted = True + _pdh_query = _open_shared_usage_query(pdh) + if _pdh_query is None: + return None return _shared_usage_by_instance(pdh, *_pdh_query).get(instance) except Exception: return None 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/tests/app/util/test_torch_cuda_allocator.py b/tests/app/util/test_torch_cuda_allocator.py index 3e2ddf8715e..a226f889d4f 100644 --- a/tests/app/util/test_torch_cuda_allocator.py +++ b/tests/app/util/test_torch_cuda_allocator.py @@ -1,5 +1,6 @@ import os import sys +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -72,6 +73,39 @@ def test_too_late_once_torch_is_imported(self, rocm_windows): 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/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 933d539c9bd..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 @@ -214,9 +214,13 @@ def test_the_windows_video_memory_budget_caps_the_measured_free_vram(monkeypatch assert cache._get_vram_available(None) == 6 * GB -def test_offloading_under_expandable_segments_stops_once_enough_is_free(monkeypatch: pytest.MonkeyPatch): +@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.""" + 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, @@ -234,9 +238,12 @@ def move_to_ram(entry, bytes_to_free): device["unmappable"] += 4 * GB return 4 * GB - def empty_cache(): + 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)) 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 d02979d0147..79c3e920157 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -201,7 +201,9 @@ def _scores(shape_pair: tuple[torch.Size, torch.Size]) -> int: 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_is_bitwise_unchunked(self, monkeypatch): + 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) From 9989d684ea4c78f28fa1eaf002ecf99cd99fb349 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 05:54:49 +0200 Subject: [PATCH 14/21] fix(vae): keep Anima tiling and price the AutoencoderKL nodes per backend - auto_tiled_decode no longer disables Anima's own rule, which is a speed optimization, not an OOM fallback - SD1/SDXL, SD3 and CogView4 take the MIOpen constant on ROCm, like the FLUX.1 autoencoder they share a stack with - Encode nodes place their tensors on the VAE's device, matching the decode nodes --- .../content/docs/configuration/low-vram-mode.mdx | 4 ++-- docs/src/generated/settings.json | 2 +- .../app/invocations/vae/anima_latents_to_image.py | 5 ++++- .../invocations/vae/cogview4_image_to_latents.py | 2 +- .../invocations/vae/cogview4_latents_to_image.py | 2 +- invokeai/app/invocations/vae/flux_vae_encode.py | 7 ++++--- invokeai/app/invocations/vae/image_to_latents.py | 1 + invokeai/app/invocations/vae/latents_to_image.py | 1 + .../app/invocations/vae/sd3_image_to_latents.py | 2 +- .../app/invocations/vae/sd3_latents_to_image.py | 2 +- .../invocations/vae/z_image_image_to_latents.py | 7 ++++--- invokeai/app/services/config/config_default.py | 4 ++-- invokeai/frontend/web/openapi.json | 4 ++-- invokeai/frontend/web/src/services/api/schema.ts | 4 ++-- tests/app/invocations/test_anima_vae.py | 14 ++++++-------- 15 files changed, 33 insertions(+), 28 deletions(-) diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 77208d4e531..bb11d4f1d1d 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -161,9 +161,9 @@ device_working_mem_gb: 4 ### Tiled VAE decoding -Turning latents into the final image (the VAE decode) needs working memory that grows with the image's pixel count. When an untiled decode would claim most of the GPU's memory, Invoke decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Krea-2, Anima and Wan decodes. A tiled decode is not pixel-identical to an untiled one, so 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. Anima tiles a little earlier, because its untiled decode pushes the Anima model itself out of VRAM on 8GB cards. +Turning latents into the final image (the VAE decode) needs working memory that grows with the image's pixel count. When an untiled decode would claim most of the memory the GPU keeps resident, Invoke decodes in tiles instead, for the FLUX.1, Z-Image, Qwen-Image, Krea-2, Anima and Wan decodes. A tiled decode is not pixel-identical to an untiled one, so 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 it off. A decode too large for the card then runs out of memory: the FLUX.1, Z-Image and Anima 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. +To always decode untiled unless you ask for tiling, turn it 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 diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index dae131d13da..96cc847a159 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -728,7 +728,7 @@ { "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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.", + "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", diff --git a/invokeai/app/invocations/vae/anima_latents_to_image.py b/invokeai/app/invocations/vae/anima_latents_to_image.py index 10433e38830..d6243d888a5 100644 --- a/invokeai/app/invocations/vae/anima_latents_to_image.py +++ b/invokeai/app/invocations/vae/anima_latents_to_image.py @@ -106,7 +106,10 @@ def invoke(self, context: InvocationContext) -> ImageOutput: vae=vae_info.model, tile_size=None, ) - use_tiling = context.config.get().auto_tiled_decode and should_pretile_vae_decode( + # 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( 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_encode.py b/invokeai/app/invocations/vae/flux_vae_encode.py index cfa471697a9..3fc1adc307f 100644 --- a/invokeai/app/invocations/vae/flux_vae_encode.py +++ b/invokeai/app/invocations/vae/flux_vae_encode.py @@ -14,7 +14,6 @@ from invokeai.backend.flux.modules.autoencoder import AutoEncoder from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor -from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux @@ -46,11 +45,13 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso estimated_working_memory = estimate_vae_working_memory_flux( operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) - 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, AutoEncoder) 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) latents = vae.encode(image_tensor, sample=True, generator=generator) return latents 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/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/z_image_image_to_latents.py b/invokeai/app/invocations/vae/z_image_image_to_latents.py index 48d65d3ee7a..6a85f62f1a6 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -19,7 +19,6 @@ from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder from invokeai.backend.model_manager.load.load_base import LoadedModel from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor -from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux # Z-Image can use either the Diffusers AutoencoderKL or the FLUX AutoEncoder @@ -64,12 +63,14 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso ) 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 torch.inference_mode(): if isinstance(vae, FluxAutoEncoder): # FLUX VAE handles scaling internally - generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0) + generator = torch.Generator(device=vae_info.compute_device).manual_seed(0) latents = vae.encode(image_tensor, sample=True, generator=generator) else: # AutoencoderKL - needs manual scaling diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index d0e9152a341..d80c3bd4806 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -125,7 +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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. + 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` @@ -250,7 +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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.") + 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/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index a31844effe9..7908b45aa76 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -55093,7 +55093,7 @@ "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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.", + "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": { @@ -55390,7 +55390,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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory.\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; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, 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. As of this release the setting also applies to FLUX.1, which previously ignored 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/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 8a8f990b7d0..810a69927bf 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -21129,7 +21129,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. As of this release the setting also applies to FLUX.1, which previously ignored 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. + * 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` @@ -21560,7 +21560,7 @@ export type components = { 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, Anima and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1, Z-Image and Anima decodes then still retry tiled after running out of memory. + * @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; diff --git a/tests/app/invocations/test_anima_vae.py b/tests/app/invocations/test_anima_vae.py index 40061116c19..766fa4319eb 100644 --- a/tests/app/invocations/test_anima_vae.py +++ b/tests/app/invocations/test_anima_vae.py @@ -331,8 +331,10 @@ def test_decode_requests_estimated_working_memory(self): 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_unless_switched_off(self, auto): - """Anima asks with its own, lower fraction (measured on 8GB cards) about the VAE's own device.""" + 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 @@ -346,12 +348,8 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se ): _build_l2i_invocation().invoke(context) - if auto: - pretile.assert_called_once_with(vae_info.compute_device, full, ANIMA_PRETILE_VRAM_FRACTION) - vae.enable_tiling.assert_called_once() - else: - pretile.assert_not_called() - vae.enable_tiling.assert_not_called() + pretile.assert_called_once_with(vae_info.compute_device, full, ANIMA_PRETILE_VRAM_FRACTION) + vae.enable_tiling.assert_called_once() class TestAnimaImageToLatentsEncode: From feb4d8e2ab20792f67c2a3ff706a8a9d8fd08b14 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 07:04:11 +0200 Subject: [PATCH 15/21] fix(vae): price up-front tiling from the measured peak and a stable ceiling - Keep the padded Qwen-Image reservations; the tiling decision reads the measured curve instead, so a decode whose real peak fits the card is no longer tiled by a reservation's headroom - Compare against the video-memory budget plus what this process can release: Windows halves the budget once the process passes ~12 of 16 GiB, and the bare figure tiled a 7.0 GiB decode that fits --- .../vae/qwen_image_latents_to_image.py | 8 +- invokeai/backend/util/vae_working_memory.py | 101 ++++++++++++++---- invokeai/backend/util/wddm.py | 9 +- .../test_qwen_image_working_memory.py | 85 ++++++++++----- tests/backend/util/test_vae_pretile.py | 15 +++ 5 files changed, 166 insertions(+), 52 deletions(-) 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 f1f83123804..62b9f5eced6 100644 --- a/invokeai/app/invocations/vae/qwen_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/qwen_image_latents_to_image.py @@ -24,6 +24,7 @@ ) from invokeai.backend.util.vae_working_memory import ( estimate_vae_working_memory_qwen_image, + qwen_image_untiled_decode_peak_bytes, should_pretile_vae_decode, ) @@ -77,10 +78,15 @@ def invoke(self, context: InvocationContext) -> ImageOutput: 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, estimated_working_memory) + 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) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 1bb90aa16c6..18ccd0ea4a4 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -10,7 +10,7 @@ from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.backend.flux.modules.autoencoder import AutoEncoder, resolve_tile_size from invokeai.backend.util import wddm -from invokeai.backend.util.attention import rocm_sdpa_chunks_math, sdpa_score_matrix_bytes +from invokeai.backend.util.attention import sdpa_score_matrix_bytes from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger @@ -31,15 +31,20 @@ def should_pretile_vae_decode( 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. - The comparison is against the Windows video-memory budget where there is one, not the card's nameplate total: - that budget is the point past which Windows pages this process out, and another GPU process lowers it within - about a second. A decode measured against the total alone would look like it fits and then crawl. + Where Windows grants a video-memory budget, that budget bounds the card's nameplate total: it is the point past + which Windows pages this process out, and another GPU process lowers it within about a second. What the decode + may claim is the budget plus whatever this process already holds, because the cache evicts models to make room + for the reservation -- measured on an RX 9060 XT, Windows keeps the budget at 15.09 GiB until the process passes + about 12 GiB and then halves it to 7.62, so comparing against the bare budget mid-session tiles a decode that + fits (a 7.0 GiB decode against a 6.9 GiB line, with output that is not pixel-identical). """ if device.type == "cuda": total_bytes = torch.cuda.get_device_properties(device).total_memory budget_bytes = wddm.video_memory_budget(device) if budget_bytes is not None: - total_bytes = min(total_bytes, budget_bytes) + # On a ROCm build under Windows torch's free figure is the total minus this process's own allocations. + free_bytes, _total = torch.cuda.mem_get_info(device) + total_bytes = min(total_bytes, budget_bytes + (total_bytes - free_bytes)) elif device.type == "xpu": total_bytes = torch.xpu.get_device_properties(device).total_memory else: @@ -121,7 +126,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). @@ -151,7 +156,9 @@ def estimate_vae_working_memory_cogview4( # 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. @@ -549,6 +556,60 @@ 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 largest of the measured ROCm cards at each point (the per-point table in +# `estimate_vae_working_memory_qwen_image`, W7900 column; the RX 9060 XT measures 2650-2772 throughout), and the CUDA +# maximum for everything else. +# +# 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. +_QWEN_VAE_MEASURED_DECODE_PEAKS: tuple[tuple[int, int], ...] = ( + (512 * 512, 5132), + (768 * 768, 4596), + (1024 * 1024, 4570), + (1536 * 1536, 3273), + (1792 * 1792, 3735), + (2048 * 2048, 4813), +) +_QWEN_VAE_MEASURED_DECODE_PEAK_CUDA = 2690 + + +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 takes the larger of the two, and beyond the largest measured area the largest of all: the + curve is not monotonic, and a decode that big is tiled on any card either way. + """ + 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 + if not is_rocm: + return h * w * element_size * _QWEN_VAE_MEASURED_DECODE_PEAK_CUDA + + area = h * w + measured = dict(_QWEN_VAE_MEASURED_DECODE_PEAKS) + if area in measured: + constant = measured[area] + elif area > _QWEN_VAE_MEASURED_DECODE_PEAKS[-1][0]: + constant = max(measured.values()) + else: + previous = _QWEN_VAE_MEASURED_DECODE_PEAKS[0][1] + constant = previous + for measured_area, peak in _QWEN_VAE_MEASURED_DECODE_PEAKS: + if measured_area > area: + constant = max(previous, peak) + break + previous = peak + return h * w * element_size * constant + + def estimate_vae_working_memory_qwen_image( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, @@ -608,20 +669,18 @@ 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. - # On ROCm the figures above were measured while a math-kernel attention call built its whole score matrix. Where - # the guard bounds that call (`rocm_sdpa_chunks_math`), the super-linear term is gone and the curve is flat and - # linear again -- measured 2026-09-23 on an RX 9060 XT (gfx1200, torch 2.12+rocm7.14, fp16, the same script): - # decode 2650-2772 across 512^2 to 2048^2, encode 1541-1552, i.e. what CUDA measures (2660 decode). Shipping the - # unbounded figures there costs two things at once: eviction of models the decode does not need out of the way, - # and -- since the same figure decides it -- tiling a decode that fits, which is not pixel-identical. + # - 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. device = device if device is not None else TorchDevice.choose_torch_device() is_rocm = device.type == "cuda" and torch.version.hip is not None - if is_rocm and rocm_sdpa_chunks_math(device.type): - scaling_constant = 3000 if operation == "decode" else 1700 - elif is_rocm: - scaling_constant = 5500 if operation == "decode" else 6300 - else: - scaling_constant = 2900 if operation == "decode" else 1600 + if operation == "decode": + scaling_constant = 5500 if is_rocm else 2900 + else: # encode + scaling_constant = 6300 if is_rocm else 1600 if tile_size is not None and tile_size > 0: # Bounded by one tile (plus overlap) rather than the full frame. @@ -667,4 +726,6 @@ def estimate_vae_working_memory_sd3( # 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 index fc042211cd8..376c9003997 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -220,8 +220,10 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: return None matches: list[_AdapterInfo] = [] + pending = list(infos[: enum.NumAdapters]) try: - for info in infos[: enum.NumAdapters]: + while pending: + info = pending.pop() address = _AdapterAddress() query = _QueryAdapterInfo( info.hAdapter, @@ -237,8 +239,9 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: else: _close(lib, info.hAdapter) except Exception: - # Every handle opened by the enumeration belongs to this process until it is closed. - for info in matches: + # Every handle the enumeration opened belongs to this process until it is closed -- the ones matched so far + # and the ones this loop never reached. + for info in matches + pending: _close(lib, info.hAdapter) raise diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index b476221672e..0a019b7d919 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: @@ -28,27 +31,15 @@ class TestQwenImageWorkingMemoryEstimate: # (operation, latent_h, latent_w) -> the estimator scales pixel area (latent * 8 for decode, # raw for encode) by element_size and the constant. @pytest.mark.parametrize( - "operation, is_rocm, chunked, expected_constant", + "operation, is_rocm, expected_constant", [ - ("decode", True, True, 3000), - ("decode", True, False, 5500), - ("decode", False, False, 2900), - ("encode", True, True, 1700), - ("encode", True, False, 6300), - ("encode", False, False, 1600), - ], - ids=[ - "rocm-decode-chunked", - "rocm-decode-whole", - "cuda-decode", - "rocm-encode-chunked", - "rocm-encode-whole", - "cuda-encode", + ("decode", True, 5500), + ("decode", False, 2900), + ("encode", True, 6300), + ("encode", False, 1600), ], ) - def test_constant_selected_per_backend(self, operation, is_rocm, chunked, expected_constant): - """The ROCm figures split by whether math attention is chunked: the larger pair was measured while a single - call built its whole score matrix, which is the super-linear term the guard removes.""" + def test_constant_selected_per_backend(self, operation, is_rocm, expected_constant): mock_vae = MagicMock(spec=AutoencoderKLQwenImage) mock_vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.float16)]) # element_size == 2 @@ -61,10 +52,7 @@ def test_constant_selected_per_backend(self, operation, is_rocm, chunked, expect h = w = 512 hip_value = "7.1.0" if is_rocm else None - with ( - patch("torch.version.hip", hip_value), - patch("invokeai.backend.util.vae_working_memory.rocm_sdpa_chunks_math", return_value=chunked), - ): + with patch("torch.version.hip", hip_value): result = estimate_vae_working_memory_qwen_image( operation=operation, image_tensor=image_tensor, vae=mock_vae, device=torch.device("cuda") ) @@ -449,6 +437,7 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se 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()), ): @@ -462,11 +451,14 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se invocation.invoke(mock_context) if auto: - pretile.assert_called_once_with(mock_vae_info.compute_device, 20 * 2**30) + # 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() @@ -479,6 +471,46 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se 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(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_takes_the_larger(self): + # 1280^2 sits between the 1024^2 (4570) and 1536^2 (3273) measurements. + assert self._peak(1280) == 1280 * 1280 * 2 * 4570 + + def test_past_the_measured_range_takes_the_largest(self): + assert self._peak(3072) == 3072 * 3072 * 2 * 5132 + + def test_a_cuda_card_is_priced_from_the_cuda_measurements(self): + assert self._peak(1536, hip=None) == 1536 * 1536 * 2 * 2690 + + 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. @@ -496,11 +528,8 @@ def test_a_decode_whose_measured_peak_fits_the_card_is_not_tiled(self, px, total 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. A running app - # always has the ROCm attention guard installed (`apply_monkeypatches`), and the ROCm constants differ by - # whether it is: measured 2650-2772 with math attention chunked against 5500 with the whole score matrix. + # 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("invokeai.backend.util.vae_working_memory.rocm_sdpa_chunks_math", lambda device: True) 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, diff --git a/tests/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py index a0213b28089..76fd578fff1 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -72,3 +72,18 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa 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_the_gate_counts_memory_this_process_can_still_release(monkeypatch): + """Windows halves the budget once the process passes about 12 of 16 GiB, i.e. below what it already holds. The + cache evicts models for the reservation, so the ceiling is the budget plus this process's own allocations -- + measured: comparing against the bare budget tiled a 7.0 GiB Z-Image decode against a 6.9 GiB line mid-session, + changing 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)) + # Mid-session: 12 GiB of models resident, and Windows has trimmed the budget to 7.6 GiB in response. + monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (4 * 2**30, total_bytes)) + monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(7.6 * 2**30)) + + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is False From 1a3e8104a2b2180b9d4ff75df3af40ef3d3be37a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 07:17:54 +0200 Subject: [PATCH 16/21] fix(tests): point the merged FLUX.2 constant test at the shared table main added a case reading _FLUX2_VAE_SCALING_CONSTANTS, which this branch merged into _FLUX_VAE_SCALING_CONSTANTS. --- tests/app/invocations/test_flux2_working_memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 58ffa0abd31..286b6b44511 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -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" From f8d7dfb7f35abfe91294a9ef265685bedd05a7e9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 23 Sep 2026 07:46:24 +0200 Subject: [PATCH 17/21] fix(vae): keep the build-wide constant when no VAE device is named The new device argument made a device-less estimate resolve the session device, so the same call answered MIOpen on a GPU box and cuDNN on CI. A named device still answers for that device. Pin the convolution backend separately from the attention guard in the ROCm estimator tests, which only diverged on a ROCm build. --- invokeai/backend/util/vae_working_memory.py | 6 ++++-- tests/backend/util/test_rocm_sdpa_head_dim_guard.py | 4 ++++ tests/backend/util/test_vae_pretile.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 48b594596f7..9c0daa81bb1 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -795,8 +795,10 @@ def estimate_vae_working_memory_qwen_image( # 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. - device = device if device is not None else TorchDevice.choose_torch_device() - is_rocm = device.type == "cuda" and torch.version.hip is not None + # 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") if operation == "decode": scaling_constant = 5500 if is_rocm else 2900 else: # encode 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 79c3e920157..662dd13e8e6 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -486,9 +486,13 @@ class TestClassicVaeEstimators: @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 index 76fd578fff1..79e208b833d 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -73,6 +73,7 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True + def test_the_gate_counts_memory_this_process_can_still_release(monkeypatch): """Windows halves the budget once the process passes about 12 of 16 GiB, i.e. below what it already holds. The cache evicts models for the reservation, so the ceiling is the budget plus this process's own allocations -- From 355595ffe3a1085966abd8272aec2f4ce6f5259f Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 27 Sep 2026 18:50:28 +0200 Subject: [PATCH 18/21] fix(vae): decide tiling on what other processes hold, not on the budget - The ceiling is the card minus foreign residency (new wddm.other_process_local_bytes): measured on an RX 9060 XT, the budget reads 15.09 GiB next to an 8 GiB holder with ~7.7 GiB available, and 7.62 GiB when Invoke itself holds 12 GiB and eviction would free the card. The budget stays as the fallback - Both backends get their measured decode curve; past the last measured area the decision uses the reservation constant instead of extrapolating through a calibration point that ran out of memory - Make the allocator and pre-tile tests hermetic on a ROCm host, rename the guard call in the LTX-2 slow test, and cover the MIOpen branch of the SD/SD3/CogView4 estimators --- .../load/model_cache/model_cache.py | 8 +- invokeai/backend/util/vae_working_memory.py | 84 ++++++++++------- invokeai/backend/util/wddm.py | 94 +++++++++++++++---- .../test_qwen_image_working_memory.py | 15 ++- .../test_vae_working_memory_backend.py | 63 +++++++++++++ tests/app/util/test_torch_cuda_allocator.py | 17 +++- tests/backend/ltx2/test_ltx2_real_files.py | 4 +- tests/backend/util/test_vae_pretile.py | 48 ++++++++-- tests/backend/util/test_wddm.py | 52 ++++++++-- 9 files changed, 309 insertions(+), 76 deletions(-) create mode 100644 tests/app/invocations/test_vae_working_memory_backend.py 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 fa049ef6593..c64bc9fe69e 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2557,8 +2557,12 @@ def _offload_unlocked_models( # 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 -- they sit in - # this process's allocator, which reuses them for the load this is making room for. + # 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. diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 9c0daa81bb1..1eb19c55f30 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -31,20 +31,21 @@ def should_pretile_vae_decode( 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. - Where Windows grants a video-memory budget, that budget bounds the card's nameplate total: it is the point past - which Windows pages this process out, and another GPU process lowers it within about a second. What the decode - may claim is the budget plus whatever this process already holds, because the cache evicts models to make room - for the reservation -- measured on an RX 9060 XT, Windows keeps the budget at 15.09 GiB until the process passes - about 12 GiB and then halves it to 7.62, so comparing against the bare budget mid-session tiles a decode that - fits (a 7.0 GiB decode against a 6.9 GiB line, with output that is not pixel-identical). + On Windows the card's nameplate total is not what this decode can have: what the rest of the card is holding is. + The cache evicts Invoke's own models to honour the reservation, so the ceiling is the total minus what *other* + processes hold (`wddm.other_process_local_bytes`). Measured on an RX 9060 XT (16 GB), neither figure Windows + reports answers that on its own: with another program holding 8 GiB and Invoke idle the budget still read + 15.09 GiB while only ~7.7 GiB could be had, and with Invoke holding 12 GiB and nothing else on the card the + budget read 7.62 GiB although releasing its own models would have freed almost all of it. Where the counters + cannot answer, the budget is the fallback: too pessimistic mid-session, but never too optimistic. """ if device.type == "cuda": total_bytes = torch.cuda.get_device_properties(device).total_memory budget_bytes = wddm.video_memory_budget(device) if budget_bytes is not None: - # On a ROCm build under Windows torch's free figure is the total minus this process's own allocations. - free_bytes, _total = torch.cuda.mem_get_info(device) - total_bytes = min(total_bytes, budget_bytes + (total_bytes - free_bytes)) + other_bytes = wddm.other_process_local_bytes(device) + ceiling = total_bytes - other_bytes if other_bytes is not None else budget_bytes + total_bytes = min(total_bytes, ceiling) elif device.type == "xpu": total_bytes = torch.xpu.get_device_properties(device).total_memory else: @@ -677,24 +678,43 @@ def estimate_vae_working_memory_wan( # 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 largest of the measured ROCm cards at each point (the per-point table in -# `estimate_vae_working_memory_qwen_image`, W7900 column; the RX 9060 XT measures 2650-2772 throughout), and the CUDA -# maximum for everything else. +# 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. -_QWEN_VAE_MEASURED_DECODE_PEAKS: tuple[tuple[int, int], ...] = ( - (512 * 512, 5132), - (768 * 768, 4596), - (1024 * 1024, 4570), - (1536 * 1536, 3273), - (1792 * 1792, 3735), - (2048 * 2048, 4813), -) -_QWEN_VAE_MEASURED_DECODE_PEAK_CUDA = 2690 +# +# 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( @@ -703,26 +723,25 @@ def qwen_image_untiled_decode_peak_bytes( """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 takes the larger of the two, and beyond the largest measured area the largest of all: the - curve is not monotonic, and a decode that big is tiled on any card either way. + measured points this takes the larger of the two; past the last one it takes that 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 - if not is_rocm: - return h * w * element_size * _QWEN_VAE_MEASURED_DECODE_PEAK_CUDA + points = _QWEN_VAE_MEASURED_DECODE_PEAKS["rocm" if is_rocm else "cuda"] area = h * w - measured = dict(_QWEN_VAE_MEASURED_DECODE_PEAKS) + measured = dict(points) if area in measured: constant = measured[area] - elif area > _QWEN_VAE_MEASURED_DECODE_PEAKS[-1][0]: - constant = max(measured.values()) + elif area > points[-1][0]: + constant = _QWEN_VAE_RESERVATION_CONSTANTS["rocm" if is_rocm else "cuda"]["decode"] else: - previous = _QWEN_VAE_MEASURED_DECODE_PEAKS[0][1] + previous = points[0][1] constant = previous - for measured_area, peak in _QWEN_VAE_MEASURED_DECODE_PEAKS: + for measured_area, peak in points: if measured_area > area: constant = max(previous, peak) break @@ -799,10 +818,7 @@ def estimate_vae_working_memory_qwen_image( # 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") - if operation == "decode": - scaling_constant = 5500 if is_rocm else 2900 - else: # encode - scaling_constant = 6300 if is_rocm else 1600 + 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. diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index 376c9003997..89f9d1010b6 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -38,6 +38,10 @@ _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" +# Adapter-wide, i.e. every process on the adapter. Per-process attribution from `GPU Process Memory` was +# measured to mix up its instances while two processes shared the card, so what "other processes hold" is +# derived from this figure and torch's own-usage view instead. +_PDH_ADAPTER_USAGE_COUNTER = r"\GPU Adapter Memory(*)\Dedicated Usage" # Win32 types spelled as fixed-width ctypes, so this module imports on every platform. @@ -138,8 +142,8 @@ class _Adapter: _adapters: dict[int, Optional[_Adapter]] = {} _pdh_lock = threading.Lock() -# One PDH query over every process's shared usage, as (query, counter), opened on first use; None if that failed. -_pdh_query: Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]] = None +# The PDH query and its counters, opened on first use: (query, {counter name: handle}). +_pdh_query: Optional[tuple[ctypes.c_void_p, dict[str, ctypes.c_void_p]]] = None _pdh_attempted = False @@ -221,9 +225,13 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: matches: list[_AdapterInfo] = [] pending = list(infos[: enum.NumAdapters]) + current: list[_AdapterInfo] = [] try: while pending: - info = pending.pop() + # 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, @@ -238,10 +246,11 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: 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 - # and the ones this loop never reached. - for info in matches + pending: + # 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 @@ -296,19 +305,22 @@ def video_memory_budget(device: torch.device) -> Optional[int]: return None -def _open_shared_usage_query(pdh: ctypes.CDLL) -> Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]: +def _open_gpu_memory_query(pdh: ctypes.CDLL) -> Optional[tuple[ctypes.c_void_p, dict[str, 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, _PDH_SHARED_USAGE_COUNTER, None, ctypes.byref(counter)) != 0: - pdh.PdhCloseQuery(query) - return None - return query, counter + counters: dict[str, ctypes.c_void_p] = {} + for name, path in (("shared", _PDH_SHARED_USAGE_COUNTER), ("adapter", _PDH_ADAPTER_USAGE_COUNTER)): + counter = ctypes.c_void_p() + if pdh.PdhAddEnglishCounterW(query, path, None, ctypes.byref(counter)) != 0: + pdh.PdhCloseQuery(query) + return None + counters[name] = counter + return query, counters -def _shared_usage_by_instance(pdh: ctypes.CDLL, query: ctypes.c_void_p, counter: ctypes.c_void_p) -> dict[str, int]: - """Collect the query once; every ``GPU Process Memory`` instance's shared usage, keyed by its lower-case name.""" +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) @@ -340,20 +352,66 @@ def paged_bytes(device: torch.device) -> Optional[int]: pdh = _load_pdh() if resolved is None or pdh is None: return None - luid = resolved[1].luid - instance = f"pid_{os.getpid()}_luid_0x{(luid >> 32) & 0xFFFFFFFF:08x}_0x{luid & 0xFFFFFFFF:08x}_phys_0" + instance = f"pid_{os.getpid()}_{_adapter_instance(resolved[1].luid)}" with _pdh_lock: try: if not _pdh_attempted: _pdh_attempted = True - _pdh_query = _open_shared_usage_query(pdh) + _pdh_query = _open_gpu_memory_query(pdh) if _pdh_query is None: return None - return _shared_usage_by_instance(pdh, *_pdh_query).get(instance) + query, counters = _pdh_query + return _usage_by_instance(pdh, query, counters["shared"]).get(instance) except Exception: return None +def _adapter_instance(luid: int) -> str: + return f"luid_0x{(luid >> 32) & 0xFFFFFFFF:08x}_0x{luid & 0xFFFFFFFF:08x}_phys_0" + + +def other_process_local_bytes(device: torch.device) -> Optional[int]: + """Dedicated video memory on the device's adapter held by processes other than this one. + + This is what bounds how large an allocation this process could still make after releasing everything of its own, + and it is the figure the budget does not answer. Measured on an RX 9060 XT (16 GB): with another program holding + 8 GiB and Invoke idle, the budget still reported 15.09 GiB while only ~7.7 GiB could be had; with Invoke itself + holding 12 GiB and nothing else on the card, the budget read 7.62 GiB while releasing its own memory would have + freed almost the whole card. + + ``None`` off Windows ROCm and whenever the counters cannot answer. One PDH sample, so a process that takes memory + in the same instant is missed; the model cache's own per-load measurement is what catches that. + """ + global _pdh_query, _pdh_attempted + + if not _supported(device): + return None + resolved = _adapter_for(device) + pdh = _load_pdh() + if resolved is None or pdh is None: + return None + instance = _adapter_instance(resolved[1].luid) + with _pdh_lock: + try: + if not _pdh_attempted: + _pdh_attempted = True + _pdh_query = _open_gpu_memory_query(pdh) + if _pdh_query is None: + return None + query, counters = _pdh_query + adapter_bytes = _usage_by_instance(pdh, query, counters["adapter"]).get(instance) + except Exception: + return None + if adapter_bytes is None: + return None + try: + free_bytes, total_bytes = torch.cuda.mem_get_info(device) + except Exception: + return None + # On a ROCm build under Windows torch's free figure is the total minus this process's own allocations. + return max(0, adapter_bytes - (total_bytes - free_bytes)) + + def reset_cache() -> None: """Close the cached adapter handles and the PDH query, and forget the loaders (tests only).""" global _pdh_query, _pdh_attempted diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index 9a1ae9ad53a..d12c44b19f8 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -567,11 +567,16 @@ def test_between_two_points_takes_the_larger(self): # 1280^2 sits between the 1024^2 (4570) and 1536^2 (3273) measurements. assert self._peak(1280) == 1280 * 1280 * 2 * 4570 - def test_past_the_measured_range_takes_the_largest(self): - assert self._peak(3072) == 3072 * 3072 * 2 * 5132 - - def test_a_cuda_card_is_priced_from_the_cuda_measurements(self): - assert self._peak(1536, hip=None) == 1536 * 1536 * 2 * 2690 + 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.""" 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/util/test_torch_cuda_allocator.py b/tests/app/util/test_torch_cuda_allocator.py index a226f889d4f..d3461983ec4 100644 --- a/tests/app/util/test_torch_cuda_allocator.py +++ b/tests/app/util/test_torch_cuda_allocator.py @@ -14,10 +14,23 @@ @pytest.fixture -def rocm_windows(monkeypatch: pytest.MonkeyPatch): - """A Windows ROCm install with no allocator configuration in the environment.""" +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 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/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py index 79e208b833d..ae9404d21a4 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -53,7 +53,8 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa 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. + so it is not tiled. With no adapter-wide reading available this is the budget's own case; the + tests below cover what the gate does when it can tell whose memory it is. """ total_bytes = 16 * 2**30 budget_bytes = int(12.5 * 2**30) @@ -67,6 +68,8 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa # `TorchDevice.cuda_mem_get_info` and one calling `wddm.video_memory_budget` are both covered. monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: budget_bytes) monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: budget_bytes) + # Without this the gate reads the real card here, and any foreign usage below ~1.2 GiB flips the case. + monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: None) 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" @@ -74,17 +77,48 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True -def test_the_gate_counts_memory_this_process_can_still_release(monkeypatch): - """Windows halves the budget once the process passes about 12 of 16 GiB, i.e. below what it already holds. The - cache evicts models for the reservation, so the ceiling is the budget plus this process's own allocations -- - measured: comparing against the bare budget tiled a 7.0 GiB Z-Image decode against a 6.9 GiB line mid-session, - changing output that had been pixel-identical.""" +def test_a_budget_trimmed_by_our_own_residency_does_not_tile_a_decode_that_fits(monkeypatch): + """Windows trims the budget when the process itself grows -- measured on an RX 9060 XT, 15.09 GiB until it passes + about 12 of 16 GiB, then 7.62. The cache evicts those models for the reservation, so a decode that fits must not + be tiled on that reading: it tiled a 7.0 GiB Z-Image decode against a 6.9 GiB line mid-session, which 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)) - # Mid-session: 12 GiB of models resident, and Windows has trimmed the budget to 7.6 GiB in response. + # 12 GiB of our own models resident, and nothing else on the card. monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (4 * 2**30, total_bytes)) monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(7.6 * 2**30)) + monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: 0) assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is False + + +def test_another_process_holding_the_card_tiles_even_while_the_budget_looks_fine(monkeypatch): + """The case the budget does not answer: measured with another program holding 8 GiB and Invoke idle, the budget + still read 15.09 GiB although only ~7.7 GiB could be had. Evicting Invoke's own models cannot recover memory it + does not hold, so a 8.9 GiB Qwen-Image decode has to be tiled here.""" + total_bytes = 16 * 2**30 + estimate = int(8.9 * 2**30) + + monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) + # Invoke holds almost nothing; torch's free figure does not see the other process at all. + monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (total_bytes - 2**28, total_bytes)) + monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(15.09 * 2**30)) + monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: 8 * 2**30) + + assert estimate < VAE_PRETILE_VRAM_FRACTION * total_bytes, "the card total and the budget both look fine" + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True + + +def test_the_budget_is_the_fallback_when_the_counters_cannot_answer(monkeypatch): + """No adapter-wide figure means no way to tell whose memory it is; the budget is then the conservative answer.""" + 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("torch.cuda.mem_get_info", lambda device: (4 * 2**30, total_bytes)) + monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(7.6 * 2**30)) + monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: None) + + assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py index 5db03b5f52c..de7209c72e1 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -26,13 +26,15 @@ def _luid_low(handle: int) -> int: class _FakePdh: - """A PDH that reports the given ``GPU Process Memory`` instances' shared usage.""" + """A PDH with the module's two counters: per-process shared usage, and adapter-wide dedicated usage.""" - def __init__(self, instances: dict[str, int]): + def __init__(self, instances: dict[str, int], adapter_instances: dict[str, int] | None = None): self.instances = instances + self.adapter_instances = adapter_instances or {} self.counter_paths: list[str] = [] self.closed = 0 self._alive: list[object] = [] + self._by_handle: dict[int, dict[str, int]] = {} def PdhOpenQueryW(self, source, user_data, query_ref): query_ref._obj.value = 1 @@ -40,20 +42,23 @@ def PdhOpenQueryW(self, source, user_data, query_ref): def PdhAddEnglishCounterW(self, query, path, user_data, counter_ref): self.counter_paths.append(path) - counter_ref._obj.value = 2 + handle = 2 + len(self.counter_paths) + counter_ref._obj.value = handle + self._by_handle[handle] = self.adapter_instances if "Adapter" in path else self.instances return 0 def PdhCollectQueryData(self, query): return 0 def PdhGetFormattedCounterArrayW(self, counter, fmt, size_ref, count_ref, buffer): - count = len(self.instances) + instances = self._by_handle[counter.value if hasattr(counter, "value") else counter] + 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, self.instances.items(), strict=True): + for item, (name, value) in zip(items, instances.items(), strict=True): item.szName = name item.FmtValue.largeValue = value self._alive.append(items) # the name buffers live as long as the array object @@ -205,12 +210,47 @@ def test_paged_bytes_reads_this_process_on_the_devices_adapter(windows_rocm, mon assert wddm.paged_bytes(DEVICE) == 3 * GIB assert wddm.paged_bytes(DEVICE) == 3 * GIB - assert pdh.counter_paths == [r"\GPU Process Memory(*)\Shared Usage"], "one query, opened on first use" + assert pdh.counter_paths == [ + r"\GPU Process Memory(*)\Shared Usage", + r"\GPU Adapter Memory(*)\Dedicated Usage", + ], "one query with both counters, opened on first use" wddm.reset_cache() assert pdh.closed == 1 +def _adapter_instance(handle: int) -> str: + return f"luid_0x00000000_0x{_luid_low(handle):08x}_phys_0" + + +@pytest.mark.parametrize( + ("adapter_gib", "ours_gib", "expected_gib"), + [(8.4, 0.2, 8.2), (12.2, 12.1, 0.0), (0.2, 6.0, 0.0)], + ids=["another-process-holds-the-card", "all-of-it-is-ours", "adapter-reads-below-our-own"], +) +def test_other_process_local_bytes_subtracts_our_own_usage( + windows_rocm, monkeypatch, adapter_gib, ours_gib, expected_gib +): + """What bounds a decode is what OTHER processes hold: the budget does not track them (measured 15.09 GiB next to + an 8 GiB holder), and the per-process PDH counter swapped its instances while two processes shared the card, so + this comes from the adapter-wide figure minus torch's own-usage view. A reading below our own clamps to zero.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + total = 16 * GIB + pdh = _FakePdh({}, adapter_instances={_adapter_instance(20): int(adapter_gib * GIB)}) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (total - int(ours_gib * GIB), total)) + + assert wddm.other_process_local_bytes(DEVICE) == pytest.approx(int(expected_gib * GIB), abs=GIB // 10) + + +def test_other_process_local_bytes_is_unknown_without_an_adapter_reading(windows_rocm, monkeypatch): + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + pdh = _FakePdh({}, adapter_instances={_adapter_instance(21): 8 * GIB}) # a different adapter + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.other_process_local_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}) From bbd71a1470c08e024a563a8d8100ed204e7c12f2 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 28 Sep 2026 05:51:25 +0200 Subject: [PATCH 19/21] fix(vae): tile on the budget and our own residency, not on an adapter counter - The adapter-wide PDH counter cannot be used: measured with an 8 GiB holder, the holding process read 15.66 GiB and another read 0.176 GiB for the same instance at the same moment. It also shares the driver's basis, so subtracting torch's live figure invented ~2 GiB of foreign residency after every load/unload cycle - The ceiling is min(total, max(budget, our CurrentUsage)): a budget below our own residency asks this process to trim itself, which the cache does for the reservation. A foreign holder that leaves the budget alone is documented as undetectable - One PDH query per counter, per-item CStatus checked, usage clamped rather than dropped with the budget - Interpolate decode peaks in bytes so a smaller image is never priced above a larger one --- invokeai/backend/util/vae_working_memory.py | 61 ++++---- invokeai/backend/util/wddm.py | 144 +++++++----------- .../test_qwen_image_working_memory.py | 26 +++- tests/backend/util/test_vae_pretile.py | 55 +++---- tests/backend/util/test_wddm.py | 76 ++++----- 5 files changed, 179 insertions(+), 183 deletions(-) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 1eb19c55f30..335889948cf 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -31,21 +31,25 @@ def should_pretile_vae_decode( 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 card's nameplate total is not what this decode can have: what the rest of the card is holding is. - The cache evicts Invoke's own models to honour the reservation, so the ceiling is the total minus what *other* - processes hold (`wddm.other_process_local_bytes`). Measured on an RX 9060 XT (16 GB), neither figure Windows - reports answers that on its own: with another program holding 8 GiB and Invoke idle the budget still read - 15.09 GiB while only ~7.7 GiB could be had, and with Invoke holding 12 GiB and nothing else on the card the - budget read 7.62 GiB although releasing its own models would have freed almost all of it. Where the counters - cannot answer, the budget is the fallback: too pessimistic mid-session, but never too optimistic. + 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 already holds is + Windows asking it to trim itself, and the cache does exactly that to honour the reservation, so the ceiling is then + what it holds. 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. + + Known limit: a foreign process that takes memory without the budget moving is invisible here. Windows keeps + reporting 15.09 GiB next to an 8 GiB holder, 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), so there is no reading to base it on. Such a decode pages, 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 - budget_bytes = wddm.video_memory_budget(device) - if budget_bytes is not None: - other_bytes = wddm.other_process_local_bytes(device) - ceiling = total_bytes - other_bytes if other_bytes is not None else budget_bytes - total_bytes = min(total_bytes, ceiling) + local = wddm.local_video_memory(device) + if local is not None: + budget_bytes, usage_bytes = local + total_bytes = min(total_bytes, max(budget_bytes, usage_bytes)) elif device.type == "xpu": total_bytes = torch.xpu.get_device_properties(device).total_memory else: @@ -723,7 +727,9 @@ def qwen_image_untiled_decode_peak_bytes( """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 takes the larger of the two; past the last one it takes that reservation figure, since a + 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] @@ -733,20 +739,23 @@ def qwen_image_untiled_decode_peak_bytes( points = _QWEN_VAE_MEASURED_DECODE_PEAKS["rocm" if is_rocm else "cuda"] area = h * w - measured = dict(points) - if area in measured: - constant = measured[area] - elif area > points[-1][0]: + if area > points[-1][0]: constant = _QWEN_VAE_RESERVATION_CONSTANTS["rocm" if is_rocm else "cuda"]["decode"] - else: - previous = points[0][1] - constant = previous - for measured_area, peak in points: - if measured_area > area: - constant = max(previous, peak) - break - previous = peak - return h * w * element_size * constant + 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( diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index 89f9d1010b6..64405508ace 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -38,10 +38,8 @@ _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" -# Adapter-wide, i.e. every process on the adapter. Per-process attribution from `GPU Process Memory` was -# measured to mix up its instances while two processes shared the card, so what "other processes hold" is -# derived from this figure and torch's own-usage view instead. -_PDH_ADAPTER_USAGE_COUNTER = r"\GPU Adapter Memory(*)\Dedicated 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. @@ -142,9 +140,9 @@ class _Adapter: _adapters: dict[int, Optional[_Adapter]] = {} _pdh_lock = threading.Lock() -# The PDH query and its counters, opened on first use: (query, {counter name: handle}). -_pdh_query: Optional[tuple[ctypes.c_void_p, dict[str, ctypes.c_void_p]]] = None -_pdh_attempted = False +# One PDH query per counter path, opened on first use: path -> (query, counter) or None when that counter is absent. +# Per counter so that a counter this build does not have disables only what reads it. +_pdh_queries: dict[str, Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]] = {} def _supported(device: torch.device) -> bool: @@ -281,11 +279,11 @@ def _adapter_for(device: torch.device) -> Optional[tuple[ctypes.CDLL, _Adapter]] return None if adapter is None else (lib, adapter) -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. +def local_video_memory(device: torch.device) -> Optional[tuple[int, 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 something implausible. It is called - per generation and per decode, so every failure answers "unknown" here rather than reaching those callers. + ``None`` off Windows ROCm, and whenever the driver cannot answer or answers something implausible. Called per + generation and per decode, so every failure answers "unknown" here rather than reaching those callers. """ if not _supported(device): return None @@ -299,24 +297,47 @@ def video_memory_budget(device: torch.device) -> Optional[int]: return None if not 0 < info.Budget <= adapter.total_bytes: return None - return int(info.Budget) + # Clamped, not rejected: `CurrentUsage` counts bytes Windows has already paged out, so it can exceed the + # adapter's own size exactly when the process over-commits -- and dropping the pair there would also take the + # budget away from the model cache's free-VRAM cap, which is the caller that needs it most in that state. + return int(info.Budget), max(0, min(int(info.CurrentUsage), adapter.total_bytes)) except Exception: - InvokeAILogger.get_logger(__name__).debug(f"WDDM budget lookup for {device} failed", exc_info=True) + InvokeAILogger.get_logger(__name__).debug(f"WDDM memory lookup for {device} failed", exc_info=True) return None -def _open_gpu_memory_query(pdh: ctypes.CDLL) -> Optional[tuple[ctypes.c_void_p, dict[str, ctypes.c_void_p]]]: +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 - counters: dict[str, ctypes.c_void_p] = {} - for name, path in (("shared", _PDH_SHARED_USAGE_COUNTER), ("adapter", _PDH_ADAPTER_USAGE_COUNTER)): - counter = ctypes.c_void_p() - if pdh.PdhAddEnglishCounterW(query, path, None, ctypes.byref(counter)) != 0: - pdh.PdhCloseQuery(query) - return None - counters[name] = counter - return query, counters + 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]: @@ -332,7 +353,11 @@ def _usage_by_instance(pdh: ctypes.CDLL, query: ctypes.c_void_p, counter: ctypes 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} + 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]: @@ -342,80 +367,23 @@ def paged_bytes(device: torch.device) -> Optional[int]: 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 the PDH query (~0.3 s); each later call collects it once (~0.3 ms). + The first call opens this counter's PDH query (~0.3 s); each later call collects it once (~0.3 ms). """ - global _pdh_query, _pdh_attempted - if not _supported(device): return None resolved = _adapter_for(device) - pdh = _load_pdh() - if resolved is None or pdh is None: + if resolved is None: return None instance = f"pid_{os.getpid()}_{_adapter_instance(resolved[1].luid)}" - with _pdh_lock: - try: - if not _pdh_attempted: - _pdh_attempted = True - _pdh_query = _open_gpu_memory_query(pdh) - if _pdh_query is None: - return None - query, counters = _pdh_query - return _usage_by_instance(pdh, query, counters["shared"]).get(instance) - except Exception: - return None + 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 other_process_local_bytes(device: torch.device) -> Optional[int]: - """Dedicated video memory on the device's adapter held by processes other than this one. - - This is what bounds how large an allocation this process could still make after releasing everything of its own, - and it is the figure the budget does not answer. Measured on an RX 9060 XT (16 GB): with another program holding - 8 GiB and Invoke idle, the budget still reported 15.09 GiB while only ~7.7 GiB could be had; with Invoke itself - holding 12 GiB and nothing else on the card, the budget read 7.62 GiB while releasing its own memory would have - freed almost the whole card. - - ``None`` off Windows ROCm and whenever the counters cannot answer. One PDH sample, so a process that takes memory - in the same instant is missed; the model cache's own per-load measurement is what catches that. - """ - global _pdh_query, _pdh_attempted - - if not _supported(device): - return None - resolved = _adapter_for(device) - pdh = _load_pdh() - if resolved is None or pdh is None: - return None - instance = _adapter_instance(resolved[1].luid) - with _pdh_lock: - try: - if not _pdh_attempted: - _pdh_attempted = True - _pdh_query = _open_gpu_memory_query(pdh) - if _pdh_query is None: - return None - query, counters = _pdh_query - adapter_bytes = _usage_by_instance(pdh, query, counters["adapter"]).get(instance) - except Exception: - return None - if adapter_bytes is None: - return None - try: - free_bytes, total_bytes = torch.cuda.mem_get_info(device) - except Exception: - return None - # On a ROCm build under Windows torch's free figure is the total minus this process's own allocations. - return max(0, adapter_bytes - (total_bytes - free_bytes)) - - def reset_cache() -> None: """Close the cached adapter handles and the PDH query, and forget the loaders (tests only).""" - global _pdh_query, _pdh_attempted - with _lock: opened = [adapter for adapter in _adapters.values() if adapter is not None] _adapters.clear() @@ -426,8 +394,10 @@ def reset_cache() -> None: _load_gdi32.cache_clear() with _pdh_lock: - query, _pdh_query, _pdh_attempted = _pdh_query, None, False - pdh = _load_pdh() if query is not None else None - if pdh is not None and query is not None: - pdh.PdhCloseQuery(query[0]) + 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/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index d12c44b19f8..49fd219fa0f 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -551,6 +551,14 @@ 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)]) @@ -563,9 +571,21 @@ def _peak(self, px: int, device: str = "cuda", hip: str | None = "7.2.0") -> int 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_takes_the_larger(self): - # 1280^2 sits between the 1024^2 (4570) and 1536^2 (3273) measurements. - assert self._peak(1280) == 1280 * 1280 * 2 * 4570 + 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.""" + assert self._peak_hw(1728, 1856, hip=None) <= self._peak_hw(1792, 1792, hip=None) 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 diff --git a/tests/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py index ae9404d21a4..0900aea2090 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -53,8 +53,7 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa 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. With no adapter-wide reading available this is the budget's own case; the - tests below cover what the gate does when it can tell whose memory it is. + so it is not tiled. """ total_bytes = 16 * 2**30 budget_bytes = int(12.5 * 2**30) @@ -64,12 +63,9 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa # torch cannot see the shortfall: on Windows ROCm its free figure is the device total minus this # process's own live allocations, and this process has allocated nothing yet. monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (total_bytes, total_bytes)) - # Patch both binding sites: `devices` imports the name directly, so a fix routing through - # `TorchDevice.cuda_mem_get_info` and one calling `wddm.video_memory_budget` are both covered. + # Patch both binding sites: `devices` imports its own name, and the gate reads the budget/usage pair. monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: budget_bytes) - monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: budget_bytes) - # Without this the gate reads the real card here, and any foreign usage below ~1.2 GiB flips the case. - monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: None) + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: (budget_bytes, 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" @@ -77,48 +73,43 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True -def test_a_budget_trimmed_by_our_own_residency_does_not_tile_a_decode_that_fits(monkeypatch): - """Windows trims the budget when the process itself grows -- measured on an RX 9060 XT, 15.09 GiB until it passes - about 12 of 16 GiB, then 7.62. The cache evicts those models for the reservation, so a decode that fits must not - be tiled on that reading: it tiled a 7.0 GiB Z-Image decode against a 6.9 GiB line mid-session, which changed - output that had been pixel-identical.""" +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)) - # 12 GiB of our own models resident, and nothing else on the card. monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (4 * 2**30, total_bytes)) - monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(7.6 * 2**30)) - monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: 0) + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: (int(7.6 * 2**30), 12 * 2**30)) assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is False -def test_another_process_holding_the_card_tiles_even_while_the_budget_looks_fine(monkeypatch): - """The case the budget does not answer: measured with another program holding 8 GiB and Invoke idle, the budget - still read 15.09 GiB although only ~7.7 GiB could be had. Evicting Invoke's own models cannot recover memory it - does not hold, so a 8.9 GiB Qwen-Image decode has to be tiled here.""" +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.9 * 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)) - # Invoke holds almost nothing; torch's free figure does not see the other process at all. - monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (total_bytes - 2**28, total_bytes)) - monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(15.09 * 2**30)) - monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: 8 * 2**30) + monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (int(7.8 * 2**30), total_bytes)) + monkeypatch.setattr( + "invokeai.backend.util.wddm.local_video_memory", lambda device: (int(7.55 * 2**30), int(8.2 * 2**30)) + ) - assert estimate < VAE_PRETILE_VRAM_FRACTION * total_bytes, "the card total and the budget both look fine" + 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 -def test_the_budget_is_the_fallback_when_the_counters_cannot_answer(monkeypatch): - """No adapter-wide figure means no way to tell whose memory it is; the budget is then the conservative answer.""" +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 - estimate = 7 * 2**30 monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) - monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (4 * 2**30, total_bytes)) - monkeypatch.setattr("invokeai.backend.util.wddm.video_memory_budget", lambda device: int(7.6 * 2**30)) - monkeypatch.setattr("invokeai.backend.util.wddm.other_process_local_bytes", lambda device: None) + monkeypatch.setattr("invokeai.backend.util.wddm.local_video_memory", lambda device: None) - assert should_pretile_vae_decode(torch.device("cuda", 0), estimate) is True + 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 index de7209c72e1..2ec16a9ae1a 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -26,32 +26,38 @@ def _luid_low(handle: int) -> int: class _FakePdh: - """A PDH with the module's two counters: per-process shared usage, and adapter-wide dedicated usage.""" - - def __init__(self, instances: dict[str, int], adapter_instances: dict[str, int] | None = None): + """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.adapter_instances = adapter_instances or {} + self.unavailable = unavailable + self.invalid = invalid self.counter_paths: list[str] = [] self.closed = 0 self._alive: list[object] = [] - self._by_handle: dict[int, dict[str, int]] = {} 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 - self._by_handle[handle] = self.adapter_instances if "Adapter" in path else self.instances return 0 def PdhCollectQueryData(self, query): return 0 def PdhGetFormattedCounterArrayW(self, counter, fmt, size_ref, count_ref, buffer): - instances = self._by_handle[counter.value if hasattr(counter, "value") else counter] + instances = self.instances count = len(instances) if buffer is None: size_ref._obj.value = ctypes.sizeof(wddm._PdhCounterValueItem) * count @@ -61,6 +67,8 @@ def PdhGetFormattedCounterArrayW(self, counter, fmt, size_ref, count_ref, 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 @@ -73,9 +81,10 @@ def PdhCloseQuery(self, query): class _FakeGdi32: """Adapters as (handle, bus, device); every thunk returns NTSTATUS 0 unless told otherwise.""" - def __init__(self, adapters, budget=15 * GIB, query_status=0): + 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] = [] @@ -101,7 +110,7 @@ def D3DKMTQueryAdapterInfo(self, ref): def D3DKMTQueryVideoMemoryInfo(self, ref): info = ref._obj - info.Budget = self.budget + info.Budget, info.CurrentUsage = self.budget, self.usage return self.query_status def D3DKMTCloseAdapter(self, ref): @@ -142,6 +151,16 @@ def test_resolves_the_adapter_once(windows_rocm): assert gdi32.enum_calls == 2, "one count call and one fill call, for the first query only" +def test_usage_beyond_the_card_is_clamped_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. Dropping the pair there would take the budget away from the model cache's free-VRAM cap in + exactly that state, so the usage is clamped instead.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)], budget=15 * GIB, usage=17 * GIB)) + + assert wddm.local_video_memory(DEVICE) == (15 * GIB, 16 * GIB) # the adapter's total + assert wddm.video_memory_budget(DEVICE) == 15 * GIB + + @pytest.mark.parametrize( ("adapters", "kwargs"), [ @@ -210,45 +229,32 @@ def test_paged_bytes_reads_this_process_on_the_devices_adapter(windows_rocm, mon assert wddm.paged_bytes(DEVICE) == 3 * GIB assert wddm.paged_bytes(DEVICE) == 3 * GIB - assert pdh.counter_paths == [ - r"\GPU Process Memory(*)\Shared Usage", - r"\GPU Adapter Memory(*)\Dedicated Usage", - ], "one query with both counters, opened on first use" + 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 _adapter_instance(handle: int) -> str: - return f"luid_0x00000000_0x{_luid_low(handle):08x}_phys_0" - - -@pytest.mark.parametrize( - ("adapter_gib", "ours_gib", "expected_gib"), - [(8.4, 0.2, 8.2), (12.2, 12.1, 0.0), (0.2, 6.0, 0.0)], - ids=["another-process-holds-the-card", "all-of-it-is-ours", "adapter-reads-below-our-own"], -) -def test_other_process_local_bytes_subtracts_our_own_usage( - windows_rocm, monkeypatch, adapter_gib, ours_gib, expected_gib -): - """What bounds a decode is what OTHER processes hold: the budget does not track them (measured 15.09 GiB next to - an 8 GiB holder), and the per-process PDH counter swapped its instances while two processes shared the card, so - this comes from the adapter-wide figure minus torch's own-usage view. A reading below our own clamps to zero.""" +def test_a_counter_this_build_does_not_have_only_disables_what_reads_it(windows_rocm, monkeypatch): + """One query per counter: sharing one meant any missing counter took every PDH-backed answer with it.""" windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) - total = 16 * GIB - pdh = _FakePdh({}, adapter_instances={_adapter_instance(20): int(adapter_gib * GIB)}) + 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)) - monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (total - int(ours_gib * GIB), total)) - assert wddm.other_process_local_bytes(DEVICE) == pytest.approx(int(expected_gib * GIB), abs=GIB // 10) + assert wddm.paged_bytes(DEVICE) is None + assert wddm.local_video_memory(DEVICE) is not None, "a D3DKMT answer does not depend on the counters" -def test_other_process_local_bytes_is_unknown_without_an_adapter_reading(windows_rocm, monkeypatch): +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)])) - pdh = _FakePdh({}, adapter_instances={_adapter_instance(21): 8 * GIB}) # a different adapter + 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.other_process_local_bytes(DEVICE) is None + assert wddm.paged_bytes(DEVICE) is None def test_paged_bytes_is_unknown_before_this_process_used_the_adapter(windows_rocm, monkeypatch): From 5878d9d102c5e2937cea7e577e9ebd8057be1b20 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Tue, 29 Sep 2026 05:23:43 +0200 Subject: [PATCH 20/21] fix(vae): bound the pre-tiling ceiling by residency, not by usage Windows has paged out Only what this process keeps in VRAM (CurrentUsage minus its paged bytes) may lift the ceiling above the WDDM budget; usage beyond the card or an unknown paged amount leaves the budget in charge. The measured budget behaviour next to a foreign holder now reads the same in wddm.py and the gate docstring. Tests cover the foreign-holder, over-commit and unknown-paged states and area monotonicity of the Qwen pricing on both backends. --- invokeai/backend/util/vae_working_memory.py | 26 ++++++--- invokeai/backend/util/wddm.py | 32 ++++++----- .../test_qwen_image_working_memory.py | 13 ++++- tests/backend/util/test_vae_pretile.py | 54 +++++++++++++++---- tests/backend/util/test_wddm.py | 13 +++-- 5 files changed, 98 insertions(+), 40 deletions(-) diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 335889948cf..f1b321a20a1 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -32,24 +32,34 @@ def should_pretile_vae_decode( 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 already holds is + 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 - what it holds. Measured on an RX 9060 XT (16 GB): with Invoke holding 12 GiB and nothing else on the card the + 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. - Known limit: a foreign process that takes memory without the budget moving is invisible here. Windows keeps - reporting 15.09 GiB next to an 8 GiB holder, 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), so there is no reading to base it on. Such a decode pages, and the - end-of-session paging warning is what surfaces it. + 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 - total_bytes = min(total_bytes, max(budget_bytes, usage_bytes)) + 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: diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index 64405508ace..ec39f7b7389 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -7,11 +7,16 @@ ``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; next to another GPU process, Windows lowers the budget within -about a second. 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. +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. @@ -141,7 +146,6 @@ class _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. -# Per counter so that a counter this build does not have disables only what reads it. _pdh_queries: dict[str, Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]] = {} @@ -279,11 +283,12 @@ def _adapter_for(device: torch.device) -> Optional[tuple[ctypes.CDLL, _Adapter]] return None if adapter is None else (lib, adapter) -def local_video_memory(device: torch.device) -> Optional[tuple[int, int]]: +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 something implausible. Called per - generation and per decode, so every failure answers "unknown" here rather than reaching those callers. + ``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 @@ -297,10 +302,11 @@ def local_video_memory(device: torch.device) -> Optional[tuple[int, int]]: return None if not 0 < info.Budget <= adapter.total_bytes: return None - # Clamped, not rejected: `CurrentUsage` counts bytes Windows has already paged out, so it can exceed the - # adapter's own size exactly when the process over-commits -- and dropping the pair there would also take the - # budget away from the model cache's free-VRAM cap, which is the caller that needs it most in that state. - return int(info.Budget), max(0, min(int(info.CurrentUsage), adapter.total_bytes)) + # `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 diff --git a/tests/app/invocations/test_qwen_image_working_memory.py b/tests/app/invocations/test_qwen_image_working_memory.py index 49fd219fa0f..246477618fd 100644 --- a/tests/app/invocations/test_qwen_image_working_memory.py +++ b/tests/app/invocations/test_qwen_image_working_memory.py @@ -584,8 +584,19 @@ def test_between_two_points_interpolates_the_bytes(self): 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.""" + """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 diff --git a/tests/backend/util/test_vae_pretile.py b/tests/backend/util/test_vae_pretile.py index 0900aea2090..4aba3c32e7b 100644 --- a/tests/backend/util/test_vae_pretile.py +++ b/tests/backend/util/test_vae_pretile.py @@ -45,10 +45,8 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa 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, and another GPU - process lowers that budget within about a second -- which is exactly why this PR added - `TorchDevice.cuda_mem_get_info` to cap the model cache's free figure by `wddm.video_memory_budget`. - The tiling gate consults neither that budget nor free memory, only `total_memory`. + 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 @@ -60,12 +58,9 @@ def test_the_pretile_gate_uses_the_memory_the_device_will_keep_resident(monkeypa 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)) - # torch cannot see the shortfall: on Windows ROCm its free figure is the device total minus this - # process's own live allocations, and this process has allocated nothing yet. - monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (total_bytes, total_bytes)) - # Patch both binding sites: `devices` imports its own name, and the gate reads the budget/usage pair. - monkeypatch.setattr("invokeai.backend.util.devices.video_memory_budget", lambda device: budget_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" @@ -82,8 +77,8 @@ def test_a_budget_below_our_own_residency_does_not_tile_a_decode_that_fits(monke estimate = 7 * 2**30 monkeypatch.setattr("torch.cuda.get_device_properties", lambda device: MagicMock(total_memory=total_bytes)) - monkeypatch.setattr("torch.cuda.mem_get_info", lambda device: (4 * 2**30, 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 @@ -95,15 +90,52 @@ def test_a_decode_beyond_both_the_budget_and_our_residency_is_tiled(monkeypatch) 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("torch.cuda.mem_get_info", lambda device: (int(7.8 * 2**30), 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 diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py index 2ec16a9ae1a..a2c53b65ea3 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -151,13 +151,13 @@ def test_resolves_the_adapter_once(windows_rocm): assert gdi32.enum_calls == 2, "one count call and one fill call, for the first query only" -def test_usage_beyond_the_card_is_clamped_and_keeps_the_budget(windows_rocm): +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. Dropping the pair there would take the budget away from the model cache's free-VRAM cap in - exactly that state, so the usage is clamped instead.""" + 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, 16 * GIB) # the adapter's total + assert wddm.local_video_memory(DEVICE) == (15 * GIB, None) assert wddm.video_memory_budget(DEVICE) == 15 * GIB @@ -237,14 +237,13 @@ def test_paged_bytes_reads_this_process_on_the_devices_adapter(windows_rocm, mon assert pdh.closed == 1 -def test_a_counter_this_build_does_not_have_only_disables_what_reads_it(windows_rocm, monkeypatch): - """One query per counter: sharing one meant any missing counter took every PDH-backed answer with it.""" +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 - assert wddm.local_video_memory(DEVICE) is not None, "a D3DKMT answer does not depend on the counters" def test_an_instance_with_invalid_data_is_skipped(windows_rocm, monkeypatch): From 40fb6750480afb9c9e9dcf7ab95958fd95551a4e Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 30 Sep 2026 03:18:54 +0200 Subject: [PATCH 21/21] style(tests): reflow a call the merge left wrapped --- tests/app/invocations/test_z_image_tiled_decode.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index eb1bd744a01..2f8a46a894d 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -389,9 +389,7 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se 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) - ) + _, 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,