[diffusion] refactor: introduce component residency manager (#23771)

python/sglang/multimodal_gen/runtime/loader/weights_updater.py

@@ -53,9 +53,9 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

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

@@ -0,0 +1,673 @@
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Mapping, MutableMapping, Protocol, Sequence, TypeVar
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.managers.component_resident_strategies import (
+    ComponentResidencyStrategy,
+    LayerwiseOffloadStrategy,
+    ResidentStrategy,
+    VanillaD2HStrategy,
+)
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_T = TypeVar("_T")
+
+
+@dataclass(slots=True)
+class ComponentUse:
+    """Describes one stage/use-site access to a pipeline component."""
+
+    stage_name: str
+    # Pipeline module key: transformer / video_dit / text_encoder / ...
+    component_name: str
+    # Model-specific phase for sequential components, e.g. stage1 or stage2.
+    # TODO: Replace this with ordered timeline identity. In an all-sequential
+    # pipeline, use-site identity should come from the declared ComponentUse
+    # order instead of a per-use `phase` field.
+    phase: str | None = None
+    # Whether the manager may prepare this component for the next request.
+    preferred_ready_after_request: bool = False
+    # Whether cross-stage prefetch may prepare this use before the use-site.
+    allow_prefetch: bool = True
+    # Whether this use is expensive enough that earlier timeline prefetch matters.
+    # TODO: Replace this boolean hint with a budget-aware lookahead planner:
+    # estimate memory/load cost and reuse distance, keep small and early-request
+    # components resident within budget, prefetch as soon as VRAM slack appears,
+    # and release completed components only when the budget requires it.
+    memory_intensive: bool = False
+    # Optional module dtype required by this use-site.
+    target_dtype: torch.dtype | None = None
+    # Some components are intentionally kept ready between warmup and the first
+    # real request to avoid measuring a cold H2D in the user-visible request.
+    keep_ready_after_warmup: bool = False
+
+
+@dataclass(slots=True)
+class ResidencyState:
+    """
+    Necessary internal runtime info of ComponentResidencyManager
+    """
+
+    stages: Sequence["ComponentResidencyStage"] = ()
+    stage_index: int = -1
+    stage_name: str | None = None
+    next_stage_name: str | None = None
+    current_use: ComponentUse | None = None
+    # the ComponentUses of the preceding stages
+    future_uses: tuple[ComponentUse, ...] = ()
+    batch_is_warmup: bool = False
+    manager_mode: str = "static"
+    trace_enabled: bool = False
+
+
+class ResidencyBatch(Protocol):
+    is_warmup: bool
+
+
+class ComponentResidencyStage(Protocol):
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]: ...
+
+
+class ComponentResidencyPipeline(Protocol):
+    modules: Mapping[str, object]
+    _stage_name_mapping: Mapping[str, ComponentResidencyStage]
+    component_residency_strategies: MutableMapping[str, "ComponentResidencyStrategy"]
+
+
+def build_dit_residency_strategy(
+    module: nn.Module,
+    server_args: ServerArgs,
+) -> ComponentResidencyStrategy:
+    if (
+        isinstance(module, OffloadableDiTMixin)
+        and module.layerwise_offload_managers
+        and any(manager.enabled for manager in module.layerwise_offload_managers)
+    ):
+        # only if dit_layerwise_offload is enabled
+        return LayerwiseOffloadStrategy()
+    if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
+        # handles offload by vanalla D2H
+        return VanillaD2HStrategy()
+    return ResidentStrategy()
+
+
+def is_fsdp_managed_module(module: nn.Module) -> bool:
+    return module.__class__.__name__.startswith("FSDP")
+
+
+def build_component_residency_strategy(
+    component_name: str,
+    module: nn.Module,
+    server_args: ServerArgs,
+) -> ComponentResidencyStrategy:
+    if component_name in {
+        "transformer",
+        "transformer_2",
+        "video_dit",
+        "video_dit_2",
+        "audio_dit",
+        "dual_tower_bridge",
+    }:
+        return build_dit_residency_strategy(module, server_args)
+
+    if component_name.startswith("text_encoder") or component_name.endswith(
+        "text_encoder"
+    ):
+        if (
+            server_args.text_encoder_cpu_offload
+            and not server_args.use_fsdp_inference
+            and not is_fsdp_managed_module(module)
+        ):
+            return VanillaD2HStrategy()
+        return ResidentStrategy()
+
+    if component_name == "image_encoder":
+        if server_args.image_encoder_cpu_offload and not server_args.use_fsdp_inference:
+            return VanillaD2HStrategy()
+        return ResidentStrategy()
+
+    if component_name in {
+        "vae",
+        "video_vae",
+        "audio_vae",
+        "vocoder",
+        "spatial_upsampler",
+        "condition_image_encoder",
+    }:
+        if server_args.vae_cpu_offload and not server_args.use_fsdp_inference:
+            return VanillaD2HStrategy()
+        return ResidentStrategy()
+
+    return ResidentStrategy()
+
+
+class ComponentResidencyManager:
+    """Executor-owned component lifecycle coordinator. Provide hooks for a PipelineExecutor
+
+    Hooks are called around executor progress:
+        before request: collect a flat ordered ComponentUse timeline.
+        before stage: update current/next stage context only.
+        begin use: finish previous active use, prepare current use, wait until ready.
+        end use: finish or keep current use, then prefetch the next heavy timeline use.
+        finish request: finish active use and schedule preferred next-request prefetch.
+
+    The manager instance is global and rebound to the active pipeline before request execution.
+    This manager is designed only for sequential execution order for now
+    """
+
+    def __init__(
+        self, pipeline: ComponentResidencyPipeline, server_args: ServerArgs
+    ) -> None:
+        self.pipeline = pipeline
+        self.server_args = server_args
+        self.state = ResidencyState(trace_enabled=False)
+        self._stage_names_by_id: dict[int, str] = {}
+        self._stage_uses_by_index: list[tuple[ComponentUse, ...]] = []
+        self._ordered_uses: tuple[ComponentUse, ...] = ()
+        self._current_use_index: int = -1
+        self._active_use: ComponentUse | None = None
+        self._active_use_module: nn.Module | None = None
+        self._prefetched_use_keys: set[tuple[str, str, str | None]] = set()
+        self._custom_strategies: dict[str, ComponentResidencyStrategy] = dict(
+            pipeline.component_residency_strategies
+        )
+        self._uses_seen: dict[str, ComponentUse] = {}
+
+    @property
+    def enabled(self) -> bool:
+        return True
+
+    def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None:
+        custom_strategies = dict(pipeline.component_residency_strategies)
+        if pipeline is not self.pipeline:
+            self.strategy_for.cache_clear()
+            self._should_keep_single_dit.cache_clear()
+            self._active_use = None
+            self._active_use_module = None
+            self._uses_seen.clear()
+            self._prefetched_use_keys.clear()
+        elif custom_strategies != self._custom_strategies:
+            self.strategy_for.cache_clear()
+        self.pipeline = pipeline
+        self._custom_strategies = custom_strategies
+        self._stage_names_by_id = {
+            id(stage): name for name, stage in pipeline._stage_name_mapping.items()
+        }
+
+    def refresh_server_args(self, server_args: ServerArgs) -> None:
+        if server_args is not self.server_args:
+            self.strategy_for.cache_clear()
+        self.server_args = server_args
+
+    def register_strategy(
+        self, component_name: str, strategy: ComponentResidencyStrategy
+    ) -> None:
+        self.pipeline.component_residency_strategies[component_name] = strategy
+        self._custom_strategies[component_name] = strategy
+        self.strategy_for.cache_clear()
+
+    def begin_request(
+        self,
+        stages: Sequence[ComponentResidencyStage],
+        batch: ResidencyBatch,
+        server_args: ServerArgs,
+    ) -> None:
+        """A hook called before processing an actual request"""
+        self.refresh_server_args(server_args)
+        self.state = ResidencyState(
+            stages=stages, batch_is_warmup=batch.is_warmup, trace_enabled=False
+        )
+        self._active_use = None
+        self._active_use_module = None
+        self._current_use_index = -1
+        self._prefetched_use_keys.clear()
+        self._uses_seen.clear()
+        if self.enabled:
+            self._stage_uses_by_index = [
+                tuple(stage.component_uses(server_args, self.stage_name(stage)))
+                for stage in stages
+            ]
+            self._ordered_uses = tuple(
+                use for uses in self._stage_uses_by_index for use in uses
+            )
+        else:
+            self._stage_uses_by_index = []
+            self._ordered_uses = ()
+        self._trace(
+            "request_start",
+            detail=f"stages={len(stages)} uses={len(self._ordered_uses)}",
+        )
+
+    def before_stage(
+        self,
+        stage: ComponentResidencyStage,
+        stage_index: int,
+        batch: ResidencyBatch,
+        server_args: ServerArgs,
+    ) -> None:
+        """called after stage starts"""
+        if not self.enabled:
+            return
+        # update state before entering the stage
+        self.state.stage_index = stage_index
+        self.state.stage_name = self.stage_name(stage)
+        self.state.next_stage_name = self._next_stage_name(stage_index)
+        self._trace("stage_enter", detail=f"index={stage_index}")
+
+    def after_stage(self, stage_index: int) -> None:
+        """called after stage exits"""
+        if not self.enabled:
+            return
+        self._trace("stage_exit", detail=f"index={stage_index}")
+
+    def before_use(self, use: ComponentUse) -> None:
+        """component use-site starts"""
+        if not self.enabled:
+            return
+        self.begin_use(use)
+
+    def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
+        """Begin one sequential component use interval. this is idempotent
+
+        1. Finish the previous active use if this is a different timeline use.
+        2. Prepare the current component.
+        3. Wait until the current component is ready, then prefetch the next heavy use.
+        """
+        if self._active_use is not None and self._same_use(self._active_use, use):
+            return
+        if self._active_use is not None:
+            # finish previous active use
+            self._finish_use(
+                self._active_use,
+                module=self._active_use_module,
+                keep_on_warmup=self._active_use.keep_ready_after_warmup,
+            )
+            self._active_use = None
+            self._active_use_module = None
+            self.state.current_use = None
+        self._mark_current_use(use)
+        self._prepare_forward_use(use, module=module)
+        self._active_use = use
+        self._active_use_module = module
+        self._prefetch_next_memory_intensive_use()
+
+    def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
+        """End one sequential component use interval.
+
+        1. Finish or keep the current component.
+        2. Clear it as the active use.
+        3. Prefetch the next memory-intensive use without waiting.
+        """
+        if self._active_use is None or not self._same_use(self._active_use, use):
+            return
+        self._finish_use(
+            self._active_use,
+            module=self._active_use_module or module,
+            keep_on_warmup=self._active_use.keep_ready_after_warmup,
+        )
+        self._active_use = None
+        self._active_use_module = None
+        self.state.current_use = None
+        self._prefetch_next_memory_intensive_use()
+
+    @contextmanager
+    def use_component(
+        self, use: ComponentUse, module: nn.Module | None = None
+    ) -> Iterator[nn.Module | None]:
+        self.begin_use(use, module=module)
+        try:
+            yield module if module is not None else self.get_module(use.component_name)
+        finally:
+            self.end_use(use, module=module)
+
+    def call_component(
+        self,
+        use: ComponentUse,
+        module: Callable[..., _T],
+        *args,
+        **kwargs,
+    ) -> _T:
+        with self.use_component(use):
+            return module(*args, **kwargs)
+
+    def prefetch_use(self, use: ComponentUse) -> None:
+        """Prepare a future use without blocking the current use."""
+        if not self.enabled:
+            return
+        self._prefetch_use(use)
+
+    def ensure_ready(self, use: ComponentUse, module: nn.Module | None = None) -> None:
+        """Prepare a shared component and wait without making it the active use."""
+        if not self.enabled:
+            return
+        self._prepare_forward_use(use, module=module)
+
+    def prefetch_checkpoint(self, anchor: ComponentUse | None = None) -> None:
+        """Give the manager a timeline overlap point.
+
+        1. Locate the anchor or current use in the ordered timeline.
+        2. Find the next prefetchable memory-intensive use.
+        3. Prepare it opportunistically without waiting.
+        """
+        if not self.enabled:
+            return
+        if anchor is not None:
+            self._mark_current_use(anchor)
+        self._prefetch_next_memory_intensive_use()
+
+    def finish_active_use(self, *, prefetch_next: bool = True) -> None:
+        """Finish the currently active sequential use, if any."""
+        if self._active_use is None:
+            return
+        active_use = self._active_use
+        self._finish_use(
+            active_use,
+            module=self._active_use_module,
+            keep_on_warmup=active_use.keep_ready_after_warmup,
+        )
+        self._active_use = None
+        self._active_use_module = None
+        self.state.current_use = None
+        if prefetch_next:
+            self._prefetch_next_memory_intensive_use()
+
+    def _prepare_forward_use(
+        self, use: ComponentUse, module: nn.Module | None = None
+    ) -> None:
+        """Prepare a component that is about to run and wait until it is ready."""
+        module = module or self.get_module(use.component_name)
+        if module is None:
+            self._trace("skip_missing", use)
+            return
+        strategy = self.strategy_for(use.component_name, module)
+        self._uses_seen[use.component_name] = use
+        self.state.current_use = use
+        self._trace("prepare", use, strategy, module)
+        strategy.prepare_for_use(module, use, self.state)
+        self._trace("wait", use, strategy, module)
+        strategy.wait_for_use(module, use, self.state)
+
+    def _prefetch_use(self, use: ComponentUse) -> None:
+        """Prepare a future component opportunistically without waiting.
+
+        This is called when the component is memory-intensive so it may takes a long time to prefetch.
+
+        manager will perform the prefetch at some checkpoints, if necessary
+        """
+        if not use.allow_prefetch:
+            return
+        module = self.get_module(use.component_name)
+        if module is None:
+            self._trace("skip_missing", use)
+            return
+        strategy = self.strategy_for(use.component_name, module)
+        if isinstance(strategy, VanillaD2HStrategy) and self._active_use is not None:
+            # Avoid making two vanilla-offloaded heavy components resident before
+            # a budget-aware planner can prove the overlap is safe.
+            self._trace("prefetch_skip_active_vanilla", use, strategy, module)
+            return
+
+        self._uses_seen[use.component_name] = use
+        self._trace("prefetch", use, strategy, module)
+        if strategy.prefetch_for_use(module, use, self.state):
+            self._prefetched_use_keys.add(self._use_key(use))
+
+    def after_use(self, use: ComponentUse) -> None:
+        if not self.enabled:
+            return
+        self.end_use(use)
+
+    def _finish_use(
+        self,
+        use: ComponentUse,
+        *,
+        module: nn.Module | None = None,
+        keep_on_warmup: bool,
+    ) -> None:
+        """finish a specific use by keeping them resident or call finish_use hook"""
+        module = module or self.get_module(use.component_name)
+        if module is None:
+            self._trace("skip_missing", use)
+            return
+        should_keep = (
+            keep_on_warmup and self.state.batch_is_warmup
+        ) or self._should_keep_after_use(use)
+        if should_keep:
+            self._trace(
+                "keep",
+                use,
+                self.strategy_for(use.component_name, module),
+                module,
+            )
+            return
+        strategy = self.strategy_for(use.component_name, module)
+        self._trace("finish", use, strategy, module)
+        strategy.finish_use(module, use, self.state)
+
+    def finish_request(self) -> None:
+        if not self.enabled and not self._uses_seen and self._active_use is None:
+            return
+        # 1. Close the currently active sequential use.
+        self.finish_active_use(prefetch_next=False)
+        # 2. Pick components that should be ready for the next request.
+        preferred_uses = self._preferred_request_end_uses()
+        # 3. Finish everything else, or prepare preferred uses for request tail.
+        for component_name, use in list(self._uses_seen.items()):
+            module = self.get_module(component_name)
+            if module is None:
+                continue
+            if self.state.batch_is_warmup and use.keep_ready_after_warmup:
+                self._trace(
+                    "request_keep_warmup",
+                    use,
+                    self.strategy_for(component_name, module),
+                    module,
+                )
+                continue
+            preferred = component_name in preferred_uses
+            if not preferred and self._should_keep_single_dit(component_name):
+                self._trace(
+                    "keep",
+                    use,
+                    self.strategy_for(component_name, module),
+                    module,
+                    detail="single_dit",
+                )
+                continue
+            strategy = self.strategy_for(component_name, module)
+            if preferred and not self.state.batch_is_warmup:
+                self._trace("request_prefetch", use, strategy, module)
+                strategy.prepare_after_request(module, use, self.state)
+            else:
+                action = "request_resident" if preferred else "request_finish"
+                self._trace(action, use, strategy, module)
+                strategy.finish_request(module, use, self.state, preferred=preferred)
+        self._trace("request_end")
+
+    def stage_name(self, stage: ComponentResidencyStage) -> str:
+        return self._stage_names_by_id.get(id(stage), stage.__class__.__name__)
+
+    def component_name_for_module(self, module: nn.Module | None, default: str) -> str:
+        if module is None:
+            return default
+        for name, candidate in self.pipeline.modules.items():
+            if candidate is module:
+                return name
+        return default
+
+    def get_module(self, component_name: str) -> nn.Module | None:
+        module = self.pipeline.modules.get(component_name)
+        return module if isinstance(module, nn.Module) else None
+
+    @lru_cache(maxsize=None)
+    def strategy_for(
+        self, component_name: str, module: nn.Module
+    ) -> ComponentResidencyStrategy:
+        """Return the pre-registered strategy for a specific component"""
+        custom_strategy = self._custom_strategies.get(component_name)
+        if custom_strategy is not None:
+            return custom_strategy
+        return build_component_residency_strategy(
+            component_name, module, self.server_args
+        )
+
+    def _stage_uses(self, stage_index: int) -> tuple[ComponentUse, ...]:
+        """Returns the ComponentUse(s) of a specific stage"""
+        if stage_index < 0 or stage_index >= len(self._stage_uses_by_index):
+            return ()
+        return self._stage_uses_by_index[stage_index]
+
+    def _next_stage_name(self, stage_index: int) -> str | None:
+        next_index = stage_index + 1
+        if next_index < 0 or next_index >= len(self.state.stages):
+            return None
+        return self.stage_name(self.state.stages[next_index])
+
+    def _mark_current_use(self, use: ComponentUse) -> None:
+        index = self._locate_use_index(use)
+        if index is None:
+            self._current_use_index = len(self._ordered_uses)
+            self.state.future_uses = ()
+            return
+        self._current_use_index = index
+        self.state.future_uses = self._ordered_uses[index + 1 :]
+
+    def _locate_use_index(self, use: ComponentUse) -> int | None:
+        for index in range(self._current_use_index + 1, len(self._ordered_uses)):
+            if self._same_use(self._ordered_uses[index], use):
+                return index
+        for index, candidate in enumerate(self._ordered_uses):
+            if self._same_use(candidate, use):
+                return index
+        return None
+
+    def _prefetch_next_memory_intensive_use(self) -> None:
+        for use in self._ordered_uses[self._current_use_index + 1 :]:
+            if not use.memory_intensive:
+                continue
+            if self._use_key(use) in self._prefetched_use_keys:
+                return
+            self.prefetch_use(use)
+            return
+
+    def _should_keep_after_use(self, use: ComponentUse) -> bool:
+        future_component_names = {
+            future.component_name for future in self.state.future_uses
+        }
+        if use.component_name in future_component_names:
+            return True
+        if self._should_keep_single_dit(use.component_name):
+            return True
+        return False
+
+    @lru_cache(maxsize=None)
+    def _should_keep_single_dit(self, component_name: str) -> bool:
+        modules = self.pipeline.modules
+        return (component_name == "transformer" and "transformer_2" not in modules) or (
+            component_name == "video_dit" and "video_dit_2" not in modules
+        )
+
+    def _preferred_request_end_use(self) -> ComponentUse | None:
+        """Returns a ComponentUse preferred to be resident after a request finishes, to prepare for next request"""
+        for uses in self._stage_uses_by_index:
+            for use in uses:
+                if use.preferred_ready_after_request:
+                    return use
+        for uses in self._stage_uses_by_index:
+            if uses:
+                return uses[0]
+        return None
+
+    def _preferred_request_end_uses(self) -> dict[str, ComponentUse]:
+        preferred_uses: dict[str, ComponentUse] = {}
+        for uses in self._stage_uses_by_index:
+            for use in uses:
+                if use.preferred_ready_after_request:
+                    preferred_uses[use.component_name] = use
+        for use in self._uses_seen.values():
+            if use.preferred_ready_after_request:
+                preferred_uses[use.component_name] = use
+        if preferred_uses:
+            return preferred_uses
+        preferred_use = self._preferred_request_end_use()
+        if preferred_use is None:
+            return {}
+        return {preferred_use.component_name: preferred_use}
+
+    @staticmethod
+    def _same_use(lhs: ComponentUse, rhs: ComponentUse) -> bool:
+        return lhs.component_name == rhs.component_name and lhs.phase == rhs.phase
+
+    @staticmethod
+    def _use_key(use: ComponentUse) -> tuple[str, str, str | None]:
+        return (use.stage_name, use.component_name, use.phase)
+
+    def _trace(
+        self,
+        action: str,
+        use: ComponentUse | None = None,
+        strategy: ComponentResidencyStrategy | None = None,
+        module: nn.Module | None = None,
+        *,
+        component_name: str | None = None,
+        detail: str = "",
+    ) -> None:
+        if not self.state.trace_enabled:
+            return
+        if use is not None:
+            component_name = use.component_name
+        device = self._module_device(module)
+        logger.info(
+            "[component_residency] action=%s stage=%s next_stage=%s component=%s "
+            "strategy=%s phase=%s device=%s warmup=%s mode=%s %s",
+            action,
+            self.state.stage_name,
+            self.state.next_stage_name,
+            component_name,
+            strategy.name if strategy is not None else None,
+            use.phase if use is not None else None,
+            device,
+            self.state.batch_is_warmup,
+            self.state.manager_mode,
+            detail,
+        )
+
+    def _module_device(self, module: nn.Module | None) -> str | None:
+        if module is None:
+            return None
+        param = next(module.parameters(), None)
+        if param is not None:
+            return param.device.type
+        buffer = next(module.buffers(), None)
+        return buffer.device.type if buffer is not None else None
+
+
+_GLOBAL_COMPONENT_RESIDENCY_MANAGER: ComponentResidencyManager | None = None
+
+
+def get_global_component_residency_manager(
+    pipeline: ComponentResidencyPipeline,
+    server_args: ServerArgs,
+) -> ComponentResidencyManager:
+    global _GLOBAL_COMPONENT_RESIDENCY_MANAGER
+
+    if _GLOBAL_COMPONENT_RESIDENCY_MANAGER is None:
+        _GLOBAL_COMPONENT_RESIDENCY_MANAGER = ComponentResidencyManager(
+            pipeline, server_args
+        )
+    else:
+        _GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_server_args(server_args)
+    _GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_pipeline(pipeline)
+
+    return _GLOBAL_COMPONENT_RESIDENCY_MANAGER

