[Model] Derive MiniMax-H3 DLO component cache budget from staged bytes

Replace BoundedAllocatorCache's device-capacity fraction bound (25% of
total memory) with a budget derived from the model topology: each attached
PinnedModuleStager reports its unique staged storage bytes once, and the
retention ceiling is that total times a fixed 2.0 headroom multiplier
covering allocator fragmentation and transient storages freed at component
boundaries. The 5% physically-free floor is unchanged.

The budget therefore scales with the staged components instead of device
capacity, so no user-chosen fraction is needed and behavior is identical
across differently sized devices with the same model. Re-attaching the
same cache does not double count; a second stager grows the budget by its
own bytes.

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -177,8 +177,12 @@ 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 is at most 25% of device capacity and at least 5% of device capacity
-remains physically free. Crossing either bound releases the allocator cache.
+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.
 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

@@ -210,6 +210,7 @@ def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
         lambda _device: (free, 100),
     )
     retention = BoundedAllocatorCache(torch.device("cpu"))
+    retention.grow_retention_budget(10)
 
     released = retention.release_if_needed()
 
@@ -217,6 +218,31 @@ 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):
+    left, right = _aliased_modules()
+    retention = BoundedAllocatorCache(torch.device("cpu"))
+    assert retention.budget_bytes == 0
+
+    stager = PinnedModuleStager(
+        [left, right],
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    expected_bytes = 12 * 4 + 8 * 8  # float32 parameter + int64 buffer storages
+    assert stager.staged_bytes == expected_bytes
+    assert retention.budget_bytes == 2 * expected_bytes
+
+    # Re-attaching the same cache does not double count.
+    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)
+
+
 def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime):
     empty_cache, _ = patched_runtime
 

vllm_omni/diffusion/offloader/module_residency.py

@@ -20,6 +20,13 @@ from .tensor_utils import is_dtensor, set_tensor_storage
 logger = init_logger(__name__)
 
 
+# Cache budget 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.
+_RETENTION_HEADROOM_MULTIPLIER = 2.0
+
+
 class BoundedAllocatorCache:
     """Retain reusable allocator blocks without monopolizing device memory.
 
@@ -28,35 +35,48 @@ class BoundedAllocatorCache:
     cached blocks are immediately reusable. This policy keeps the cache while
     both of these bounds hold:
 
-    * cached-but-unallocated memory is at most 25% of device capacity; and
+    * 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
     * at least 5% of device capacity is physically free.
 
-    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 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.
     """
 
     def __init__(
         self,
         device: torch.device,
         *,
-        max_cached_fraction: float = 0.25,
         min_free_fraction: float = 0.05,
     ) -> None:
-        if not 0.0 <= max_cached_fraction <= 1.0:
-            raise ValueError(f"max_cached_fraction must be in [0, 1], got {max_cached_fraction}")
         if not 0.0 <= min_free_fraction <= 1.0:
             raise ValueError(f"min_free_fraction must be in [0, 1], got {min_free_fraction}")
         self.device = device
-        self.max_cached_fraction = max_cached_fraction
         self.min_free_fraction = min_free_fraction
+        self._staged_bytes = 0
+
+    def grow_retention_budget(self, staged_bytes: int) -> None:
+        """Grow the derived cache budget by a component'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)
+
+    @property
+    def budget_bytes(self) -> int:
+        """Reusable-cache ceiling: staged bytes times the headroom multiplier."""
+        return int(self._staged_bytes * _RETENTION_HEADROOM_MULTIPLIER)
 
     def _should_release(self) -> bool:
         reserved = int(torch.accelerator.memory_reserved(self.device))
         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 > int(total * self.max_cached_fraction) or free < int(total * self.min_free_fraction)
+        return cached > self.budget_bytes or free < int(total * self.min_free_fraction)
 
     def release_if_needed(self, *, force: bool = False) -> bool:
         """Release cached blocks when a bound is crossed or release is forced."""
@@ -116,11 +136,12 @@ class PinnedModuleStager:
         self.device = device
         self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream()
         self._ready_event = current_omni_platform.Event()
-        self.cache_retention = cache_retention
+        self.cache_retention = None
         self.loaded = False
         self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
         self._device_storages: list[torch.Tensor] = []
         self._restore_masters()
+        self.set_cache_retention(cache_retention)
 
     @staticmethod
     def _local_tensor(target: torch.Tensor) -> torch.Tensor:
@@ -198,8 +219,17 @@ class PinnedModuleStager:
     def _restore_masters(self) -> None:
         self._bind([group.master for group in self._groups])
 
+    @property
+    def staged_bytes(self) -> int:
+        """Total unique staged storage bytes (masters are uint8 byte views)."""
+        return sum(group.master.numel() for group in self._groups)
+
     def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None:
+        if cache_retention is self.cache_retention:
+            return
         self.cache_retention = cache_retention
+        if cache_retention is not None:
+            cache_retention.grow_retention_budget(self.staged_bytes)
 
     def _release_cache(self, *, force: bool = False) -> None:
         if self.cache_retention is None: