[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>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -177,18 +177,19 @@ component can be reused by the next. The DiT prefetch path and its two shared
 device buffers are unchanged.
 
 After a component is offloaded, cached-but-unallocated memory is retained only
-while it stays within a budget derived from the workload itself and at least
-5% of device capacity remains physically free. The budget starts from the
-model topology — the staged component byte total reported by the attached
-`PinnedModuleStager` instances times a seed multiplier — and then tracks the
-largest cached value observed at any component 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
-a margin, and later boundaries retain. The budget is capped at a quarter of
-device capacity so neither a large staged total nor a spurious observation
-can monopolize a small device, and a detached stager returns its bytes to the
-seed. No user-chosen fraction and no warmup flag are needed. Each decision
-logs the staged byte total, observed peak, effective budget and cap,
+while it stays within a budget learned from the running workload and at least
+5% of device capacity remains physically free. The budget is the largest
+cached value observed at any component boundary times a margin, so it tracks
+the weight-reuse payload plus the transient footprint whatever the
+resolution. It starts at zero: a cold start pays one release per observation
+step while the peak climbs, after which boundaries retain steadily, and a
+workload change whose cached footprint grows by more than the margin pays one
+further release while the new peak lands. The budget is capped at a quarter
+of device capacity so a spurious observation cannot monopolize a small
+device, and detaching the last attached stager resets the learned peak so one
+workload never permanently influences another. No user-chosen fraction, no
+multiplier constant, and no warmup flag are needed. Each decision logs the
+staged byte total, observed peak, effective budget and cap,
 cached-but-unallocated bytes, and free memory.
 Missing allocator telemetry also releases it conservatively. Component or
 staging failure forces a release, and an out-of-memory allocation gets one

tests/diffusion/offloader/test_module_residency.py

@@ -202,15 +202,23 @@ def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
     expected_release,
 ):
     empty_cache, _ = patched_runtime
-    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved)
-    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated)
+    telemetry = {"reserved": 20, "allocated": 10, "free": 80}
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: telemetry["reserved"])
+    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: telemetry["allocated"])
     monkeypatch.setattr(
         residency_module.current_omni_platform,
         "get_device_memory",
-        lambda _device: (free, 100),
+        lambda _device: (telemetry["free"], 400),
     )
     retention = BoundedAllocatorCache(torch.device("cpu"))
-    retention.grow_retention_budget(10)
+
+    # Prime the learned budget: the cold-start check releases (budget 0) and
+    # records cached = 10, so later checks retain at budget 12.5.
+    assert retention.release_if_needed()
+    assert retention.observed_cached_peak == 10
+    empty_cache.reset_mock()
+
+    telemetry.update(reserved=reserved, allocated=allocated, free=free)
 
     released = retention.release_if_needed()
 
@@ -218,7 +226,7 @@ def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
     assert empty_cache.call_count == int(expected_release)
 
 
-def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_runtime):
+def test_attached_stagers_form_a_detach_ledger_that_resets_the_observed_peak(patched_runtime):
     left, right = _aliased_modules()
     retention = BoundedAllocatorCache(torch.device("cpu"))
     assert retention.budget_bytes == 0
@@ -231,57 +239,57 @@ def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_ru
     )
     expected_bytes = 12 * 4 + 8 * 8  # float32 parameter + int64 buffer storages
     assert stager.staged_bytes == expected_bytes
-    assert retention.budget_bytes == 2 * expected_bytes
+    assert retention._staged_bytes == expected_bytes
 
     # Re-attaching the same cache does not double count.
     stager.set_cache_retention(retention)
-    assert retention.budget_bytes == 2 * expected_bytes
+    assert retention._staged_bytes == 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.
+    # A second stager adds its own staged bytes to the ledger.
     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)
+    assert retention._staged_bytes == expected_bytes + extra.staged_bytes
+
+    # A learned peak survives a partial detach...
+    retention._observed_cached_peak = 30
+    extra.set_cache_retention(None)
+    assert retention.observed_cached_peak == 30
+    assert retention.budget_bytes == 37
+
+    # ...but the last detach resets it so one workload cannot influence the next.
+    stager.set_cache_retention(None)
+    assert retention.observed_cached_peak == 0
+    assert retention.budget_bytes == 0
 
 
-@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,
-):
+def test_allocator_cache_budget_is_capped_by_device_capacity(monkeypatch, patched_runtime):
     empty_cache, _ = patched_runtime
-    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved)
-    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated)
+    telemetry = {"reserved": 40, "allocated": 10, "free": 80}
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: telemetry["reserved"])
+    monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: telemetry["allocated"])
     monkeypatch.setattr(
         residency_module.current_omni_platform,
+        # Capacity 100 keeps the 25% cap at 25, below the learned budget of
+        # 30 * 1.25 = 37.5, so the cap must govern.
         "get_device_memory",
-        lambda _device: (free, 100),
+        lambda _device: (telemetry["free"], 100),
     )
     retention = BoundedAllocatorCache(torch.device("cpu"))
-    retention.grow_retention_budget(1000)
 
-    released = retention.release_if_needed()
+    # Cold start: budget 0 releases and records the peak of 30.
+    assert retention.release_if_needed()
+    assert retention.observed_cached_peak == 30
+    empty_cache.reset_mock()
 
-    assert released is expected_release
-    assert empty_cache.call_count == int(expected_release)
+    # The learned budget 37.5 is capped to 25, so cached 30 keeps releasing.
+    assert retention.release_if_needed()
+    assert retention.release_if_needed()
+    empty_cache.reset_mock()
+
+    # Within the cap (cached 20 <= 25) the policy retains.
+    telemetry.update(reserved=30, allocated=10)
+    assert not retention.release_if_needed()
+    assert empty_cache.call_count == 0
 
 
 def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime):
@@ -310,10 +318,9 @@ def test_allocator_cache_budget_adapts_to_observed_peak(monkeypatch, patched_run
         lambda _device: (telemetry["free"], 400),
     )
     retention = BoundedAllocatorCache(torch.device("cpu"))
-    retention.grow_retention_budget(10)
 
-    # First boundary: cached 30 exceeds the seed budget 20, so the policy
-    # releases once and records the observed peak.
+    # Cold-start boundary: budget 0, so the policy releases and records the
+    # observed peak.
     assert retention.release_if_needed()
     assert retention.observed_cached_peak == 30
 

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))