RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix pin registration priority (#15226) 

This priority scheme was broken in the case where you have pin
registration exhaustion while loading a VBAR that gets a big evicition.
The weight would stay in the loaded set but inherit the MRU priority
against other workflow models WRT pin registration which leads to async
offload without pinning.

Fix by universally promiting active pin registration above workflow
pins without concern for the weights/weights-loaded split. This diverges
from the actual budgeting where the split still makes sense.
This commit is contained in:
rattus 2026-08-03 01:16:06 +10:00 committed by GitHub
parent 364081170f
commit 611f2a4e0f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
comfy/model_management.py
View File

@ -671,6 +671,19 @@ def pin_eviction_tiers(loaded, evict_active):
tiers.append((PIN_SUBSETS, True, True))
return tiers
def registration_eviction_tiers(evict_active):
subsets = PIN_SUBSETS + LOADED_PIN_SUBSETS
tiers = [
(subsets, False, False),
(subsets, True, False),
]
if evict_active:
tiers.extend([
(subsets, False, True),
(subsets, True, True),
])
return tiers
def free_pins(size, evict_active=False, loaded=False):
freed = 0
for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
@ -703,19 +716,19 @@ def ensure_pin_budget(size, evict_active=False, loaded=False):
to_free = shortfall + PIN_PRESSURE_HYSTERESIS
return free_pins(to_free, evict_active=evict_active, loaded=loaded) >= shortfall
def free_registrations(shortfall, evict_active=True, loaded=False):
def free_registrations(shortfall, evict_active=True):
if MAX_PINNED_MEMORY <= 0:
return False
if shortfall <= 0:
return True
shortfall += REGISTERABLE_PIN_HYSTERESIS
for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
for subsets, current_prompt, active in registration_eviction_tiers(evict_active):
shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True)
return shortfall <= REGISTERABLE_PIN_HYSTERESIS
def ensure_pin_registerable(size, evict_active=True, loaded=False):
return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active, loaded=loaded)
def ensure_pin_registerable(size, evict_active=True):
return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active)
class LoadedModel:
def __init__(self, model: ModelPatcher):

6
comfy/pinned_memory.py
View File

@ -56,7 +56,7 @@ def get_pin(module, subset="weights"):
_, _, stack_split, pinned_size, *_ = module._pin_state[subset]
size = pin.nbytes
comfy.model_management.ensure_pin_registerable(size, loaded=subset.endswith("-loaded"))
comfy.model_management.ensure_pin_registerable(size)
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
@ -93,7 +93,7 @@ def pin_memory(module, subset="weights", size=None):
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or
not comfy.model_management.ensure_pin_registerable(registerable_size, loaded=loaded)):
not comfy.model_management.ensure_pin_registerable(registerable_size)):
return _steal_pin(module, stack, buckets, size, priority, subset)
offset = hostbuf.size
@ -105,7 +105,7 @@ def pin_memory(module, subset="weights", size=None):
pin.untyped_storage()._comfy_hostbuf = hostbuf
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
comfy.model_management.free_registrations(size, loaded=loaded)
comfy.model_management.free_registrations(size)
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
del pin