chore: import upstream snapshot with attribution

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

@@ -0,0 +1,174 @@
+"""Memory-aware ordering for pipeline component weight loads to avoid OOM while loading.
+
+Load the VRAM-intensive components earlier than others
+
+The pipeline owns component selection, path resolution, and actual loading; this
+module only ranks already-selected load specs.
+"""
+
+import glob
+import json
+import os
+from dataclasses import dataclass
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
+    is_dit_component_name,
+    is_image_encoder_component_name,
+    is_text_encoder_component_name,
+    is_vae_component_name,
+)
+
+
+@dataclass(frozen=True)
+class ComponentLoadSpec:
+    """One pipeline component that still needs a real weight load."""
+
+    module_name: str
+    load_module_name: str
+    component_model_path: str
+    transformers_or_diffusers: str
+    architecture: str | None
+    index: int
+
+
+_WEIGHT_FILE_SUFFIXES = (".bin", ".pt", ".pth")
+
+
+def _component_base_name(component_name: str) -> str:
+    prefix, separator, suffix = component_name.rpartition("_")
+    if separator and suffix.isdigit():
+        return prefix
+    return component_name
+
+
+def _component_variant_priority(component_name: str) -> int:
+    _, separator, suffix = component_name.rpartition("_")
+    if separator and suffix.isdigit():
+        return -int(suffix)
+    return 0
+
+
+def component_load_risk_rank(component_name: str) -> int:
+    """Fallback type rank when checkpoint size cannot be inferred."""
+    candidate_names = (component_name, _component_base_name(component_name))
+    if any(is_dit_component_name(name) for name in candidate_names):
+        return 0
+    if any(is_text_encoder_component_name(name) for name in candidate_names):
+        return 1
+    if any(is_image_encoder_component_name(name) for name in candidate_names):
+        return 2
+    if any(is_vae_component_name(name) for name in candidate_names):
+        return 3
+    return 10
+
+
+def _safe_file_size(file_path: str) -> int | None:
+    try:
+        return os.path.getsize(file_path)
+    except OSError:
+        return None
+
+
+def _safetensors_payload_size_bytes(file_path: str) -> int | None:
+    try:
+        with open(file_path, "rb") as f:
+            header_size_bytes = f.read(8)
+            if len(header_size_bytes) != 8:
+                return _safe_file_size(file_path)
+            header_size = int.from_bytes(header_size_bytes, "little")
+            header = json.loads(f.read(header_size))
+    except (OSError, json.JSONDecodeError, ValueError):
+        return _safe_file_size(file_path)
+
+    payload_size = 0
+    for tensor_name, tensor_info in header.items():
+        if tensor_name == "__metadata__":
+            continue
+        offsets = tensor_info.get("data_offsets")
+        if not isinstance(offsets, list) or len(offsets) != 2:
+            return _safe_file_size(file_path)
+        payload_size += offsets[1] - offsets[0]
+    return payload_size
+
+
+def _safetensors_files_from_index(component_model_path: str) -> list[str]:
+    indexed_files: set[str] = set()
+    index_paths = sorted(
+        glob.glob(os.path.join(component_model_path, "*.safetensors.index.json"))
+    )
+    for index_path in index_paths:
+        try:
+            with open(index_path) as f:
+                weight_map = json.load(f).get("weight_map", {})
+        except (OSError, json.JSONDecodeError):
+            continue
+        for shard_name in weight_map.values():
+            shard_path = os.path.join(component_model_path, shard_name)
+            if os.path.isfile(shard_path):
+                indexed_files.add(shard_path)
+    return sorted(indexed_files)
+
+
+def _list_component_safetensors_files(component_model_path: str) -> list[str]:
+    if os.path.isfile(component_model_path):
+        if component_model_path.endswith(".safetensors"):
+            return [component_model_path]
+        return []
+    if not os.path.isdir(component_model_path):
+        return []
+
+    indexed_files = _safetensors_files_from_index(component_model_path)
+    if indexed_files:
+        return indexed_files
+    return sorted(glob.glob(os.path.join(component_model_path, "*.safetensors")))
+
+
+def infer_component_weight_size_bytes(component_model_path: str) -> int | None:
+    """Infer checkpoint payload size from safetensors without materializing tensors."""
+    safetensors_files = _list_component_safetensors_files(component_model_path)
+    if safetensors_files:
+        sizes = [
+            size
+            for size in (
+                _safetensors_payload_size_bytes(file_path)
+                for file_path in safetensors_files
+            )
+            if size is not None
+        ]
+        return sum(sizes) if sizes else None
+
+    if os.path.isfile(component_model_path):
+        if component_model_path.endswith(_WEIGHT_FILE_SUFFIXES):
+            return _safe_file_size(component_model_path)
+        return None
+    if not os.path.isdir(component_model_path):
+        return None
+
+    weight_files = []
+    for suffix in _WEIGHT_FILE_SUFFIXES:
+        weight_files.extend(glob.glob(os.path.join(component_model_path, f"*{suffix}")))
+    if not weight_files:
+        return None
+    sizes = [
+        size
+        for size in (_safe_file_size(file_path) for file_path in weight_files)
+        if size is not None
+    ]
+    return sum(sizes) if sizes else None
+
+
+def order_component_load_specs(
+    component_specs: list[ComponentLoadSpec],
+) -> list[ComponentLoadSpec]:
+    # load larger weight payloads before small helpers to reduce startup peak OOMs
+    return sorted(
+        component_specs,
+        key=lambda spec: (
+            # 1. model size inferred from checkpoints
+            -(infer_component_weight_size_bytes(spec.component_model_path) or 0),
+            # 2. infer from component name
+            component_load_risk_rank(spec.load_module_name),
+            _component_variant_priority(spec.load_module_name),
+            spec.index,
+        ),
+    )

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

