[Model] Cap derived cache budget and account for stager detach

Address review on #6629:

- Cap the derived retention budget at a quarter of device capacity so a
  large staged total cannot monopolize a small device; the multiplier
  itself stays pending an instrumented target run.
- Log every release decision (reason, cached bytes, effective budget and
  cap, staged total, free/total memory) and every staging OOM retry with
  its staged byte total, so a single target run validates the headroom
  multiplier directly.
- Make stager detach symmetric: set_cache_retention now shrinks the
  previous cache's budget before growing the new one, so cache -> None ->
  same cache counts the staged bytes exactly once (regression test added).

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -180,9 +180,14 @@ After a component is offloaded, cached-but-unallocated memory is retained only
 while it stays within a budget derived from the model topology — the staged
 component byte total reported by the attached `PinnedModuleStager` instances
 times a fixed headroom multiplier — and at least 5% of device capacity
-remains physically free. Crossing either bound releases the allocator cache,
-so the budget scales with the staged components rather than with device
-capacity and needs no user-chosen fraction.
+remains physically free. The derived budget is capped at a quarter of device
+capacity so a large staged total cannot monopolize a small device, and a
+detached stager returns its bytes to the budget. Crossing either bound
+releases the allocator cache, so the budget scales with the staged components
+rather than with device capacity and needs no user-chosen fraction. Each
+release decision logs the staged byte total, the effective budget and cap,
+cached-but-unallocated bytes, and free memory, so a target run can validate
+the headroom multiplier directly.
 Missing allocator telemetry also releases it conservatively. Component or
 staging failure forces a release, and an out-of-memory allocation gets one
 retry after release. This is a component-local policy rather than a global

tests/diffusion/offloader/test_module_residency.py

@@ -237,12 +237,53 @@ def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_ru
     stager.set_cache_retention(retention)
     assert retention.budget_bytes == 2 * expected_bytes
 
+    # Detaching and re-attaching the same cache also counts the bytes once.
+    stager.set_cache_retention(None)
+    assert retention.budget_bytes == 0
+    stager.set_cache_retention(retention)
+    assert retention.budget_bytes == 2 * expected_bytes
+
     # A second stager grows the budget by its own staged bytes.
     extra = PinnedModuleStager(torch.nn.Linear(2, 2), torch.device("cpu"), pin_memory=False)
     extra.set_cache_retention(retention)
     assert retention.budget_bytes == 2 * (expected_bytes + extra.staged_bytes)
 
 
+@pytest.mark.parametrize(
+    ("reserved", "allocated", "free", "expected_release"),
+    [
+        # Derived budget 2 * 1000 = 2000 uncapped, but capacity cap = 25 and
+        # cached = 30 exceeds it, so the cap must trigger the release.
+        (40, 10, 80, True),
+        # Within the cap (cached = 20 <= 25) and the free floor holds.
+        (30, 10, 80, False),
+    ],
+)
+def test_allocator_cache_budget_is_capped_by_device_capacity(
+    monkeypatch,
+    patched_runtime,
+    reserved,
+    allocated,
+    free,
+    expected_release,
+):
+    empty_cache, _ = patched_runtime
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved)
+    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated)
+    monkeypatch.setattr(
+        residency_module.current_omni_platform,
+        "get_device_memory",
+        lambda _device: (free, 100),
+    )
+    retention = BoundedAllocatorCache(torch.device("cpu"))
+    retention.grow_retention_budget(1000)
+
+    released = retention.release_if_needed()
+
+    assert released is expected_release
+    assert empty_cache.call_count == int(expected_release)
+
+
 def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime):
     empty_cache, _ = patched_runtime
 

vllm_omni/diffusion/offloader/module_residency.py

@@ -25,6 +25,10 @@ logger = init_logger(__name__)
 # weights at component boundaries, so the reusable payload alone under-bounds
 # the working set the policy must keep to be effective.
 _RETENTION_HEADROOM_MULTIPLIER = 2.0
