diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 8d94d51411ed..82990a6c9734 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py index 07530ca0e744..63a5d52939d2 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index c569eff34368..41e490bd3281 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/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