python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py

@@ -0,0 +1,502 @@
+"""
+Basic Component Resident Strategy Utilities for defining usage of components, to let ComponentResidencyManager to coordinate
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.managers.component_manager import (
+        ComponentUse,
+        ResidencyState,
+    )
+
+logger = init_logger(__name__)
+
+
+def _module_to_local_device(
+    module: nn.Module, *, dtype: torch.dtype | None = None
+) -> None:
+    device = get_local_torch_device()
+    tensor = _module_reference_tensor(module)
+    if tensor is not None and tensor.device == device:
+        if dtype is None or tensor.dtype == dtype:
+            return
+    if dtype is None:
+        module.to(device, non_blocking=True)
+    else:
+        module.to(device, dtype=dtype, non_blocking=True)
+
+
+def _module_reference_tensor(module: nn.Module) -> torch.Tensor | None:
+    tensor = next(module.parameters(), None)
+    if tensor is None:
+        tensor = next(module.buffers(), None)
+    return tensor
+
+
+def _module_ready_on_local_device(
+    module: nn.Module, *, dtype: torch.dtype | None = None
+) -> bool:
+    tensor = _module_reference_tensor(module)
+    if tensor is None:
+        return True
+    if tensor.device != get_local_torch_device():
+        return False
+    return dtype is None or tensor.dtype == dtype
+
+
+class ComponentResidencyStrategy:
+    """Baseclass for describing how a component should be treated (regarding where its weights locates)
+
+    e.g., a LayerwiseOffloadStrategy would override:
+        enter: to prefetch some layers before DiT is used, and
+        exits: to release GPU weight snapshot after DiT is used
+    to achieve desired behavior
+
+    """
+
+    name = "resident"
+
+    def prepare_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        """hook called"""
+        self.enter(module)
+
+    def wait_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        """Wait for the preparation to be ready, only applicable for async device syncs"""
+        pass
+
+    def finish_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        """Finish a specific component use"""
+        self.exit(module)
+
+    def prepare_after_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        """Called after a request is finished, to prepare for the upcoming request"""
+        pass
+
+    def finish_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+        *,
+        preferred: bool,
+    ) -> None:
+        if preferred:
+            self.prepare_for_use(module, use, state)
+            self.wait_for_use(module, use, state)
+        else:
+            self.finish_use(module, use, state)
+
+    def prefetch_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> bool:
+        self.prepare_for_use(module, use, state)
+        return True
+
+    def enter(self, module: nn.Module) -> None:
+        pass
+
+    def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
+        pass
+
+
+class ResidentStrategy(ComponentResidencyStrategy):
+    name = "resident"
+
+    def prepare_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        if use.target_dtype is not None:
+            _module_to_local_device(module, dtype=use.target_dtype)
+
+
+class SnapshotModuleResidency:
+    """Reusable snapshot-based module residency primitive.
+
+    This helper only knows how to:
+    - keep CPU parameter/buffer snapshots,
+    - prefetch a module (H2D) to the local device on a CUDA side stream
+    - release a module by rebinding tensors to those snapshots,
+    - track and wait for readiness events.
+
+    It deliberately does not know about pipeline stages, phases, or model-specific
+    ordering. Strategy subclasses decide when each primitive is called.
+    """
+
+    def __init__(self, *, pin_cpu_memory: bool, enable_async_prefetch: bool) -> None:
+        self.pin_cpu_memory = pin_cpu_memory
+        self.enable_async_prefetch = enable_async_prefetch
+        self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {}
+        self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {}
+        self._prefetch_stream: object | None = None
+        self._ready_events: dict[str, object] = {}
+
+    @staticmethod
+    def is_on_gpu(module: nn.Module | None) -> bool:
+        if module is None:
+            return False
+        param = next(module.parameters(), None)
+        return param is not None and param.device.type == "cuda"
+
+    def is_ready(self, component_name: str) -> bool:
+        return component_name in self._ready_events
+
+    def wait_ready(self, component_name: str) -> None:
+        """wait for the (H2D) stream to be ready"""
+        ready_event = self._ready_events.get(component_name)
+        if ready_event is None or not current_platform.is_cuda():
+            return
+        torch.get_device_module().current_stream().wait_event(ready_event)
+
+    def record_ready(self, component_name: str, module: nn.Module | None) -> None:
+        if not current_platform.is_cuda():
+            self._ready_events.pop(component_name, None)
+            return
+        if not self.is_on_gpu(module):
+            self._ready_events.pop(component_name, None)
+            return
+        event = torch.get_device_module().Event()
+        event.record(torch.get_device_module().current_stream())
+        self._ready_events[component_name] = event
+
+    @staticmethod
+    def _clone_cpu_tensor_snapshot(
+        tensor: torch.Tensor, *, pin_memory: bool
+    ) -> torch.Tensor:
+        snapshot = tensor.detach()
+        if snapshot.device.type == "cpu":
+            if pin_memory and not snapshot.is_pinned():
+                return snapshot.pin_memory()
+            return snapshot
+
+        cpu_tensor = snapshot.to("cpu")
+        if pin_memory:
+            return cpu_tensor.pin_memory()
+        return cpu_tensor
+
+    def _should_pin_memory(self) -> bool:
+        return bool(self.pin_cpu_memory and torch.get_device_module().is_available())
+
+    def capture(self, component_name: str, module: nn.Module) -> None:
+        """Capture a CPU snapshot for a component"""
+        if component_name in self._cpu_param_snapshots:
+            return
+
+        pin_memory = self._should_pin_memory()
+        self._cpu_param_snapshots[component_name] = {
+            name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory)
+            for name, param in module.named_parameters()
+        }
+        self._cpu_buffer_snapshots[component_name] = {
+            name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory)
+            for name, buffer in module.named_buffers()
+        }
+
+    def release_to_snapshot(
+        self,
+        component_name: str,
+        module: nn.Module,
+        *,
+        copy_runtime_buffers: bool = False,
+    ) -> None:
+        """Release CUDA storages by rebinding tensors to cached CPU snapshots.
+
+        This does not call `module.to("cpu")`. Instead, parameter and buffer
+        storages are rebound to pre-captured CPU tensors so CUDA storages can be
+        released by the allocator without an explicit D2H transfer.
+        """
+        param_snapshots = self._cpu_param_snapshots.get(component_name)
+        buffer_snapshots = self._cpu_buffer_snapshots.get(component_name)
+        if param_snapshots is None or buffer_snapshots is None:
+            module.to("cpu")
+            self._ready_events.pop(component_name, None)
+            return
+
+        pin_memory = self._should_pin_memory()
+        for name, param in module.named_parameters():
+            snapshot = param_snapshots.get(name)
+            if snapshot is None:
+                snapshot = self._clone_cpu_tensor_snapshot(
+                    param.data, pin_memory=pin_memory
+                )
+                param_snapshots[name] = snapshot
+            param.data = snapshot
+
+        for name, buffer in module.named_buffers():
+            snapshot = buffer_snapshots.get(name)
+            if snapshot is None:
+                snapshot = self._clone_cpu_tensor_snapshot(
+                    buffer.data, pin_memory=pin_memory
+                )
+                buffer_snapshots[name] = snapshot
+            if copy_runtime_buffers:
+                # Preserve runtime-updated buffers (e.g., lazily built caches) when
+                # releasing back to CPU snapshots.
+                if buffer.device.type == "cuda":
+                    snapshot.copy_(
+                        buffer.detach().to(device="cpu", dtype=snapshot.dtype)
+                    )
+                elif buffer.device.type == "cpu":
+                    snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype))
+            buffer.data = snapshot
+
+        self._ready_events.pop(component_name, None)
+
+    def _supports_async_prefetch(self) -> bool:
+        return self.enable_async_prefetch and current_platform.is_cuda()
+
+    def _get_prefetch_stream(self):
+        """returns a stream is async-prefetch is enabled"""
+        if not self._supports_async_prefetch():
+            return None
+        if self._prefetch_stream is None:
+            self._prefetch_stream = torch.get_device_module().Stream(
+                device=get_local_torch_device()
+            )
+        return self._prefetch_stream
+
+    def prefetch_to_device(self, component_name: str, module: nn.Module | None) -> None:
+        if module is None:
+            self._ready_events.pop(component_name, None)
+            return
+        prefetch_stream = self._get_prefetch_stream()
+        if prefetch_stream is None:
+            # if the async prefetching is disabled
+            module.to(get_local_torch_device(), non_blocking=True)
+            self.record_ready(component_name, module)
+            return
+        with torch.get_device_module().stream(prefetch_stream):
+            module.to(get_local_torch_device(), non_blocking=True)
+            event = torch.get_device_module().Event()
+            event.record(prefetch_stream)
+        self._ready_events[component_name] = event
+
+
+class SnapshotStrategy(ComponentResidencyStrategy):
+    """Snapshot residency: async H2D before use and light snapshot release after use."""
+
+    name = "snapshot"
+
+    def __init__(
+        self,
+        *,
+        pin_cpu_memory: bool,
+        enable_async_prefetch: bool,
+        copy_runtime_buffers_on_release: bool = False,
+    ) -> None:
+        self._snapshot_residency = SnapshotModuleResidency(
+            pin_cpu_memory=pin_cpu_memory,
+            enable_async_prefetch=enable_async_prefetch,
+        )
+        self._copy_runtime_buffers_on_release = copy_runtime_buffers_on_release
+
+    def capture(self, component_name: str, module: nn.Module) -> None:
+        self._snapshot_residency.capture(component_name, module)
+
+    def is_ready(self, component_name: str) -> bool:
+        return self._snapshot_residency.is_ready(component_name)
+
+    def record_ready(self, component_name: str, module: nn.Module | None) -> None:
+        self._snapshot_residency.record_ready(component_name, module)
+
+    def prefetch_component(self, component_name: str, module: nn.Module | None) -> None:
+        if SnapshotModuleResidency.is_on_gpu(module):
+            self._snapshot_residency.record_ready(component_name, module)
+            return
+        self._snapshot_residency.prefetch_to_device(component_name, module)
+
+    def wait_component_ready(self, component_name: str) -> None:
+        self._snapshot_residency.wait_ready(component_name)
+
+    def release_component(self, component_name: str, module: nn.Module) -> None:
+        self._snapshot_residency.release_to_snapshot(
+            component_name,
+            module,
+            copy_runtime_buffers=self._copy_runtime_buffers_on_release,
+        )
+
+    def prepare_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.prefetch_component(use.component_name, module)
+
+    def wait_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.wait_component_ready(use.component_name)
+
+    def finish_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.release_component(use.component_name, module)
+
+    def prepare_after_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.prepare_for_use(module, use, state)
+
+
+class VanillaD2HStrategy(ComponentResidencyStrategy):
+    """A strategy that performs native torch D2H and H2D for a component"""
+
+    name = "vanilla"
+
+    def __init__(self) -> None:
+        self._prefetch_stream: object | None = None
+        self._ready_events: dict[str, object] = {}
+
+    def prepare_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        _module_to_local_device(module, dtype=use.target_dtype)
+
+    def wait_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        ready_event = self._ready_events.get(use.component_name)
+        if ready_event is None or not current_platform.is_cuda():
+            return
+        torch.get_device_module().current_stream().wait_event(ready_event)
+
+    def prefetch_for_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> bool:
+        if not current_platform.is_cuda():
+            self.prepare_for_use(module, use, state)
+            return True
+        if _module_ready_on_local_device(module, dtype=use.target_dtype):
+            return True
+        if self._prefetch_stream is None:
+            self._prefetch_stream = torch.get_device_module().Stream(
+                device=get_local_torch_device()
+            )
+        with torch.get_device_module().stream(self._prefetch_stream):
+            _module_to_local_device(module, dtype=use.target_dtype)
+            event = torch.get_device_module().Event()
+            event.record(self._prefetch_stream)
+        self._ready_events[use.component_name] = event
+        return True
+
+    def enter(self, module: nn.Module) -> None:
+        param = next(module.parameters(), None)
+        if param is not None and param.device.type == "cpu":
+            _module_to_local_device(module)
+
+    def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
+        param = next(module.parameters(), None)
+        if param is not None and param.device.type == "cuda":
+            module.to("cpu", non_blocking=True)
+
+    def finish_use(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.wait_for_use(module, use, state)
+        self.exit(module)
+        self._ready_events.pop(use.component_name, None)
+
+    def prepare_after_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.prefetch_for_use(module, use, state)
+
+    def finish_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+        *,
+        preferred: bool,
+    ) -> None:
+        if preferred and state.batch_is_warmup:
+            self.prepare_for_use(module, use, state)
+            self.wait_for_use(module, use, state)
+            return
+        if not preferred:
+            self.finish_use(module, use, state)
+
+
+class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
+    """A wrapper around LayerwiseOffloadManager to fit in a ComponentResidencyStrategy"""
+
+    name = "layerwise"
+
+    def enter(self, module: nn.Module) -> None:
+        if isinstance(module, OffloadableDiTMixin):
+            module.prepare_for_next_req()
+
+    def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
+        if not isinstance(module, OffloadableDiTMixin):
+            return
+        for manager in module.layerwise_offload_managers:
+            manager.release_all()
+
+    def prepare_after_request(
+        self,
+        module: nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.prepare_for_use(module, use, state)

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -37,6 +37,10 @@ from sglang.multimodal_gen.runtime.loader.weights_updater import (
     WeightsUpdater,
     get_updatable_modules,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import (
+    OffloadableDiTMixin,
+    iter_materialized_weights,
+)
 from sglang.multimodal_gen.runtime.pipelines_core import (
     ComposedPipelineBase,
     LoRAPipeline,
@@ -47,10 +51,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
 from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
-    OffloadableDiTMixin,
-    iter_materialized_weights,
-)
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     configure_logger,
     globally_suppress_loggers,

python/sglang/multimodal_gen/runtime/managers/layerwise_offload.py

@@ -199,6 +199,11 @@ class LayerwiseOffloadManager:
 
                 self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer
 
+        # Keep non-layer parameters resident on GPU. Layer tensors have already
+        # been replaced by tiny device placeholders, so this does not reload the
+        # offloaded layer weights.
+        self.model.to(self.device)
+
         # prefetch the first layer for warm-up
         self.prepare_for_next_req(non_blocking=False)
 

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -467,22 +467,20 @@ class Scheduler(SchedulerDisaggMixin):
             reqs = [item[1] for item in items]
 
             try:
-                processed_req = reqs[0]
+                first_req = reqs[0]
-                if isinstance(processed_req, list) and processed_req:
+                if isinstance(first_req, list) and first_req:
-                    is_warmup = processed_req[0].is_warmup
+                    is_warmup = first_req[0].is_warmup
                 else:
                     is_warmup = (
-                        processed_req.is_warmup
+                        first_req.is_warmup if isinstance(first_req, Req) else False
-                        if isinstance(processed_req, Req)
-                        else False
                     )
 
-                handler = self.request_handlers.get(type(processed_req))
+                handler = self.request_handlers.get(type(first_req))
                 if handler:
                     output_batch = handler(reqs)
                 else:
                     output_batch = OutputBatch(
-                        error=f"Unknown request type: {type(processed_req)}"
+                        error=f"Unknown request type: {type(first_req)}"
                     )
             except Exception as e:
                 logger.error(
@@ -493,13 +491,11 @@ class Scheduler(SchedulerDisaggMixin):
 
             # 3. return results
             try:
-                if isinstance(processed_req, list) and processed_req:
+                if isinstance(first_req, list) and first_req:
-                    is_warmup = processed_req[0].is_warmup
+                    is_warmup = first_req[0].is_warmup
                 else:
                     is_warmup = (
-                        processed_req.is_warmup
+                        first_req.is_warmup if isinstance(first_req, Req) else False
-                        if isinstance(processed_req, Req)
-                        else False
                     )
                 if is_warmup:
                     if output_batch.error is None:

python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py

@@ -26,8 +26,8 @@ from sglang.multimodal_gen.runtime.layers.linear import (
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -13,7 +13,7 @@ from torch.nn.attention.flex_attention import (
     flex_attention,
 )
 
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 
 # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
 # see https://github.com/pytorch/pytorch/issues/133254

python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py

@@ -33,8 +33,8 @@ from sglang.multimodal_gen.runtime.layers.linear import (
     RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 
 
 def _rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -52,9 +52,9 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepTextProjEmbeddings,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import current_platform
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -42,12 +42,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     NDRotaryEmbedding,
     apply_flashinfer_rope_qk_inplace,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name

python/sglang/multimodal_gen/runtime/models/dits/glm_image.py

@@ -37,12 +37,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
 from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/helios.py

@@ -46,8 +46,8 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
     TimestepEmbedder,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py

@@ -22,9 +22,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
     RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py

@@ -37,13 +37,13 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
     unpatchify,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.utils import modulate
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 
 
 class MMDoubleStreamBlock(nn.Module):

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -31,9 +31,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
     QuantizationConfig,
 )
 from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py

@@ -18,8 +18,8 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 
 # Reuse common functions and classes from mova_video_dit
 from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d

python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py

@@ -33,9 +33,9 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import current_platform
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -46,12 +46,12 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name

python/sglang/multimodal_gen/runtime/models/dits/sana.py

@@ -8,8 +8,8 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbed
 from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -49,6 +49,7 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
     TimestepEmbedder,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.utils import (
     _use_aiter,
@@ -58,7 +59,6 @@ from sglang.multimodal_gen.runtime.platforms import (
     current_platform,
 )
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.utils import add_prefix
 

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -41,9 +41,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     _apply_rotary_emb,
     apply_flashinfer_rope_qk_inplace,
 )
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import current_platform
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 try:

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -16,6 +16,15 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
     PipelineComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB
+from sglang.multimodal_gen.runtime.managers.component_manager import (
+    ComponentResidencyStrategy,
+    ComponentUse,
+    ResidencyState,
+)
+from sglang.multimodal_gen.runtime.managers.component_resident_strategies import (
+    SnapshotModuleResidency,
+    SnapshotStrategy,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
@@ -321,44 +330,150 @@ class LTX2Pipeline(_BaseLTX2Pipeline):
         _add_ltx2_decoding_stage(self)
 
 
-class LTX2TwoStageDeviceManager:
+class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy):
-    """
+    name = "ltx2_original"
-    Device residency manager for LTX-2.3 two-stage DiT switching.
 