+# Hard ceiling for the derived budget so a large staged total can never
+# monopolize a small device; the #6526 L20X retention plateau (~23.8 GB on a
+# ~140 GB card) stayed well below this fraction.
+_MAX_BUDGET_CAPACITY_FRACTION = 0.25
 
 
 class BoundedAllocatorCache:
@@ -37,15 +41,16 @@ class BoundedAllocatorCache:
 
     * cached-but-unallocated memory stays within the derived budget: the
       staged component byte total reported by attached ``PinnedModuleStager``
-      instances times a fixed headroom multiplier; and
+      instances times a fixed headroom multiplier, capped at a quarter of
+      device capacity; and
     * at least 5% of device capacity is physically free.
 
     The budget derives from the model topology instead of a user-chosen
     fraction, so it scales with the staged components rather than with device
-    capacity. Missing memory telemetry is handled conservatively by releasing
-    the cache. Failure paths can force release; normal executor shutdown keeps
-    its own unconditional device-cache cleanup because this policy is not
-    global.
+    capacity. Detaching a stager returns its bytes to the budget. Missing
+    memory telemetry is handled conservatively by releasing the cache.
+    Failure paths can force release; normal executor shutdown keeps its own
+    unconditional device-cache cleanup because this policy is not global.
     """
 
     def __init__(
@@ -66,6 +71,12 @@ class BoundedAllocatorCache:
             raise ValueError(f"staged_bytes must be >= 0, got {staged_bytes}")
         self._staged_bytes += int(staged_bytes)
 
+    def shrink_retention_budget(self, staged_bytes: int) -> None:
+        """Return bytes previously grown, e.g. when a stager detaches."""
+        if staged_bytes < 0:
+            raise ValueError(f"staged_bytes must be >= 0, got {staged_bytes}")
+        self._staged_bytes = max(0, self._staged_bytes - int(staged_bytes))
+
     @property
     def budget_bytes(self) -> int:
         """Reusable-cache ceiling: staged bytes times the headroom multiplier."""
@@ -76,7 +87,32 @@ class BoundedAllocatorCache:
         allocated = int(torch.accelerator.memory_allocated(self.device))
         free, total = current_omni_platform.get_device_memory(self.device)
         cached = max(0, reserved - allocated)
-        return cached > self.budget_bytes or free < int(total * self.min_free_fraction)
+        capacity_cap = int(total * _MAX_BUDGET_CAPACITY_FRACTION)
+        budget = min(self.budget_bytes, capacity_cap)
+        if cached > budget or free < int(total * self.min_free_fraction):
+            reason = "cached-over-budget" if cached > budget else "low-free-memory"
+            logger.debug(
+                "Releasing retained allocator cache (%s): cached=%d budget=%d staged=%d "
+                "capacity_cap=%d free=%d total=%d",
+                reason,
+                cached,
+                budget,
+                self._staged_bytes,
+                capacity_cap,
+                free,
+                total,
+            )
+            return True
+        logger.debug(
+            "Retaining allocator cache: cached=%d budget=%d staged=%d capacity_cap=%d free=%d total=%d",
+            cached,
+            budget,
+            self._staged_bytes,
+            capacity_cap,
+            free,
+            total,
+        )
+        return False
 
     def release_if_needed(self, *, force: bool = False) -> bool:
         """Release cached blocks when a bound is crossed or release is forced."""
@@ -227,6 +263,8 @@ class PinnedModuleStager:
     def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None:
         if cache_retention is self.cache_retention:
             return
+        if self.cache_retention is not None:
+            self.cache_retention.shrink_retention_budget(self.staged_bytes)
         self.cache_retention = cache_retention
         if cache_retention is not None:
             cache_retention.grow_retention_budget(self.staged_bytes)
@@ -269,6 +307,10 @@ class PinnedModuleStager:
         except torch.OutOfMemoryError:
             # A retained cache is normally reusable by this process, but an
             # explicit flush gives external memory pressure one bounded retry.
+            logger.info(
+                "Staging out of memory (%d staged bytes); releasing retained cache and retrying once",
+                self.staged_bytes,
+            )
             self._cleanup_failed_load()
             try:
                 self._load_once()