@@ -0,0 +1,619 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Mapping, MutableMapping, Protocol, Sequence
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
+    ComponentResidencyStrategy,
+    LayerwiseOffloadStrategy,
+    ResidentStrategy,
+    VanillaD2HStrategy,
+    is_fsdp_managed_module,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    is_layerwise_offloaded_module,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
+    is_dit_component_name,
+    is_image_encoder_component_name,
+    is_text_encoder_component_name,
+    is_vae_component_name,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import DiffusionNvtxHooks
+
+logger = init_logger(__name__)
+
+
+@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
+
+
+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 should_cpu_offload_component(
+    component_name: str, module: nn.Module, server_args: ServerArgs
+) -> bool:
+    if server_args.use_fsdp_inference or is_fsdp_managed_module(module):
+        return False
+    if is_dit_component_name(component_name):
+        return bool(server_args.dit_cpu_offload)
+    if is_text_encoder_component_name(component_name):
+        return bool(server_args.text_encoder_cpu_offload)
+    if is_image_encoder_component_name(component_name):
+        return bool(server_args.image_encoder_cpu_offload)
+    if is_vae_component_name(component_name):
+        return bool(server_args.vae_cpu_offload)
+    return False
+
+
+def build_component_residency_strategy(
+    component_name: str,
+    module: nn.Module,
+    server_args: ServerArgs,
+) -> ComponentResidencyStrategy:
+    if is_layerwise_offloaded_module(module):
+        return LayerwiseOffloadStrategy()
+    if should_cpu_offload_component(component_name, module, server_args):
+        return VanillaD2HStrategy()
+    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()
+        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._active_nvtx_key: tuple[str, str, str | None] | None = None
+        self._nvtx_hooks_by_use_key: dict[
+            tuple[str, str, str | None], tuple[int, DiffusionNvtxHooks]
+        ] = {}
+        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] = {}
+
+    def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None:
+        custom_strategies = dict(pipeline.component_residency_strategies)
+        if pipeline is not self.pipeline:
+            self._remove_nvtx_hooks()
+            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 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)
+        self._active_use = None
+        self._active_use_module = None
+        self._disable_active_nvtx()
+        self._current_use_index = -1
+        self._prefetched_use_keys.clear()
+        self._uses_seen.clear()
+        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
+        )
+
+    def before_stage(
+        self,
+        stage: ComponentResidencyStage,
+        stage_index: int,
+        batch: ResidencyBatch,
+        server_args: ServerArgs,
+    ) -> None:
+        """called after stage starts"""
+        # 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)
+
+    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):
+            if self._use_key(self._active_use) != self._use_key(use):
+                self._mark_current_use(use)
+                self._active_use = use
+                self.state.current_use = use
+            self._enable_nvtx_for_use(
+                use,
+                module
+                or self._active_use_module
+                or self.get_module(use.component_name),
+            )
+            return
+        if self._active_use is not None:
+            self._disable_active_nvtx()
+            # 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)
+        module = self._prepare_forward_use(use, module=module)
+        self._active_use = use
+        self._active_use_module = module
+        self._enable_nvtx_for_use(use, 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._disable_active_nvtx()
+        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 ensure_ready(self, use: ComponentUse, module: nn.Module | None = None) -> None:
+        """Prepare a shared component and wait without making it the active use."""
+        self._prepare_forward_use(use, module=module)
+
+    def remove_nvtx_hooks_for_module(self, module: nn.Module | None) -> None:
+        """Detach NVTX hooks before a component object is deleted or replaced."""
+        if module is None:
+            return
+        module_id = id(module)
+        for key, (registered_id, hooks) in list(self._nvtx_hooks_by_use_key.items()):
+            if registered_id != module_id:
+                continue
+            if self._active_nvtx_key == key:
+                hooks.set_enabled(False)
+                self._active_nvtx_key = None
+            hooks.remove_hooks()
+            del self._nvtx_hooks_by_use_key[key]
+
+    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._disable_active_nvtx()
+        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
+    ) -> nn.Module | 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:
+            return None
+        strategy = self.strategy_for(use.component_name, module)
+        self._uses_seen[use.component_name] = use
+        self.state.current_use = use
+        strategy.prepare_for_use(module, use, self.state)
+        strategy.wait_for_use(module, use, self.state)
+        return module
+
+    def _enable_nvtx_for_use(
+        self, use: ComponentUse, module: nn.Module | None = None
+    ) -> None:
+        if (
+            not self.server_args.enable_layerwise_nvtx_marker
+            or self.state.batch_is_warmup
+            or not isinstance(module, nn.Module)
+        ):
+            self._disable_active_nvtx()
+            return
+
+        key = self._use_key(use)
+        if self._active_nvtx_key != key:
+            self._disable_active_nvtx()
+
+        module_id = id(module)
+        existing = self._nvtx_hooks_by_use_key.get(key)
+        if existing is None or existing[0] != module_id:
+            if existing is not None:
+                existing[1].remove_hooks()
+                self._nvtx_hooks_by_use_key.pop(key, None)
+            hooks = DiffusionNvtxHooks()
+            prefix = self._nvtx_prefix_for_use(use)
+            total = hooks.register_hooks(module, prefix=prefix)
+            if total == 0:
+                return
+            logger.debug(
+                "[component_residency] Registered NVTX hooks for %s on %d submodules",
+                prefix,
+                total,
+            )
+            self._nvtx_hooks_by_use_key[key] = (module_id, hooks)
+        else:
+            hooks = existing[1]
+
+        hooks.set_enabled(True)
+        self._active_nvtx_key = key
+
+    def _disable_active_nvtx(self) -> None:
+        if self._active_nvtx_key is None:
+            return
+        existing = self._nvtx_hooks_by_use_key.get(self._active_nvtx_key)
+        if existing is not None:
+            existing[1].set_enabled(False)
+        self._active_nvtx_key = None
+
+    def _remove_nvtx_hooks(self) -> None:
+        self._disable_active_nvtx()
+        for _, hooks in self._nvtx_hooks_by_use_key.values():
+            hooks.remove_hooks()
+        self._nvtx_hooks_by_use_key.clear()
+
+    @staticmethod
+    def _nvtx_prefix_for_use(use: ComponentUse) -> str:
+        parts = [use.stage_name, use.component_name]
+        if use.phase is not None and use.phase != use.component_name:
+            parts.append(use.phase)
+        return ".".join(parts)
+
+    def _prefetch_use(self, use: ComponentUse) -> None:
+        """Prepare a future component opportunistically without waiting.
+
+        This is called for memory-intensive future uses where H2D placement can
+        overlap with the current stage.
+        """
+        if not use.allow_prefetch:
+            return
+        module = self.get_module(use.component_name)
+        if module is None:
+            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.
+            return
+
+        self._uses_seen[use.component_name] = use
+        if strategy.prefetch_for_use(module, use, self.state):
+            self._prefetched_use_keys.add(self._use_key(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:
+            return
+        should_keep = (
+            keep_on_warmup and self.state.batch_is_warmup
+        ) or self._should_keep_after_use(use)
+        if should_keep:
+            return
+        strategy = self.strategy_for(use.component_name, module)
+        was_on_cuda = self._module_on_cuda(module)
+        strategy.finish_use(module, use, self.state)
+        self._empty_cache_after_large_release(use, strategy, module, was_on_cuda)
+
+    def finish_request(self) -> None:
+        # 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:
+                continue
+            preferred = component_name in preferred_uses
+            if not preferred and self._should_keep_single_dit(component_name):
+                continue
+            strategy = self.strategy_for(component_name, module)
+            if preferred and not self.state.batch_is_warmup:
+                strategy.prepare_after_request(module, use, self.state)
+            else:
+                was_on_cuda = self._module_on_cuda(module)
+                strategy.finish_request(module, use, self.state, preferred=preferred)
+                self._empty_cache_after_large_release(
+                    use, strategy, module, was_on_cuda
+                )
+
+    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 _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 _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
+
+    def _module_on_cuda(self, module: nn.Module | None) -> bool:
+        return self._module_device(module) == "cuda"
+
+    def _empty_cache_after_large_release(
+        self,
+        use: ComponentUse,
+        strategy: ComponentResidencyStrategy,
+        module: nn.Module,
+        was_on_cuda: bool,
+    ) -> None:
+        """explicitly empty cache after potential release of large component"""
+        if not use.memory_intensive:
+            return
+        released_cuda_storage = was_on_cuda and not self._module_on_cuda(module)
+        released_layerwise_storage = isinstance(strategy, LayerwiseOffloadStrategy)
+        if not (released_cuda_storage or released_layerwise_storage):
+            return
+        if not torch.get_device_module().is_available():
+            return
+        torch.get_device_module().empty_cache()
+
+
+_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/memory_managers/component_resident_strategies.py

@@ -0,0 +1,508 @@
+"""
+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.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+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.memory_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
+
+
+def is_fsdp_managed_module(module: nn.Module) -> bool:
+    return module.__class__.__name__.startswith("FSDP")
+
+
+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:
+        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 is_fsdp_managed_module(module):
+            return
+        _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, LayerwiseOffloadableModuleMixin):
+            module.prepare_for_next_req()
+
+    def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
+        if not isinstance(module, LayerwiseOffloadableModuleMixin):
+            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/memory_managers/layerwise_offload.py

@@ -0,0 +1,816 @@
+import re
+from collections.abc import Mapping, Sequence
+from typing import Any, Dict, List, Set, Tuple
+
+import torch
+from torch.distributed.tensor import DTensor
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
+    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
+    LAYERWISE_OFFLOAD_DIT_GROUP,
+    layerwise_component_matches_any_selection,
+    normalize_layerwise_offload_components,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# Adapted from skywork AI Infra diffusion optimize
+class LayerwiseOffloadManager:
+    """A lightweight layerwise CPU offload manager.
+
+    This utility offloads per-layer parameters/buffers from GPU to CPU, and
+    supports async H2D prefetch using a dedicated CUDA stream.
+
+    Typical usage:
+    - Construct the manager with the target model and the list-like module
+      attribute that represents transformer blocks (e.g. ``blocks``).
+    - Call :meth:`initialize` once to offload weights and prefetch layer 0.
+    - During forward, call :meth:`prefetch_layer` for the next layer and
+      :meth:`release_layer` for the finished layer.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        *,
+        layers_attr_str: str,
+        num_layers: int,
+        enabled: bool,
+        pin_cpu_memory: bool = True,
+        prefetch_size: int = 1,
+    ) -> None:
+        self.model = model
+        self.layers_attr_str = layers_attr_str
+        self.num_layers = num_layers
+        self.pin_cpu_memory = pin_cpu_memory
+        self.prefetch_size = min(max(1, prefetch_size), self.num_layers)
+        self.enabled = bool(enabled and torch.get_device_module().is_available())
+        if not self.enabled:
+            return
+        self.device = torch.device(
+            current_platform.device_type, torch.get_device_module().current_device()
+        )
+        self.copy_stream = torch.get_device_module().Stream()
+
+        self._layer_name_re = re.compile(
+            rf"(^|\.){re.escape(layers_attr_str)}\.(\d+)(\.|$)"
+        )
+
+        # layer_idx -> {dtype: consolidated_pinned_cpu_tensor}
+        # stores the consolidated weight from a same layer, of same dtype
+        self._consolidated_cpu_weights: Dict[int, Dict[torch.dtype, torch.Tensor]] = {}
+        # layer_idx -> {name: pinned_cpu_tensor_with_original_stride}
+        # stores tensors whose original non-contiguous stride/layout must be preserved
+        self._strided_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
+        # layer_idx -> {name: {dtype, offset, numel, shape}}
+        # stores the offset and numel of each weight from a same layer, of same dtype
+        self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
+        # layer indices that are already in gpu
+        self._gpu_layers: Set[int] = set()
+        # layer_idx -> torch.get_device_module().Event for fine-grained sync, to make sure the weight is resident in pre-hook
+        self._prefetch_events: Dict[int, torch.get_device_module().Event] = {}
+
+        self._named_parameters: Dict[str, torch.nn.Parameter] = {}
+        self._named_buffers: Dict[str, torch.Tensor] = {}
+        self._offload_placeholders: Dict[torch.dtype, torch.Tensor] = {}
+        self._has_dtensor_weights = False
+        # Store forward hooks for removal
+        self._forward_hooks: List[Any] = []
+
+        self._initialize()
+
+    def _match_layer_idx(self, name: str) -> int | None:
+        m = self._layer_name_re.search(name)
+        if not m:
+            return None
+        try:
+            return int(m.group(2))
+        except Exception:
+            return None
+
+    def _get_shared_empty_tensor(self, dtype: torch.dtype) -> torch.Tensor:
+        placeholder = self._offload_placeholders.get(dtype)
+        if placeholder is None:
+            placeholder = torch.empty((1,), device=self.device, dtype=dtype)
+            self._offload_placeholders[dtype] = placeholder
+        return placeholder
+
+    @staticmethod
+    def _to_local_tensor(tensor: torch.Tensor) -> torch.Tensor:
+        if isinstance(tensor, DTensor):
+            return tensor.to_local()
+        return tensor
+
+    def _wrap_for_target(
+        self, target: torch.Tensor, local_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        if isinstance(target, DTensor):
+            return DTensor.from_local(
+                local_tensor, target.device_mesh, target.placements
+            )
+        return local_tensor
+
+    def _get_shared_empty_tensor_for_target(
+        self, target: torch.Tensor, dtype: torch.dtype
+    ) -> torch.Tensor:
+        return self._wrap_for_target(target, self._get_shared_empty_tensor(dtype))
+
+    @staticmethod
+    def _get_alignment_numel(dtype: torch.dtype, alignment_bytes: int = 32) -> int:
+        element_size = torch.empty((), dtype=dtype).element_size()
+        return max(1, alignment_bytes // element_size)
+
+    @classmethod
+    def _align_numel_offset(
+        cls, offset: int, dtype: torch.dtype, alignment_bytes: int = 32
+    ) -> int:
+        alignment_numel = cls._get_alignment_numel(dtype, alignment_bytes)
+        remainder = offset % alignment_numel
+        if remainder == 0:
+            return offset
+        return offset + alignment_numel - remainder
+
+    @torch.compiler.disable
+    def _initialize(self) -> None:
+        if not self.enabled:
+            return
+
+        self._named_parameters = dict(self.model.named_parameters())
+        self._named_buffers = dict(self.model.named_buffers())
+
+        # 1. collect and group layer parameters by dtype. Keep buffers resident:
+        # shared buffers such as RoPE caches may be referenced by many layers.
+        layer_groups: Dict[int, Dict[torch.dtype, List[Tuple[str, torch.Tensor]]]] = {}
+        for name, tensor in self._named_parameters.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            self._has_dtensor_weights = self._has_dtensor_weights or isinstance(
+                tensor, DTensor
+            )
+            local_tensor = self._to_local_tensor(tensor)
+            layer_groups.setdefault(layer_idx, {}).setdefault(
+                local_tensor.dtype, []
+            ).append((name, tensor))
+
+        # 2. concat and offload (in pinned memory)
+        for layer_idx, dtype_to_params in layer_groups.items():
+            self._consolidated_cpu_weights[layer_idx] = {}
+            self._strided_cpu_weights[layer_idx] = {}
+            self._weight_metadata[layer_idx] = {}
+
+            for dtype, weights in dtype_to_params.items():
+                contiguous_weights: List[Tuple[str, torch.Tensor, torch.Tensor]] = []
+                for name, weight in weights:
+                    local_weight = self._to_local_tensor(weight)
+                    if local_weight.is_contiguous():
+                        contiguous_weights.append((name, weight, local_weight))
+                        continue
+
+                    # Preserve non-contiguous layouts such as the transposed FP8
+                    # weight views expected by CUTLASS kernels.
+                    cpu_tensor = torch.empty_strided(
+                        size=local_weight.shape,
+                        stride=local_weight.stride(),
+                        dtype=dtype,
+                        pin_memory=self.pin_cpu_memory,
+                    )
+                    cpu_tensor.copy_(local_weight)
+                    self._strided_cpu_weights[layer_idx][name] = cpu_tensor
+                    self._weight_metadata[layer_idx][name] = {
+                        "dtype": dtype,
+                        "shape": local_weight.shape,
+                        "stride": local_weight.stride(),
+                        "preserve_strides": True,
+                    }
+                    weight.data = self._get_shared_empty_tensor_for_target(
+                        weight, dtype
+                    )
+
+                if not contiguous_weights:
+                    continue
+
+                current_offset = 0
+                aligned_offsets: Dict[str, int] = {}
+                for name, weight, local_weight in contiguous_weights:
+                    # Some fused diffusion kernels require tensor base pointers to
+                    # satisfy a 32-byte alignment contract. Reusing one flat buffer
+                    # is still fine, but each logical tensor slice must start on an
+                    # aligned offset inside that buffer.
+                    current_offset = self._align_numel_offset(current_offset, dtype)
+                    aligned_offsets[name] = current_offset
+                    current_offset += local_weight.numel()
+
+                total_numel = current_offset
+
+                # create concatenated CPU buffer (in pinned memory)
+                cpu_buffer = torch.empty(
+                    total_numel, dtype=dtype, pin_memory=self.pin_cpu_memory
+                )
+
+                # offload weights to the buffer
+                for name, weight, local_weight in contiguous_weights:
+                    current_offset = aligned_offsets[name]
+                    numel = local_weight.numel()
+                    cpu_buffer[current_offset : current_offset + numel].copy_(
+                        local_weight.flatten()
+                    )
+                    self._weight_metadata[layer_idx][name] = {
+                        "dtype": dtype,
+                        "offset": current_offset,
+                        "numel": numel,
+                        "shape": local_weight.shape,
+                        "stride": local_weight.stride(),
+                        "preserve_strides": False,
+                    }
+
+                    weight.data = self._get_shared_empty_tensor_for_target(
+                        weight, dtype
+                    )
+
+                    current_offset += numel
+
+                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.
+        if not self._has_dtensor_weights:
+            self.model.to(self.device)
+
+        # prefetch the first layer for warm-up
+        self.prepare_for_next_req(non_blocking=False)
+
+        self.register_forward_hooks()
+        logger.info(
+            f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, total num layers: {self.num_layers}"
+        )
+
+    def prepare_for_next_req(self, non_blocking=True):
+        """
+        Prepare for the next round of denoising loop with prefetching the necessary layers
+        """
+        for i in range(self.prefetch_size):
+            self.prefetch_layer(i, non_blocking=non_blocking)
+        if not non_blocking and self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+    def get_target_with_name(self, name: str) -> torch.Tensor:
+        """get the target model weight/buffer to be replaced"""
+        if name in self._named_parameters:
+            target = self._named_parameters[name]
+        else:
+            target = self._named_buffers[name]
+        return target
+
+    @torch.compiler.disable
+    def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
+        """
+        idempotent
+        """
+        if not self.enabled or self.device is None or self.copy_stream is None:
+            return
+        if layer_idx < 0 or layer_idx >= self.num_layers:
+            return
+        if layer_idx in self._gpu_layers:
+            return
+        if layer_idx not in self._consolidated_cpu_weights:
+            return
+        self.copy_stream.wait_stream(torch.get_device_module().current_stream())
+
+        # create gpu buffer and load from CPU buffer
+        gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
+        with (
+            torch.inference_mode(False),
+            torch.no_grad(),
+            torch.get_device_module().stream(self.copy_stream),
+        ):
+            for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items():
+                gpu_buffer = torch.empty(
+                    cpu_buffer.shape, dtype=dtype, device=self.device
+                )
+                gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking)
+                gpu_buffers[dtype] = gpu_buffer
+
+            # restore model's weights by their metadata using the same copy stream
+            # so the recorded event covers both flat-buffer and stride-preserving copies.
+            for name, meta in self._weight_metadata[layer_idx].items():
+                target = self.get_target_with_name(name)
+                if meta.get("preserve_strides", False):
+                    # Recreate the original view layout instead of flatten+view.
+                    # ModelOpt FP8 relies on a transposed runtime weight layout,
+                    # so preserving stride is part of correctness, not just an
+                    # optimization detail.
+                    cpu_tensor = self._strided_cpu_weights[layer_idx][name]
+                    gpu_tensor = torch.empty_strided(
+                        size=meta["shape"],
+                        stride=meta["stride"],
+                        dtype=meta["dtype"],
+                        device=self.device,
+                    )
+                    gpu_tensor.copy_(cpu_tensor, non_blocking=non_blocking)
+                    target.data = self._wrap_for_target(target, gpu_tensor)
+                    continue
+
+                dtype = meta["dtype"]
+                gpu_buffer = gpu_buffers[dtype]
+
+                # map the parameter's data to the correct slice of the GPU buffer
+                local_tensor = gpu_buffer[
+                    meta["offset"] : meta["offset"] + meta["numel"]
+                ].view(meta["shape"])
+                target.data = self._wrap_for_target(target, local_tensor)
+
+        # record the prefetch event of this layer after all copies are enqueued
+        event = torch.get_device_module().Event()
+        event.record(self.copy_stream)
+        self._prefetch_events[layer_idx] = event
+
+        self._gpu_layers.add(layer_idx)
+
+    @torch.compiler.disable
+    def release_layer(self, layer_idx: int) -> None:
+        """
+        lightweight release layer weights
+        Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched
+        """
+        if not self.enabled or self.device is None:
+            return
+
+        # clear prefetch event, since it's useless and needs to be reset
+        self._prefetch_events.pop(layer_idx, None)
+
+        if layer_idx not in self._gpu_layers:
+            return
+
+        with torch.inference_mode(False), torch.no_grad():
+            for name, meta in self._weight_metadata.get(layer_idx, {}).items():
+                target = self.get_target_with_name(name)
+                # Wraparound prefetch will reload the layer when it is needed again
+                target.data = self._get_shared_empty_tensor_for_target(
+                    target, meta["dtype"]
+                )
+
+        self._gpu_layers.discard(layer_idx)
+
+    @torch.compiler.disable
+    def release_all(self) -> None:
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in list(self._gpu_layers):
+            self.release_layer(layer_idx)
+
+    @torch.compiler.disable
+    def load_all_layers(self) -> None:
+        """Load all layers from CPU to GPU."""
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in range(self.num_layers):
+            if layer_idx not in self._gpu_layers:
+                self.prefetch_layer(layer_idx, non_blocking=False)
+
+    @torch.compiler.disable
+    def sync_layer_to_cpu(self, layer_idx: int) -> None:
+        """Sync a layer's weights from GPU back to CPU."""
+        if not self.enabled or layer_idx not in self._gpu_layers:
+            return
+        if layer_idx not in self._consolidated_cpu_weights:
+            return
+
+        if self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+        # Collect current GPU weights and write back to CPU buffer
+        for name, meta in self._weight_metadata.get(layer_idx, {}).items():
+            target = self.get_target_with_name(name)
+            target_local = self._to_local_tensor(target)
+            if meta.get("preserve_strides", False):
+                self._strided_cpu_weights[layer_idx][name].copy_(target_local.cpu())
+                continue
+
+            gpu_weight = target_local.flatten().cpu()
+
+            dtype = meta["dtype"]
+            cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype]
+            offset = meta["offset"]
+            numel = meta["numel"]
+            cpu_buffer[offset : offset + numel].copy_(gpu_weight)
+
+    @torch.compiler.disable
+    def sync_all_layers_to_cpu(self) -> None:
+        """Sync all loaded layers' weights from GPU back to CPU."""
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.get_device_module().current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in list(self._gpu_layers):
+            self.sync_layer_to_cpu(layer_idx)
+
+    @torch.compiler.disable
+    def update_cpu_weights(
+        self, weight_dict: Dict[str, torch.Tensor]
+    ) -> Set[str] | None:
+        """Update consolidated CPU buffers with new weights.
+
+        When layerwise offload (--dit-layerwise-offload) is enabled, the
+        offload manager replaces GPU parameters with small torch.empty((1,))
+        placeholders while real weights live in consolidated pinned CPU
+        buffers.
+
+        The refit process writes new weights directly into the CPU buffers,
+        bypassing the placeholders.  For any layer that happens to be resident
+        on the GPU at update time, the live GPU tensor is also updated.
+
+        Args:
+            weight_dict: Mapping of parameter name to new weight tensor.
+
+        Returns:
+            Set of parameter names that were successfully updated.
+
+        Raises:
+            ValueError: If a weight's shape does not match the recorded
+                metadata (i.e., the real shape, not the placeholder shape).
+        """
+        if not self.enabled:
+            return None
+
+        updated_names: Set[str] = set()
+        for name, loaded_weight in weight_dict.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None:
+                continue
+            meta_layer = self._weight_metadata.get(layer_idx)
+            if meta_layer is None or name not in meta_layer:
+                continue
+
+            meta = meta_layer[name]
+            local_loaded_weight = self._to_local_tensor(loaded_weight)
+            if tuple(meta["shape"]) != tuple(local_loaded_weight.shape):
+                raise ValueError(
+                    f"Shape mismatch for {name}: "
+                    f"expected={tuple(meta['shape'])}, "
+                    f"loaded={tuple(local_loaded_weight.shape)}"
+                )
+
+            dtype = meta["dtype"]
+            if meta.get("preserve_strides", False):
+                self._strided_cpu_weights[layer_idx][name].copy_(
+                    local_loaded_weight.to(dtype=dtype)
+                )
+            else:
+                offset = meta["offset"]
+                numel = meta["numel"]
+                cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype]
+                cpu_buffer[offset : offset + numel].copy_(
+                    local_loaded_weight.to(dtype=dtype).flatten()
+                )
+
+            # If this layer is currently on GPU, update the live parameter.
+            if layer_idx in self._gpu_layers:
+                target = self.get_target_with_name(name)
+                target_local = self._to_local_tensor(target)
+                target_local.copy_(local_loaded_weight.to(dtype=target_local.dtype))
+
+            updated_names.add(name)
+
+        return updated_names
+
+    def iter_cpu_weights(self):
+        """Yield (name, tensor) pairs from consolidated CPU buffers.
+
+        This reconstructs the original weight tensors (with correct shapes)
+        from the flat CPU buffers using stored metadata.  Unlike
+        model.named_parameters(), which returns (1,) placeholders
+        when offload is enabled, this method returns the real weights and
+        can be used for checksum computation.
+        """
+        for layer_idx in sorted(self._weight_metadata):
+            for name, meta in self._weight_metadata[layer_idx].items():
+                if meta.get("preserve_strides", False):
+                    # Some quantized weights rely on a non-contiguous layout.
+                    # Yield the strided tensor directly instead of rebuilding it
+                    # from the flat buffer, which would silently lose the
+                    # original stride information.
+                    yield name, self._strided_cpu_weights[layer_idx][name]
+                    continue
+
+                dtype = meta["dtype"]
+                offset = meta["offset"]
+                numel = meta["numel"]
+                shape = meta["shape"]
+                cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype]
+                yield name, cpu_buffer[offset : offset + numel].reshape(shape)
+
+    def register_forward_hooks(self) -> None:
+        if not self.enabled:
+            return
+
+        layers = dict(self.model.named_modules())[self.layers_attr_str]
+
+        def make_pre_hook(i):
+            def hook(module, input):
+                if i == 0:
+                    self.prepare_for_next_req(non_blocking=False)
+                if i not in self._gpu_layers:
+                    # LTX audio VAE traverses decoder.up in reverse order
+                    self.prefetch_layer(i, non_blocking=False)
+                if i in self._prefetch_events:
+                    torch.get_device_module().current_stream().wait_event(
+                        self._prefetch_events[i]
+                    )
+
+                # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed
+                if i % self.prefetch_size == 0:
+                    for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size):
+                        layer_to_prefetch = j % self.num_layers
+                        self.prefetch_layer(layer_to_prefetch, non_blocking=True)
+
+            return hook
+
+        def make_post_hook(i):
+            def hook(module, input, output):
+                # previous, we wait here, until the copy stream for next layer is finished,
+                # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer
+                self.release_layer(i)
+
+            return hook
+
+        # register prefetch & release hooks for each layer
+        self._forward_hooks.clear()
+        for i, layer in enumerate(layers):
+            pre_hook_handle = layer.register_forward_pre_hook(make_pre_hook(i))
+            post_hook_handle = layer.register_forward_hook(make_post_hook(i))
+            self._forward_hooks.extend([pre_hook_handle, post_hook_handle])
+
+    def remove_forward_hooks(self) -> None:
+        """Remove all registered forward hooks."""
+        for hook_handle in self._forward_hooks:
+            hook_handle.remove()
+        self._forward_hooks.clear()
+
+
+class LayerwiseOffloadableModuleMixin:
+    """A mixin that registers forward hooks to enable layerwise offload."""
+
+    # whether the current module is selected by the `dit` group
+    layerwise_offload_dit_group_enabled: bool = True
+
+    # The list of names of this module's layer/block ModuleList or Sequential attributes.
+    layer_names: List[str] = []
+    layerwise_offload_managers: list[LayerwiseOffloadManager] = []
+
+    def configure_layerwise_offload(self, server_args: ServerArgs):
+        self.layerwise_offload_managers = []
+        named_modules = dict(self.named_modules())
+        configured_layer_names = []
+        for layer_name in self.layer_names:
+            module_list = named_modules.get(layer_name)
+            if not isinstance(module_list, (torch.nn.ModuleList, torch.nn.Sequential)):
+                continue
+            if len(module_list) == 0:
+                continue
+
+            num_layers = len(module_list)
+            if server_args.dit_offload_prefetch_size < 1.0:
+                prefetch_size = 1 + int(
+                    round(server_args.dit_offload_prefetch_size * (num_layers - 1))
+                )
+            else:
+                prefetch_size = int(server_args.dit_offload_prefetch_size)
+
+            manager = LayerwiseOffloadManager(
+                model=self,
+                layers_attr_str=layer_name,
+                num_layers=num_layers,
+                enabled=True,
+                pin_cpu_memory=server_args.pin_cpu_memory,
+                prefetch_size=prefetch_size,
+            )
+            self.layerwise_offload_managers.append(manager)
+            configured_layer_names.append(layer_name)
+
+        if configured_layer_names:
+            logger.info(
+                "Enabled layerwise offload for %s on modules: %s",
+                self.__class__.__name__,
+                configured_layer_names,
+            )
+        else:
+            logger.info(
+                "No layerwise-offloadable ModuleList found for %s. Candidates: %s",
+                self.__class__.__name__,
+                self.layer_names,
+            )
+
+    def prepare_for_next_req(self):
+        if self.layerwise_offload_managers is None:
+            return
+        for manager in self.layerwise_offload_managers:
+            manager.prepare_for_next_req(non_blocking=True)
+
+    def disable_offload(self) -> None:
+        """Disable layerwise offload: load all layers to GPU and remove hooks."""
+        if self.layerwise_offload_managers is None:
+            return
+        for manager in self.layerwise_offload_managers:
+            if manager.enabled:
+                manager.remove_forward_hooks()
+                manager.load_all_layers()
+
+    def enable_offload(self) -> None:
+        """Re-enable layerwise offload: sync weights to CPU, release layers, and restore hooks."""
+        if self.layerwise_offload_managers is None:
+            return
+        for manager in self.layerwise_offload_managers:
+            if manager.enabled:
+                manager.sync_all_layers_to_cpu()
+                manager.release_all()
+                manager.register_forward_hooks()
+
+
+def iter_materialized_weights(module: torch.nn.Module):
+    """Yield (name, tensor) pairs with materialized weights, even under offload.
+
+    When layerwise offload is active, module.named_parameters() returns
+    (1,) placeholders for offloaded layers.  This function reads the
+    actual data from the offload manager's CPU buffers and chains it with
+    the non-offloaded parameters.
+    """
+    offload_managers: list = []
+    if is_layerwise_offloaded_module(module):
+        offload_managers = [m for m in module.layerwise_offload_managers if m.enabled]
+
+    if not offload_managers:
+        yield from module.named_parameters()
+        return
+
+    # Collect offloaded names and their real tensors from CPU buffers.
+    offloaded_names: set[str] = set()
+    for manager in offload_managers:
+        for name, tensor in manager.iter_cpu_weights():
+            offloaded_names.add(name)
+            yield name, tensor
+
+    # Yield non-offloaded parameters (e.g. final norms, embeddings).
+    for name, param in module.named_parameters():
+        if name not in offloaded_names:
+            yield name, param
+
+
+def is_layerwise_offloaded_module(module: torch.nn.Module) -> bool:
+    return isinstance(module, LayerwiseOffloadableModuleMixin) and any(
+        manager.enabled for manager in module.layerwise_offload_managers
+    )
+
+
+def get_layerwise_offload_component_names_for_pipeline(
+    modules: Mapping[str, object],
+    component_names: Sequence[str] | None = None,
+) -> list[str]:
+    """Resolve layerwise selectors against the current pipeline modules.
+
+    Explicit unsupported component names are kept so callers can report them.
+    """
+    normalized_component_names = normalize_layerwise_offload_components(component_names)
+    selected_component_names = (
+        set(normalized_component_names)
+        if normalized_component_names is not None
+        else None
+    )
+
+    if selected_component_names is None:
+        return [
+            component_name
+            for component_name, module in modules.items()
+            if isinstance(module, LayerwiseOffloadableModuleMixin)
+            and module.layerwise_offload_dit_group_enabled
+        ]
+
+    if LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names:
+        return [
+            component_name
+            for component_name, module in modules.items()
+            if isinstance(module, LayerwiseOffloadableModuleMixin)
+        ]
+
+    explicit_component_names = selected_component_names - {LAYERWISE_OFFLOAD_DIT_GROUP}
+    select_dit_group = LAYERWISE_OFFLOAD_DIT_GROUP in selected_component_names
+    selected_pipeline_component_names: list[str] = []
+    for component_name, module in modules.items():
+        if layerwise_component_matches_any_selection(
+            component_name, explicit_component_names
+        ):
+            selected_pipeline_component_names.append(component_name)
+            continue
+        if (
+            select_dit_group
+            and isinstance(module, LayerwiseOffloadableModuleMixin)
+            and module.layerwise_offload_dit_group_enabled
+        ):
+            selected_pipeline_component_names.append(component_name)
+    return selected_pipeline_component_names
+
+
+def configure_layerwise_offload_modules(
+    modules: Mapping[str, object],
+    server_args: ServerArgs,
+    component_names: Sequence[str] | None = None,
+    warn_missing: bool = True,
+) -> list[str]:
+    """Configure layerwise offload for the given modules, from the given component_names
+
+    Args:
+        modules: the dict of {component_name: component}, containing the components to be chosen from
+        component_names: list of component names. component with names not in this list shouldn't be configured
+
+    Returns a list of component names of modules configured to be layerwise-offload
+    """
+
+    # components which has already been configured to be layerwise-offload
+    configured_component_names: list[str] = []
+    configured_module_ids: set[int] = set()
+    normalized_component_names = normalize_layerwise_offload_components(component_names)
+    selected_component_names = (
+        set(normalized_component_names)
+        if normalized_component_names is not None
+        else None
+    )
+    select_all = (
+        selected_component_names is not None
+        and LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names
+    )
+    selected_pipeline_component_names = (
+        get_layerwise_offload_component_names_for_pipeline(
+            modules,
+            normalized_component_names,
+        )
+    )
+
+    if warn_missing and selected_component_names is not None and not select_all:
+        explicit_component_names = selected_component_names - {
+            LAYERWISE_OFFLOAD_DIT_GROUP
+        }
+        missing_component_names = [
+            selected_component_name
+            for selected_component_name in explicit_component_names
+            if not any(
+                layerwise_component_matches_any_selection(
+                    component_name, [selected_component_name]
+                )
+                for component_name in modules
+            )
+        ]
+        if missing_component_names:
+            logger.warning(
+                "Layerwise offload components are not currently loaded: %s. "
+                "Available pipeline components: %s",
+                sorted(missing_component_names),
+                sorted(modules),
+            )
+
+        unsupported_component_names = [
+            component_name
+            for component_name in selected_pipeline_component_names
+            if not isinstance(modules[component_name], LayerwiseOffloadableModuleMixin)
+        ]
+        if unsupported_component_names:
+            logger.warning(
+                "Layerwise offload components do not support layerwise offload: %s",
+                sorted(unsupported_component_names),
+            )
+
+    for component_name in selected_pipeline_component_names:
+        module = modules[component_name]
+        if not isinstance(module, LayerwiseOffloadableModuleMixin):
+            continue
+        module_id = id(module)
+        if module_id in configured_module_ids:
+            # avoid duplicated configures on a same module
+            continue
+
+        configured_module_ids.add(module_id)
+
+        if not is_layerwise_offloaded_module(module):
+            module.configure_layerwise_offload(server_args)
+        if is_layerwise_offloaded_module(module):
+            configured_component_names.append(component_name)
+
+    if configured_component_names:
+        logger.info(
+            "Enabled layerwise offload for pipeline components: %s",
+            configured_component_names,
+        )
+    else:
+        logger.info("No pipeline component supports layerwise offload.")
+    return configured_component_names

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