-    Modes:
+    def __init__(self, manager: "LTX2TwoStageResidencyController") -> None:
-    - resident: keep both DiTs on GPU; phase switch is pointer rebinding only.
+        self.manager = manager
-    - snapshot: keep CPU snapshots and prefetch the target DiT (DiT2 with pre-merged LoRA) with async H2D (similar to dit layerwise offload).
+
-      The DiT_1 will always be kept a replica in CPU.
+    @property
-      - default snapshot behavior: allow stage1/stage2 overlap by prefetching
+    def pipeline(self) -> "LTX2TwoStagePipeline":
-        stage2 while stage1 is still running.
+        return self.manager.pipeline
-      - snapshot low-VRAM behavior (`_snapshot_low_vram_mode=True`): evict
+
-        stage1 before stage2 prefetch and disable early overlap prefetch to
+    @property
-        reduce peak VRAM, at the cost of higher phase-switch latency.
+    def server_args(self) -> ServerArgs:
-      - default toggle: low-VRAM auto-enables on H100-like (<130 GiB) CUDA
+        return self.manager.server_args
-        GPUs, and stays disabled by default on higher-memory GPUs. It can be
+
-        overridden with `SGLANG_LTX2_SNAPSHOT_LOW_VRAM_MODE`.
+    def _phase(self, use: ComponentUse) -> str:
-    - original: official two-stage semantics without premerged stage-2.
+        if use.phase in ("stage1", "stage2"):
+            return use.phase
+        return "stage2" if use.component_name == "transformer_2" else "stage1"
+
+    def initialize(self) -> None:
+        pass
+
+    def prepare_for_use(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        phase = self._phase(use)
+        if phase != self.manager._active_phase:
+            self.enter_phase(phase)
+
+    def wait_for_use(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.ensure_phase_ready(self._phase(use))
+
+    def finish_use(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        self.exit_phase(self._phase(use))
+
+    def prepare_after_request(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        phase = self._phase(use)
+        if phase != self.manager._active_phase:
+            self.enter_phase(phase)
+
+    def enter_phase(self, phase: str) -> bool:
+        return False
+
+    def exit_phase(self, phase: str | None, next_phase: str | None = None) -> None:
+        pass
+
+    def ensure_phase_ready(self, phase: str | None) -> None:
+        """wait for the preparation to be ready"""
+        pass
+
+    def _ensure_on_gpu(self, module_name: str) -> None:
+        module = self.pipeline.get_module(module_name)
+        if module is None:
+            return
+        param = next(module.parameters(), None)
+        if param is not None and param.device.type == "cpu":
+            module.to(get_local_torch_device(), non_blocking=True)
+
+    @staticmethod
+    def _module_is_on_gpu(module: torch.nn.Module | None) -> bool:
+        return SnapshotModuleResidency.is_on_gpu(module)
+
+
+class LTX2OriginalResidencyStrategy(LTX2TwoStageResidencyStrategy):
+    pass
+
+
+class LTX2ResidentResidencyStrategy(LTX2TwoStageResidencyStrategy):
+    """A residency strategy for ltx two-stage pipeline with pre-merged lora, that keep both dits always resident"""
+
+    name = "ltx2_resident"
+
+    def initialize(self) -> None:
+        self._ensure_on_gpu("transformer")
+        self._ensure_on_gpu("transformer_2")
+        logger.info(
+            "Using resident LTX-2.3 two-stage transformers mode (both DiTs stay on GPU)"
+        )
+        self.manager._active_phase = "stage1"
+        self.manager._sync_refinement_stage_transformer("stage1")
+
+    def enter_phase(self, phase: str) -> bool:
+        self.manager._sync_refinement_stage_transformer(phase)
+        self.manager._active_phase = phase
+        return True
+
+
+class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
+    """
+    Snapshot mode keeps CPU snapshots and prefetches the target DiT with async H2D. (only with pre-merged lora enabled)
+
+    The DiT_1 will always be kept a replica in CPU.
+    - default snapshot behavior: allow stage1/stage2 overlap by prefetching
+      stage2 while stage1 is still running.
+    - snapshot low-VRAM behavior (`_snapshot_low_vram_mode=True`): evict
+      stage1 before stage2 prefetch and disable early overlap prefetch to
+      reduce peak VRAM, at the cost of higher phase-switch latency.
+    - default toggle: low-VRAM auto-enables on H100-like (<130 GiB) CUDA
+      GPUs, and stays disabled by default on higher-memory GPUs. It can be
+      overridden with `SGLANG_LTX2_SNAPSHOT_LOW_VRAM_MODE`.
     """
 
-    VALID_MODES = ("original", "snapshot", "resident")
+    name = "ltx2_snapshot"
 
-    def __init__(self, pipeline: "LTX2TwoStagePipeline", server_args: ServerArgs):
+    def __init__(self, manager: "LTX2TwoStageResidencyController") -> None:
-        self.pipeline = pipeline
+        super().__init__(manager)
-        self.server_args = server_args
+        self._snapshot_strategy = SnapshotStrategy(
-        self.mode = self._resolve_mode(server_args)
+            pin_cpu_memory=manager.server_args.pin_cpu_memory,
-        self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {}
+            enable_async_prefetch=manager.server_args.dit_cpu_offload,
-        self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {}
+        )
-        self._active_phase: str | None = None
-        self._prefetch_stream: object | None = None
-        self._phase_ready_events: dict[str, object] = {}
         self._snapshot_low_vram_mode = self._resolve_snapshot_low_vram_mode()
         self._snapshot_release_empty_cache = get_bool_env_var(
             "SGLANG_LTX2_SNAPSHOT_RELEASE_EMPTY_CACHE",
             default="false",
         )
 
+    @staticmethod
+    def _module_name_for_phase(phase: str | None) -> str | None:
+        if phase == "stage1":
+            return "transformer"
+        if phase == "stage2":
+            return "transformer_2"
+        return None
+
     def _resolve_snapshot_low_vram_mode(self) -> bool:
-        if self.mode != "snapshot" or not current_platform.is_cuda():
+        if not current_platform.is_cuda():
             return False
         device_name = str(current_platform.get_device_name(0)).upper()
         device_total_memory_gb = (
@@ -384,243 +499,100 @@ class LTX2TwoStageDeviceManager:
             )
         return enabled
 
-    @classmethod
-    def _resolve_mode(cls, server_args: ServerArgs) -> str:
-        mode = getattr(server_args, "ltx2_two_stage_device_mode", None)
-        if mode is None:
-            env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE")
-            mode = env_mode.lower() if env_mode else "snapshot"
-        if mode not in cls.VALID_MODES:
-            raise ValueError(
-                f"Invalid ltx2_two_stage_device_mode={mode!r}. "
-                f"Expected one of {cls.VALID_MODES}."
-            )
-        return mode
-
-    @property
-    def should_use_premerged(self) -> bool:
-        """Whether to keep a pre-merged stage-2 DiT for LTX-2.3 two-stage.
-
-        We only enable this optimization for native LTX-2.3 two-stage and when
-        users did not explicitly provide a stage-1 LoRA path
-        """
-        return (
-            self.mode != "original"
-            and self.pipeline._should_merge_stage2_distilled_lora(self.server_args)
-            and getattr(self.pipeline, "_stage1_lora_path", None) is None
-        )
-
     def initialize(self) -> None:
-        if not self.should_use_premerged:
+        # Snapshot mode keeps both DiT CPU snapshots for cheap GPU release
-            return
+        # and re-hydrates stage-2 with async H2D when stage-1 finishes.
-
+        self._capture_module_cpu_snapshot("transformer")
-        self.pipeline._initialize_premerged_stage2_transformer(self.server_args)
+        self._capture_module_cpu_snapshot("transformer_2")
-        if self.mode == "snapshot":
+        self._pin_stage1_transformer_if_beneficial()
-            # Snapshot mode keeps both DiT CPU snapshots for cheap GPU release
+        self.manager._sync_refinement_stage_transformer("stage1")
-            # and re-hydrates stage-2 with async H2D when stage-1 finishes.
+        self._record_component_ready("transformer")
-            self._capture_module_cpu_snapshot("transformer")
-            self._capture_module_cpu_snapshot("transformer_2")
-            self._pin_stage1_transformer_if_beneficial()
-        elif self.mode == "resident":
-            self._ensure_on_gpu("transformer")
-            self._ensure_on_gpu("transformer_2")
-            logger.info(
-                "Using resident LTX-2.3 two-stage transformers mode (both DiTs stay on GPU)"
-            )
-            self._active_phase = "stage1"
-
-        self._sync_refinement_stage_transformer("stage1")
-        self._record_phase_ready_event("stage1")
-
-    def switch_phase(self, phase: str) -> bool:
-        """Switch active two-stage DiT with minimal transfer/sync overhead."""
-        if not self.should_use_premerged:
-            return False
-        if phase == self._active_phase:
-            return True
-
-        if self.mode == "resident":
-            self._sync_refinement_stage_transformer(phase)
-            self._active_phase = phase
-            return True
 
+    def enter_phase(self, phase: str) -> bool:
         if self.server_args.dit_cpu_offload:
-            target_name = "transformer_2" if phase == "stage2" else "transformer"
+            target_module_name = self._module_name_for_phase(phase)
-            target_module = self.pipeline.get_module(target_name)
+            if target_module_name is None:
-            if self.mode == "snapshot" and self._snapshot_low_vram_mode:
+                return False
+            target_module = self.pipeline.get_module(target_module_name)
+            if self._snapshot_low_vram_mode:
                 # Trade a bit of phase-switch latency for lower peak VRAM:
                 # evict stage-1 before stage-2 H2D.
-                if phase == "stage2" and phase not in self._phase_ready_events:
+                if phase == "stage2" and not self._snapshot_strategy.is_ready(
-                    stage1_module = self.pipeline.get_module("transformer")
+                    target_module_name
-                    stage1_param = (
+                ):
-                        next(stage1_module.parameters(), None)
+                    self._release_stage1_for_low_vram()
-                        if stage1_module is not None
+
-                        else None
+            # make sure the component is pre-fetched
-                    )
+            if not self._snapshot_strategy.is_ready(target_module_name):
-                    if stage1_param is not None and stage1_param.device.type == "cuda":
-                        self._release_module_to_cpu_snapshot("transformer")
-            if phase not in self._phase_ready_events:
                 if self._module_is_on_gpu(target_module):
-                    self._record_phase_ready_event(phase)
+                    self._record_component_ready(target_module_name)
                 else:
-                    self._schedule_phase_prefetch(phase, target_module)
+                    self._snapshot_strategy.prefetch_component(
-
+                        target_module_name, target_module
-            # Stage-2 is only consumed after stage-1 denoising + upsample.
+                    )
-            # Kick off the H2D early in stage-1 to overlap transfer with compute.
-            if (
-                phase == "stage1"
-                and "stage2" not in self._phase_ready_events
-                and not self._snapshot_low_vram_mode
-            ):
-                self._schedule_phase_prefetch(
-                    "stage2", self.pipeline.get_module("transformer_2")
-                )
         else:
-            self._record_phase_ready_event(phase)
+            component_name = self._module_name_for_phase(phase)
+            if component_name is not None:
+                self._record_component_ready(component_name)
 
-        self._sync_refinement_stage_transformer(phase)
+        self.manager._sync_refinement_stage_transformer(phase)
-        self._active_phase = phase
+        self.manager._active_phase = phase
         return True
 
-    def prefetch_stage2_after_stage1(self) -> None:
+    def prepare_after_request(
-        """Kick off stage-2 H2D right after stage-1 denoising to hide switch latency."""
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        phase = self._phase(use)
+        if phase != "stage1":
+            return
+        if self.server_args.dit_cpu_offload:
+            target_module = self.pipeline.get_module("transformer")
+            if self._module_is_on_gpu(target_module):
+                self._record_component_ready("transformer")
+            elif not self._snapshot_strategy.is_ready("transformer"):
+                self._snapshot_strategy.prefetch_component("transformer", target_module)
+        else:
+            self._record_component_ready("transformer")
+        self.manager._sync_refinement_stage_transformer("stage1")
+        self.manager._active_phase = "stage1"
+
+    def finish_use(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> None:
+        phase = self._phase(use)
+        if self.server_args.dit_cpu_offload:
+            # release cuda storage
+            self._snapshot_strategy.release_component(use.component_name, module)
         if (
-            not self.should_use_premerged
+            phase == "stage2"
-            or self.mode != "snapshot"
+            and self._snapshot_release_empty_cache
-            or not self.server_args.dit_cpu_offload
-        ):
-            return
-
-        if "stage2" in self._phase_ready_events:
-            return
-        if self._snapshot_low_vram_mode:
-            self._release_stage1_for_low_vram()
-
-        self._schedule_phase_prefetch(
-            "stage2", self.pipeline.get_module("transformer_2")
-        )
-
-    def prepare_upsample_after_stage1(self) -> bool:
-        if (
-            not self.should_use_premerged
-            or self.mode != "snapshot"
-            or not self.server_args.dit_cpu_offload
-            or not self._snapshot_low_vram_mode
-        ):
-            return False
-        if "stage2" in self._phase_ready_events:
-            return False
-        self._release_stage1_for_low_vram()
-        return True
-
-    def ensure_phase_ready(self, phase: str | None) -> None:
-        if not self.should_use_premerged or phase not in ("stage1", "stage2"):
-            return
-        if self.mode == "resident":
-            return
-        ready_event = self._phase_ready_events.get(phase)
-        if ready_event is None or not current_platform.is_cuda():
-            return
-        torch.get_device_module().current_stream().wait_event(ready_event)
-
-    def release_premerged_transformers(self) -> None:
-        if not self.should_use_premerged or self.mode != "snapshot":
-            return
-        # Keep stage-1 resident across requests so the next request can start
-        # denoising immediately while stage-2 is prefetched in the background.
-        for module_name in ("transformer_2",):
-            module = self.pipeline.get_module(module_name)
-            param = next(module.parameters(), None) if module is not None else None
-            if param is not None and param.device.type == "cuda":
-                self._release_module_to_cpu_snapshot(module_name)
-        if (
-            self._snapshot_release_empty_cache
             and torch.get_device_module().is_available()
         ):
             torch.get_device_module().empty_cache()
-        self._record_phase_ready_event("stage1")
 
-    @staticmethod
+    def ensure_phase_ready(self, phase: str | None) -> None:
-    def _clone_cpu_tensor_snapshot(
+        component_name = self._module_name_for_phase(phase)
-        tensor: torch.Tensor, *, pin_memory: bool
+        if component_name is None:
-    ) -> torch.Tensor:
+            return
-        snapshot = tensor.detach()
+        self._snapshot_strategy.wait_component_ready(component_name)
-        if snapshot.device.type == "cpu":
-            if pin_memory and not snapshot.is_pinned():
-                return snapshot.pin_memory()
-            return snapshot
-
-        cpu_tensor = snapshot.to("cpu")
-        if pin_memory:
-            return cpu_tensor.pin_memory()
-        return cpu_tensor
 
     def _capture_module_cpu_snapshot(self, module_name: str) -> None:
-        if module_name in self._cpu_param_snapshots:
-            return
-
         module = self.pipeline.get_module(module_name)
         if module is None:
             raise ValueError(f"Module {module_name} is not available.")
-
+        self._snapshot_strategy.capture(module_name, module)
-        pin_memory = bool(
-            self.server_args.pin_cpu_memory and torch.get_device_module().is_available()
-        )
-        self._cpu_param_snapshots[module_name] = {
-            name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory)
-            for name, param in module.named_parameters()
-        }
-        self._cpu_buffer_snapshots[module_name] = {
-            name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory)
-            for name, buffer in module.named_buffers()
-        }
 
     def _release_module_to_cpu_snapshot(self, module_name: str) -> None:
-        """Replace module tensors with cached CPU snapshots to avoid D2H copies.
-
-        This does not call `module.to("cpu")`. Instead, parameter and buffer storages
-        are rebound to pre-captured CPU tensors so CUDA storages can be released by
-        the allocator without an explicit D2H transfer.
-        """
         module = self.pipeline.get_module(module_name)
         if module is None:
             return
-
+        self._snapshot_strategy.release_component(module_name, module)
-        param_snapshots = self._cpu_param_snapshots.get(module_name)
-        buffer_snapshots = self._cpu_buffer_snapshots.get(module_name)
-        if param_snapshots is None or buffer_snapshots is None:
-            module.to("cpu")
-            return
-
-        pin_memory = bool(
-            self.server_args.pin_cpu_memory and torch.get_device_module().is_available()
-        )
-        for name, param in module.named_parameters():
-            snapshot = param_snapshots.get(name)
-            if snapshot is None:
-                snapshot = self._clone_cpu_tensor_snapshot(
-                    param.data, pin_memory=pin_memory
-                )
-                param_snapshots[name] = snapshot
-            param.data = snapshot
-
-        for name, buffer in module.named_buffers():
-            snapshot = buffer_snapshots.get(name)
-            if snapshot is None:
-                snapshot = self._clone_cpu_tensor_snapshot(
-                    buffer.data, pin_memory=pin_memory
-                )
-                buffer_snapshots[name] = snapshot
-            # Preserve runtime-updated buffers (e.g., lazily built caches) when
-            # releasing back to CPU snapshots.
-            if buffer.device.type == "cuda":
-                snapshot.copy_(buffer.detach().to(device="cpu", dtype=snapshot.dtype))
-            elif buffer.device.type == "cpu":
-                snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype))
-            buffer.data = snapshot
-
-        phase = "stage2" if module_name == "transformer_2" else "stage1"
-        self._phase_ready_events.pop(phase, None)
 
     def _release_stage1_for_low_vram(self) -> None:
         stage1_module = self.pipeline.get_module("transformer")
@@ -632,66 +604,29 @@ class LTX2TwoStageDeviceManager:
         if stage1_param is not None and stage1_param.device.type == "cuda":
             self._release_module_to_cpu_snapshot("transformer")
 
-    def _ensure_on_gpu(self, module_name: str) -> None:
+    def _record_component_ready(self, module_name: str) -> None:
-        module = self.pipeline.get_module(module_name)
+        self._snapshot_strategy.record_ready(
-        if module is None:
+            module_name, self.pipeline.get_module(module_name)
-            return
-        param = next(module.parameters(), None)
-        if param is not None and param.device.type == "cpu":
-            module.to(get_local_torch_device(), non_blocking=True)
-
-    @staticmethod
-    def _module_is_on_gpu(module: torch.nn.Module | None) -> bool:
-        if module is None:
-            return False
-        param = next(module.parameters(), None)
-        return param is not None and param.device.type == "cuda"
-
-    def _supports_async_phase_prefetch(self) -> bool:
-        return (
-            self.mode == "snapshot"
-            and self.server_args.dit_cpu_offload
-            and current_platform.is_cuda()
         )
 
-    def _get_prefetch_stream(self):
+    def prefetch_for_use(
-        if not self._supports_async_phase_prefetch():
+        self,
-            return None
+        module: torch.nn.Module,
-        if self._prefetch_stream is None:
+        use: ComponentUse,
-            self._prefetch_stream = torch.get_device_module().Stream(
+        state: ResidencyState,
-                device=get_local_torch_device()
+    ) -> bool:
-            )
+        if not self.server_args.dit_cpu_offload:
-        return self._prefetch_stream
+            return True
-
+        phase = self._phase(use)
-    def _record_phase_ready_event(self, phase: str) -> None:
+        if phase == "stage2":
-        if not current_platform.is_cuda():
+            if self._snapshot_strategy.is_ready("transformer_2"):
-            self._phase_ready_events.pop(phase, None)
+                return True
-            return
+            if self._snapshot_low_vram_mode and state.current_use is not None:
-        module_name = "transformer_2" if phase == "stage2" else "transformer"
+                return False
-        module = self.pipeline.get_module(module_name)
+            if self._snapshot_low_vram_mode:
-        if not self._module_is_on_gpu(module):
+                self._release_stage1_for_low_vram()
-            self._phase_ready_events.pop(phase, None)
+        self._snapshot_strategy.prefetch_component(use.component_name, module)
-            return
+        return True
-        event = torch.get_device_module().Event()
-        event.record(torch.get_device_module().current_stream())
-        self._phase_ready_events[phase] = event
-
-    def _schedule_phase_prefetch(
-        self, phase: str, module: torch.nn.Module | None
-    ) -> None:
-        if module is None:
-            self._phase_ready_events.pop(phase, None)
-            return
-        prefetch_stream = self._get_prefetch_stream()
-        if prefetch_stream is None:
-            module.to(get_local_torch_device(), non_blocking=True)
-            self._record_phase_ready_event(phase)
-            return
-        with torch.get_device_module().stream(prefetch_stream):
-            module.to(get_local_torch_device(), non_blocking=True)
-            event = torch.get_device_module().Event()
-            event.record(prefetch_stream)
-        self._phase_ready_events[phase] = event
 
     def _pin_stage1_transformer_if_beneficial(self) -> None:
         """Optionally pin stage-1 DiT on GPU to remove first-stage cold H2D stall.
@@ -717,7 +652,80 @@ class LTX2TwoStageDeviceManager:
             logger.info(
                 "Pinned stage1 transformer on GPU for LTX-2.3 two-stage startup"
             )
-        self._active_phase = "stage1"
+        self.manager._active_phase = "stage1"
+
+
+class LTX2TwoStageResidencyController:
+    """
+    LTX-2.3 two-stage residency controller.
+    It builds the selected LTX2 ComponentResidencyStrategy and keeps the
+    thin stage adapter methods that are specific to two-stage LoRA flow.
+
+    Modes:
+    - resident: keep both DiTs on GPU; phase switch is pointer rebinding only.
+    - snapshot: keep CPU snapshots and prefetch the target DiT.
+    - original: official two-stage semantics without premerged stage-2.
+    """
+
+    VALID_MODES = ("original", "snapshot", "resident")
+
+    def __init__(self, pipeline: "LTX2TwoStagePipeline", server_args: ServerArgs):
+        self.pipeline = pipeline
+        self.server_args = server_args
+        self.mode = self._resolve_mode(server_args)
+        self._active_phase: str | None = None
+        self._strategy = self._build_strategy()
+
+    @classmethod
+    def _resolve_mode(cls, server_args: ServerArgs) -> str:
+        mode = server_args.ltx2_two_stage_device_mode
+        if mode is None:
+            env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE")
+            mode = env_mode.lower() if env_mode else "snapshot"
+        if mode not in cls.VALID_MODES:
+            raise ValueError(
+                f"Invalid ltx2_two_stage_device_mode={mode!r}. "
+                f"Expected one of {cls.VALID_MODES}."
+            )
+        return mode
+
+    def _build_strategy(self) -> LTX2TwoStageResidencyStrategy:
+        if self.mode == "snapshot":
+            return LTX2SnapshotResidencyStrategy(self)
+        if self.mode == "resident":
+            return LTX2ResidentResidencyStrategy(self)
+        return LTX2OriginalResidencyStrategy(self)
+
+    @property
+    def strategy(self) -> ComponentResidencyStrategy:
+        return self._strategy
+
+    @property
+    def should_use_premerged(self) -> bool:
+        """Whether to keep a pre-merged stage-2 DiT for LTX-2.3 two-stage.
+
+        We only enable this optimization for native LTX-2.3 two-stage and when
+        users did not explicitly provide a stage-1 LoRA path
+        """
+        return (
+            self.mode != "original"
+            and self.pipeline._should_merge_stage2_distilled_lora(self.server_args)
+            and self.pipeline._stage1_lora_path is None
+        )
+
+    def initialize(self) -> None:
+        if not self.should_use_premerged:
+            return
+        self.pipeline._initialize_premerged_stage2_transformer(self.server_args)
+        self._strategy.initialize()
+
+    def enter_phase(self, phase: str) -> bool:
+        """Switch active two-stage DiT with minimal transfer/sync overhead."""
+        if not self.should_use_premerged:
+            return False
+        if phase == self._active_phase:
+            return True
+        return self._strategy.enter_phase(phase)
 
     def _sync_refinement_stage_transformer(self, phase: str) -> None:
         """Keep stage-2 refinement bound to the expected DiT for current phase."""
@@ -740,11 +748,18 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._device_manager = LTX2TwoStageDeviceManager(self, self.server_args)
+        self._ltx2_residency = LTX2TwoStageResidencyController(self, self.server_args)
         self._use_premerged_stage2_transformer = (
-            self._device_manager.should_use_premerged
+            self._ltx2_residency.should_use_premerged
         )
-        self._device_manager.initialize()
+        self._ltx2_residency.initialize()
+        if self._use_premerged_stage2_transformer:
+            self.component_residency_strategies["transformer"] = (
+                self._ltx2_residency.strategy
+            )
+            self.component_residency_strategies["transformer_2"] = (
+                self._ltx2_residency.strategy
+            )
 
     @staticmethod
     def _should_merge_stage2_distilled_lora(server_args: ServerArgs) -> bool:
@@ -809,25 +824,8 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
             merge_weights=True,
         )
 
-    def release_premerged_transformers_to_cpu_snapshots(self) -> None:
-        """Release inactive premerged DiTs according to the selected device mode."""
-        self._device_manager.release_premerged_transformers()
-
-    def release_ltx2_phase_state(self, phase: str | None) -> None:
-        if phase == "stage2":
-            self.release_premerged_transformers_to_cpu_snapshots()
-
-    def ensure_ltx2_phase_ready(self, phase: str | None) -> None:
-        self._device_manager.ensure_phase_ready(phase)
-
-    def prefetch_ltx2_stage2_after_stage1(self) -> None:
-        self._device_manager.prefetch_stage2_after_stage1()
-
-    def prepare_ltx2_upsample_after_stage1(self) -> bool:
-        return self._device_manager.prepare_upsample_after_stage1()
-
     def should_skip_ltx2_lora_switch_stage(self) -> bool:
-        return self._use_premerged_stage2_transformer and self._device_manager.mode in (
+        return self._use_premerged_stage2_transformer and self._ltx2_residency.mode in (
             "snapshot",
             "resident",
         )
@@ -912,7 +910,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
         if phase_signature == self._active_lora_signature:
             return
 
-        if self._device_manager.switch_phase(
+        if self._ltx2_residency.enter_phase(
             phase
         ) and self._can_short_circuit_lora_switch(phase, batch):
             self._active_lora_phase = phase

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -21,6 +21,11 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import (
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     PipelineComponentLoader,
 )
+from sglang.multimodal_gen.runtime.managers.component_manager import (
+    ComponentResidencyManager,
+    ComponentResidencyStrategy,
+    get_global_component_residency_manager,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
     PipelineExecutor,
 )
@@ -91,7 +96,9 @@ class ComposedPipelineBase(ABC):
         self.model_path: str = model_path
         self._stages: list[PipelineStage] = []
         self._stage_name_mapping: dict[str, PipelineStage] = {}
+        self.component_residency_strategies: dict[str, ComponentResidencyStrategy] = {}
         self.executor = executor or self.build_executor(server_args=server_args)
+        self.component_residency_manager: ComponentResidencyManager | None = None
 
         if required_config_modules is not None:
             self._required_config_modules = required_config_modules
@@ -738,6 +745,11 @@ class ComposedPipelineBase(ABC):
                 main_process_only=True,
             )
 
+        self.component_residency_manager = get_global_component_residency_manager(
+            self, server_args
+        )
+        self.executor.component_residency_manager = self.component_residency_manager
+
         return self.executor.execute_with_profiling(self.stages, batch, server_args)
 
     @torch.no_grad()

python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py

@@ -55,7 +55,7 @@ class ParallelExecutor(PipelineExecutor):
     def _execute_stages(
         self,
         stages: List[PipelineStage],
-        payload: Any,
+        batch: Any,
         server_args: ServerArgs,
         run_stage: Callable[[PipelineStage, Any], Any],
     ) -> Any:
@@ -66,30 +66,40 @@ class ParallelExecutor(PipelineExecutor):
             rank = get_world_rank()
         cfg_group = get_cfg_group()
 
-        # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY
+        self.begin_component_residency_request(stages, batch, server_args)
-        for stage in stages:
+        try:
-            paradigm = stage.parallelism_type
+            # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY
+            for stage_index, stage in enumerate(stages):
+                paradigm = stage.parallelism_type
 
-            if paradigm == StageParallelismType.MAIN_RANK_ONLY:
+                if paradigm == StageParallelismType.MAIN_RANK_ONLY:
-                if rank == 0:
+                    if rank == 0:
-                    # Only main rank executes, others just wait
+                        # Only main rank executes, others just wait
-                    payload = run_stage(stage, payload)
+                        self.before_stage(stage, stage_index, batch, server_args)
-                torch.distributed.barrier()
+                        batch = stage(batch, server_args)
+                        self.after_stage(stage_index)
+                    torch.distributed.barrier()
 
-            elif paradigm == StageParallelismType.CFG_PARALLEL:
+                elif paradigm == StageParallelismType.CFG_PARALLEL:
-                obj_list = [payload] if rank == 0 else []
+                    obj_list = [batch] if rank == 0 else []
-                broadcasted_list = broadcast_pyobj(
+                    broadcasted_list = broadcast_pyobj(
-                    obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0
+                        obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0
-                )
+                    )
-                if rank != 0:
+                    if rank != 0:
-                    payload = broadcasted_list[0]
+                        batch = broadcasted_list[0]
-                payload = run_stage(stage, payload)
+                    self.before_stage(stage, stage_index, batch, server_args)
+                    batch = stage(batch, server_args)
+                    self.after_stage(stage_index)
 
-                torch.distributed.barrier()
+                    torch.distributed.barrier()
 
-            elif paradigm == StageParallelismType.REPLICATED:
+                elif paradigm == StageParallelismType.REPLICATED:
-                payload = run_stage(stage, payload)
+                    self.before_stage(stage, stage_index, batch, server_args)
-        return payload
+                    batch = stage(batch, server_args)
+                    self.after_stage(stage_index)
+        finally:
+            self.finish_component_residency_request()
+        return batch
 
     def execute(
         self,

python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py

@@ -45,6 +45,33 @@ class PipelineExecutor(ABC):
 
     def __init__(self, server_args):
         self.server_args = server_args
+        self.component_residency_manager = None
+
+    def begin_component_residency_request(
+        self,
+        stages: List["PipelineStage"],
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> None:
+        self.component_residency_manager.begin_request(stages, batch, server_args)
+
+    def before_stage(
+        self,
+        stage: "PipelineStage",
+        stage_index: int,
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> None:
+        stage.set_component_residency_manager(self.component_residency_manager)
+        self.component_residency_manager.before_stage(
+            stage, stage_index, batch, server_args
+        )
+
+    def after_stage(self, stage_index: int) -> None:
+        self.component_residency_manager.after_stage(stage_index)
+
+    def finish_component_residency_request(self) -> None:
+        self.component_residency_manager.finish_request()
 
     def execute_with_profiling(
         self,

python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py

@@ -29,11 +29,17 @@ class SyncExecutor(PipelineExecutor):
         run_stage: Callable[[PipelineStage, Any], Any],
     ) -> Any:
         """Execute all pipeline stages sequentially and step the profiler."""
-        for stage in stages:
+        self.begin_component_residency_request(stages, payload, server_args)
-            payload = run_stage(stage, payload)
+        try:
-            profiler = SGLDiffusionProfiler.get_instance()
+            for stage_index, stage in enumerate(stages):
-            if profiler:
+                self.before_stage(stage, stage_index, payload, server_args)
-                profiler.step_stage()
+                payload = run_stage(stage, payload)
+                self.after_stage(stage_index)
+                profiler = SGLDiffusionProfiler.get_instance()
+                if profiler:
+                    profiler.step_stage()
+        finally:
+            self.finish_component_residency_request()
         return payload
 
     def run_profile_all_stages(

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -169,7 +169,7 @@ class LoRAPipeline(ComposedPipelineBase):
         Yields:
             List of modules that had offload disabled.
         """
-        from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+        from sglang.multimodal_gen.runtime.managers.layerwise_offload import (
             OffloadableDiTMixin,
         )
 

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -69,8 +69,8 @@ class Req:
     # Primary encoder embeddings
     prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list)
     negative_prompt_embeds: list[torch.Tensor] | None = None
-    prompt_attention_mask: list[torch.Tensor] | None = None
+    prompt_attention_mask: list[torch.Tensor | None] | None = None
-    negative_attention_mask: list[torch.Tensor] | None = None
+    negative_attention_mask: list[torch.Tensor | None] | None = None
     clip_embedding_pos: list[torch.Tensor] | None = None
     clip_embedding_neg: list[torch.Tensor] | None = None
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -9,11 +9,15 @@ composed to create complete diffusion pipelines.
 """
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import replace
 from enum import Enum, auto
 
 import torch
 
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.dedup import StageDedupMixin
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
@@ -53,6 +57,7 @@ class PipelineStage(StageDedupMixin, ABC):
 
     def __init__(self):
         self.server_args = get_global_server_args()
+        self._component_residency_manager = None
 
     def log_info(self, msg, *args):
         """Logs an informational message with the stage name as a prefix."""
@@ -105,6 +110,81 @@ class PipelineStage(StageDedupMixin, ABC):
         """
         pass
 
+    def set_component_residency_manager(self, manager) -> None:
+        self._component_residency_manager = manager
+
+    def _component_stage_name(self, stage_name: str | None = None) -> str:
+        return stage_name or self.__class__.__name__
+
+    def _active_component_stage_name(self) -> str:
+        manager = self._component_residency_manager
+        if manager is not None and manager.state.stage_name is not None:
+            return manager.state.stage_name
+        return self.__class__.__name__
+
+    def _finish_active_component_use(self) -> None:
+        if self._component_residency_manager is not None:
+            self._component_residency_manager.finish_active_use()
+
+    @contextmanager
+    def _use_component(
+        self,
+        use: ComponentUse,
+        module=None,
+    ) -> Iterator[object | None]:
+        if self._component_residency_manager is None:
+            yield module
+            return
+        with self._component_residency_manager.use_component(use, module) as component:
+            yield component
+
+    def _declared_component_use(
+        self,
+        *,
+        component_name: str,
+        phase: str | None = None,
+        target_dtype: torch.dtype | None = None,
+    ) -> ComponentUse:
+        manager = self._component_residency_manager
+        stage_name = self._active_component_stage_name()
+        server_args = manager.server_args if manager is not None else self.server_args
+        for use in self.component_uses(server_args, stage_name):
+            if use.component_name != component_name:
+                continue
+            if phase is not None and use.phase != phase:
+                continue
+            if target_dtype is not None:
+                return replace(use, target_dtype=target_dtype)
+            return use
+        raise ValueError(
+            f"{self.__class__.__name__} did not declare component use: "
+            f"{component_name}"
+        )
+
+    @contextmanager
+    def use_declared_component(
+        self,
+        *,
+        component_name: str,
+        module=None,
+        phase: str | None = None,
+        target_dtype: torch.dtype | None = None,
+    ) -> Iterator[object | None]:
+        """reference a component already declared in `component_uses`"""
+        use = self._declared_component_use(
+            component_name=component_name,
+            phase=phase,
+            target_dtype=target_dtype,
+        )
+        with self._use_component(use, module) as component:
+            yield component
+
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        """Declares component uses of current stage for unified residency scheduling."""
+        return []
+
     # Default role affinity: ENCODER. Override in subclasses for DENOISING/DECODER.
     @property
     def role_affinity(self) -> RoleType:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -11,6 +11,7 @@ import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import VAELoader
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
@@ -68,6 +69,20 @@ class DecodingStage(PipelineStage):
         self.pipeline = weakref.ref(pipeline) if pipeline else None
         self.component_name = component_name
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name,
+                self.component_name,
+                target_dtype=vae_dtype,
+                keep_ready_after_warmup=True,
+            )
+        ]
+
     @property
     def parallelism_type(self) -> StageParallelismType:
         if get_global_server_args().enable_cfg_parallel:
