fix(multi-gpu): disable torch's post-conv-algorithm-search global emptyCache

invokeai/app/services/session_processor/session_processor_default.py

@@ -42,7 +42,7 @@
 from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context
 from invokeai.app.util.profiler import Profiler
 from invokeai.backend.util.device_pool import GENERATION_DEVICE_POOL
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import TorchDevice, disable_conv_benchmark_empty_cache
 
 # 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 —
@@ -618,6 +618,12 @@ def start(self, invoker: Invoker) -> None:
         # encoders on (see offload_text_encoders_to_idle_gpus). None means legacy single-device mode.
         GENERATION_DEVICE_POOL.set_generation_devices([d for d in devices if d is not None])
 
+        # With more than one CUDA/HIP generation device, torch's post-conv-algorithm-search global
+        # emptyCache() convoys the peer GPU's in-flight step from C++, where the peer-aware
+        # empty_cache wrapper cannot intercept it. Trade it for cached workspace blocks instead.
+        if sum(1 for d in devices if d is not None and d.type == "cuda") > 1:
+            disable_conv_benchmark_empty_cache()
+
         # If profiling is enabled, create a profiler. The same profiler will be used for all sessions. Internally,
         # the profiler will create a new profile for each session. Profiling uses a process-global cProfile, which
         # cannot cleanly attribute work when multiple sessions run concurrently, so it is disabled in multi-GPU mode.

invokeai/backend/util/devices.py

@@ -498,3 +498,34 @@ def peer_aware_empty_cache() -> None:
 
     setattr(peer_aware_empty_cache, _PEER_AWARE_SENTINEL, True)
     torch.cuda.empty_cache = peer_aware_empty_cache
+
+
+def disable_conv_benchmark_empty_cache() -> None:
+    """Stop torch's conv algorithm search from calling global emptyCache() after each find.
+
+    When torch searches for a convolution algorithm (cuDNN benchmark mode; MIOpen on ROCm on
+    every algo-cache miss, benchmark flag or not), ``findAlgorithm`` ends the search with a
+    process-global ``CUDACachingAllocator::emptyCache()`` to release benchmarking workspace
+    (``aten/src/ATen/native/miopen/Conv_miopen.cpp``, likewise the cuDNN path). That call takes
+    every CUDA/HIP device's allocator mutex and frees their cached blocks — freezing a peer
+    GPU's worker mid-step exactly like the Python-level ``torch.cuda.empty_cache`` calls
+    handled by ``install_peer_aware_empty_cache``, but from C++, out of reach of that wrapper
+    (observed via py-spy on a dual-GPU ROCm rig: one worker inside
+    ``chooseAlgorithm -> emptyCache -> hipFree`` waiting out the other worker's 40-100 s
+    denoise step; each new conv shape in the process re-triggers it).
+
+    The call is gated on ``_cudnn_get_conv_benchmark_empty_cache()``, which torch exposes a
+    setter for. Disabling it leaves the benchmarking workspace blocks cached in the allocator
+    for reuse instead of returning them to the driver — the same trade the peer-aware skips
+    already make everywhere else. Called only on multi-GPU installs (see
+    ``DefaultSessionProcessor.start``); single-GPU installs keep torch's default behavior.
+
+    No-op on torch builds that lack the flag (e.g. CPU-only builds).
+    """
+    setter = getattr(torch._C, "_cudnn_set_conv_benchmark_empty_cache", None)
+    if setter is None:
+        return
+    setter(False)
+    InvokeAILogger.get_logger("TorchDevice").debug(
+        "Disabled torch's post-conv-algorithm-search global emptyCache (multi-GPU install)."
+    )

tests/backend/util/test_devices.py

@@ -714,3 +714,27 @@ def test_install_peer_aware_empty_cache_wraps_torch_entry_point(monkeypatch):
         # monkeypatch restores the attribute we set; make sure the true original is back for
         # other tests regardless of ordering.
         torch_mod.cuda.empty_cache = original
+
+
+def test_disable_conv_benchmark_empty_cache_flips_torch_flag():
+    """The multi-GPU startup path must clear torch's post-conv-find emptyCache flag (and no-op
+    gracefully on builds that lack it)."""
+    from invokeai.backend.util.devices import disable_conv_benchmark_empty_cache
+
+    getter = getattr(torch._C, "_cuda_get_conv_benchmark_empty_cache", None)
+    setter = getattr(torch._C, "_cudnn_set_conv_benchmark_empty_cache", None)
+    if getter is None or setter is None:
+        # CPU-only torch builds lack the flag; the function must still be a safe no-op.
+        disable_conv_benchmark_empty_cache()
+        return
+
+    original = getter()
+    try:
+        setter(True)
+        disable_conv_benchmark_empty_cache()
+        assert getter() is False
+        # Idempotent: a second call keeps it disabled without raising.
+        disable_conv_benchmark_empty_cache()
+        assert getter() is False
+    finally:
+        setter(original)