[diffusion] feat: allow a component use retain its layerwise resident set

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py

@@ -49,6 +49,12 @@ class ComponentUse:
     target_dtype: torch.dtype | None = None
     keep_ready_after_warmup: bool = False
     start_at_stage_entry: bool = True
+    # Layerwise components release their resident set when the use ends, which
+    # is right for a DiT (one use spans every denoise step) and buys nothing for
+    # a component used once per forward. Set this when something that knows the
+    # pipeline's per-phase headroom has decided the room is better spent holding
+    # the set until this component runs again.
+    retain_resident_layers: bool = False
 
 
 @dataclass(slots=True)

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py

@@ -278,14 +278,19 @@ def finish_use(
     ) -> None:
         if not isinstance(module, LayerwiseOffloadableModuleMixin):
             return
+        # Not release_all: this is a use ending, not a reset. Whether the
+        # resident set outlives the use is declared on the use, by whoever has
+        # the pipeline's per-phase headroom in view; the default is off, so
+        # this stays the long-standing behaviour until something sets it.
+        keep_resident = use.retain_resident_layers
         for manager in module.layerwise_offload_managers:
-            # Not release_all: this is a use ending, not a reset. The default
-            # still drops the resident set, so behaviour is unchanged here.
-            manager.release_after_use()
+            manager.release_after_use(keep_resident=keep_resident)
         # The layers are gone; the rest of this component is dead weight on the
         # device until it is used again, and the stage that follows may be the
-        # one that needs the room.
-        module.park_non_layer_weights()
+        # one that needs the room. That reasoning does not hold when the room
+        # was just judged available: parking would undo the transfer we kept.
+        if not keep_resident:
+            module.park_non_layer_weights()
         if current_platform.is_mps():
             torch.mps.synchronize()
             module.restore_mps_cpu_non_layer_weights()

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -1381,6 +1381,60 @@ def test_disable_offload_short_circuits_residency_release(monkeypatch):
         assert tuple(param.shape) != (1,), name
 
 
+def test_finish_use_drops_residents_unless_the_use_says_otherwise(monkeypatch):
+    """`retain_resident_layers` defaults off, so finish_use behaves as before.
+
+    Every existing pipeline builds its uses without the flag, so this is the
+    path they all take and it must keep releasing the resident set.
+    """
+    model = _configure_mixin_model(monkeypatch)
+    released = []
+    parked = []
+    for manager in model.layerwise_offload_managers:
+        manager.release_after_use = lambda *, keep_resident=False: released.append(
+            keep_resident
+        )
+    model.park_non_layer_weights = lambda: parked.append(True)
+
+    LayerwiseOffloadStrategy().finish_use(
+        model,
+        ComponentUse(stage_name="test", component_name="transformer"),
+        SimpleNamespace(),
+    )
+
+    assert released and all(keep is False for keep in released)
+    assert parked == [True]
+
+
+def test_finish_use_keeps_residents_when_the_use_declares_it(monkeypatch):
+    """The declaration reaches the manager, and parking is skipped with it.
+
+    Parking pushes the component's non-layer weights to host; doing that right
+    after deciding the room is available would undo the transfer being kept.
+    """
+    model = _configure_mixin_model(monkeypatch)
+    released = []
+    parked = []
+    for manager in model.layerwise_offload_managers:
+        manager.release_after_use = lambda *, keep_resident=False: released.append(
+            keep_resident
+        )
+    model.park_non_layer_weights = lambda: parked.append(True)
+
+    LayerwiseOffloadStrategy().finish_use(
+        model,
+        ComponentUse(
+            stage_name="test",
+            component_name="transformer",
+            retain_resident_layers=True,
+        ),
+        SimpleNamespace(),
+    )
+
+    assert released and all(keep is True for keep in released)
+    assert parked == []
+
+
 def test_enable_offload_rearms_after_disable(monkeypatch):
     model = _configure_mixin_model(monkeypatch)
     # blocks[2] holds a placeholder right after configure; the real values are