feat: multi-GPU parallel session execution

[lstein]

Summary

This PR adds multi-GPU parallel generation: on a machine with more than one GPU, InvokeAI runs several generation sessions concurrently — one per GPU — instead of draining the queue one job at a time. Jobs are distributed fairly across users so a single user's large batch can't monopolize every GPU while others wait.

It's controlled by a new generation_devices config setting (defaults to auto = use every available CUDA GPU). Setting it to a single device, or leaving CUDA out of the picture, preserves the previous serial behavior exactly. The choice of GPUs can also be controlled via a new section of the Settings dialogue (restart required to take effect).

Demo (turn on the sound!)

invoke-mgpu.mp4

How it works — the change is built around five small backend seams plus a frontend update, rather than per-node edits:

• Device context (invokeai/backend/util/devices.py): a thread-local set/get/clear_session_device on TorchDevice; choose_torch_device() consults it first. This is the lynchpin — the ~79 existing call sites resolve to the worker's GPU with no per-node changes.
• Per-device model caches (model_manager_default, model_load_default): one ModelCache per device, resolved by the current thread's device, with fan-out for clear/drop/shutdown. Model construction is serialized against VRAM moves to prevent meta-device corruption. A single global RAM budget is shared across the per-device caches, and identical CPU weights are deduplicated across devices (see RAM management below).
• Atomic dequeue (session_queue_sqlite.dequeue): a lock makes select+claim atomic so concurrent workers never grab the same queue item.
• Worker pool (session_processor_default): one _SessionWorker per device, each pinning torch.cuda.set_device + the session device, with its own runner and cancel event; cancellation is routed per item. Profiling is disabled when more than one worker is active.
• Concurrency hardening: added a Lock to ObjectSerializerForwardCache and made DiskImageFileStorage thread-safe for parallel sessions.

Frontend: during parallel generation the progress display stacks one progress bar per active session (each disappears as its session finishes), and the image viewer tiles per-session progress previews when ≥2 sessions are active.

Idle-GPU text-encoder offload

When more than one generation device is configured and a GPU is idle, a session's text/prompt encoder runs on the idle GPU instead of the one running its denoise pipeline. This avoids evicting the denoise model from VRAM to make room for the encoder, and lets a cached encoder be reused across generations. Under full load (no idle GPU) behavior is unchanged. Controlled by offload_text_encoders_to_idle_gpus (default on); inspired by #9310.

• Device-pool arbiter (backend/util/device_pool.py): GENERATION_DEVICE_POOL gives each generation device one exclusive-use lock. A native session blocking-acquires its own device's lock for the whole run; an encoder node try-borrows an idle device's lock for the duration of that node. A borrowed encoder and a native session are therefore mutually exclusive on a GPU — preventing the shared-encoder corruption that produced garbled images — and the design is deadlock-free (borrows are non-blocking; a session only ever blocks on its own device).
• Opt-in marker: nodes declare support via @invocation(idle_gpu_offloadable=True), mirroring the existing bottleneck ClassVar. Applied to the text/prompt-encoder nodes (compel + sdxl/refiner, flux, sd3, qwen-image, anima, cogview4, flux2 klein, z-image, flux_redux). The runner re-pins the worker thread to the borrowed device for the node; conditioning is stored on the CPU so the denoiser picks it up on its own GPU afterward.

RAM management for parallel sessions

Running N sessions in parallel multiplies memory pressure, so this PR also makes the model cache parallel-aware:

• Shared global RAM budget, clamped to a safe fraction of system RAM, so summing per-device cache sizes across GPUs can't claim ~N× RAM and drive the box into swap.
• Cross-device weight de-duplication: when a second GPU loads a model another GPU already holds, it adopts the resident CPU weights (a meta-weight structural clone + load_state_dict(assign=True)) instead of re-reading from disk and materializing a second copy. This is loader-agnostic and now also covers GGUF models — GGMLTensor doesn't implement aten.empty_like, which previously made the largest quantized models (e.g. a Q8_0 transformer) silently re-load on every device and spike RAM; the adopted GGMLTensor shares the quantized storage, so it's one copy across devices.

