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

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