@@ -0,0 +1,157 @@
+from collections.abc import Collection, Sequence
+
+LAYERWISE_OFFLOAD_ALL_COMPONENTS = "all"
+LAYERWISE_OFFLOAD_DIT_GROUP = "dit"
+LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP = "text_encoder"
+LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP = "image_encoder"
+LAYERWISE_OFFLOAD_VAE_GROUP = "vae"
+LAYERWISE_OFFLOAD_DEFAULT_GROUP = "default"
+
+# Components whose layerwise policy has been validated as a better default than
+# component-level CPU offload when the user has not pinned their placement.
+LAYERWISE_OFFLOAD_DEFAULT_GROUP_COMPONENTS = (
+    LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
+    LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
+    LAYERWISE_OFFLOAD_VAE_GROUP,
+)
+DIT_COMPONENT_NAMES = frozenset(
+    {
+        "transformer",
+        "transformer_2",
+        "video_dit",
+        "video_dit_2",
+        "audio_dit",
+        "dual_tower_bridge",
+    }
+)
+VAE_COMPONENT_NAMES = frozenset(
+    {
+        "vae",
+        "video_vae",
+        "audio_vae",
+        "vocoder",
+        "spatial_upsampler",
+        "condition_image_encoder",
+    }
+)
+DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
+    {
+        "vae",
+        "video_vae",
+        "condition_image_encoder",
+    }
+)
+CPU_OFFLOAD_FLAG_NAMES = (
+    "dit_cpu_offload",
+    "text_encoder_cpu_offload",
+    "image_encoder_cpu_offload",
+    "vae_cpu_offload",
+)
+
+
+def is_dit_component_name(component_name: str) -> bool:
+    return component_name in DIT_COMPONENT_NAMES
+
+
+def is_text_encoder_component_name(component_name: str) -> bool:
+    return component_name.startswith("text_encoder") or component_name.endswith(
+        "text_encoder"
+    )
+
+
+def is_image_encoder_component_name(component_name: str) -> bool:
+    return component_name == "image_encoder"
+
+
+def is_vae_component_name(component_name: str) -> bool:
+    return component_name in VAE_COMPONENT_NAMES
+
+
+def layerwise_component_matches_selection(
+    component_name: str,
+    selected_component_name: str,
+) -> bool:
+    """if the provided component_name (unnormalized, e.g., text_encoder_2)  matches with the selected_component_name (normalized)"""
+    if selected_component_name == LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP:
+        return is_text_encoder_component_name(component_name)
+    if selected_component_name == LAYERWISE_OFFLOAD_VAE_GROUP:
+        # `vae` is a default-policy selector; AV-side decoders remain explicit-only
+        return component_name in DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES
+    return component_name == selected_component_name
+
+
+def layerwise_component_matches_any_selection(
+    component_name: str,
+    selected_component_names: Collection[str],
+) -> bool:
+    return any(
+        layerwise_component_matches_selection(component_name, selected_component_name)
+        for selected_component_name in selected_component_names
+    )
+
+
+def cpu_offload_flags_for_layerwise_components(
+    component_names: Sequence[str],
+) -> tuple[str, ...]:
+    component_names = normalize_layerwise_offload_components(component_names) or []
+    if LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names:
+        return CPU_OFFLOAD_FLAG_NAMES
+
+    flag_names: list[str] = []
+    if LAYERWISE_OFFLOAD_DIT_GROUP in component_names:
+        flag_names.append("dit_cpu_offload")
+
+    for component_name in component_names:
+        if component_name == LAYERWISE_OFFLOAD_DIT_GROUP:
+            continue
+        if is_dit_component_name(component_name):
+            flag_name = "dit_cpu_offload"
+        elif is_text_encoder_component_name(component_name):
+            flag_name = "text_encoder_cpu_offload"
+        elif is_image_encoder_component_name(component_name):
+            flag_name = "image_encoder_cpu_offload"
+        elif is_vae_component_name(component_name):
+            flag_name = "vae_cpu_offload"
+        else:
+            continue
+
+        if flag_name not in flag_names:
+            flag_names.append(flag_name)
+
+    return tuple(flag_names)
+
+
+def expand_layerwise_offload_component_group(component_name: str) -> tuple[str, ...]:
+    if component_name == LAYERWISE_OFFLOAD_DEFAULT_GROUP:
+        return LAYERWISE_OFFLOAD_DEFAULT_GROUP_COMPONENTS
+    return (component_name,)
+
+
+def normalize_layerwise_offload_components(
+    component_names: str | Sequence[str] | None,
+) -> list[str] | None:
+    if component_names is None:
+        return None
+
+    raw_components = (
+        [component_names] if isinstance(component_names, str) else component_names
+    )
+    normalized_components: list[str] = []
+    for raw_component in raw_components:
+        if not isinstance(raw_component, str):
+            raise ValueError(
+                f"Invalid layerwise offload component name: {raw_component}."
+            )
+        for component_name in raw_component.split(","):
+            component_name = component_name.strip().replace("-", "_").lower()
+            if not component_name:
+                continue
+            for expanded_component_name in expand_layerwise_offload_component_group(
+                component_name
+            ):
+                if expanded_component_name == LAYERWISE_OFFLOAD_ALL_COMPONENTS:
+                    return [LAYERWISE_OFFLOAD_ALL_COMPONENTS]
+                if expanded_component_name not in normalized_components:
+                    normalized_components.append(expanded_component_name)
+
+    return normalized_components or None

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

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+
+import torch
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    is_layerwise_offloaded_module,
+)
+from sglang.multimodal_gen.runtime.pipelines_core import ComposedPipelineBase
+from sglang.multimodal_gen.runtime.post_training.weights_updater import (
+    get_updatable_modules,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def _get_module_device(module: torch.nn.Module) -> str:
+    """Return best-effort device string for a module."""
+    param = next(module.parameters(), None)
+    if param is not None:
+        return str(param.device)
+    buffer = next(module.buffers(), None)
+    if buffer is not None:
+        return str(buffer.device)
+
+    for key, val in vars(module).items():
+        if key.startswith("_"):
+            continue
+        if isinstance(val, torch.Tensor):
+            return str(val.device)
+
+    return "cpu"
+
+
+def _move_unregistered_tensors(module: torch.nn.Module, device: str) -> None:
+    """Move tensor attributes that are not covered by `module.to(device)`."""
+
+    def move_tensors(obj):
+        if torch.is_tensor(obj):
+            return obj.to(device)
+        if isinstance(obj, dict):
+            return {k: move_tensors(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [move_tensors(v) for v in obj]
+        if isinstance(obj, tuple):
+            return tuple(move_tensors(v) for v in obj)
+        return obj
+
+    attrs = module.__dict__
+    for attr_name, attr_value in list(attrs.items()):
+        if attr_name.startswith("_"):
+            continue
+        if attr_name in {"_parameters", "_buffers", "_modules"}:
+            continue
+
+        moved_value = move_tensors(attr_value)
+        if moved_value is not attr_value:
+            attrs[attr_name] = moved_value
+
+
+def _is_layerwise_offload_managed(module: torch.nn.Module) -> bool:
+    return is_layerwise_offloaded_module(module)
+
+
+class MemoryOccupationController:
+    def __init__(
+        self,
+        pipeline: ComposedPipelineBase | None,
+        rank: int,
+        use_fsdp_inference: bool,
+    ):
+        self.pipeline = pipeline
+        self.rank = rank
+        self.use_fsdp_inference = use_fsdp_inference
+        self._sleeping = False
+        self._sleep_restore_map: dict[str, str] = {}
+
+    def is_sleeping(self) -> bool:
+        return self._sleeping
+
+    def _memory_occupation_result(
+        self, success: bool, message: str
+    ) -> dict[str, bool | str]:
+        return {
+            "success": success,
+            "sleeping": self._sleeping,
+            "message": message,
+        }
+
+    @staticmethod
+    def _clear_torch_device_cache() -> None:
+        device = torch.get_device_module()
+        device.synchronize()
+        gc.collect()
+        device.empty_cache()
+
+    def _move_modules(self, names: list[str], device: str) -> None:
+        """
+        Move selected modules to device.
+
+        This function has all-or-nothing semantics:
+        - Stop on first failure (device query / move / sanitize).
+        - Roll back modules already moved in this call.
+        - Raise RuntimeError to caller after rollback.
+        """
+        modules = get_updatable_modules(self.pipeline)
+        moved: list[str] = []
+        src_device_map: dict[str, str] = {}
+
+        try:
+            for name in names:
+                module = modules[name]
+                src_device_map[name] = _get_module_device(module)
+                module.to(device)
+                moved.append(name)
+                _move_unregistered_tensors(module, device)
+        except Exception as e:
+            logger.warning(
+                f"[_move_modules] move failed, rollback started: target={device} moved={moved} error={e}",
+            )
+            for name in moved:
+                module = modules.get(name)
+                src_dev = src_device_map.get(name)
+                module.to(src_dev)
+                _move_unregistered_tensors(module, src_dev)
+            raise RuntimeError(
+                f"failed to move modules to {device}; rollback finished: error={e}"
+            ) from e
+
+    def _offload_active_modules_to_cpu(self) -> dict[str, str]:
+        restore_map: dict[str, str] = {}
+        for name, module in get_updatable_modules(self.pipeline).items():
+            if _is_layerwise_offload_managed(module):
+                continue
+            device = _get_module_device(module)
+            if not device.startswith("cpu"):
+                restore_map[name] = device
+
+        self._move_modules(list(restore_map.keys()), "cpu")
+        self._clear_torch_device_cache()
+        return restore_map
+
+    def _restore_modules_to_original_devices(
+        self, module_device_map: dict[str, str]
+    ) -> None:
+        grouped: dict[str, list[str]] = {}
+        for name, device in module_device_map.items():
+            grouped.setdefault(device, []).append(name)
+
+        for device, names in grouped.items():
+            self._move_modules(names, device)
+
+    def release_memory_occupation(self) -> dict[str, bool | str]:
+        logger.info(f"[SLEEP] release_memory_occupation rank={self.rank}")
+        if self._sleeping:
+            return self._memory_occupation_result(
+                success=True,
+                message="already sleeping",
+            )
+        if self.use_fsdp_inference:
+            raise RuntimeError("sleep/wake does not support FSDP inference")
+        if self.pipeline is None:
+            return self._memory_occupation_result(
+                success=False,
+                message="pipeline not initialized",
+            )
+
+        self._sleep_restore_map = self._offload_active_modules_to_cpu()
+        self._sleeping = True
+        return self._memory_occupation_result(
+            success=True,
+            message="released GPU memory (moved active modules to CPU)",
+        )
+
+    def resume_memory_occupation(self) -> dict[str, bool | str]:
+        logger.info(f"[WAKE] resume_memory_occupation rank={self.rank}")
+        if not self._sleeping:
+            return self._memory_occupation_result(
+                success=True,
+                message="already awake",
+            )
+        if self.pipeline is None:
+            return self._memory_occupation_result(
+                success=False,
+                message="pipeline not initialized",
+            )
+
+        if not self._sleep_restore_map:
+            self._sleeping = False
+            return self._memory_occupation_result(
+                success=True,
+                message="no restore map; marked awake",
+            )
+
+        self._restore_modules_to_original_devices(self._sleep_restore_map)
+        self._sleep_restore_map = {}
+        self._sleeping = False
+        return self._memory_occupation_result(
+            success=True,
+            message="resumed GPU memory (restored modules to original devices)",
+        )