@@ -110,7 +125,13 @@ class DecodingStage(PipelineStage):
         return latents
 
     @torch.no_grad()
-    def decode(self, latents: torch.Tensor, server_args: ServerArgs) -> torch.Tensor:
+    def decode(
+        self,
+        latents: torch.Tensor,
+        server_args: ServerArgs,
+        *,
+        vae_dtype: torch.dtype,
+    ) -> torch.Tensor:
         """
         Decode latent representations into pixel space using VAE.
 
@@ -125,8 +146,6 @@ class DecodingStage(PipelineStage):
             Decoded video tensor with shape (batch, channels, frames, height, width),
             normalized to [0, 1] range and moved to CPU as float32
         """
-        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
-        self.vae = self.vae.to(device=get_local_torch_device(), dtype=vae_dtype)
         latents = latents.to(get_local_torch_device())
         vae_autocast_enabled = (
             vae_dtype != torch.float32
@@ -175,22 +194,6 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module(self.component_name, self.vae)
             self.server_args.model_loaded[self.component_name] = True
 
-    def offload_model(self):
-        # Offload models if needed
-        self.maybe_free_model_hooks()
-
-        if self.server_args.vae_cpu_offload:
-            self.vae.to("cpu", non_blocking=True)
-
-        if torch.backends.mps.is_available():
-            # Flush lazy MPS kernels before freeing weights to avoid hangs.
-            torch.mps.synchronize()
-            del self.vae
-            pipeline = self.pipeline() if self.pipeline else None
-            if pipeline is not None and self.component_name in pipeline.modules:
-                del pipeline.modules[self.component_name]
-            self.server_args.model_loaded[self.component_name] = False
-
     @torch.no_grad()
     def forward(
         self,
@@ -208,32 +211,42 @@ class DecodingStage(PipelineStage):
         # load vae if not already loaded (used for memory constrained devices)
         self.load_model()
 
-        frames = self.decode(batch.latents, server_args)
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        with self.use_declared_component(
+            component_name=self.component_name,
+            module=self.vae,
+        ) as vae:
+            assert vae is not None
+            self.vae = vae
 
-        # decode trajectory latents if needed
+            frames = self.decode(batch.latents, server_args, vae_dtype=vae_dtype)
-        if batch.return_trajectory_decoded:
-            assert (
-                batch.trajectory_latents is not None
-            ), "batch should have trajectory latents"
 
-            # 1. Batch trajectory decoding to improve GPU utilization
+            # decode trajectory latents if needed
-            # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+            if batch.return_trajectory_decoded:
-            B, T, C, F, H, W = batch.trajectory_latents.shape
+                assert (
-            flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W)
+                    batch.trajectory_latents is not None
+                ), "batch should have trajectory latents"
 
-            logger.info("decoding %s trajectory latents in batch", B * T)
+                # 1. Batch trajectory decoding to improve GPU utilization
-            # Use the optimized batch decode
+                # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
-            all_decoded = self.decode(flat_latents, server_args)
+                B, T, C, F, H, W = batch.trajectory_latents.shape
+                flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W)
 
