From 611f2a4e0f30ea7f50451fd34ab25b8f9365ff2f Mon Sep 17 00:00:00 2001
From: rattus <46076784+rattus128@users.noreply.github.com>
Date: Mon, 3 Aug 2026 01:16:06 +1000
Subject: [PATCH] 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.
---
 comfy/model_management.py | 21 +++++++++++++++++----
 comfy/pinned_memory.py    |  6 +++---
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index e3c94c15a..1000f69e1 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -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):
diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py
index d78ab3c76..e9a9a70e2 100644
--- a/comfy/pinned_memory.py
+++ b/comfy/pinned_memory.py
@@ -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