Generation Devices settings refinements

A few small fixes to the Generation Devices selector and its logging:

• Stable device numbering: the disambiguating #N suffix on identically-named GPUs is now tied to each device's cuda index (its position in the full available-device set) rather than its position in the possibly-filtered generation_devices list. Previously, disabling e.g. cuda:1 renumbered the survivors in the backend startup log (cuda:2 became #2), disagreeing with the frontend, which always labels over the full set. Now both stay consistent — cuda:2 remains #3.
• Restart reminder: reworded the Settings caption to "Restart InvokeAI for changes to take effect." and flash that same warning as a toast on every successful change, since generation_devices only takes effect after a restart.

Related Issues / Discussions

• Builds conceptually alongside [feat] Round robin job scheduling in multiuser mode #9086 (round-robin per-user dequeue). This PR does not depend on [feat] Round robin job scheduling in multiuser mode #9086 being merged — the atomic claim works on the current FIFO dequeue, and [feat] Round robin job scheduling in multiuser mode #9086's round-robin CTE can slot in cleanly when it lands.
• Idle-GPU text-encoder offload was inspired by Add split GPU text encoder cache #9310 (now folded into this PR; supersedes the earlier private-repo prototype).

QA Instructions

On a multi-GPU machine:

• With default config (generation_devices: auto), enqueue a batch larger than the GPU count and confirm multiple sessions run simultaneously (one per GPU), with stacked progress bars and tiled previews in the viewer.
• Set generation_devices: [cuda:0] and confirm generation runs serially, exactly as before this PR.
• Set generation_devices: [cuda:0, cuda:2] and confirm only those devices are used.
• Cancel an in-flight item and confirm only that session stops.
• On a single-GPU / CPU / MPS machine, confirm auto resolves to the one best device and behavior is unchanged.
• Idle-GPU offload: with ≥2 GPUs and a single running session, confirm the text encoder runs on a different (idle) GPU than the denoiser (debug log: Running ... on idle device ...), and that the denoise model is not evicted to load the encoder. Set offload_text_encoders_to_idle_gpus: false and confirm the encoder runs on the session's own GPU.
• Parallel RAM: run several parallel sessions with the same model (ideally a GGUF/quantized one) and confirm process RAM stays bounded — the transformer/text-encoder show an Adopted shared CPU weights ... log on the second device rather than a second disk load.

New automated tests cover device routing (test_model_load_device_routing.py), dequeue concurrency (test_session_queue_dequeue_concurrency.py), device resolution (test_devices.py), the device-pool lock semantics and offload mutual-exclusion (test_device_pool.py, test_encoder_offload.py), and cross-device weight adoption incl. GGUF (test_shared_weight_adoption.py).

Merge Plan

Standard merge. No DB schema or redux migrations. Touches the session processor and model cache, so worth a careful look from those areas' owners.

The idle-GPU text-encoder offload (originally prototyped as a follow-on PR) is now included in this branch, along with the cross-device GGUF weight de-duplication that keeps parallel-session RAM bounded.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — N/A, no slice changes
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[JPPhoto]

• Stale shared CPU weights can still survive a load-affecting model settings change. Affected paths: invokeai/app/api/routers/model_manager.py:445, invokeai/backend/model_manager/load/model_cache/model_cache.py:1144, invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py:59, invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py:52, invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py:61. drop_model() only marks locked entries stale, so their shared-store reference stays live until unlock. Another device cache can reload the same key before that unlock and adopt the old canonical tensors. To expose this issue, add a test that locks a shared cached model, calls drop_model() across caches, then reloads the same key in another cache with different weights/settings and asserts the new entry cannot adopt the old canonical state dict.