-            # 2. Reshape back
+                logger.info("decoding %s trajectory latents in batch", B * T)
-            # Keep on GPU to allow faster vectorized post-processing
+                # Use the optimized batch decode
-            decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:])
+                all_decoded = self.decode(
+                    flat_latents, server_args, vae_dtype=vae_dtype
+                )
 
-            # Convert to list of tensors (per timestep) as expected by OutputBatch
+                # 2. Reshape back
-            # Each element in list is [B, channels, frames, H_out, W_out]
+                # Keep on GPU to allow faster vectorized post-processing
-            trajectory_decoded = [decoded_tensor[:, i] for i in range(T)]
+                decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:])
-        else:
+
-            trajectory_decoded = None
+                # Convert to list of tensors (per timestep) as expected by OutputBatch
+                # Each element in list is [B, channels, frames, H_out, W_out]
+                trajectory_decoded = [decoded_tensor[:, i] for i in range(T)]
+            else:
+                trajectory_decoded = None
 
         frames = server_args.pipeline_config.post_decoding(frames, server_args)
 
@@ -248,8 +261,4 @@ class DecodingStage(PipelineStage):
             noise_pred=None,
         )
 
-        # Keep VAE resident during warmup; the real request needs it next.
-        if not getattr(batch, "is_warmup", False):
-            self.offload_model()
-
         return output_batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py

@@ -1,6 +1,7 @@
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
 from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -25,6 +26,16 @@ class LTX2AVDecodingStage(DecodingStage):
 
         self.video_processor = VideoProcessor(vae_scale_factor=32)
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(stage_name, "vae", target_dtype=torch.bfloat16),
+            ComponentUse(stage_name, "audio_vae"),
+            ComponentUse(stage_name, "vocoder"),
+        ]
+
     @staticmethod
     def _ltx2_should_externally_denorm_video_latents(server_args: ServerArgs) -> bool:
         arch_config = server_args.pipeline_config.vae_config.arch_config
@@ -33,45 +44,44 @@ class LTX2AVDecodingStage(DecodingStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
         self.load_model()
 
-        self.vae = self.vae.to(get_local_torch_device())
-        self.vae.eval()
-        latents = batch.latents.to(get_local_torch_device())
-
         vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (
             vae_dtype != torch.float32
         ) and not server_args.disable_autocast
 
         original_dtype = vae_dtype
-        self.vae.to(torch.bfloat16)
+        with self.use_declared_component(component_name="vae", module=self.vae) as vae:
-        latents = latents.to(torch.bfloat16)
+            assert vae is not None
-        if self._ltx2_should_externally_denorm_video_latents(server_args):
+            self.vae = vae
-            std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents)
+            self.vae.eval()
-            mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents)
+            latents = batch.latents.to(get_local_torch_device(), dtype=torch.bfloat16)
-            latents = latents * std + mean
+            if self._ltx2_should_externally_denorm_video_latents(server_args):
-        latents = server_args.pipeline_config.preprocess_decoding(
+                std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents)
-            latents, server_args, vae=self.vae
+                mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents)
-        )
+                latents = latents * std + mean
+            latents = server_args.pipeline_config.preprocess_decoding(
+                latents, server_args, vae=self.vae
+            )
 
-        with torch.autocast(
+            with torch.autocast(
-            device_type=current_platform.device_type,
+                device_type=current_platform.device_type,
-            dtype=vae_dtype,
+                dtype=vae_dtype,
-            enabled=vae_autocast_enabled,
+                enabled=vae_autocast_enabled,
-        ):
+            ):
-            try:
+                try:
-                if server_args.pipeline_config.vae_tiling:
+                    if server_args.pipeline_config.vae_tiling:
-                    self.vae.enable_tiling()
+                        self.vae.enable_tiling()
-            except Exception:
+                except Exception:
-                pass
+                    pass
-            decode_output = self.vae.decode(latents)
+                decode_output = self.vae.decode(latents)
-            if isinstance(decode_output, tuple):
+                if isinstance(decode_output, tuple):
-                video = decode_output[0]
+                    video = decode_output[0]
-            elif hasattr(decode_output, "sample"):
+                elif hasattr(decode_output, "sample"):
-                video = decode_output.sample
+                    video = decode_output.sample
-            else:
+                else:
-                video = decode_output
+                    video = decode_output
 
-        self.vae.to(original_dtype)
+            self.vae.to(original_dtype)
         video = self.video_processor.postprocess_video(video, output_type="np")
 
         output_batch = OutputBatch(
@@ -90,49 +100,64 @@ class LTX2AVDecodingStage(DecodingStage):
         if audio_latents is not None:
             # Ensure device/dtype
             device = get_local_torch_device()
-            self.audio_vae = self.audio_vae.to(device)
+            with self.use_declared_component(
-            self.vocoder = self.vocoder.to(device)
+                component_name="audio_vae",
-            self.audio_vae.eval()
+                module=self.audio_vae,
-            self.vocoder.eval()
+            ) as audio_vae:
-            try:
+                assert audio_vae is not None
-                dtype = self.audio_vae.dtype
+                self.audio_vae = audio_vae
-            except AttributeError:
+                self.audio_vae.eval()
-                dtype = None
-            if dtype is None:
                 try:
-                    dtype = next(self.audio_vae.parameters()).dtype
+                    dtype = self.audio_vae.dtype
-                except StopIteration:
+                except AttributeError:
-                    dtype = torch.float32
+                    dtype = None
-            audio_latents = audio_latents.to(device, dtype=dtype)
+                if dtype is None:
-            try:
+                    try:
-                latents_std = self.audio_vae.latents_std
+                        dtype = next(self.audio_vae.parameters()).dtype
-            except AttributeError:
+                    except StopIteration:
-                latents_std = None
+                        dtype = torch.float32
-            if isinstance(latents_std, torch.Tensor) and torch.all(latents_std == 0):
+                audio_latents = audio_latents.to(device, dtype=dtype)
-                logger.warning(
+                try:
-                    "audio_vae.latents_std is all zeros; audio denorm may be incorrect."
+                    latents_std = self.audio_vae.latents_std
-                )
+                except AttributeError:
-            try:
+                    latents_std = None
-                latents_mean = self.audio_vae.latents_mean
+                if isinstance(latents_std, torch.Tensor) and torch.all(
-            except AttributeError:
+                    latents_std == 0
-                latents_mean = None
+                ):
-            if isinstance(latents_mean, torch.Tensor) and isinstance(
+                    logger.warning(
-                latents_std, torch.Tensor
+                        "audio_vae.latents_std is all zeros; audio denorm may be incorrect."
-            ):
-                latents_mean = latents_mean.to(device=device, dtype=dtype)
-                latents_std = latents_std.to(device=device, dtype=dtype)
-                if audio_latents.ndim == 4:
-                    latents_mean = latents_mean.view(
-                        1, audio_latents.shape[1], 1, audio_latents.shape[3]
                     )
-                    latents_std = latents_std.view(
+                try:
-                        1, audio_latents.shape[1], 1, audio_latents.shape[3]
+                    latents_mean = self.audio_vae.latents_mean
-                    )
+                except AttributeError:
-                audio_latents = audio_latents * latents_std + latents_mean
+                    latents_mean = None
+                if isinstance(latents_mean, torch.Tensor) and isinstance(
+                    latents_std, torch.Tensor
+                ):
+                    latents_mean = latents_mean.to(device=device, dtype=dtype)
+                    latents_std = latents_std.to(device=device, dtype=dtype)
+                    if audio_latents.ndim == 4:
+                        latents_mean = latents_mean.view(
+                            1, audio_latents.shape[1], 1, audio_latents.shape[3]
+                        )
+                        latents_std = latents_std.view(
+                            1, audio_latents.shape[1], 1, audio_latents.shape[3]
+                        )
+                    audio_latents = audio_latents * latents_std + latents_mean
 
-            with torch.no_grad():
+                with torch.no_grad():
-                # Decode latents to spectrogram
+                    # Decode latents to spectrogram
-                spectrogram = self.audio_vae.decode(audio_latents, return_dict=False)[0]
+                    spectrogram = self.audio_vae.decode(
+                        audio_latents, return_dict=False
+                    )[0]
+
+            with self.use_declared_component(
+                component_name="vocoder",
+                module=self.vocoder,
+            ) as vocoder:
+                assert vocoder is not None
+                self.vocoder = vocoder
+                self.vocoder.eval()
                 if hasattr(self.vocoder, "conv_in") and hasattr(
                     self.vocoder.conv_in, "in_channels"
                 ):
@@ -143,7 +168,8 @@ class LTX2AVDecodingStage(DecodingStage):
                             f"Vocoder expects channels*mel_bins={expected_in}, got {actual_in} from spectrogram shape {tuple(spectrogram.shape)}"
                         )
                 # Decode spectrogram to waveform
-                waveform = self.vocoder(spectrogram)
+                with torch.no_grad():
+                    waveform = self.vocoder(spectrogram)
             output_batch.audio = waveform.cpu().float()
             try:
                 pipeline_audio_cfg = server_args.pipeline_config.audio_vae_config
@@ -170,5 +196,4 @@ class LTX2AVDecodingStage(DecodingStage):
                 vocoder_sr or audio_vae_sr or pipeline_audio_sr
             )
 
-        self.offload_model()
         return output_batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -60,6 +60,7 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
 from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
     TransformerLoader,
 )
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
@@ -87,11 +88,10 @@ from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import
     RolloutDenoisingMixin,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
-from sglang.multimodal_gen.utils import dict_to_3d_list
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, dict_to_3d_list
 from sglang.srt.utils.common import get_compiler_backend
 
 logger = init_logger(__name__)
