Multi-GPU support: per-device VRAM management

[Kosinkadink]

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.

[rattus128]

this appears to leave vbars_free behind as dead code:

(venv) rattus@rattus-box2:~/ComfyUI-MuiMui$ git grep "vbars_free[^_]"
src/model-vbar.c:size_t vbars_free(size_t size) {
src/plat.h:size_t vbars_free(size_t size);

If you support multi-gpu the global API doesnt make sense any more. This should just be a signature change of vbars_free to add the require the new argument.

[rattus128]

src/pyt-cu-plug-alloc-async.c has the same. Should be generalize if we still need this lock (more to come).

[rattus128]

I think we are better off with cleaning up the C and calling device_init() per device rather than the assymetry on the init paths. If we want to keep comfy 0-change we can do the lazy in python layer, although how hard would a device_init() call per device be in comfy main @Kosinkadink ?

[rattus128]

Amp is seeing a cross GPU state sharing of the last_check timestamp with is a local static. I think this might need to be per GPU state.

[rattus128]

I think this might be a regression. amp might have misunderstood the intent of this code.

The original bug here was the poll loop in ComfyUI-Manager for the repositories could call a python gc() from an alien thread. This in turn would cause the cudaFree to be completely detached from contexts. Since we interpose between the rt and driver APIs we need to "guess" the context. aimdo_cuda_ctx is the best guess as that is the users elected GPU (--cuda-device X).

Defaulting to 0 as top priority would there break the case of:

--cuda-device 1
ComfyUI-Manager installed, online and still doing its poll while doing intensive workflows.

Theres a little chaos agent thread I added to comfy main.py to make it more reproducible in the original PR which is more reliable than trying to use Manager to repro.

[rattus128]

the longer term solution is to stop hooking the RT API in linux and instead do it the windows way. @Kosinkadink if you hard needed this change for functionality we may want to prioritize and bite that bullet. For 100% solution we need to adopt a detours framework for Linux. Was has a PR with PLT/GOT based solution that double hooks but that doesnt solve for this unfortunately so we would have to go ultimate solution of detours.

[rattus128]

Test case and chaos agent here:

Comfy-Org/ComfyUI#12840

[rattus128]

ModelPatcher in ComfyUI needs vbars_analyze to return the VRAM per device so that get_free_memory does the right thing. Currently this will aggregate all gpus and over-report free-able memory.

From amp:

Short answer: PR 31 did not make vbars_analyze itself device-scoped. It mainly made it thread-safe, while the multi-GPU changes around it caused its effective meaning to become “global VBAR residency across all devices,” not “the current device’s VBAR residency.”

What PR 31 Changed

From the PR diff, vbars_analyze itself kept the same core algorithm:

It still walks the VBAR list.
It still recomputes resident VRAM by counting rp->handle pages instead of trusting resident_count.
It still logs per-VBAR residency and returns the summed byte count.

What PR 31 directly changed in that function was:

It wrapped vbars_analyze in the new global vbar_list_lock.
It kept vbars_dirty as a single global dirty bit, now protected by that same lock.

So semantically, the function became safer to call concurrently, but its accounting logic did not change.

The Important Semantic Shift After PR 31

The real semantic change came from the surrounding PR 31 refactor:

PR 31 introduced multi-GPU behavior but still kept one global VBAR linked list and one global vbars_dirty.
Because of that, vbars_analyze() in PR 31 no longer implicitly meant “analyze the one device we have.”
It now meant “analyze every VBAR in the global list,” which in practice is “sum resident VBAR pages across all devices.”

That means the return value changed in meaning:

Before multi-GPU, it was effectively single-device VBAR residency.
After PR 31, it became aggregate cross-device VBAR residency.