[PR #31] Multi-GPU support: per-device VRAM management #71

Open
opened 2026-10-26 22:24:18 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-aimdo/pull/31
Author: @Kosinkadink
Created: 4/9/2026
Status: 🔄 Open

Base: master ← Head: multigpu-thread-safety


📝 Commits (4)

  • 607f103 Add mutex synchronization and device-aware eviction for multi-GPU
  • 7593759 multi-GPU phases 1-5: per-device state, context safety, hybrid budget, device-aware allocator, atomic counters
  • 718e01b fix code review issues: thread-safe init, atomic reads, per-device ctx sync
  • ac7f22e fix Windows build: include <intrin.h> for Interlocked intrinsics

📊 Changes

7 files changed (+543 additions, -31 deletions)

View changed files

➕ MULTIGPU_PLAN.md (+139 -0)
📝 build-linux-docker (+1 -0)
📝 src/control.c (+125 -0)
📝 src/model-vbar.c (+122 -19)
📝 src/plat.h (+122 -1)
📝 src/pyt-cu-plug-alloc-async.c (+31 -8)
📝 src/vrambuf.c (+3 -3)

📄 Description

Summary

Makes aimdo work properly with multiple GPUs. Previously, aimdo assumed a single GPU — budget_deficit() used global total_vram_usage (sum of ALL GPUs) against one GPU's vram_capacity, creating a phantom deficit that triggered constant unnecessary eviction. With 2×RTX 4090, this made 2-GPU slower than 1-GPU.

Changes

Commit 1: Mutex + device-aware eviction (foundation)

  • Added mutex synchronization (pthread/CRITICAL_SECTION) to protect the VBAR priority linked list
  • Made eviction device-aware: vbars_free_locked_dev() with device_filter parameter
  • Added per-device VRAM accounting: dev_vram_usage[], dev_vram_add/sub, current_cuda_device()

Commit 2: Phases 1–5 (core multi-GPU logic)

  • Phase 1 — Per-device state: AimdoDeviceState struct with per-device vram_capacity, ctx, deficit_sync, etc. Lazy ensure_device_init() called from vbar_allocate and allocator hooks. vbar_allocate uses per-device capacity.
  • Phase 2 — Context safety: with_device_ctx()/restore_ctx() helpers. poll_budget_deficit_dev() switches to correct device context for cuMemGetInfo. Fixed ensure_ctx() to not force device 0 context on multi-GPU.
  • Phase 3 — Per-device hybrid budget: budget_deficit_dev() uses per-device usage + per-device cuMemGetInfo backstop (not just simple accounting). vbar_fault uses this instead of global budget_deficit(), eliminating phantom cross-GPU deficit.
  • Phase 4 — Device-aware allocator hooks: aimdo_cuda_malloc/aimdo_cuda_malloc_async determine current device, use vbars_free_dev + budget_deficit_dev. vrambuf_grow also device-aware.
  • Phase 5 — Atomic counters: dev_vram_add/sub use __atomic_add_fetch (Linux) / InterlockedExchangeAdd64 (Windows).

API Compatibility

Fully backward compatible — no changes to exported function signatures. ComfyUI does not need any modifications. init(device_id) still works; other devices are initialized lazily.

Benchmarks (2×RTX 4090, Qwen-Image 38GB, CFG=7, 20 steps)

Config Time vs 1-GPU Status
1-GPU + aimdo 56.39s 1.00× ✅ stable
2-GPU + aimdo (before, mutex only) 95.39s 0.59× ✅ stable
2-GPU + aimdo (before, dev-aware evict) 78.85s 0.72× ✅ stable
2-GPU + aimdo (this PR) 41.29s 1.37× ✅ stable (7/7)
2-GPU no aimdo 37.32s 1.51× ❌ crashes after ~2 runs
1-GPU + aimdo (this PR) 56.41s 1.00× ✅ no regression

Key design decisions

  1. Must keep cuMemGetInfo hybrid backstop — pure per-device accounting OOM'd because cudaFreeAsync decrements counters before memory is actually reusable.
  2. Lazy device init — ComfyUI only calls init() for device 0. Other devices are initialized on first use via ensure_device_init().
  3. Global budget_deficit() kept as fallback — budget_deficit_dev() falls back to global when device state isn't initialized.

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-aimdo/pull/31 **Author:** [@Kosinkadink](https://github.com/Kosinkadink) **Created:** 4/9/2026 **Status:** 🔄 Open **Base:** `master` ← **Head:** `multigpu-thread-safety` --- ### 📝 Commits (4) - [`607f103`](https://github.com/Comfy-Org/comfy-aimdo/commit/607f10373ac9376994a7b79c4c6814370eb65482) Add mutex synchronization and device-aware eviction for multi-GPU - [`7593759`](https://github.com/Comfy-Org/comfy-aimdo/commit/75937591e210f04bb1a0d6878bd1d95c95644b9c) multi-GPU phases 1-5: per-device state, context safety, hybrid budget, device-aware allocator, atomic counters - [`718e01b`](https://github.com/Comfy-Org/comfy-aimdo/commit/718e01b646308e5de18dc61389eb9e9f48fa3e0f) fix code review issues: thread-safe init, atomic reads, per-device ctx sync - [`ac7f22e`](https://github.com/Comfy-Org/comfy-aimdo/commit/ac7f22ee092bbb2d107495aa3b4eb2c580871b6c) fix Windows build: include <intrin.h> for Interlocked intrinsics ### 📊 Changes **7 files changed** (+543 additions, -31 deletions) <details> <summary>View changed files</summary> ➕ `MULTIGPU_PLAN.md` (+139 -0) 📝 `build-linux-docker` (+1 -0) 📝 `src/control.c` (+125 -0) 📝 `src/model-vbar.c` (+122 -19) 📝 `src/plat.h` (+122 -1) 📝 `src/pyt-cu-plug-alloc-async.c` (+31 -8) 📝 `src/vrambuf.c` (+3 -3) </details> ### 📄 Description ## Summary Makes aimdo work properly with multiple GPUs. Previously, aimdo assumed a single GPU — `budget_deficit()` used global `total_vram_usage` (sum of ALL GPUs) against one GPU's `vram_capacity`, creating a phantom deficit that triggered constant unnecessary eviction. With 2×RTX 4090, this made 2-GPU **slower** than 1-GPU. ## Changes ### Commit 1: Mutex + device-aware eviction (foundation) - Added mutex synchronization (pthread/CRITICAL_SECTION) to protect the VBAR priority linked list - Made eviction device-aware: `vbars_free_locked_dev()` with device_filter parameter - Added per-device VRAM accounting: `dev_vram_usage[]`, `dev_vram_add/sub`, `current_cuda_device()` ### Commit 2: Phases 1–5 (core multi-GPU logic) - **Phase 1 — Per-device state**: `AimdoDeviceState` struct with per-device `vram_capacity`, `ctx`, `deficit_sync`, etc. Lazy `ensure_device_init()` called from `vbar_allocate` and allocator hooks. `vbar_allocate` uses per-device capacity. - **Phase 2 — Context safety**: `with_device_ctx()`/`restore_ctx()` helpers. `poll_budget_deficit_dev()` switches to correct device context for `cuMemGetInfo`. Fixed `ensure_ctx()` to not force device 0 context on multi-GPU. - **Phase 3 — Per-device hybrid budget**: `budget_deficit_dev()` uses per-device usage + per-device `cuMemGetInfo` backstop (not just simple accounting). `vbar_fault` uses this instead of global `budget_deficit()`, eliminating phantom cross-GPU deficit. - **Phase 4 — Device-aware allocator hooks**: `aimdo_cuda_malloc`/`aimdo_cuda_malloc_async` determine current device, use `vbars_free_dev` + `budget_deficit_dev`. `vrambuf_grow` also device-aware. - **Phase 5 — Atomic counters**: `dev_vram_add/sub` use `__atomic_add_fetch` (Linux) / `InterlockedExchangeAdd64` (Windows). ## API Compatibility Fully backward compatible — no changes to exported function signatures. ComfyUI does not need any modifications. `init(device_id)` still works; other devices are initialized lazily. ## Benchmarks (2×RTX 4090, Qwen-Image 38GB, CFG=7, 20 steps) | Config | Time | vs 1-GPU | Status | |--------|------|----------|--------| | 1-GPU + aimdo | 56.39s | 1.00× | ✅ stable | | 2-GPU + aimdo (before, mutex only) | 95.39s | 0.59× | ✅ stable | | 2-GPU + aimdo (before, dev-aware evict) | 78.85s | 0.72× | ✅ stable | | **2-GPU + aimdo (this PR)** | **41.29s** | **1.37×** | ✅ stable (7/7) | | 2-GPU no aimdo | 37.32s | 1.51× | ❌ crashes after ~2 runs | | 1-GPU + aimdo (this PR) | 56.41s | 1.00× | ✅ no regression | ## Key design decisions 1. **Must keep cuMemGetInfo hybrid backstop** — pure per-device accounting OOM'd because `cudaFreeAsync` decrements counters before memory is actually reusable. 2. **Lazy device init** — ComfyUI only calls `init()` for device 0. Other devices are initialized on first use via `ensure_device_init()`. 3. **Global budget_deficit() kept as fallback** — `budget_deficit_dev()` falls back to global when device state isn't initialized. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/comfy-aimdo#71
No description provided.