@@ -195,6 +195,49 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self._cached_num_steps = None
         self._is_warmed_up = False
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        uses: list[ComponentUse] = []
+        if self.vae is not None:
+            vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+            uses.append(
+                ComponentUse(
+                    stage_name=stage_name,
+                    component_name="vae",
+                    target_dtype=vae_dtype,
+                )
+            )
+        for default_name, module in (
+            ("transformer", self.transformer),
+            ("transformer_2", self.transformer_2),
+        ):
+            if module is None:
+                continue
+            component_name = self._component_name_for_stage_module(module, default_name)
+            uses.append(
+                ComponentUse(
+                    stage_name=stage_name,
+                    component_name=component_name,
+                    phase=component_name,
+                    preferred_ready_after_request=component_name == "transformer",
+                    memory_intensive=True,
+                )
+            )
+        return uses
+
+    def _component_name_for_stage_module(
+        self, module: nn.Module | None, default_name: str
+    ) -> str:
+        pipeline = self.pipeline() if self.pipeline else None
+        if pipeline is None or module is None:
+            return default_name
+        for name, candidate in pipeline.modules.items():
+            if candidate is module:
+                return name
+        return default_name
+
     def _maybe_enable_torch_compile(self, module: object) -> None:
         """
         Compile a module with torch.compile, and enable inductor overlap tweak if available.
@@ -572,13 +615,22 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
 
         # TI2V specific preparations - before SP sharding
         if should_preprocess_for_wan_ti2v:
-            seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents(
+            vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
-                self.vae,
+            with self.use_declared_component(
-                latents,
+                component_name="vae",
-                target_dtype,
+                module=self.vae,
-                batch,
+                target_dtype=vae_dtype,
-                server_args,
+            ) as vae:
-            )
+                assert vae is not None
+                self.vae = vae
+                seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents(
+                    self.vae,
+                    latents,
+                    target_dtype,
+                    vae_dtype,
+                    batch,
+                    server_args,
+                )
         else:
             seq_len, z, reserved_frames_masks = (
                 None,
@@ -765,7 +817,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self, batch: Req
     ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
         """Return the prompt-embedding validator used by verify_input."""
-        del batch
         return V.list_not_empty
 
     def _get_negative_prompt_embeds_validator(
@@ -928,9 +979,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         ):
             self.save_sta_search_results(batch)
 
-        # Capture references before potential deletion on MPS
-        dits = list(filter(None, [self.transformer, self.transformer_2]))
-
         # deallocate transformer if on mps
         pipeline = self.pipeline() if self.pipeline else None
         if torch.backends.mps.is_available() and not is_warmup:
@@ -947,14 +995,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 torch.mps.current_allocated_memory(),
             )
 
-        # reset offload managers with prefetching first layer for next forward
-        for dit in dits:
-            if isinstance(dit, OffloadableDiTMixin):
-                # release all DiT weights to avoid peak VRAM usage, which may increasing the latency for next req
-                # TODO: should be make this an option?
-                for manager in dit.layerwise_offload_managers:
-                    manager.release_all()
-
     def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
         """Shard latents for Sequence Parallelism if applicable."""
         if get_sp_world_size() <= 1:
@@ -1026,35 +1066,32 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         if profiler:
             profiler.step_denoising_step()
 
-    def _manage_device_placement(
+    def _manage_dit_use_site(
         self,
-        model_to_use: nn.Module,
+        current_model: nn.Module,
-        model_to_offload: nn.Module | None,
+        current_phase: str,
-        server_args: ServerArgs,
+        batch: Req,
-    ):
+    ) -> None:
         """
-        Manages the offload / load behavior of dit
+        manage dit's residency by reporting the active sequential use
+
+        only applicable for dual-dit architecture like Wan
+
+        Args:
+            current_model: the next active dit, transformer_1 or transformer_2
         """
-        if not server_args.dit_cpu_offload:
+        manager = self._component_residency_manager
-            return
 
-        # FSDP manages offloading internally
+        component_name = manager.component_name_for_module(current_model, current_phase)
-        if server_args.use_fsdp_inference:
+        phase = str(batch.extra.get("ltx2_phase", current_phase))
-            return
+        use = ComponentUse(
-
+            stage_name=self._active_component_stage_name(),
-        # Offload the unused model if it's on CUDA
+            component_name=component_name,
-        if (
+            phase=phase,
-            model_to_offload is not None
+            preferred_ready_after_request=component_name == "transformer",
-            and next(model_to_offload.parameters()).device.type == "cuda"
+            memory_intensive=True,
-        ):
+        )
-            model_to_offload.to("cpu")
+        manager.begin_use(use)
-
-        # Load the model to use if it's on CPU
-        if (
-            model_to_use is not None
-            and next(model_to_use.parameters()).device.type == "cpu"
-        ):
-            model_to_use.to(get_local_torch_device())
 
     def _select_and_manage_model(
         self,
@@ -1066,15 +1103,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         if boundary_timestep is None or t_int >= boundary_timestep:
             # High-noise stage
             current_model = self.transformer
-            model_to_offload = self.transformer_2
             current_guidance_scale = batch.guidance_scale
+            current_phase = "transformer"
         else:
             # Low-noise stage
             current_model = self.transformer_2
-            model_to_offload = self.transformer
             current_guidance_scale = batch.guidance_scale_2
+            current_phase = "transformer_2"
 
-        self._manage_device_placement(current_model, model_to_offload, server_args)
+        self._manage_dit_use_site(current_model, current_phase, batch)
 
         assert current_model is not None, "The model for the current step is not set."
         return current_model, current_guidance_scale
@@ -1195,6 +1232,8 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 (denoising_end_time - denoising_start_time) / len(ctx.timesteps),
             )
 
+        self._finish_active_component_use()
+
         # Rollout postprocessing must run BEFORE _finalize_denoising_loop so
         # the final scheduler.step output (ctx.latents) is still SP-sharded and
         # can be gathered uniformly alongside the per-step dit_trajectory via

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -2,6 +2,7 @@ import torch
 from diffusers.utils.torch_utils import randn_tensor
 
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import is_ltx23_native_variant
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
     clone_scheduler_runtime,
 )
@@ -10,7 +11,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import
     LTX2DenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -28,6 +28,20 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
         )
         self.audio_vae = audio_vae
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="transformer",
+                phase="stage1",
+                preferred_ready_after_request=True,
+                memory_intensive=True,
+            )
+        ]
+
     def _post_denoising_loop(
         self,
         batch: Req,
@@ -83,24 +97,6 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
             batch.latents = latents
             batch.audio_latents = audio_latents
 
-        pipeline = self.pipeline() if self.pipeline else None
-        current_phase = (
-            str(getattr(batch, "extra", {}).get("ltx2_phase", ""))
-            if hasattr(batch, "extra")
-            else ""
-        )
-        release_phase_state = (
-            getattr(pipeline, "release_ltx2_phase_state", None)
-            if pipeline is not None
-            else None
-        )
-        if callable(release_phase_state):
-            release_phase_state(current_phase)
-
-        if isinstance(self.transformer, OffloadableDiTMixin):
-            for manager in self.transformer.layerwise_offload_managers:
-                manager.release_all()
-
 
 class LTX2RefinementStage(LTX2AVDenoisingStage):
     """Stage-2 refinement wrapper that re-noises distilled LTX latents once."""
@@ -125,6 +121,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
         )
         self.distilled_sigmas = torch.tensor(distilled_sigmas)
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        component_name = "transformer_2"
+        pipeline = self.pipeline() if self.pipeline else None
+        if pipeline is not None and "transformer_2" not in pipeline.modules:
+            component_name = "transformer"
+        return [
+            ComponentUse(
+                stage_name=stage_name,
+                component_name=component_name,
+                phase="stage2",
+                memory_intensive=True,
+            )
+        ]
+
     @staticmethod
     def _randn_like_with_batch_generators(
         reference_tensor: torch.Tensor, batch: Req
@@ -218,14 +231,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Run the distilled refinement schedule on top of the shared AV denoiser."""
         batch.extra["ltx2_phase"] = "stage2"
-        pipeline = self.pipeline() if self.pipeline else None
-        ensure_phase_ready = (
-            getattr(pipeline, "ensure_ltx2_phase_ready", None)
-            if pipeline is not None
-            else None
-        )
-        if callable(ensure_phase_ready):
-            ensure_phase_ready("stage2")
         original_clean_latent_background = getattr(
             batch, "ltx2_ti2v_clean_latent_background", None
         )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -121,7 +121,7 @@ class DmdDenoisingStage(DenoisingStage):
                         )
                     else:
                         current_model = self.transformer
-                        self._manage_device_placement(current_model, None, server_args)
+                        self._manage_dit_use_site(current_model, "transformer", batch)
                     # Expand latents for I2V
                     noise_latents = latents.clone()
                     latent_model_input = latents.to(target_dtype)
@@ -232,45 +232,19 @@ class DmdDenoisingStage(DenoisingStage):
         if boundary_timestep is None or t_int >= boundary_timestep:
             # High-noise stage
             current_model = self.transformer
-            model_to_offload = self.transformer_2
             current_guidance_scale = batch.guidance_scale
+            current_phase = "transformer"
         else:
             # Low-noise stage
             current_model = self.transformer_2
-            model_to_offload = self.transformer
             current_guidance_scale = batch.guidance_scale_2
+            current_phase = "transformer_2"
 
-        self._manage_device_placement(current_model, model_to_offload, server_args)
+        self._manage_dit_use_site(current_model, current_phase, batch)
 
         assert current_model is not None, "The model for the current step is not set."
         return current_model, current_guidance_scale
 
-    def _manage_device_placement(
-        self,
-        model_to_use: torch.nn.Module,
-        model_to_offload: torch.nn.Module | None,
-        server_args: ServerArgs,
-    ):
-        """
-        Manages the offload / load behavior of dit
-        """
-        if not server_args.dit_cpu_offload:
-            return
-
-        # Offload the unused model if it's on CUDA
-        if (
-            model_to_offload is not None
-            and next(model_to_offload.parameters()).device.type == "cuda"
-        ):
-            model_to_offload.to("cpu")
-
-        # Load the model to use if it's on CPU
-        if (
-            model_to_use is not None
-            and next(model_to_use.parameters()).device.type == "cpu"
-        ):
-            model_to_use.to(get_local_torch_device())
-
     def _handle_boundary_ratio(
         self,
         server_args,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py

@@ -8,6 +8,7 @@ Encoding stage for diffusion pipelines.
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
@@ -37,6 +38,19 @@ class EncodingStage(PipelineStage):
         super().__init__()
         self.vae: ParallelTiledVAE = vae
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name,
+                "vae",
+                target_dtype=vae_dtype,
+            )
+        ]
+
     @torch.no_grad()
     def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
         """Verify encoding stage inputs."""
