[Model] Optimize MiniMax-H3 DLO component lifecycle

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -169,6 +169,31 @@ This mode means:
   copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
+### Component allocator-cache retention
+
+MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged
+encoder, DiT, and VAE component boundaries so allocations released by one
+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.
+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
+`empty_cache` override, so the executor's unconditional shutdown cleanup is
+preserved.
+
+Retaining allocator blocks trades higher reserved and peak device memory for
+lower lifecycle latency. The benefit is device- and workload-dependent, and a
+smaller or externally loaded device may cross a bound and return to the
+ordinary per-component cache-release behavior. Current performance validation
+covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer
+topologies and platforms receive the same bounded policy but are not claimed
+performance targets; other DLO models remain unchanged unless they explicitly
+adopt it.
+
 ### Final-layout Host Weight Runtime consumer
 
 > **Status:** PR2 implementation. The representation-independent identity,
@@ -390,6 +415,8 @@ Current source-level validation includes:
   transforms without parameter-side flags;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
+- bounded component allocator-cache retention, conservative/forced release,
+  OOM retry, and immutable encoder non-block staging;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
   tests.
 

docs/design/module/diffusion/offloader.md

@@ -18,7 +18,7 @@ validation_paths:
   - tests/diffusion/offloader/**
 upstream_refs:
   - torch.nn.Module.to
-last_reviewed: 2026-07-16
+last_reviewed: 2026-08-25
 ---
 
 # Diffusion offloader
@@ -48,6 +48,11 @@ and correctness required by the selected execution mode.
 **Rule:** Retained host and device copies MUST respect configured limits and
 provide deterministic teardown.
 
+Allocator-cache retention MUST remain local to an explicit component owner,
+declare both cache and physical-free-memory bounds, release conservatively
+when telemetry is unavailable, and force release on failure or memory
+pressure. It MUST NOT replace unconditional executor shutdown cleanup.
+
 ## Safe-change guide
 
 Test enabled and disabled paths, repeated execution, asynchronous transfer,

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _OffloadAbort(BaseException):
+    pass
+
+
+def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+    class VisionStack(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.patch_embed = torch.nn.Linear(2, 2)
+            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+    class TextStack(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed_tokens = torch.nn.Embedding(4, 2)
+            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+    encoder = object.__new__(encoder_module.MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    encoder.device_target = torch.device("cpu")
+    encoder.vision = VisionStack()
+    encoder.text_model = TextStack()
+    hook = mocker.Mock()
+    hook.pin_memory = False
+    encoder._omni_layerwise_hooks = [hook]
+    encoder._omni_layerwise_enabled = True
+    cache = mocker.Mock()
+    encoder.set_omni_component_cache(cache)
+    stager = mocker.Mock()
+    stager_cls = mocker.Mock(return_value=stager)
+    monkeypatch.setattr(encoder_module, "PinnedModuleStager", stager_cls)
+
+    encoder.load_to_device()
+    encoder.offload_to_cpu()
+
+    stager_cls.assert_called_once_with(
+        [encoder.vision.patch_embed, encoder.text_model.embed_tokens],
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=cache,
+    )
+    stager.load.assert_called_once_with()
+    stager.offload.assert_called_once_with()
+    hook.offload_layer.assert_called_once_with()
+
+
+def test_manual_component_failure_forces_retained_cache_release(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    pipeline._dlo_component_cache = mocker.Mock()
+    component = mocker.Mock()
+    component.offload_to_cpu.side_effect = _OffloadAbort("offload failed")
+
+    with pytest.raises(RuntimeError, match="component failed"):
+        with pipeline._component_on_device(component):
+            raise RuntimeError("component failed")
+
+    component.load_to_device.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+    pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    pipeline._dlo_component_cache = mocker.Mock()
+    component = mocker.Mock()
+    component.offload_to_cpu.side_effect = [None, _OffloadAbort("offload failed")]
+
+    with pipeline._component_on_device(component):
+        pass
+    with pytest.raises(_OffloadAbort, match="offload failed"):
+        with pipeline._component_on_device(component):
+            pass
+
+    assert component.load_to_device.call_count == 2
+    assert component.offload_to_cpu.call_count == 2
+    pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)

tests/diffusion/offloader/test_module_residency.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+from torch import nn
+
+import vllm_omni.diffusion.offloader.module_residency as residency_module
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class _DummyEvent:
+    def __init__(self) -> None:
+        self.record_count = 0
+
+    def record(self, _stream) -> None:
+        self.record_count += 1
+
+
+@contextmanager
+def _dummy_stream(_stream):
+    yield
+
+
+@pytest.fixture
+def patched_runtime(monkeypatch, mocker):
+    empty_cache = mocker.Mock()
+    synchronize = mocker.Mock()
+    monkeypatch.setattr(residency_module.current_omni_platform, "Stream", _DummyStream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "Event", _DummyEvent)
+    monkeypatch.setattr(residency_module.current_omni_platform, "current_stream", _DummyStream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "stream", _dummy_stream)
+    monkeypatch.setattr(residency_module.current_omni_platform, "empty_cache", empty_cache)
+    monkeypatch.setattr(residency_module.current_omni_platform, "synchronize", synchronize)
+    return empty_cache, synchronize
+
+
+def _aliased_modules() -> tuple[nn.Module, nn.Module]:
+    parameter_storage = torch.arange(12, dtype=torch.float32)
+    buffer_storage = torch.arange(8, dtype=torch.int64)
+    left = nn.Module()
+    right = nn.Module()
+    left.weight = nn.Parameter(parameter_storage[:8].view(2, 4))
+    right.weight = nn.Parameter(parameter_storage[4:].view(2, 4))
+    left.register_buffer("state", buffer_storage[:6])
+    right.register_buffer("state", buffer_storage[2:])
+    return left, right
+
+
+def test_module_group_preserves_parameter_and_buffer_storage_aliases(patched_runtime, mocker):
+    left, right = _aliased_modules()
+    expected_left_weight = left.weight.detach().clone()
+    expected_right_weight = right.weight.detach().clone()
+    expected_left_state = left.state.clone()
+    expected_right_state = right.state.clone()
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        [left, right],
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+
+    master_weight_ptr = left.weight.untyped_storage().data_ptr()
+    master_state_ptr = left.state.untyped_storage().data_ptr()
+    assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert right.state.untyped_storage().data_ptr() == master_state_ptr
+    assert (left.weight.storage_offset(), right.weight.storage_offset()) == (0, 4)
+    assert (left.state.storage_offset(), right.state.storage_offset()) == (0, 2)
+
+    stager.load()
+
+    assert left.weight.untyped_storage().data_ptr() != master_weight_ptr
+    assert right.weight.untyped_storage().data_ptr() == left.weight.untyped_storage().data_ptr()
+    assert right.state.untyped_storage().data_ptr() == left.state.untyped_storage().data_ptr()
+    left.weight.data[1, 0] = -1
+    left.state[2] = -1
+    assert right.weight[0, 0].item() == -1
+    assert right.state[0].item() == -1
+
+    stager.offload()
+
+    assert left.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+    assert left.state.untyped_storage().data_ptr() == master_state_ptr
+    assert right.state.untyped_storage().data_ptr() == master_state_ptr
+    assert torch.equal(left.weight, expected_left_weight)
+    assert torch.equal(right.weight, expected_right_weight)
+    assert torch.equal(left.state, expected_left_state)
+    assert torch.equal(right.state, expected_right_state)
+    retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_staging_transitions_are_idempotent_and_reuse_one_event(patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    ready_event = stager._ready_event
+
+    stager.load()
+    first_storage = module.weight.untyped_storage().data_ptr()
+    stager.load()
+    assert module.weight.untyped_storage().data_ptr() == first_storage
+    assert stager._ready_event is ready_event
+    assert ready_event.record_count == 1
+
+    stager.offload()
+    stager.offload()
+    assert not stager.loaded
+    retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_failed_load_restores_master_and_forces_cache_release(monkeypatch, patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    expected = module.weight.detach().clone()
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    master_ptr = module.weight.untyped_storage().data_ptr()
+
+    def fail_after_rebind() -> None:
+        device_storages = [torch.empty_like(group.master) for group in stager._groups]
+        stager._bind(device_storages)
+        module.weight.data.zero_()
+        raise RuntimeError("copy failed")
+
+    monkeypatch.setattr(stager, "_load_once", fail_after_rebind)
+    with pytest.raises(RuntimeError, match="copy failed"):
+        stager.load()
+
+    assert not stager.loaded
+    assert module.weight.untyped_storage().data_ptr() == master_ptr
+    assert torch.equal(module.weight, expected)
+    retention.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_out_of_memory_releases_cache_and_retries_once(monkeypatch, patched_runtime, mocker):
+    module = nn.Linear(2, 2)
+    retention = mocker.Mock(spec=BoundedAllocatorCache)
+    stager = PinnedModuleStager(
+        module,
+        torch.device("cpu"),
+        pin_memory=False,
+        cache_retention=retention,
+    )
+    load_once = stager._load_once
+    attempts = 0
+
+    def fail_once() -> None:
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise torch.OutOfMemoryError("memory pressure")
+        load_once()
+
+    monkeypatch.setattr(stager, "_load_once", fail_once)
+    stager.load()
+
+    assert stager.loaded
+    assert attempts == 2
+    retention.release_if_needed.assert_called_once_with(force=True)
+
+
+@pytest.mark.parametrize(
+    ("reserved", "allocated", "free", "expected_release"),
+    [
+        (20, 10, 80, False),
+        (40, 10, 80, True),
+        (20, 10, 4, True),
+    ],
+)
+def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
+    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"))
+
+    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
+
+    def unavailable(_device):
+        raise NotImplementedError
+
+    monkeypatch.setattr(torch.accelerator, "memory_reserved", unavailable)
+    retention = BoundedAllocatorCache(torch.device("cpu"))
+
+    assert retention.release_if_needed()
+    empty_cache.assert_called_once_with()

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 Qwen3-VL layer-50 text/vision encoder.
 
 The encoder is reimplemented on top of vLLM-style tensor-parallel building
@@ -42,6 +43,10 @@ from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMetho
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
 
 MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
@@ -1211,12 +1216,18 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         if not self.is_loaded:
             return
         if getattr(self, "_omni_layerwise_enabled", False):
-            for name, child in self.vision.named_children():
-                if name != "blocks":
-                    child.to(self.device_target)
-            for name, child in self.text_model.named_children():
-                if name != "layers":
-                    child.to(self.device_target)
+            stager = getattr(self, "_omni_non_block_stager", None)
+            if stager is None:
+                hooks = getattr(self, "_omni_layerwise_hooks", ())
+                pin_memory = bool(getattr(hooks[0], "pin_memory", True)) if hooks else True
+                stager = PinnedModuleStager(
+                    self._omni_non_block_modules(),
+                    self.device_target,
+                    pin_memory=pin_memory,
+                    cache_retention=getattr(self, "_omni_component_cache", None),
+                )
+                self._omni_non_block_stager = stager
+            stager.load()
             return
         self.vision.to(self.device_target)
         self.text_model.to(self.device_target)
@@ -1227,18 +1238,39 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         if getattr(self, "_omni_layerwise_enabled", False):
             for hook in self._omni_layerwise_hooks:
                 hook.offload_layer()
-            for name, child in self.vision.named_children():
-                if name != "blocks":
+            stager = getattr(self, "_omni_non_block_stager", None)
+            if stager is not None:
+                stager.offload()
+            else:
+                # The initial DLO setup reaches this branch before the first
+                # load. These modules were constructed on CPU, so this is not
+                # a device-to-host rematerialization.
+                for child in self._omni_non_block_modules():
                     child.to("cpu")
-            for name, child in self.text_model.named_children():
-                if name != "layers":
-                    child.to("cpu")
-            torch.accelerator.empty_cache()
+                self._release_omni_component_cache()
             return
         self.vision.to("cpu")
         self.text_model.to("cpu")
         torch.accelerator.empty_cache()
 
+    def _omni_non_block_modules(self) -> list[nn.Module]:
+        modules = [child for name, child in self.vision.named_children() if name != "blocks"]
+        modules.extend(child for name, child in self.text_model.named_children() if name != "layers")
+        return modules
+
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        stager = getattr(self, "_omni_non_block_stager", None)
+        if stager is not None:
+            stager.set_cache_retention(cache)
+
+    def _release_omni_component_cache(self) -> None:
+        cache = getattr(self, "_omni_component_cache", None)
+        if cache is None:
+            torch.accelerator.empty_cache()
+        else:
+            cache.release_if_needed()
+
     def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None:
         """Stream the TP-local Qwen vision/text blocks for low-HBM serving.
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -46,6 +46,7 @@ from vllm_omni.diffusion.models.interface import (
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.offloader import (
+    BoundedAllocatorCache,
     OffloadPlan,
     apply_sequential_offload,
     remove_sequential_offload,
@@ -952,6 +953,12 @@ class MiniMaxH3Pipeline(
         # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
         self.vae = self.video_vae
 
+        self._dlo_component_cache = None
+        if getattr(od_config, "enable_distributed_layerwise_offload", False):
+            self._dlo_component_cache = BoundedAllocatorCache(self.device)
+            for component in (self.text_encoder, self.video_vae, self.audio_vae):
+                component.set_omni_component_cache(self._dlo_component_cache)
+
         self._quality_policy = MiniMaxH3QualityPolicy(od_config)
         self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
 
@@ -1371,11 +1378,8 @@ class MiniMaxH3Pipeline(
         ):
             # Layerwise DiT offload already provides the low-residency encoder
             # phase used by the checkpoint reference.
-            self.text_encoder.load_to_device()
-            try:
+            with self._component_on_device(self.text_encoder):
                 return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
-            finally:
-                self.text_encoder.offload_to_cpu()
 
         # Keep both Qwen and DiT resident across requests. Moving either model
         # here makes encoder latency include a tens-of-gigabytes PCIe transfer,
@@ -1433,13 +1437,35 @@ class MiniMaxH3Pipeline(
                 yield
             return
         staged = self._uses_manual_component_offload()
-        if staged:
-            component.load_to_device()
         try:
-            yield
-        finally:
             if staged:
-                component.offload_to_cpu()
+                component.load_to_device()
+            yield
+        except BaseException:
+            if staged:
+                try:
+                    component.offload_to_cpu()
+                except BaseException:
+                    logger.exception("Failed to release %s after component failure", component.__class__.__name__)
+                cache = getattr(self, "_dlo_component_cache", None)
+                if cache is not None:
+                    try:
+                        cache.release_if_needed(force=True)
+                    except BaseException:
+                        logger.exception("Failed to release retained allocator cache after component failure")
+            raise
+        else:
+            if staged:
+                try:
+                    component.offload_to_cpu()
+                except BaseException:
+                    cache = getattr(self, "_dlo_component_cache", None)
+                    if cache is not None:
+                        try:
+                            cache.release_if_needed(force=True)
+                        except BaseException:
+                            logger.exception("Failed to release retained allocator cache after offload failure")
+                    raise
 
     def _encode_visual_conditions(
         self,

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 remote-code VAE adapters and exact latent contracts."""
 
 from __future__ import annotations
@@ -21,7 +22,10 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     DistributedVaeMixin,
 )
 from vllm_omni.diffusion.distributed.parallel_state import get_world_group
-from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
+from vllm_omni.diffusion.offloader.module_residency import (
+    BoundedAllocatorCache,
+    PinnedModuleStager,
+)
 
 from .packed_tokens import minimax_h3_patchify_video_latent
 
@@ -148,12 +152,21 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         else:
             self.remote.to(self._device_target)
 
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        if self._stager is not None:
+            self._stager.set_cache_retention(cache)
+
     def offload_to_cpu(self) -> None:
         if self._stager is not None:
             self._stager.offload()
         else:
             self.remote.to("cpu")
-            torch.accelerator.empty_cache()
+            cache = getattr(self, "_omni_component_cache", None)
+            if cache is None:
+                torch.accelerator.empty_cache()
+            else:
+                cache.release_if_needed()
 
     def set_parallel_size(
         self,
@@ -411,12 +424,21 @@ class MiniMaxH3AudioVAE(nn.Module):
         else:
             self.remote.to(self._device_target)
 
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._omni_component_cache = cache
+        if self._stager is not None:
+            self._stager.set_cache_retention(cache)
+
     def offload_to_cpu(self) -> None:
         if self._stager is not None:
             self._stager.offload()
         else:
             self.remote.to("cpu")
-            torch.accelerator.empty_cache()
+            cache = getattr(self, "_omni_component_cache", None)
+            if cache is None:
+                torch.accelerator.empty_cache()
+            else:
+                cache.release_if_needed()
 
     @torch.inference_mode()
     def encode_waveform(

vllm_omni/diffusion/offloader/__init__.py

@@ -17,7 +17,7 @@ from .distributed_layerwise_backend import (
     remove_distributed_block_hook,
 )
 from .layerwise_backend import LayerWiseOffloadBackend
-from .module_residency import PinnedModuleStager
+from .module_residency import BoundedAllocatorCache, PinnedModuleStager
 from .offload_plan import OffloadPlan, get_offload_plan
 from .sequential_backend import (
     ModelLevelOffloadBackend,
@@ -46,6 +46,7 @@ __all__ = [
     "DistributedLayerwiseOffloadBackend",
     "DistributedLayerwiseOffloadHook",
     "ModelLevelOffloadBackend",
+    "BoundedAllocatorCache",
     "PinnedModuleStager",
     "apply_sequential_offload",
     "remove_sequential_offload",

vllm_omni/diffusion/offloader/module_residency.py

@@ -1,89 +1,271 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""On-demand module staging backed by an immutable pinned CPU snapshot."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""On-demand module staging backed by immutable pinned CPU storage."""
 
 from __future__ import annotations
 
+from collections.abc import Iterable
+from dataclasses import dataclass
 from itertools import chain
 from typing import Any
 
 import torch
 from torch import nn
+from vllm.logger import init_logger
 
 from vllm_omni.platforms import current_omni_platform
 
-from .tensor_utils import set_tensor_storage
+from .tensor_utils import is_dtensor, set_tensor_storage
+
+logger = init_logger(__name__)
 
 
-class PinnedModuleStager:
-    """Stage an immutable module without copying device weights back to CPU.
+class BoundedAllocatorCache:
+    """Retain reusable allocator blocks without monopolizing device memory.
 
-    ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter
-    after each forward. In inference the weights are immutable, so retain one
-    pinned CPU master instead. ``load`` materializes fresh device tensors from
-    that master; ``offload`` only rebinds the Parameters/Buffers to it and
-    releases the device allocations.
+    Component offload normally calls ``empty_cache`` after every stage. That
+    makes the next stage return to the device allocator even though PyTorch's
+    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
+    * 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.
     """
 
     def __init__(
         self,
-        module: nn.Module,
+        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
+
+    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)
+
+    def release_if_needed(self, *, force: bool = False) -> bool:
+        """Release cached blocks when a bound is crossed or release is forced."""
+        if not force:
+            try:
+                if not self._should_release():
+                    return False
+            except Exception as exc:
+                # Preserve the pre-retention behavior on platforms that do not
+                # expose allocator telemetry through torch.accelerator.
+                logger.debug("Allocator cache telemetry unavailable; releasing cache: %s", exc)
+        current_omni_platform.empty_cache()
+        return True
+
+
+@dataclass(frozen=True)
+class _TensorBinding:
+    target: torch.Tensor
+    dtype: torch.dtype
+    shape: tuple[int, ...]
+    stride: tuple[int, ...]
+    storage_offset: int
+
+
+@dataclass
+class _StorageGroup:
+    master: torch.Tensor
+    bindings: list[_TensorBinding]
+
+
+class PinnedModuleStager:
+    """Stage immutable module groups without copying device weights to CPU.
+
+    ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter
+    after each forward. In inference the weights are immutable, so retain one
+    pinned CPU master instead. ``load`` materializes device storage from that
+    master; ``offload`` only rebinds Parameters and buffers to the master.
+
+    A module iterable is treated as one staging group. It uses one copy stream
+    and one reusable completion event. Tensors sharing storage keep their
+    shapes, strides, offsets, dtypes, and aliases across every transition.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module | Iterable[nn.Module],
         device: torch.device,
         *,
         pin_memory: bool = True,
         copy_stream: Any | None = None,
+        cache_retention: BoundedAllocatorCache | None = None,
     ) -> None:
-        self.device = device
-        self.copy_stream = copy_stream or current_omni_platform.Stream()
-        self.loaded = False
-        self._entries: list[tuple[torch.Tensor, torch.Tensor]] = []
-        self._device_tensors: list[torch.Tensor] = []
+        modules = (module,) if isinstance(module, nn.Module) else tuple(module)
+        if not modules or not all(isinstance(item, nn.Module) for item in modules):
+            raise ValueError("PinnedModuleStager requires at least one nn.Module")
 
-        # parameters()/buffers() remove duplicate objects by default, which
-        # preserves tied weights and buffer aliases when their storage moves.
-        for target in chain(module.parameters(), module.buffers()):
-            master = target.detach()
-            if master.device.type != "cpu":
-                master = master.to("cpu")
+        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.loaded = False
+        self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
+        self._device_storages: list[torch.Tensor] = []
+        self._restore_masters()
+
+    @staticmethod
+    def _local_tensor(target: torch.Tensor) -> torch.Tensor:
+        return target.to_local() if is_dtensor(target) else target
+
+    @classmethod
+    def _snapshot_groups(
+        cls,
+        modules: tuple[nn.Module, ...],
+        *,
+        pin_memory: bool,
+    ) -> list[_StorageGroup]:
+        targets: list[torch.Tensor] = []
+        seen_targets: set[int] = set()
+        for target in chain.from_iterable(chain(item.parameters(), item.buffers()) for item in modules):
+            if id(target) not in seen_targets:
+                seen_targets.add(id(target))
+                targets.append(target)
+
+        grouped: dict[tuple[Any, ...], tuple[torch.Tensor, list[_TensorBinding]]] = {}
+        for target in targets:
+            local = cls._local_tensor(target)
+            if local.is_meta:
+                raise ValueError("PinnedModuleStager cannot snapshot a meta tensor")
+            storage = local.untyped_storage()
+            if storage.nbytes() == 0:
+                storage_key: tuple[Any, ...] = ("empty", id(target))
+            else:
+                storage_key = (
+                    local.device.type,
+                    local.device.index,
+                    storage.data_ptr(),
+                    storage.nbytes(),
+                )
+            binding = _TensorBinding(
+                target=target,
+                dtype=local.dtype,
+                shape=tuple(local.shape),
+                stride=tuple(local.stride()),
+                storage_offset=local.storage_offset(),
+            )
+            if storage_key not in grouped:
+                grouped[storage_key] = (local, [])
+            grouped[storage_key][1].append(binding)
+
+        groups: list[_StorageGroup] = []
+        for source, bindings in grouped.values():
+            storage = source.untyped_storage()
+            storage_view = torch.empty(0, dtype=torch.uint8, device=source.device).set_(
+                storage,
+                0,
+                (storage.nbytes(),),
+                (1,),
+            )
+            master = storage_view.detach() if storage_view.device.type == "cpu" else storage_view.to("cpu")
             if pin_memory and not master.is_pinned():
                 master = master.pin_memory()
-            set_tensor_storage(target, master)
-            self._entries.append((target, master))
+            groups.append(_StorageGroup(master=master, bindings=bindings))
+        return groups
+
+    @staticmethod
+    def _view(backing: torch.Tensor, binding: _TensorBinding) -> torch.Tensor:
+        return torch.empty(0, dtype=binding.dtype, device=backing.device).set_(
+            backing.untyped_storage(),
+            binding.storage_offset,
+            binding.shape,
+            binding.stride,
+        )
+
+    def _bind(self, storages: list[torch.Tensor]) -> None:
+        for storage, group in zip(storages, self._groups):
+            for binding in group.bindings:
+                set_tensor_storage(binding.target, self._view(storage, binding))
+
+    def _restore_masters(self) -> None:
+        self._bind([group.master for group in self._groups])
+
+    def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None:
+        self.cache_retention = cache_retention
+
+    def _release_cache(self, *, force: bool = False) -> None:
+        if self.cache_retention is None:
+            current_omni_platform.empty_cache()
+        else:
+            self.cache_retention.release_if_needed(force=force)
+
+    def _load_once(self) -> None:
+        device_storages = [torch.empty_like(group.master, device=self.device) for group in self._groups]
+        compute_stream = current_omni_platform.current_stream()
+        self.copy_stream.wait_stream(compute_stream)
+        with current_omni_platform.stream(self.copy_stream):
+            for device_storage, group in zip(device_storages, self._groups):
+                device_storage.copy_(group.master, non_blocking=group.master.is_pinned())
+            self._ready_event.record(self.copy_stream)
+
+        self._bind(device_storages)
+        compute_stream.wait_event(self._ready_event)
+        self._device_storages = device_storages
+        self.loaded = True
+
+    def _cleanup_failed_load(self) -> None:
+        try:
+            current_omni_platform.synchronize()
+        except Exception:
+            logger.debug("Device synchronization failed while cleaning up module staging", exc_info=True)
+        self._restore_masters()
+        self._device_storages.clear()
+        self.loaded = False
+        self._release_cache(force=True)
 
     def load(self) -> None:
         if self.loaded:
             return
-
-        # Allocate on the compute stream so the caching allocator can reuse
-        # memory released by the preceding encoder/DiT stage.
-        device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries]
-        compute_stream = current_omni_platform.current_stream()
-        self.copy_stream.wait_stream(compute_stream)
-        ready = current_omni_platform.Event()
-        with current_omni_platform.stream(self.copy_stream):
-            for device_tensor, (_, master) in zip(device_tensors, self._entries):
-                device_tensor.copy_(master, non_blocking=master.is_pinned())
-            ready.record(self.copy_stream)
-
-        for device_tensor, (target, _) in zip(device_tensors, self._entries):
-            set_tensor_storage(target, device_tensor)
-        compute_stream.wait_event(ready)
-        self._device_tensors = device_tensors
-        self.loaded = True
+        try:
+            self._load_once()
+        except torch.OutOfMemoryError:
+            # A retained cache is normally reusable by this process, but an
+            # explicit flush gives external memory pressure one bounded retry.
+            self._cleanup_failed_load()
+            try:
+                self._load_once()
+            except BaseException:
+                self._cleanup_failed_load()
+                raise
+        except BaseException:
+            self._cleanup_failed_load()
+            raise
 
     def offload(self) -> None:
         if not self.loaded:
             return
 
         # The module has completed on the compute stream. Synchronize once at
-        # the stage boundary, then release HBM by rebinding to the unchanged
-        # host masters. No device-to-host transfer is necessary.
-        current_omni_platform.synchronize()
-        for target, master in self._entries:
-            set_tensor_storage(target, master)
-        self._device_tensors.clear()
+        # the stage boundary, then discard device storage without any D2H copy.
+        try:
+            current_omni_platform.synchronize()
+            self._restore_masters()
+        except BaseException:
+            self._device_storages.clear()
+            self.loaded = False
+            self._release_cache(force=True)
+            raise
+        self._device_storages.clear()
         self.loaded = False
-        current_omni_platform.empty_cache()
+        self._release_cache()
 
 
-__all__ = ["PinnedModuleStager"]
+__all__ = ["BoundedAllocatorCache", "PinnedModuleStager"]