• Runtime config accepts invalid generation_devices values at the API boundary. Affected paths: invokeai/app/api/routers/app_info.py:151, invokeai/app/services/config/config_default.py:276, invokeai/backend/util/devices.py:218. The route validator only checks the string pattern, so it accepts [] until later config validation and accepts syntactically valid but unavailable devices like cuda:99 until later device resolution/startup. To expose this issue, add route tests for [] and an out-of-range CUDA index under mocked CUDA availability/device count, asserting a 422 and no config-file mutation.

• Single-user mode is documented/configured as FIFO, but default device affinity reorders it. Affected paths: invokeai/app/services/config/config_default.py:221, invokeai/app/services/session_queue/session_queue_sqlite.py:304, invokeai/app/services/session_queue/session_queue_sqlite.py:320, tests/app/services/session_queue/test_session_queue_dequeue.py:390. The config says single-user mode always uses FIFO, but default session_queue_mode="round_robin" still enables affinity, allowing a later warm item to jump ahead of an older cold item. Either preserve strict FIFO in single-user mode or update config/docs to say same-priority jobs may be reordered for device affinity.

• Multi-GPU docs still describe invalid and stale behavior. Affected path: docs/src/content/docs/configuration/invokeai-yaml.mdx:132, docs/src/content/docs/configuration/invokeai-yaml.mdx:144. The docs say generation_devices: [] is valid, but config validation rejects empty lists. They also say model weights are duplicated in system RAM per active GPU, while this PR adds shared CPU-weight deduplication. Update the docs to match the implemented behavior.

• Cache stats still report only the API thread's default cache. Affected paths: invokeai/app/api/routers/model_manager.py:1298, invokeai/app/services/model_load/model_load_default.py:61. get_stats() returns model_manager.load.ram_cache.stats, which resolves to the default cache for API request threads, while PR 9263 now has one cache per generation device. To expose this issue, add a test with distinct stats on two ram_caches and assert the endpoint returns aggregate or per-device stats instead of only the default cache.

• Open question: should queue priority remain global in round-robin mode? ROUND_ROBIN_DEQUEUE_QUERY chooses each user's best pending item, then orders users by least-recently served. That means a recently served user's high-priority/prepended job can wait behind another user's lower-priority job. If that is intentional fairness behavior, document it; otherwise the query needs another priority tier above user rotation.

[JPPhoto]

See prior comments in this PR.

[lstein]

@JPPhoto Thanks for the thorough pass — all five findings were real. Everything is addressed in aee12c3:

Stale shared CPU weights surviving a settings change — Confirmed. drop_model() now also calls a new SharedCpuWeightsStore.invalidate(model_key), which forgets the canonical entries (and shells) for the model and all its submodels, so no future acquire()/peek() can adopt them even while a locked, stale-marked entry still aliases them. Two supporting changes: release() is now identity-checked (a stale holder releasing after a new canonical was registered under the same key is a no-op rather than a decrement of the new entry's refcount), and update_model_record holds MODEL_LOAD_LOCK.write_lock() across the multi-cache drop (via asyncio.to_thread so the event loop isn't blocked) — that closes the narrower race where an in-flight load peeks the old canonical before the drop and re-registers it after. Added the test you described (test_settings_change_rebuild_does_not_adopt_stale_shared_weights: two caches, one locked, drop on both, rebuild in the unlocked cache, assert fresh canonical + identity-safe release on unlock) plus store-level unit tests for invalidate() and mismatched release.

Runtime config accepts invalid generation_devices — Confirmed. Empty lists are now rejected by the request model (422 instead of the eventual 500 from config validation), and the route resolves the requested devices through TorchDevice.get_generation_devices() — the same code the startup path uses — before anything is mutated or persisted, so cuda:99 on a 2-GPU box gets a 422 with the same message startup would have produced. Route tests cover [], out-of-range CUDA, and CUDA-unavailable under mocked device counts, asserting 422 and unchanged config.

Single-user FIFO vs. device affinity — The affinity reordering in single-user multi-GPU mode is intentional (a single user mixing models across 2 GPUs is exactly the thrash case that motivated it), so I took the document-it option: the session_queue_mode description no longer claims "FIFO is always used in single-user mode"; it now says 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 so a freed GPU prefers models it already has loaded, and that explicit FIFO disables the reordering. The multi-GPU docs got a matching note.