@@ -67,8 +81,6 @@ class EncodingStage(PipelineStage):
         """
         assert batch.latents is not None and isinstance(batch.latents, torch.Tensor)
 
-        self.vae = self.vae.to(get_local_torch_device())
-
         # Setup VAE precision
         vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (
@@ -81,27 +93,25 @@ class EncodingStage(PipelineStage):
         # Move to appropriate device and dtype
         latents = latents.to(get_local_torch_device())
 
-        # Encode image to latents
+        with self.use_declared_component(component_name="vae", module=self.vae) as vae:
-        with torch.autocast(
+            assert vae is not None
-            device_type=current_platform.device_type,
+            self.vae = vae
-            dtype=vae_dtype,
+
-            enabled=vae_autocast_enabled,
+            # Encode image to latents
-        ):
+            with torch.autocast(
-            if server_args.pipeline_config.vae_tiling:
+                device_type=current_platform.device_type,
-                self.vae.enable_tiling()
+                dtype=vae_dtype,
-            # if server_args.vae_sp:
+                enabled=vae_autocast_enabled,
-            #     self.vae.enable_parallel()
+            ):
-            if not vae_autocast_enabled:
+                if server_args.pipeline_config.vae_tiling:
-                latents = latents.to(vae_dtype)
+                    self.vae.enable_tiling()
-            latents = self.vae.encode(latents).mean
+                # if server_args.vae_sp:
+                #     self.vae.enable_parallel()
+                if not vae_autocast_enabled:
+                    latents = latents.to(vae_dtype)
+                latents = self.vae.encode(latents).mean
 
         # Update batch with encoded latents
         batch.latents = latents
 
-        # Offload models if needed
-        self.maybe_free_model_hooks()
-
-        if server_args.vae_cpu_offload:
-            self.vae.to("cpu")
-
         return batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -22,6 +22,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     qwen_image_postprocess_text,
 )
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.models.vision_utils import (
@@ -137,26 +138,16 @@ class ImageEncodingStage(PipelineStage):
         self.image_encoder = image_encoder
         self.text_encoder = text_encoder
 
-    def load_model(self):
+    def component_uses(
-        if self.server_args.image_encoder_cpu_offload:
+        self, server_args: ServerArgs, stage_name: str | None = None
-            device = get_local_torch_device()
+    ) -> list[ComponentUse]:
-            self.move_to_device(device)
+        stage_name = self._component_stage_name(stage_name)
-
+        uses = []
-    def offload_model(self):
+        if self.image_encoder is not None:
-        if self.server_args.image_encoder_cpu_offload:
+            uses.append(ComponentUse(stage_name, "image_encoder"))
-            self.move_to_device("cpu")
+        if self.text_encoder is not None:
-
+            uses.append(ComponentUse(stage_name, "text_encoder"))
-    def move_to_device(self, device):
+        return uses
-        if self.server_args.use_fsdp_inference:
-            return
-        fields = [
-            "image_processor",
-            "image_encoder",
-        ]
-        for field in fields:
-            processor = getattr(self, field, None)
-            if processor and hasattr(processor, "to"):
-                setattr(self, field, processor.to(device))
 
     def encoding_qwen_image_edit(self, outputs, image_inputs):
         # encoder hidden state
@@ -177,8 +168,6 @@ class ImageEncodingStage(PipelineStage):
             return batch
         cuda_device = get_local_torch_device()
 
-        self.load_model()
-
         image_processor_kwargs = (
             server_args.pipeline_config.prepare_image_processor_kwargs(batch)
         )
@@ -215,14 +204,20 @@ class ImageEncodingStage(PipelineStage):
 
             if self.image_encoder:
                 # if an image encoder is provided
-                with set_forward_context(current_timestep=0, attn_metadata=None):
+                with self.use_declared_component(
-                    outputs = self.image_encoder(
+                    component_name="image_encoder",
-                        **image_inputs,
+                    module=self.image_encoder,
-                        **server_args.pipeline_config.image_encoder_extra_args,
+                ) as image_encoder:
-                    )
+                    assert image_encoder is not None
-                    image_embeds = server_args.pipeline_config.postprocess_image(
+                    self.image_encoder = image_encoder
-                        outputs
+                    with set_forward_context(current_timestep=0, attn_metadata=None):
-                    )
+                        outputs = self.image_encoder(
+                            **image_inputs,
+                            **server_args.pipeline_config.image_encoder_extra_args,
+                        )
+                        image_embeds = server_args.pipeline_config.postprocess_image(
+                            outputs
+                        )
                 batch.image_embeds.append(image_embeds)
             elif self.text_encoder:
                 # if a text encoder is provided, e.g. Qwen-Image-Edit
@@ -243,22 +238,28 @@ class ImageEncodingStage(PipelineStage):
                         **neg_image_processor_kwargs,
                     ).to(cuda_device)
 
-                with set_forward_context(current_timestep=0, attn_metadata=None):
+                with self.use_declared_component(
-                    outputs = self.text_encoder(
+                    component_name="text_encoder",
-                        input_ids=image_inputs.input_ids,
+                    module=self.text_encoder,
-                        attention_mask=image_inputs.attention_mask,
+                ) as text_encoder:
-                        pixel_values=image_inputs.pixel_values,
+                    assert text_encoder is not None
-                        image_grid_thw=image_inputs.image_grid_thw,
+                    self.text_encoder = text_encoder
-                        output_hidden_states=True,
+                    with set_forward_context(current_timestep=0, attn_metadata=None):
-                    )
+                        outputs = self.text_encoder(
-                    if batch.do_classifier_free_guidance:
+                            input_ids=image_inputs.input_ids,
-                        neg_outputs = self.text_encoder(
+                            attention_mask=image_inputs.attention_mask,
-                            input_ids=neg_image_inputs.input_ids,
+                            pixel_values=image_inputs.pixel_values,
-                            attention_mask=neg_image_inputs.attention_mask,
+                            image_grid_thw=image_inputs.image_grid_thw,
-                            pixel_values=neg_image_inputs.pixel_values,
-                            image_grid_thw=neg_image_inputs.image_grid_thw,
                             output_hidden_states=True,
                         )
+                        if batch.do_classifier_free_guidance:
+                            neg_outputs = self.text_encoder(
+                                input_ids=neg_image_inputs.input_ids,
+                                attention_mask=neg_image_inputs.attention_mask,
+                                pixel_values=neg_image_inputs.pixel_values,
+                                image_grid_thw=neg_image_inputs.image_grid_thw,
+                                output_hidden_states=True,
+                            )
 
                 all_prompt_embeds.append(
                     self.encoding_qwen_image_edit(outputs, image_inputs)
@@ -273,8 +274,6 @@ class ImageEncodingStage(PipelineStage):
         if all_neg_prompt_embeds:
             batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0))
 
-        self.offload_model()
-
         return batch
 
     def build_dedup_fingerprint(
@@ -328,20 +327,22 @@ class LTX2ImageEncodingStage(PipelineStage):
         self._condition_image_encoder = None
         self._condition_image_encoder_dir = None
 
-    # -- device management (mirrors ImageVAEEncodingStage) ---------------
+    def component_uses(
-
+        self, server_args: ServerArgs, stage_name: str | None = None
-    def load_model(self):
+    ) -> list[ComponentUse]:
-        device = get_local_torch_device()
+        arch_config = server_args.pipeline_config.vae_config.arch_config
-        if self._condition_image_encoder is not None:
+        encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", ""))
-            self._condition_image_encoder = self._condition_image_encoder.to(device)
+        stage_name = self._component_stage_name(stage_name)
-        else:
+        if encoder_subdir:
-            self.vae = self.vae.to(device)
+            return [
-
+                ComponentUse(
-    def offload_model(self):
+                    stage_name,
-        if self.server_args.vae_cpu_offload:
+                    "condition_image_encoder",
-            self.vae = self.vae.to("cpu")
+                )
-            if self._condition_image_encoder is not None:
+            ]
-                self._condition_image_encoder = self._condition_image_encoder.to("cpu")
+        if self.vae is None:
+            return []
+        return [ComponentUse(stage_name, "vae")]
 
     # -- lazy condition encoder (LTX-2.3) --------------------------------
 
@@ -576,58 +577,61 @@ class LTX2ImageEncodingStage(PipelineStage):
         if len(batch.condition_image) == 1:
             batch.condition_image = batch.condition_image[0]
 
-        # 2. Load encoder(s) to device, cast to encode_dtype
+        # 2. Select encoder(s); residency manager moves it to device and dtype.
         use_condition_encoder = self._ensure_condition_image_encoder(server_args)
-        self.load_model()
 
         device = get_local_torch_device()
         encode_dtype = batch.latents.dtype
 
-        # Cast the active encoder to the latent precision (must match original
-        # behavior — running in a different dtype shifts the encoded latents).
         if use_condition_encoder:
-            self._condition_image_encoder = self._condition_image_encoder.to(
+            component_name = "condition_image_encoder"
-                dtype=encode_dtype
+            encoder = self._condition_image_encoder
-            )
         else:
-            self.vae = self.vae.to(dtype=encode_dtype)
+            component_name = "vae"
+            encoder = self.vae
 
         packed_latents = []
-        for conditioned_img in conditioned_imgs:
+        with self.use_declared_component(
-            video_condition = self._pil_to_video_tensor(
+            component_name=component_name,
-                conditioned_img,
+            module=encoder,
-                width=int(batch.width),
+            target_dtype=encode_dtype,
-                height=int(batch.height),
+        ) as active_encoder:
-                device=device,
+            assert active_encoder is not None
-                dtype=encode_dtype,
-            )
-
-            # 3. Encode
             if use_condition_encoder:
-                latent = self._condition_encode(video_condition, server_args).to(
+                self._condition_image_encoder = active_encoder
-                    dtype=encode_dtype
-                )
             else:
-                latent = self._vae_encode(video_condition, server_args, batch.generator)
+                self.vae = active_encoder
 
-            packed = server_args.pipeline_config.maybe_pack_latents(
+            for conditioned_img in conditioned_imgs:
-                latent, latent.shape[0], batch
+                video_condition = self._pil_to_video_tensor(
-            )
+                    conditioned_img,
-            if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
+                    width=int(batch.width),
-                raise ValueError("Expected packed image latents [B, S0, D].")
+                    height=int(batch.height),
-            if int(packed.shape[1]) != expected_tokens:
+                    device=device,
-                raise ValueError(
+                    dtype=encode_dtype,
-                    f"LTX-2 conditioning token count mismatch: "
-                    f"{packed.shape[1]=} {expected_tokens=}."
                 )
-            packed_latents.append(packed)
 
-        # Restore VAE to its config dtype (shared with decoding stage)
+                # 3. Encode
-        if not use_condition_encoder:
+                if use_condition_encoder:
-            original_dtype = PRECISION_TO_TYPE[
+                    latent = self._condition_encode(video_condition, server_args).to(
-                server_args.pipeline_config.vae_precision
+                        dtype=encode_dtype
-            ]
+                    )
-            self.vae = self.vae.to(dtype=original_dtype)
+                else:
+                    latent = self._vae_encode(
+                        video_condition, server_args, batch.generator
+                    )
+
+                packed = server_args.pipeline_config.maybe_pack_latents(
+                    latent, latent.shape[0], batch
+                )
+                if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
+                    raise ValueError("Expected packed image latents [B, S0, D].")
+                if int(packed.shape[1]) != expected_tokens:
+                    raise ValueError(
+                        f"LTX-2 conditioning token count mismatch: "
+                        f"{packed.shape[1]=} {expected_tokens=}."
+                    )
+                packed_latents.append(packed)
 
         batch.image_latent = (
             packed_latents[0] if len(packed_latents) == 1 else packed_latents
@@ -643,7 +647,6 @@ class LTX2ImageEncodingStage(PipelineStage):
                 batch.height,
             )
 
-        self.offload_model()
         return batch
 
     def build_dedup_fingerprint(
@@ -688,12 +691,18 @@ class ImageVAEEncodingStage(PipelineStage):
         super().__init__()
         self.vae: ParallelTiledVAE = vae
 
-    def load_model(self):
+    def component_uses(
-        self.vae = self.vae.to(get_local_torch_device())
+        self, server_args: ServerArgs, stage_name: str | None = None
-
+    ) -> list[ComponentUse]:
-    def offload_model(self):
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
-        if self.server_args.vae_cpu_offload:
+        stage_name = self._component_stage_name(stage_name)
-            self.vae = self.vae.to("cpu")
+        return [
+            ComponentUse(
+                stage_name,
+                "vae",
+                target_dtype=vae_dtype,
+            )
+        ]
 
     def forward(
         self,
@@ -707,7 +716,6 @@ class ImageVAEEncodingStage(PipelineStage):
         if batch.condition_image is None:
             return batch
 
-        self.load_model()
         num_frames = batch.num_frames
 
         images = (
@@ -721,110 +729,115 @@ class ImageVAEEncodingStage(PipelineStage):
             server_args.pipeline_config, "prepare_condition_image_latent_ids", None
         )
         condition_latents = [] if callable(prepare_condition_image_latent_ids) else None
-        for image in images:
+        # Setup VAE precision
-            image = self.preprocess(
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
-                image,
+        vae_autocast_enabled = (
-            ).to(get_local_torch_device(), dtype=torch.float32)
+            vae_dtype != torch.float32
+        ) and not server_args.disable_autocast
 
-            # (B, C, H, W) -> (B, C, 1, H, W)
+        with self.use_declared_component(component_name="vae", module=self.vae) as vae:
-            image = image.unsqueeze(2)
+            assert vae is not None
+            self.vae = vae
 
-            if num_frames == 1:
+            for image in images:
-                video_condition = image
+                image = self.preprocess(
-            else:
+                    image,
-                video_condition = torch.cat(
+                ).to(get_local_torch_device(), dtype=torch.float32)
-                    [
+
-                        image,
+                # (B, C, H, W) -> (B, C, 1, H, W)
-                        image.new_zeros(
+                image = image.unsqueeze(2)
-                            image.shape[0],
+
-                            image.shape[1],
+                if num_frames == 1:
-                            num_frames - 1,
+                    video_condition = image
-                            image.shape[3],
+                else:
-                            image.shape[4],
+                    video_condition = torch.cat(
-                        ),
+                        [
-                    ],
+                            image,
-                    dim=2,
+                            image.new_zeros(
+                                image.shape[0],
+                                image.shape[1],
+                                num_frames - 1,
+                                image.shape[3],
+                                image.shape[4],
+                            ),
+                        ],
+                        dim=2,
+                    )
+                video_condition = video_condition.to(
+                    device=get_local_torch_device(), dtype=torch.float32
                 )
-            video_condition = video_condition.to(
-                device=get_local_torch_device(), dtype=torch.float32
-            )
 
-            # Setup VAE precision
+                # Encode Image
-            vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+                with torch.autocast(
-            vae_autocast_enabled = (
+                    device_type=current_platform.device_type,
-                vae_dtype != torch.float32
+                    dtype=vae_dtype,
-            ) and not server_args.disable_autocast
+                    enabled=vae_autocast_enabled,
+                ):
+                    if server_args.pipeline_config.vae_tiling:
+                        self.vae.enable_tiling()
+                    # if server_args.vae_sp:
+                    #     self.vae.enable_parallel()
+                    if not vae_autocast_enabled:
+                        video_condition = video_condition.to(vae_dtype)
+                    latent_dist: DiagonalGaussianDistribution = self.vae.encode(
+                        video_condition
+                    )
+                    # for auto_encoder from diffusers
+                    if isinstance(latent_dist, AutoencoderKLOutput):
+                        latent_dist = latent_dist.latent_dist
 
-            # Encode Image
+                generator = batch.generator
-            with torch.autocast(
+                if generator is None:
-                device_type=current_platform.device_type,
+                    raise ValueError("Generator must be provided")
-                dtype=vae_dtype,
+
-                enabled=vae_autocast_enabled,
+                sample_mode = (
-            ):
+                    server_args.pipeline_config.vae_config.encode_sample_mode()
-                if server_args.pipeline_config.vae_tiling:
-                    self.vae.enable_tiling()
-                # if server_args.vae_sp:
-                #     self.vae.enable_parallel()
-                if not vae_autocast_enabled:
-                    video_condition = video_condition.to(vae_dtype)
-                latent_dist: DiagonalGaussianDistribution = self.vae.encode(
-                    video_condition
                 )
-                # for auto_encoder from diffusers
-                if isinstance(latent_dist, AutoencoderKLOutput):
-                    latent_dist = latent_dist.latent_dist
 
-            generator = batch.generator
+                latent_condition = self.retrieve_latents(
-            if generator is None:
+                    latent_dist, generator, sample_mode=sample_mode
-                raise ValueError("Generator must be provided")
+                )
-
+                latent_condition = server_args.pipeline_config.postprocess_vae_encode(
-            sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode()
-
-            latent_condition = self.retrieve_latents(
-                latent_dist, generator, sample_mode=sample_mode
-            )
-            latent_condition = server_args.pipeline_config.postprocess_vae_encode(
-                latent_condition, self.vae
-            )
-            normalized_latent_condition = (
-                server_args.pipeline_config.normalize_vae_encode(
                     latent_condition, self.vae
                 )
-            )
+                normalized_latent_condition = (
-            if normalized_latent_condition is None:
+                    server_args.pipeline_config.normalize_vae_encode(
-                scaling_factor, shift_factor = (
+                        latent_condition, self.vae
-                    server_args.pipeline_config.get_decode_scale_and_shift(
-                        device=latent_condition.device,
-                        dtype=latent_condition.dtype,
-                        vae=self.vae,
                     )
                 )
+                if normalized_latent_condition is None:
+                    scaling_factor, shift_factor = (
+                        server_args.pipeline_config.get_decode_scale_and_shift(
+                            device=latent_condition.device,
+                            dtype=latent_condition.dtype,
+                            vae=self.vae,
+                        )
+                    )
 
-                # apply shift & scale if needed
+                    # apply shift & scale if needed
-                if isinstance(shift_factor, torch.Tensor):
+                    if isinstance(shift_factor, torch.Tensor):
-                    shift_factor = shift_factor.to(latent_condition.device)
+                        shift_factor = shift_factor.to(latent_condition.device)
 
-                if isinstance(scaling_factor, torch.Tensor):
+                    if isinstance(scaling_factor, torch.Tensor):
-                    scaling_factor = scaling_factor.to(latent_condition.device)
+                        scaling_factor = scaling_factor.to(latent_condition.device)
 
-                latent_condition -= shift_factor
+                    latent_condition -= shift_factor
-                latent_condition = latent_condition * scaling_factor
+                    latent_condition = latent_condition * scaling_factor
-            else:
+                else:
-                latent_condition = normalized_latent_condition
+                    latent_condition = normalized_latent_condition
 
-            if condition_latents is not None:
+                if condition_latents is not None:
-                condition_latents.append(latent_condition)
+                    condition_latents.append(latent_condition)
 
-            image_latent = server_args.pipeline_config.postprocess_image_latent(
+                image_latent = server_args.pipeline_config.postprocess_image_latent(
-                latent_condition, batch
+                    latent_condition, batch
-            )
+                )
-            all_image_latents.append(image_latent)
+                all_image_latents.append(image_latent)
 
         batch.image_latent = torch.cat(all_image_latents, dim=1)
         if condition_latents is not None:
             prepare_condition_image_latent_ids(condition_latents, batch)
 
-        self.offload_model()
         return batch
 
     def build_dedup_fingerprint(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py

@@ -1160,22 +1160,12 @@ class LTX2DenoisingStage(DenoisingStage):
         self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs
     ) -> None:
         """Reset the mirrored audio scheduler before the shared loop begins."""
-        if ctx.stage in ("stage1", "stage2"):
+        if is_ltx2_two_stage_pipeline_name(
+            server_args.pipeline_class_name
+        ) and ctx.stage in ("stage1", "stage2"):
             pipeline = self.pipeline() if self.pipeline else None
-            switch_lora_phase = (
+            if pipeline is not None:
-                getattr(pipeline, "switch_lora_phase", None)
+                pipeline.switch_lora_phase(ctx.stage, batch=batch)
-                if pipeline is not None
-                else None
-            )
-            if callable(switch_lora_phase):
-                switch_lora_phase(ctx.stage, batch=batch)
-            ensure_phase_ready = (
-                getattr(pipeline, "ensure_ltx2_phase_ready", None)
-                if pipeline is not None
-                else None
-            )
-            if callable(ensure_phase_ready):
-                ensure_phase_ready(ctx.stage)
         super()._before_denoising_loop(ctx, batch, server_args)
         if ctx.audio_scheduler is None:
             raise ValueError("LTX-2 audio scheduler was not prepared.")
@@ -1194,7 +1184,6 @@ class LTX2DenoisingStage(DenoisingStage):
     ):
         """Preserve the legacy LTX-2 attention-metadata contract."""
         # Legacy LTX-2 paths used the plain attention-metadata builder call here.
-        del ctx, t_int, timesteps_cpu
         return self._build_attn_metadata(step_index, batch, server_args)
 
     def _run_denoising_step(
@@ -1996,7 +1985,6 @@ class LTX2DenoisingStage(DenoisingStage):
 
     def _get_prompt_embeds_validator(self, batch: Req):
         """Allow either tensor or list prompt embeddings for LTX-2 prompts."""
-        del batch
         return lambda x: V.is_tensor(x) or V.list_not_empty(x)
 
     def _get_negative_prompt_embeds_validator(self, batch: Req):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py

@@ -14,6 +14,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import (
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
 
@@ -47,11 +48,20 @@ class HeliosDecodingStage(DecodingStage):
         # Load VAE if needed
         self.load_model()
 
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
         # Decode each chunk separately and concatenate in pixel space
         video_chunks = []
-        for chunk_latents in latent_chunks:
+        with self.use_declared_component(
-            chunk_video = self.decode(chunk_latents, server_args)
+            component_name=self.component_name,
-            video_chunks.append(chunk_video)
+            module=self.vae,
+        ) as vae:
+            assert vae is not None
+            self.vae = vae
+            for chunk_latents in latent_chunks:
+                chunk_video = self.decode(
+                    chunk_latents, server_args, vae_dtype=vae_dtype
+                )
+                video_chunks.append(chunk_video)
 
         frames = torch.cat(video_chunks, dim=2)
         frames = server_args.pipeline_config.post_decoding(frames, server_args)
@@ -64,5 +74,4 @@ class HeliosDecodingStage(DecodingStage):
             metrics=batch.metrics,
         )
 
-        self.offload_model()
         return output_batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py

@@ -13,7 +13,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 
-from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
     get_or_create_request_scheduler,
@@ -107,6 +107,20 @@ class HeliosChunkedDenoisingStage(PipelineStage):
     def parallelism_type(self):
         return StageParallelismType.REPLICATED
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="transformer",
+                phase="transformer",
+                preferred_ready_after_request=True,
+                memory_intensive=True,
+            )
+        ]
+
     def _denoise_one_chunk(
         self,
         latents,
@@ -483,10 +497,15 @@ class HeliosChunkedDenoisingStage(PipelineStage):
         is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk
         gamma = pipeline_config.gamma
 
-        # Move transformer to GPU if CPU-offloaded
+        transformer_use = ComponentUse(
-        if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
+            self.__class__.__name__,
-            if next(self.transformer.parameters()).device.type == "cpu":
+            "transformer",
-                self.transformer.to(get_local_torch_device())
+            phase="transformer",
+            preferred_ready_after_request=True,
+            memory_intensive=True,
+        )
+        manager = self._component_residency_manager
+        manager.begin_use(transformer_use, module=self.transformer)
 
         # Get encoder outputs (prompt_embeds is a list of tensors, one per encoder)
         prompt_embeds = batch.prompt_embeds
@@ -718,10 +737,6 @@ class HeliosChunkedDenoisingStage(PipelineStage):
             history_latents = torch.cat([history_latents, latents], dim=2)
             chunk_latents_list.append(latents)
 
-        # Move transformer back to CPU after denoising
-        if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
-            if next(self.transformer.parameters()).device.type != "cpu":
-                self.transformer.to("cpu")
         torch.cuda.empty_cache()
 
         # Store per-chunk latents for chunk-by-chunk VAE decode (matches diffusers behavior).

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py

@@ -46,6 +46,7 @@ from sglang.multimodal_gen.runtime.models.dits.mova_video_dit import (
 # Create aliases for backward compatibility
 video_sinusoidal_embedding_1d = sinusoidal_embedding_1d
 audio_sinusoidal_embedding_1d = sinusoidal_embedding_1d
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
@@ -62,7 +63,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
@@ -159,6 +159,32 @@ class MOVADenoisingStage(PipelineStage):
         self._cached_num_steps = None
         self._torch_compiled = False
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        uses = [
+            ComponentUse(stage_name, "audio_dit"),
+            ComponentUse(stage_name, "dual_tower_bridge"),
+            ComponentUse(
+                stage_name,
+                "video_dit",
+                phase="video_dit",
+                preferred_ready_after_request=True,
+                memory_intensive=True,
+            ),
+        ]
+        if self.video_dit_2 is not None:
+            uses.append(
+                ComponentUse(
+                    stage_name,
+                    "video_dit_2",
+                    phase="video_dit_2",
+                    memory_intensive=True,
+                )
+            )
+        return uses
+
     @property
     def parallelism_type(self) -> StageParallelismType:
         if get_global_server_args().enable_cfg_parallel:
@@ -329,27 +355,6 @@ class MOVADenoisingStage(PipelineStage):
     ) -> object | None:
         return None
 
-    def _manage_device_placement(
-        self,
-        model_to_use: nn.Module | None,
-        model_to_offload: nn.Module | None,
-        server_args: ServerArgs,
-    ):
-        if not server_args.dit_cpu_offload:
-            return
-
-        if (
-            model_to_offload is not None
-            and next(model_to_offload.parameters()).device.type == "cuda"
-        ):
-            model_to_offload.to("cpu")
-
-        if (
-            model_to_use is not None
-            and next(model_to_use.parameters()).device.type == "cpu"
-        ):
-            model_to_use.to(get_local_torch_device())
-
     def _select_visual_dit(
         self,
         timestep: float,
@@ -358,24 +363,52 @@ class MOVADenoisingStage(PipelineStage):
         scheduler,
     ):
         if boundary_ratio is None or self.video_dit_2 is None:
-            self._manage_device_placement(self.video_dit, None, server_args)
+            self._manage_video_dit_use(self.video_dit, "video_dit")
             return self.video_dit
 
         boundary_timestep = boundary_ratio * scheduler.num_train_timesteps
         if timestep >= boundary_timestep:
             current_model = self.video_dit
-            model_to_offload = self.video_dit_2
+            current_name = "video_dit"
         else:
             current_model = self.video_dit_2
-            model_to_offload = self.video_dit
+            current_name = "video_dit_2"
 
-        self._manage_device_placement(current_model, model_to_offload, server_args)
+        self._manage_video_dit_use(current_model, current_name)
         return current_model
 
+    def _manage_video_dit_use(
+        self, current_model: nn.Module, default_name: str
+    ) -> bool:
+        manager = self._component_residency_manager
+        if manager is None:
+            return False
+
+        component_name = manager.component_name_for_module(current_model, default_name)
+        use = ComponentUse(
+            stage_name=self._active_component_stage_name(),
+            component_name=component_name,
+            phase=component_name,
+            preferred_ready_after_request=component_name == "video_dit",
+            memory_intensive=True,
+        )
+        manager.begin_use(use, module=current_model)
+        return True
+
     def _ensure_shared_models_on_device(self, server_args: ServerArgs):
         """Ensure shared denoising modules are on the active device when cpu offload is enabled."""
-        self._manage_device_placement(self.audio_dit, None, server_args)
+        manager = self._component_residency_manager
-        self._manage_device_placement(self.dual_tower_bridge, None, server_args)
+        if manager is None:
+            return
+        stage_name = self._active_component_stage_name()
+        manager.ensure_ready(
+            ComponentUse(stage_name, "audio_dit"),
+            module=self.audio_dit,
+        )
+        manager.ensure_ready(
+            ComponentUse(stage_name, "dual_tower_bridge"),
+            module=self.dual_tower_bridge,
+        )
 
     def _apply_guidance_rescale(
         self,
@@ -598,9 +631,7 @@ class MOVADenoisingStage(PipelineStage):
                     if not is_warmup and hasattr(self, "step_profile"):
                         self.step_profile()
 
-        for dit in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]):
+        self._finish_active_component_use()
-            if isinstance(dit, OffloadableDiTMixin):
-                dit.prepare_for_next_req()
 
         return batch
 
@@ -911,6 +942,16 @@ class MOVADecodingStage(PipelineStage):
         self.video_vae = video_vae
         self.audio_vae = audio_vae
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        return [
+            ComponentUse(stage_name, "video_vae", target_dtype=vae_dtype),
+            ComponentUse(stage_name, "audio_vae"),
+        ]
+
     @property
     def parallelism_type(self) -> StageParallelismType:
         if get_global_server_args().enable_cfg_parallel:
@@ -919,36 +960,45 @@ class MOVADecodingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
-        self.video_vae = self.video_vae.to(get_local_torch_device())
-        self.audio_vae = self.audio_vae.to(get_local_torch_device())
-
-        video_latents = server_args.pipeline_config.denormalize_video_latents(
-            batch.latents, self.video_vae
-        )
-
         vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (
             vae_dtype != torch.float32
         ) and not server_args.disable_autocast
 
-        with torch.autocast(
+        with self.use_declared_component(
-            device_type=current_platform.device_type,
+            component_name="video_vae",
-            dtype=vae_dtype,
+            module=self.video_vae,
-            enabled=vae_autocast_enabled,
+        ) as video_vae:
-        ):
+            assert video_vae is not None
-            if server_args.pipeline_config.vae_tiling:
+            self.video_vae = video_vae
-                self.video_vae.enable_tiling()
+            video_latents = server_args.pipeline_config.denormalize_video_latents(
-            if not vae_autocast_enabled:
+                batch.latents, self.video_vae
-                video_latents = video_latents.to(vae_dtype)
+            )
-            decode_output = self.video_vae.decode(video_latents)
+
-            video = _ensure_tensor_decode_output(decode_output)
+            with torch.autocast(
+                device_type=current_platform.device_type,
+                dtype=vae_dtype,
+                enabled=vae_autocast_enabled,
+            ):
+                if server_args.pipeline_config.vae_tiling:
+                    self.video_vae.enable_tiling()
+                if not vae_autocast_enabled:
+                    video_latents = video_latents.to(vae_dtype)
+                decode_output = self.video_vae.decode(video_latents)
+                video = _ensure_tensor_decode_output(decode_output)
 
         video = (video / 2 + 0.5).clamp(0, 1)
 
-        with torch.autocast(
+        with self.use_declared_component(
-            device_type=current_platform.device_type, dtype=torch.float32
+            component_name="audio_vae",
-        ):
+            module=self.audio_vae,
-            audio = self.audio_vae.decode(batch.audio_latents)
+        ) as audio_vae:
+            assert audio_vae is not None
+            self.audio_vae = audio_vae
+            with torch.autocast(
+                device_type=current_platform.device_type, dtype=torch.float32
+            ):
+                audio = self.audio_vae.decode(batch.audio_latents)
         output_batch = OutputBatch(
             output=video,
             audio=audio,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py

@@ -8,6 +8,7 @@ from diffusers.image_processor import VaeImageProcessor
 from diffusers.utils.torch_utils import randn_tensor
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.models.vision_utils import load_image
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -117,13 +118,9 @@ class QwenImageLayeredBeforeDenoisingStage(PipelineStage):
         self.vae = vae.to(torch.bfloat16)
         from transformers import Qwen2_5_VLForConditionalGeneration
 
-        self.text_encoder = (
+        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_path, subfolder="text_encoder"
-                model_path, subfolder="text_encoder"
+        ).to(torch.bfloat16)
-            )
-            .to(get_local_torch_device())
-            .to(torch.bfloat16)
-        )
         self.tokenizer = tokenizer
         self.processor = processor
         self.transformer = transformer
@@ -158,6 +155,17 @@ generalizations\n - Describe all visible information in the image, while do not
 the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n"""
         self.default_sample_size = 128
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name, "qwen_layered_text_encoder", target_dtype=torch.bfloat16
+            ),
+            ComponentUse(stage_name, "vae", target_dtype=torch.bfloat16),
+        ]
+
     # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
     def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         bool_mask = mask.bool()
