[Model] Learn the retention budget directly from observed peaks

Drop the staged-bytes seed per review direction: the budget is now
min(observed_peak * 1.25, 25% of device capacity), nothing else. A cold
start pays one release per observation step while the peak climbs (two on
the validated MiniMax workload), then boundaries retain steadily; the 5%
free floor and capacity cap are unchanged. The staged-byte ledger remains
only to detect full detachment, which now resets the observed peak so one
workload cannot permanently influence another. Trade-off accepted per
review: lower reserved memory and a simpler policy over a flush-free
first request.

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

vllm_omni/diffusion/offloader/module_residency.py

@@ -20,19 +20,14 @@ from .tensor_utils import is_dtensor, set_tensor_storage
 logger = init_logger(__name__)
 
 
-# Seed headroom over the staged component bytes. Retained blocks include
-# allocator fragmentation and transient storages freed alongside the weights
-# at component boundaries, so the reusable payload alone under-bounds the
-# working set the policy must keep to be effective. Observations take over
-# from this seed at the first component boundary (see below).
-_RETENTION_HEADROOM_MULTIPLIER = 2.0
 # Margin over the largest cached-but-unallocated value observed at a component
-# boundary. The observed peak already includes the transient footprint of the
-# running workload, so this only covers check-to-check variation.
+# boundary. The observed peak already includes the weight-reuse payload and
+# the transient footprint of the running workload, so this only covers
+# check-to-check variation.
 _OBSERVED_PEAK_MARGIN = 1.25
-# Hard ceiling for the derived budget so a large staged total or a spurious
-# observation can never monopolize a small device; the #6526 L20X retention
-# plateau (~23.8 GB on a ~140 GB card) stayed well below this fraction.
+# Hard ceiling for the learned budget so a spurious observation 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
 
 
@@ -44,22 +39,22 @@ class BoundedAllocatorCache:
     cached blocks are immediately reusable. This policy keeps the cache while
     both of these bounds hold:
 
-    * cached-but-unallocated memory stays within the derived budget, capped
-      at a quarter of device capacity; and
+    * cached-but-unallocated memory stays within the learned budget — the
+      largest cached value observed at any component boundary times a margin,
+      capped at a quarter of device capacity; and
     * at least 5% of device capacity is physically free.
 
-    The budget needs no user-chosen fraction and no warmup flag: it starts
-    from the staged component byte total reported by attached
-    ``PinnedModuleStager`` instances (times the seed multiplier) and then
-    tracks the largest cached value observed at any boundary. If a workload's
-    transient footprint outgrows the seed, the first over-budget boundary
-    releases once, the observation raises the budget to the observed peak
-    times its margin, and later boundaries retain. The budget therefore never
-    exceeds real observed usage by more than the margin, whatever the
-    resolution or workload. Detaching a stager returns its bytes to the seed.
-    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.
+    The budget needs no user-chosen fraction, no multiplier constant, and no
+    warmup flag. It starts at zero, so the first boundaries of a cold start
+    release once per observation step while the peak climbs; afterwards the
+    budget tracks the observed peak and boundaries retain steadily. A
+    workload change whose cached footprint grows by more than the margin pays
+    one further release while the new peak lands. Detaching the last attached
+    ``PinnedModuleStager`` resets the observed peak, so one workload never
+    permanently influences another. 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__(
@@ -76,16 +71,18 @@ class BoundedAllocatorCache:
         self._observed_cached_peak = 0
 
     def grow_retention_budget(self, staged_bytes: int) -> None:
-        """Grow the seed cache budget by a component's staged byte total."""
+        """Record an attached stager's staged byte total."""
         if staged_bytes < 0:
             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."""
+        """Record a stager detach; the last detach resets the observed peak."""
         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))
+        if self._staged_bytes == 0:
+            self._observed_cached_peak = 0
 
     @property
     def observed_cached_peak(self) -> int:
@@ -94,10 +91,8 @@ class BoundedAllocatorCache:
 
     @property
     def budget_bytes(self) -> int:
-        """Reusable-cache ceiling: seed or observed peak, whichever is larger."""
-        seed = int(self._staged_bytes * _RETENTION_HEADROOM_MULTIPLIER)
-        learned = int(self._observed_cached_peak * _OBSERVED_PEAK_MARGIN)
-        return max(seed, learned)
+        """Reusable-cache ceiling: the observed peak times its margin."""
+        return int(self._observed_cached_peak * _OBSERVED_PEAK_MARGIN)
 
     def _should_release(self) -> bool:
         reserved = int(torch.accelerator.memory_reserved(self.device))