Stale multi-GPU docs — Fixed. The generation_devices: [] row is removed from the table (with an explicit "the list may not be empty" note), and the RAM note now describes the shared CPU-weight behavior (one RAM copy per model regardless of GPU count) instead of per-device duplication. settings.json regenerated.

Cache stats endpoint — Confirmed. /v2/models/stats now aggregates across ram_caches (deduplicating by cache object, since the default cache appears under its own device key), summing counters and merging loaded_model_sizes; single-cache installs return the same object as before. Tests cover aggregation, duplicate-object dedup, and the all-None → null case.

Open question (global priority vs. user rotation) — That ordering comes verbatim from #9086 on main (ROUND_ROBIN_DEQUEUE_QUERY is unchanged by this PR, which merely merged it), so I've treated it as intentional fairness-first design rather than changing it here: user rotation outranks priority across users; priority orders jobs within each user's queue. I've documented that in the session_queue_mode description. If we decide priority should form a tier above rotation, I'd rather that change land on main independently of this PR since it affects single-GPU multi-user installs equally.

[lstein]

@JPPhoto The fixes described above are committed and pushed (aee12c3) — this is ready for re-review whenever you have a chance. Thanks again for the careful pass.

[JPPhoto]

@lstein More:

• invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py:131 and invokeai/backend/model_manager/load/model_cache/ram_budget.py:55: invalidation undercounts RAM while stale models remain locked. invalidate() deletes the canonical entry immediately, so total_bytes_in_use() stops counting its tensors even though locked cache entries still retain them. If another device rebuilds the model before those entries unlock, both the retired and replacement weights are resident, but RamBudget counts only the replacement. Memory admission can therefore exceed max_cache_ram_gb, recreating the RAM spike or process OOM that the global budget is intended to prevent. To expose this issue, add a test that locks a shared model, invalidates it, registers replacement weights, and asserts budget usage includes both copies until the stale holder releases its reference.

• invokeai/app/api/routers/model_manager.py:1327 and invokeai/backend/model_manager/load/model_cache/model_cache.py:367: the new cache-stat aggregation multiplies system-wide values by the number of device caches. Every cache receives the same global RamBudget, so each CacheStats.cache_size is set to the same global max_bytes; _get_ram_in_use() likewise returns the shared global usage used for high_watermark. Summing these fields reports an N-GPU system's cache capacity and high-water usage approximately N times too large. The added test uses unrelated synthetic capacities and therefore codifies the incorrect behavior instead of representing the shared-budget setup. To expose this issue, add a test with two distinct caches attached to the same RamBudget and assert cache_size equals the single global limit and high_watermark is not multiplied.

• invokeai/app/api/routers/app_info.py:217 and invokeai/backend/util/devices.py:193: runtime validation still accepts an unavailable MPS device. TorchDevice.get_generation_devices() validates CUDA availability and indices but performs no equivalent check for device.type == "mps". An admin on Linux or an unsupported macOS environment can persist generation_devices: ["mps"]; startup then creates a worker and cache targeting an unusable device, deferring failure until tensor operations begin. This also contradicts the route comment claiming it prevents persistence of devices that do not exist. To expose this issue, add route and device-resolution tests with torch.backends.mps.is_available() returning false, submit ["mps"], and assert a 422 with no runtime or persisted config mutation.

• The database-facing tests changed by this PR use the shared InvokeAIAppConfig(use_memory_db=True) fixture or explicit sqlite3.connect(":memory:") connections, including tests/app/services/session_queue/test_session_queue_dequeue.py, tests/app/services/session_queue/test_session_queue_dequeue_concurrency.py, tests/app/services/session_queue/test_session_queue_multigpu_cancel.py, and tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_07_01_add_session_queue_device.py. No file-backed database dependency is apparent in those tests.