@@ -300,21 +308,23 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
 
     # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
     def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
-        self.vae = self.vae.to(get_local_torch_device())
+        with self.use_declared_component(component_name="vae", module=self.vae) as vae:
-        if isinstance(generator, list):
+            assert vae is not None
-            image_latents = [
+            self.vae = vae
-                retrieve_latents(
+            if isinstance(generator, list):
-                    self.vae.encode(image[i : i + 1]),
+                image_latents = [
-                    generator=generator[i],
+                    retrieve_latents(
-                    sample_mode="argmax",
+                        self.vae.encode(image[i : i + 1]),
+                        generator=generator[i],
+                        sample_mode="argmax",
+                    )
+                    for i in range(image.shape[0])
+                ]
+                image_latents = torch.cat(image_latents, dim=0)
+            else:
+                image_latents = retrieve_latents(
+                    self.vae.encode(image), generator=generator, sample_mode="argmax"
                 )
-                for i in range(image.shape[0])
-            ]
-            image_latents = torch.cat(image_latents, dim=0)
-        else:
-            image_latents = retrieve_latents(
-                self.vae.encode(image), generator=generator, sample_mode="argmax"
-            )
         latents_mean = (
             torch.tensor(self.vae.config.latents_mean)
             .view(1, self.latent_channels, 1, 1, 1)
@@ -326,7 +336,6 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             .to(image_latents.device, image_latents.dtype)
         )
         image_latents = (image_latents - latents_mean) / latents_std
-        self.vae.to("cpu")
         return image_latents
 
     def prepare_latents(
@@ -447,20 +456,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         image = image.to(dtype=torch.bfloat16)
 
         prompt = batch.prompt
-        if not prompt or prompt.isspace():
+        with self.use_declared_component(
-            prompt = self.get_image_caption(
+            component_name="qwen_layered_text_encoder",
-                prompt_image, use_en_prompt=use_en_prompt, device=device
+            module=self.text_encoder,
+        ) as text_encoder:
+            assert text_encoder is not None
+            self.text_encoder = text_encoder
+            if not prompt or prompt.isspace():
+                prompt = self.get_image_caption(
+                    prompt_image, use_en_prompt=use_en_prompt, device=device
+                )
+
+            prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+                prompt=prompt,
+                device=device,
             )
 
-        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-            prompt=prompt,
+                prompt=batch.negative_prompt,
-            device=device,
+                device=device,
-        )
+            )
-
-        negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-            prompt=batch.negative_prompt,
-            device=device,
-        )
 
         num_channels_latents = self.transformer.config.in_channels // 4
         latents, image_latents = self.prepare_latents(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py

@@ -33,6 +33,7 @@ def prepare_wan_ti2v_latents(
     vae: object,
     latents: torch.Tensor,
     target_dtype: torch.dtype,
+    vae_dtype: torch.dtype,
     batch: Req,
     server_args: ServerArgs,
 ) -> tuple[int, torch.Tensor, list[torch.Tensor]]:
@@ -43,10 +44,10 @@ def prepare_wan_ti2v_latents(
     assert batch.image_latent is None, "TI2V task should not have image latents"
     assert vae is not None, "VAE is not provided for TI2V task"
 
-    vae = vae.to(batch.condition_image.device)
+    condition_image = batch.condition_image.to(
-    z = vae.encode(batch.condition_image).mean.float()
+        device=get_local_torch_device(), dtype=vae_dtype
-    if getattr(vae, "device", None) != "cpu" and server_args.vae_cpu_offload:
+    )
-        vae = vae.to("cpu")
+    z = vae.encode(condition_image).mean.float()
 
     if hasattr(vae, "shift_factor") and vae.shift_factor is not None:
         if isinstance(vae.shift_factor, torch.Tensor):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -15,6 +15,7 @@ import torch
 
 from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
@@ -68,6 +69,19 @@ class TextEncodingStage(PipelineStage):
         self.tokenizers = tokenizers
         self.text_encoders = text_encoders
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name=stage_name,
+                component_name="text_encoder" if i == 0 else f"text_encoder_{i + 1}",
+                preferred_ready_after_request=i == 0,
+            )
+            for i in range(len(self.text_encoders))
+        ]
+
     @torch.no_grad()
     def forward(
         self,
@@ -166,6 +180,21 @@ class TextEncodingStage(PipelineStage):
 
         return tok_kwargs
 
+    def _manage_text_encoder_use(self, encoder_index: int) -> None:
+        manager = self._component_residency_manager
+        if manager is None:
+            return
+        component_name = (
+            "text_encoder"
+            if encoder_index == 0
+            else f"text_encoder_{encoder_index + 1}"
+        )
+        use = self._declared_component_use(component_name=component_name)
+        # TODO: Keep this begin-only interval until manager supports explicit
+        # declared-use interval grouping. Wrapping each encoder call separately
+        # can offload between positive and negative prompt encoding.
+        manager.before_use(use)
+
     def _forward_text_encoder(self, text_encoder, encoder_forward_kwargs):
         if not getattr(text_encoder, "uses_sglang_forward_context", True):
             return text_encoder(**encoder_forward_kwargs)
@@ -251,7 +280,7 @@ class TextEncodingStage(PipelineStage):
         embeds_list: list[torch.Tensor] = []
         pooled_embeds_list: list[torch.Tensor] = []
 
-        attn_masks_list: list[torch.Tensor] = []
+        attn_masks_list: list[torch.Tensor | None] = []
 
         preprocess_funcs = server_args.pipeline_config.preprocess_text_funcs
         postprocess_funcs = server_args.pipeline_config.postprocess_text_funcs
@@ -308,6 +337,7 @@ class TextEncodingStage(PipelineStage):
                 encoder_forward_kwargs["attention_mask"] = attention_mask
             if "use_cache" in inspect.signature(text_encoder.forward).parameters:
                 encoder_forward_kwargs["use_cache"] = False
+            self._manage_text_encoder_use(i)
             outputs: BaseEncoderOutput = self._forward_text_encoder(
                 text_encoder, encoder_forward_kwargs
             )
@@ -345,11 +375,19 @@ class TextEncodingStage(PipelineStage):
                         if postprocessed_attention_mask is not None
                         else None
                     )
-                elif attention_mask is not None:
+                elif attention_mask is not None and list(attention_mask.shape) == list(
+                    prompt_embeds.shape[:2]
+                ):
                     mask_to_store = attention_mask.to(device=target_device)
                 else:
                     mask_to_store = torch.ones(
-                        input_ids.shape[:2], device=target_device
+                        prompt_embeds.shape[:2],
+                        device=target_device,
+                        dtype=(
+                            attention_mask.dtype
+                            if attention_mask is not None
+                            else torch.long
+                        ),
                     )
                 attn_masks_list.append(mask_to_store)
 
@@ -379,13 +417,23 @@ class TextEncodingStage(PipelineStage):
                 )
         stacked_embeds = torch.stack(embeds_list, dim=0)
         if return_attention_mask:
-            base_mask_shape = list(attn_masks_list[0].shape)
+            stackable_masks = [
-            for m in attn_masks_list[1:]:
+                (
+                    mask
+                    if mask is not None
+                    else torch.ones(
+                        embed.shape[:2], device=embed.device, dtype=torch.long
+                    )
+                )
+                for embed, mask in zip(embeds_list, attn_masks_list, strict=True)
+            ]
+            base_mask_shape = list(stackable_masks[0].shape)
+            for m in stackable_masks[1:]:
                 if list(m.shape) != base_mask_shape:
                     raise ValueError(
-                        f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}"
+                        f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in stackable_masks]}"
                     )
-            stacked_masks = torch.stack(attn_masks_list, dim=0)
+            stacked_masks = torch.stack(stackable_masks, dim=0)
             return stacked_embeds, stacked_masks
         return stacked_embeds
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py

@@ -1,6 +1,7 @@
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
@@ -44,18 +45,10 @@ class LTX2LoRASwitchStage(PipelineStage):
         self.phase = phase
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
-        switch_fn = getattr(self.pipeline, "switch_lora_phase", None)
+        if self.pipeline.should_skip_ltx2_lora_switch_stage():
-        should_skip_switch_stage = getattr(
-            self.pipeline, "should_skip_ltx2_lora_switch_stage", None
-        )
-        if callable(should_skip_switch_stage) and should_skip_switch_stage():
             batch.extra["ltx2_phase"] = self.phase
             return batch
-        if not callable(switch_fn):
+        self.pipeline.switch_lora_phase(self.phase, batch=batch)
-            raise ValueError(
-                "LTX2LoRASwitchStage requires pipeline.switch_lora_phase()"
-            )
-        switch_fn(self.phase, batch=batch)
         batch.extra["ltx2_phase"] = self.phase
         return batch
 
@@ -63,13 +56,31 @@ class LTX2LoRASwitchStage(PipelineStage):
 class LTX2UpsampleStage(PipelineStage):
     """Upsample Stage-1 video latents and prepare Stage-2 inputs."""
 
-    def __init__(self, spatial_upsampler, vae, audio_vae=None, pipeline=None):
+    def __init__(
+        self,
+        spatial_upsampler,
+        vae,
+        audio_vae=None,
+        pipeline=None,
+    ):
         super().__init__()
         self.spatial_upsampler = spatial_upsampler
         self.vae = vae
         self.audio_vae = audio_vae
         self.pipeline = pipeline
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        uses = [
+            ComponentUse(stage_name, "spatial_upsampler"),
+            ComponentUse(stage_name, "vae"),
+        ]
+        if self.audio_vae is not None:
+            uses.append(ComponentUse(stage_name, "audio_vae"))
+        return uses
+
     def _upsample_video_latents(
         self, latents: torch.Tensor, server_args: ServerArgs, device: torch.device
     ) -> torch.Tensor:
@@ -118,25 +129,8 @@ class LTX2UpsampleStage(PipelineStage):
         )
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
-        delay_stage2_prefetch = False
-        if self.pipeline is not None:
-            prepare_upsample = getattr(
-                self.pipeline, "prepare_ltx2_upsample_after_stage1", None
-            )
-            if callable(prepare_upsample):
-                delay_stage2_prefetch = prepare_upsample()
-        prefetch_stage2 = (
-            getattr(self.pipeline, "prefetch_ltx2_stage2_after_stage1", None)
-            if self.pipeline is not None
-            else None
-        )
-        if callable(prefetch_stage2) and not delay_stage2_prefetch:
-            prefetch_stage2()
-
         device = get_local_torch_device()
         latents = self._upsample_video_latents(batch.latents, server_args, device)
-        if callable(prefetch_stage2) and delay_stage2_prefetch:
-            prefetch_stage2()
         logger.info("Upsampled video latents: %s", list(latents.shape))
         self._restore_full_resolution(batch)
         batch.image_latent = None

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

@@ -9,10 +9,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
 from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
     _ModelOptFp8OffloadAdapter,
 )
-from sglang.multimodal_gen.runtime.utils import (
+from sglang.multimodal_gen.runtime.managers import (
     layerwise_offload as layerwise_offload_mod,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+from sglang.multimodal_gen.runtime.managers.layerwise_offload import (
     LayerwiseOffloadManager,
 )