From 9d84268705a56df25842d6c7e4474b3f2ecf5230 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 1 May 2026 11:10:41 +0800 Subject: [PATCH] [diffusion] refactor: introduce component residency manager (#23771) --- .../runtime/loader/weights_updater.py | 2 +- .../runtime/managers/component_manager.py | 673 ++++++++++++++++++ .../managers/component_resident_strategies.py | 502 +++++++++++++ .../runtime/managers/gpu_worker.py | 8 +- .../{utils => managers}/layerwise_offload.py | 5 + .../runtime/managers/scheduler.py | 22 +- .../runtime/models/bridges/mova_dual_tower.py | 2 +- .../runtime/models/dits/causal_wanvideo.py | 2 +- .../runtime/models/dits/ernie_image.py | 2 +- .../runtime/models/dits/flux.py | 2 +- .../runtime/models/dits/flux_2.py | 2 +- .../runtime/models/dits/glm_image.py | 2 +- .../runtime/models/dits/helios.py | 2 +- .../runtime/models/dits/hunyuan3d.py | 2 +- .../runtime/models/dits/hunyuanvideo.py | 2 +- .../runtime/models/dits/ltx_2.py | 2 +- .../runtime/models/dits/mova_audio_dit.py | 2 +- .../runtime/models/dits/mova_video_dit.py | 2 +- .../runtime/models/dits/qwen_image.py | 2 +- .../runtime/models/dits/sana.py | 2 +- .../runtime/models/dits/wanvideo.py | 2 +- .../runtime/models/dits/zimage.py | 2 +- .../runtime/pipelines/ltx_2_pipeline.py | 636 +++++++++-------- .../pipelines_core/composed_pipeline_base.py | 12 + .../executors/parallel_executor.py | 52 +- .../executors/pipeline_executor.py | 27 + .../pipelines_core/executors/sync_executor.py | 16 +- .../runtime/pipelines_core/lora_pipeline.py | 2 +- .../runtime/pipelines_core/schedule_batch.py | 4 +- .../runtime/pipelines_core/stages/base.py | 80 +++ .../runtime/pipelines_core/stages/decoding.py | 97 +-- .../pipelines_core/stages/decoding_av.py | 171 +++-- .../pipelines_core/stages/denoising.py | 137 ++-- .../pipelines_core/stages/denoising_av.py | 59 +- .../pipelines_core/stages/denoising_dmd.py | 34 +- .../runtime/pipelines_core/stages/encoding.py | 52 +- .../pipelines_core/stages/image_encoding.py | 391 +++++----- .../pipelines_core/stages/ltx_2_denoising.py | 22 +- .../model_specific_stages/helios_decoding.py | 17 +- .../model_specific_stages/helios_denoising.py | 33 +- .../stages/model_specific_stages/mova.py | 156 ++-- .../qwen_image_layered.py | 83 ++- .../stages/model_specific_stages/wan_ti2v.py | 9 +- .../pipelines_core/stages/text_encoding.py | 62 +- .../pipelines_core/stages/upsampling.py | 50 +- .../test/unit/test_layerwise_offload.py | 4 +- 46 files changed, 2477 insertions(+), 973 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/managers/component_manager.py create mode 100644 python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py rename python/sglang/multimodal_gen/runtime/{utils => managers}/layerwise_offload.py (99%) diff --git a/python/sglang/multimodal_gen/runtime/loader/weights_updater.py b/python/sglang/multimodal_gen/runtime/loader/weights_updater.py index f170809a7..1007e9b41 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/loader/weights_updater.py @@ -53,9 +53,9 @@ from sglang.multimodal_gen.runtime.loader.utils import ( from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/component_manager.py new file mode 100644 index 000000000..b43a16ead --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/component_manager.py @@ -0,0 +1,673 @@ +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from functools import lru_cache +from typing import Mapping, MutableMapping, Protocol, Sequence, TypeVar + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.managers.component_resident_strategies import ( + ComponentResidencyStrategy, + LayerwiseOffloadStrategy, + ResidentStrategy, + VanillaD2HStrategy, +) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +_T = TypeVar("_T") + + +@dataclass(slots=True) +class ComponentUse: + """Describes one stage/use-site access to a pipeline component.""" + + stage_name: str + # Pipeline module key: transformer / video_dit / text_encoder / ... + component_name: str + # Model-specific phase for sequential components, e.g. stage1 or stage2. + # TODO: Replace this with ordered timeline identity. In an all-sequential + # pipeline, use-site identity should come from the declared ComponentUse + # order instead of a per-use `phase` field. + phase: str | None = None + # Whether the manager may prepare this component for the next request. + preferred_ready_after_request: bool = False + # Whether cross-stage prefetch may prepare this use before the use-site. + allow_prefetch: bool = True + # Whether this use is expensive enough that earlier timeline prefetch matters. + # TODO: Replace this boolean hint with a budget-aware lookahead planner: + # estimate memory/load cost and reuse distance, keep small and early-request + # components resident within budget, prefetch as soon as VRAM slack appears, + # and release completed components only when the budget requires it. + memory_intensive: bool = False + # Optional module dtype required by this use-site. + target_dtype: torch.dtype | None = None + # Some components are intentionally kept ready between warmup and the first + # real request to avoid measuring a cold H2D in the user-visible request. + keep_ready_after_warmup: bool = False + + +@dataclass(slots=True) +class ResidencyState: + """ + Necessary internal runtime info of ComponentResidencyManager + """ + + stages: Sequence["ComponentResidencyStage"] = () + stage_index: int = -1 + stage_name: str | None = None + next_stage_name: str | None = None + current_use: ComponentUse | None = None + # the ComponentUses of the preceding stages + future_uses: tuple[ComponentUse, ...] = () + batch_is_warmup: bool = False + manager_mode: str = "static" + trace_enabled: bool = False + + +class ResidencyBatch(Protocol): + is_warmup: bool + + +class ComponentResidencyStage(Protocol): + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: ... + + +class ComponentResidencyPipeline(Protocol): + modules: Mapping[str, object] + _stage_name_mapping: Mapping[str, ComponentResidencyStage] + component_residency_strategies: MutableMapping[str, "ComponentResidencyStrategy"] + + +def build_dit_residency_strategy( + module: nn.Module, + server_args: ServerArgs, +) -> ComponentResidencyStrategy: + if ( + isinstance(module, OffloadableDiTMixin) + and module.layerwise_offload_managers + and any(manager.enabled for manager in module.layerwise_offload_managers) + ): + # only if dit_layerwise_offload is enabled + return LayerwiseOffloadStrategy() + if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: + # handles offload by vanalla D2H + return VanillaD2HStrategy() + return ResidentStrategy() + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + return module.__class__.__name__.startswith("FSDP") + + +def build_component_residency_strategy( + component_name: str, + module: nn.Module, + server_args: ServerArgs, +) -> ComponentResidencyStrategy: + if component_name in { + "transformer", + "transformer_2", + "video_dit", + "video_dit_2", + "audio_dit", + "dual_tower_bridge", + }: + return build_dit_residency_strategy(module, server_args) + + if component_name.startswith("text_encoder") or component_name.endswith( + "text_encoder" + ): + if ( + server_args.text_encoder_cpu_offload + and not server_args.use_fsdp_inference + and not is_fsdp_managed_module(module) + ): + return VanillaD2HStrategy() + return ResidentStrategy() + + if component_name == "image_encoder": + if server_args.image_encoder_cpu_offload and not server_args.use_fsdp_inference: + return VanillaD2HStrategy() + return ResidentStrategy() + + if component_name in { + "vae", + "video_vae", + "audio_vae", + "vocoder", + "spatial_upsampler", + "condition_image_encoder", + }: + if server_args.vae_cpu_offload and not server_args.use_fsdp_inference: + return VanillaD2HStrategy() + return ResidentStrategy() + + return ResidentStrategy() + + +class ComponentResidencyManager: + """Executor-owned component lifecycle coordinator. Provide hooks for a PipelineExecutor + + Hooks are called around executor progress: + before request: collect a flat ordered ComponentUse timeline. + before stage: update current/next stage context only. + begin use: finish previous active use, prepare current use, wait until ready. + end use: finish or keep current use, then prefetch the next heavy timeline use. + finish request: finish active use and schedule preferred next-request prefetch. + + The manager instance is global and rebound to the active pipeline before request execution. + This manager is designed only for sequential execution order for now + """ + + def __init__( + self, pipeline: ComponentResidencyPipeline, server_args: ServerArgs + ) -> None: + self.pipeline = pipeline + self.server_args = server_args + self.state = ResidencyState(trace_enabled=False) + self._stage_names_by_id: dict[int, str] = {} + self._stage_uses_by_index: list[tuple[ComponentUse, ...]] = [] + self._ordered_uses: tuple[ComponentUse, ...] = () + self._current_use_index: int = -1 + self._active_use: ComponentUse | None = None + self._active_use_module: nn.Module | None = None + self._prefetched_use_keys: set[tuple[str, str, str | None]] = set() + self._custom_strategies: dict[str, ComponentResidencyStrategy] = dict( + pipeline.component_residency_strategies + ) + self._uses_seen: dict[str, ComponentUse] = {} + + @property + def enabled(self) -> bool: + return True + + def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None: + custom_strategies = dict(pipeline.component_residency_strategies) + if pipeline is not self.pipeline: + self.strategy_for.cache_clear() + self._should_keep_single_dit.cache_clear() + self._active_use = None + self._active_use_module = None + self._uses_seen.clear() + self._prefetched_use_keys.clear() + elif custom_strategies != self._custom_strategies: + self.strategy_for.cache_clear() + self.pipeline = pipeline + self._custom_strategies = custom_strategies + self._stage_names_by_id = { + id(stage): name for name, stage in pipeline._stage_name_mapping.items() + } + + def refresh_server_args(self, server_args: ServerArgs) -> None: + if server_args is not self.server_args: + self.strategy_for.cache_clear() + self.server_args = server_args + + def register_strategy( + self, component_name: str, strategy: ComponentResidencyStrategy + ) -> None: + self.pipeline.component_residency_strategies[component_name] = strategy + self._custom_strategies[component_name] = strategy + self.strategy_for.cache_clear() + + def begin_request( + self, + stages: Sequence[ComponentResidencyStage], + batch: ResidencyBatch, + server_args: ServerArgs, + ) -> None: + """A hook called before processing an actual request""" + self.refresh_server_args(server_args) + self.state = ResidencyState( + stages=stages, batch_is_warmup=batch.is_warmup, trace_enabled=False + ) + self._active_use = None + self._active_use_module = None + self._current_use_index = -1 + self._prefetched_use_keys.clear() + self._uses_seen.clear() + if self.enabled: + self._stage_uses_by_index = [ + tuple(stage.component_uses(server_args, self.stage_name(stage))) + for stage in stages + ] + self._ordered_uses = tuple( + use for uses in self._stage_uses_by_index for use in uses + ) + else: + self._stage_uses_by_index = [] + self._ordered_uses = () + self._trace( + "request_start", + detail=f"stages={len(stages)} uses={len(self._ordered_uses)}", + ) + + def before_stage( + self, + stage: ComponentResidencyStage, + stage_index: int, + batch: ResidencyBatch, + server_args: ServerArgs, + ) -> None: + """called after stage starts""" + if not self.enabled: + return + # update state before entering the stage + self.state.stage_index = stage_index + self.state.stage_name = self.stage_name(stage) + self.state.next_stage_name = self._next_stage_name(stage_index) + self._trace("stage_enter", detail=f"index={stage_index}") + + def after_stage(self, stage_index: int) -> None: + """called after stage exits""" + if not self.enabled: + return + self._trace("stage_exit", detail=f"index={stage_index}") + + def before_use(self, use: ComponentUse) -> None: + """component use-site starts""" + if not self.enabled: + return + self.begin_use(use) + + def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: + """Begin one sequential component use interval. this is idempotent + + 1. Finish the previous active use if this is a different timeline use. + 2. Prepare the current component. + 3. Wait until the current component is ready, then prefetch the next heavy use. + """ + if self._active_use is not None and self._same_use(self._active_use, use): + return + if self._active_use is not None: + # finish previous active use + self._finish_use( + self._active_use, + module=self._active_use_module, + keep_on_warmup=self._active_use.keep_ready_after_warmup, + ) + self._active_use = None + self._active_use_module = None + self.state.current_use = None + self._mark_current_use(use) + self._prepare_forward_use(use, module=module) + self._active_use = use + self._active_use_module = module + self._prefetch_next_memory_intensive_use() + + def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: + """End one sequential component use interval. + + 1. Finish or keep the current component. + 2. Clear it as the active use. + 3. Prefetch the next memory-intensive use without waiting. + """ + if self._active_use is None or not self._same_use(self._active_use, use): + return + self._finish_use( + self._active_use, + module=self._active_use_module or module, + keep_on_warmup=self._active_use.keep_ready_after_warmup, + ) + self._active_use = None + self._active_use_module = None + self.state.current_use = None + self._prefetch_next_memory_intensive_use() + + @contextmanager + def use_component( + self, use: ComponentUse, module: nn.Module | None = None + ) -> Iterator[nn.Module | None]: + self.begin_use(use, module=module) + try: + yield module if module is not None else self.get_module(use.component_name) + finally: + self.end_use(use, module=module) + + def call_component( + self, + use: ComponentUse, + module: Callable[..., _T], + *args, + **kwargs, + ) -> _T: + with self.use_component(use): + return module(*args, **kwargs) + + def prefetch_use(self, use: ComponentUse) -> None: + """Prepare a future use without blocking the current use.""" + if not self.enabled: + return + self._prefetch_use(use) + + def ensure_ready(self, use: ComponentUse, module: nn.Module | None = None) -> None: + """Prepare a shared component and wait without making it the active use.""" + if not self.enabled: + return + self._prepare_forward_use(use, module=module) + + def prefetch_checkpoint(self, anchor: ComponentUse | None = None) -> None: + """Give the manager a timeline overlap point. + + 1. Locate the anchor or current use in the ordered timeline. + 2. Find the next prefetchable memory-intensive use. + 3. Prepare it opportunistically without waiting. + """ + if not self.enabled: + return + if anchor is not None: + self._mark_current_use(anchor) + self._prefetch_next_memory_intensive_use() + + def finish_active_use(self, *, prefetch_next: bool = True) -> None: + """Finish the currently active sequential use, if any.""" + if self._active_use is None: + return + active_use = self._active_use + self._finish_use( + active_use, + module=self._active_use_module, + keep_on_warmup=active_use.keep_ready_after_warmup, + ) + self._active_use = None + self._active_use_module = None + self.state.current_use = None + if prefetch_next: + self._prefetch_next_memory_intensive_use() + + def _prepare_forward_use( + self, use: ComponentUse, module: nn.Module | None = None + ) -> None: + """Prepare a component that is about to run and wait until it is ready.""" + module = module or self.get_module(use.component_name) + if module is None: + self._trace("skip_missing", use) + return + strategy = self.strategy_for(use.component_name, module) + self._uses_seen[use.component_name] = use + self.state.current_use = use + self._trace("prepare", use, strategy, module) + strategy.prepare_for_use(module, use, self.state) + self._trace("wait", use, strategy, module) + strategy.wait_for_use(module, use, self.state) + + def _prefetch_use(self, use: ComponentUse) -> None: + """Prepare a future component opportunistically without waiting. + + This is called when the component is memory-intensive so it may takes a long time to prefetch. + + manager will perform the prefetch at some checkpoints, if necessary + """ + if not use.allow_prefetch: + return + module = self.get_module(use.component_name) + if module is None: + self._trace("skip_missing", use) + return + strategy = self.strategy_for(use.component_name, module) + if isinstance(strategy, VanillaD2HStrategy) and self._active_use is not None: + # Avoid making two vanilla-offloaded heavy components resident before + # a budget-aware planner can prove the overlap is safe. + self._trace("prefetch_skip_active_vanilla", use, strategy, module) + return + + self._uses_seen[use.component_name] = use + self._trace("prefetch", use, strategy, module) + if strategy.prefetch_for_use(module, use, self.state): + self._prefetched_use_keys.add(self._use_key(use)) + + def after_use(self, use: ComponentUse) -> None: + if not self.enabled: + return + self.end_use(use) + + def _finish_use( + self, + use: ComponentUse, + *, + module: nn.Module | None = None, + keep_on_warmup: bool, + ) -> None: + """finish a specific use by keeping them resident or call finish_use hook""" + module = module or self.get_module(use.component_name) + if module is None: + self._trace("skip_missing", use) + return + should_keep = ( + keep_on_warmup and self.state.batch_is_warmup + ) or self._should_keep_after_use(use) + if should_keep: + self._trace( + "keep", + use, + self.strategy_for(use.component_name, module), + module, + ) + return + strategy = self.strategy_for(use.component_name, module) + self._trace("finish", use, strategy, module) + strategy.finish_use(module, use, self.state) + + def finish_request(self) -> None: + if not self.enabled and not self._uses_seen and self._active_use is None: + return + # 1. Close the currently active sequential use. + self.finish_active_use(prefetch_next=False) + # 2. Pick components that should be ready for the next request. + preferred_uses = self._preferred_request_end_uses() + # 3. Finish everything else, or prepare preferred uses for request tail. + for component_name, use in list(self._uses_seen.items()): + module = self.get_module(component_name) + if module is None: + continue + if self.state.batch_is_warmup and use.keep_ready_after_warmup: + self._trace( + "request_keep_warmup", + use, + self.strategy_for(component_name, module), + module, + ) + continue + preferred = component_name in preferred_uses + if not preferred and self._should_keep_single_dit(component_name): + self._trace( + "keep", + use, + self.strategy_for(component_name, module), + module, + detail="single_dit", + ) + continue + strategy = self.strategy_for(component_name, module) + if preferred and not self.state.batch_is_warmup: + self._trace("request_prefetch", use, strategy, module) + strategy.prepare_after_request(module, use, self.state) + else: + action = "request_resident" if preferred else "request_finish" + self._trace(action, use, strategy, module) + strategy.finish_request(module, use, self.state, preferred=preferred) + self._trace("request_end") + + def stage_name(self, stage: ComponentResidencyStage) -> str: + return self._stage_names_by_id.get(id(stage), stage.__class__.__name__) + + def component_name_for_module(self, module: nn.Module | None, default: str) -> str: + if module is None: + return default + for name, candidate in self.pipeline.modules.items(): + if candidate is module: + return name + return default + + def get_module(self, component_name: str) -> nn.Module | None: + module = self.pipeline.modules.get(component_name) + return module if isinstance(module, nn.Module) else None + + @lru_cache(maxsize=None) + def strategy_for( + self, component_name: str, module: nn.Module + ) -> ComponentResidencyStrategy: + """Return the pre-registered strategy for a specific component""" + custom_strategy = self._custom_strategies.get(component_name) + if custom_strategy is not None: + return custom_strategy + return build_component_residency_strategy( + component_name, module, self.server_args + ) + + def _stage_uses(self, stage_index: int) -> tuple[ComponentUse, ...]: + """Returns the ComponentUse(s) of a specific stage""" + if stage_index < 0 or stage_index >= len(self._stage_uses_by_index): + return () + return self._stage_uses_by_index[stage_index] + + def _next_stage_name(self, stage_index: int) -> str | None: + next_index = stage_index + 1 + if next_index < 0 or next_index >= len(self.state.stages): + return None + return self.stage_name(self.state.stages[next_index]) + + def _mark_current_use(self, use: ComponentUse) -> None: + index = self._locate_use_index(use) + if index is None: + self._current_use_index = len(self._ordered_uses) + self.state.future_uses = () + return + self._current_use_index = index + self.state.future_uses = self._ordered_uses[index + 1 :] + + def _locate_use_index(self, use: ComponentUse) -> int | None: + for index in range(self._current_use_index + 1, len(self._ordered_uses)): + if self._same_use(self._ordered_uses[index], use): + return index + for index, candidate in enumerate(self._ordered_uses): + if self._same_use(candidate, use): + return index + return None + + def _prefetch_next_memory_intensive_use(self) -> None: + for use in self._ordered_uses[self._current_use_index + 1 :]: + if not use.memory_intensive: + continue + if self._use_key(use) in self._prefetched_use_keys: + return + self.prefetch_use(use) + return + + def _should_keep_after_use(self, use: ComponentUse) -> bool: + future_component_names = { + future.component_name for future in self.state.future_uses + } + if use.component_name in future_component_names: + return True + if self._should_keep_single_dit(use.component_name): + return True + return False + + @lru_cache(maxsize=None) + def _should_keep_single_dit(self, component_name: str) -> bool: + modules = self.pipeline.modules + return (component_name == "transformer" and "transformer_2" not in modules) or ( + component_name == "video_dit" and "video_dit_2" not in modules + ) + + def _preferred_request_end_use(self) -> ComponentUse | None: + """Returns a ComponentUse preferred to be resident after a request finishes, to prepare for next request""" + for uses in self._stage_uses_by_index: + for use in uses: + if use.preferred_ready_after_request: + return use + for uses in self._stage_uses_by_index: + if uses: + return uses[0] + return None + + def _preferred_request_end_uses(self) -> dict[str, ComponentUse]: + preferred_uses: dict[str, ComponentUse] = {} + for uses in self._stage_uses_by_index: + for use in uses: + if use.preferred_ready_after_request: + preferred_uses[use.component_name] = use + for use in self._uses_seen.values(): + if use.preferred_ready_after_request: + preferred_uses[use.component_name] = use + if preferred_uses: + return preferred_uses + preferred_use = self._preferred_request_end_use() + if preferred_use is None: + return {} + return {preferred_use.component_name: preferred_use} + + @staticmethod + def _same_use(lhs: ComponentUse, rhs: ComponentUse) -> bool: + return lhs.component_name == rhs.component_name and lhs.phase == rhs.phase + + @staticmethod + def _use_key(use: ComponentUse) -> tuple[str, str, str | None]: + return (use.stage_name, use.component_name, use.phase) + + def _trace( + self, + action: str, + use: ComponentUse | None = None, + strategy: ComponentResidencyStrategy | None = None, + module: nn.Module | None = None, + *, + component_name: str | None = None, + detail: str = "", + ) -> None: + if not self.state.trace_enabled: + return + if use is not None: + component_name = use.component_name + device = self._module_device(module) + logger.info( + "[component_residency] action=%s stage=%s next_stage=%s component=%s " + "strategy=%s phase=%s device=%s warmup=%s mode=%s %s", + action, + self.state.stage_name, + self.state.next_stage_name, + component_name, + strategy.name if strategy is not None else None, + use.phase if use is not None else None, + device, + self.state.batch_is_warmup, + self.state.manager_mode, + detail, + ) + + def _module_device(self, module: nn.Module | None) -> str | None: + if module is None: + return None + param = next(module.parameters(), None) + if param is not None: + return param.device.type + buffer = next(module.buffers(), None) + return buffer.device.type if buffer is not None else None + + +_GLOBAL_COMPONENT_RESIDENCY_MANAGER: ComponentResidencyManager | None = None + + +def get_global_component_residency_manager( + pipeline: ComponentResidencyPipeline, + server_args: ServerArgs, +) -> ComponentResidencyManager: + global _GLOBAL_COMPONENT_RESIDENCY_MANAGER + + if _GLOBAL_COMPONENT_RESIDENCY_MANAGER is None: + _GLOBAL_COMPONENT_RESIDENCY_MANAGER = ComponentResidencyManager( + pipeline, server_args + ) + else: + _GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_server_args(server_args) + _GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_pipeline(pipeline) + + return _GLOBAL_COMPONENT_RESIDENCY_MANAGER diff --git a/python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py b/python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py new file mode 100644 index 000000000..f3d3f1df2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py @@ -0,0 +1,502 @@ +""" +Basic Component Resident Strategy Utilities for defining usage of components, to let ComponentResidencyManager to coordinate +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.managers.component_manager import ( + ComponentUse, + ResidencyState, + ) + +logger = init_logger(__name__) + + +def _module_to_local_device( + module: nn.Module, *, dtype: torch.dtype | None = None +) -> None: + device = get_local_torch_device() + tensor = _module_reference_tensor(module) + if tensor is not None and tensor.device == device: + if dtype is None or tensor.dtype == dtype: + return + if dtype is None: + module.to(device, non_blocking=True) + else: + module.to(device, dtype=dtype, non_blocking=True) + + +def _module_reference_tensor(module: nn.Module) -> torch.Tensor | None: + tensor = next(module.parameters(), None) + if tensor is None: + tensor = next(module.buffers(), None) + return tensor + + +def _module_ready_on_local_device( + module: nn.Module, *, dtype: torch.dtype | None = None +) -> bool: + tensor = _module_reference_tensor(module) + if tensor is None: + return True + if tensor.device != get_local_torch_device(): + return False + return dtype is None or tensor.dtype == dtype + + +class ComponentResidencyStrategy: + """Baseclass for describing how a component should be treated (regarding where its weights locates) + + e.g., a LayerwiseOffloadStrategy would override: + enter: to prefetch some layers before DiT is used, and + exits: to release GPU weight snapshot after DiT is used + to achieve desired behavior + + """ + + name = "resident" + + def prepare_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + """hook called""" + self.enter(module) + + def wait_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + """Wait for the preparation to be ready, only applicable for async device syncs""" + pass + + def finish_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + """Finish a specific component use""" + self.exit(module) + + def prepare_after_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + """Called after a request is finished, to prepare for the upcoming request""" + pass + + def finish_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + *, + preferred: bool, + ) -> None: + if preferred: + self.prepare_for_use(module, use, state) + self.wait_for_use(module, use, state) + else: + self.finish_use(module, use, state) + + def prefetch_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> bool: + self.prepare_for_use(module, use, state) + return True + + def enter(self, module: nn.Module) -> None: + pass + + def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None: + pass + + +class ResidentStrategy(ComponentResidencyStrategy): + name = "resident" + + def prepare_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + if use.target_dtype is not None: + _module_to_local_device(module, dtype=use.target_dtype) + + +class SnapshotModuleResidency: + """Reusable snapshot-based module residency primitive. + + This helper only knows how to: + - keep CPU parameter/buffer snapshots, + - prefetch a module (H2D) to the local device on a CUDA side stream + - release a module by rebinding tensors to those snapshots, + - track and wait for readiness events. + + It deliberately does not know about pipeline stages, phases, or model-specific + ordering. Strategy subclasses decide when each primitive is called. + """ + + def __init__(self, *, pin_cpu_memory: bool, enable_async_prefetch: bool) -> None: + self.pin_cpu_memory = pin_cpu_memory + self.enable_async_prefetch = enable_async_prefetch + self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {} + self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {} + self._prefetch_stream: object | None = None + self._ready_events: dict[str, object] = {} + + @staticmethod + def is_on_gpu(module: nn.Module | None) -> bool: + if module is None: + return False + param = next(module.parameters(), None) + return param is not None and param.device.type == "cuda" + + def is_ready(self, component_name: str) -> bool: + return component_name in self._ready_events + + def wait_ready(self, component_name: str) -> None: + """wait for the (H2D) stream to be ready""" + ready_event = self._ready_events.get(component_name) + if ready_event is None or not current_platform.is_cuda(): + return + torch.get_device_module().current_stream().wait_event(ready_event) + + def record_ready(self, component_name: str, module: nn.Module | None) -> None: + if not current_platform.is_cuda(): + self._ready_events.pop(component_name, None) + return + if not self.is_on_gpu(module): + self._ready_events.pop(component_name, None) + return + event = torch.get_device_module().Event() + event.record(torch.get_device_module().current_stream()) + self._ready_events[component_name] = event + + @staticmethod + def _clone_cpu_tensor_snapshot( + tensor: torch.Tensor, *, pin_memory: bool + ) -> torch.Tensor: + snapshot = tensor.detach() + if snapshot.device.type == "cpu": + if pin_memory and not snapshot.is_pinned(): + return snapshot.pin_memory() + return snapshot + + cpu_tensor = snapshot.to("cpu") + if pin_memory: + return cpu_tensor.pin_memory() + return cpu_tensor + + def _should_pin_memory(self) -> bool: + return bool(self.pin_cpu_memory and torch.get_device_module().is_available()) + + def capture(self, component_name: str, module: nn.Module) -> None: + """Capture a CPU snapshot for a component""" + if component_name in self._cpu_param_snapshots: + return + + pin_memory = self._should_pin_memory() + self._cpu_param_snapshots[component_name] = { + name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory) + for name, param in module.named_parameters() + } + self._cpu_buffer_snapshots[component_name] = { + name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory) + for name, buffer in module.named_buffers() + } + + def release_to_snapshot( + self, + component_name: str, + module: nn.Module, + *, + copy_runtime_buffers: bool = False, + ) -> None: + """Release CUDA storages by rebinding tensors to cached CPU snapshots. + + This does not call `module.to("cpu")`. Instead, parameter and buffer + storages are rebound to pre-captured CPU tensors so CUDA storages can be + released by the allocator without an explicit D2H transfer. + """ + param_snapshots = self._cpu_param_snapshots.get(component_name) + buffer_snapshots = self._cpu_buffer_snapshots.get(component_name) + if param_snapshots is None or buffer_snapshots is None: + module.to("cpu") + self._ready_events.pop(component_name, None) + return + + pin_memory = self._should_pin_memory() + for name, param in module.named_parameters(): + snapshot = param_snapshots.get(name) + if snapshot is None: + snapshot = self._clone_cpu_tensor_snapshot( + param.data, pin_memory=pin_memory + ) + param_snapshots[name] = snapshot + param.data = snapshot + + for name, buffer in module.named_buffers(): + snapshot = buffer_snapshots.get(name) + if snapshot is None: + snapshot = self._clone_cpu_tensor_snapshot( + buffer.data, pin_memory=pin_memory + ) + buffer_snapshots[name] = snapshot + if copy_runtime_buffers: + # Preserve runtime-updated buffers (e.g., lazily built caches) when + # releasing back to CPU snapshots. + if buffer.device.type == "cuda": + snapshot.copy_( + buffer.detach().to(device="cpu", dtype=snapshot.dtype) + ) + elif buffer.device.type == "cpu": + snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype)) + buffer.data = snapshot + + self._ready_events.pop(component_name, None) + + def _supports_async_prefetch(self) -> bool: + return self.enable_async_prefetch and current_platform.is_cuda() + + def _get_prefetch_stream(self): + """returns a stream is async-prefetch is enabled""" + if not self._supports_async_prefetch(): + return None + if self._prefetch_stream is None: + self._prefetch_stream = torch.get_device_module().Stream( + device=get_local_torch_device() + ) + return self._prefetch_stream + + def prefetch_to_device(self, component_name: str, module: nn.Module | None) -> None: + if module is None: + self._ready_events.pop(component_name, None) + return + prefetch_stream = self._get_prefetch_stream() + if prefetch_stream is None: + # if the async prefetching is disabled + module.to(get_local_torch_device(), non_blocking=True) + self.record_ready(component_name, module) + return + with torch.get_device_module().stream(prefetch_stream): + module.to(get_local_torch_device(), non_blocking=True) + event = torch.get_device_module().Event() + event.record(prefetch_stream) + self._ready_events[component_name] = event + + +class SnapshotStrategy(ComponentResidencyStrategy): + """Snapshot residency: async H2D before use and light snapshot release after use.""" + + name = "snapshot" + + def __init__( + self, + *, + pin_cpu_memory: bool, + enable_async_prefetch: bool, + copy_runtime_buffers_on_release: bool = False, + ) -> None: + self._snapshot_residency = SnapshotModuleResidency( + pin_cpu_memory=pin_cpu_memory, + enable_async_prefetch=enable_async_prefetch, + ) + self._copy_runtime_buffers_on_release = copy_runtime_buffers_on_release + + def capture(self, component_name: str, module: nn.Module) -> None: + self._snapshot_residency.capture(component_name, module) + + def is_ready(self, component_name: str) -> bool: + return self._snapshot_residency.is_ready(component_name) + + def record_ready(self, component_name: str, module: nn.Module | None) -> None: + self._snapshot_residency.record_ready(component_name, module) + + def prefetch_component(self, component_name: str, module: nn.Module | None) -> None: + if SnapshotModuleResidency.is_on_gpu(module): + self._snapshot_residency.record_ready(component_name, module) + return + self._snapshot_residency.prefetch_to_device(component_name, module) + + def wait_component_ready(self, component_name: str) -> None: + self._snapshot_residency.wait_ready(component_name) + + def release_component(self, component_name: str, module: nn.Module) -> None: + self._snapshot_residency.release_to_snapshot( + component_name, + module, + copy_runtime_buffers=self._copy_runtime_buffers_on_release, + ) + + def prepare_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.prefetch_component(use.component_name, module) + + def wait_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.wait_component_ready(use.component_name) + + def finish_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.release_component(use.component_name, module) + + def prepare_after_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.prepare_for_use(module, use, state) + + +class VanillaD2HStrategy(ComponentResidencyStrategy): + """A strategy that performs native torch D2H and H2D for a component""" + + name = "vanilla" + + def __init__(self) -> None: + self._prefetch_stream: object | None = None + self._ready_events: dict[str, object] = {} + + def prepare_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + _module_to_local_device(module, dtype=use.target_dtype) + + def wait_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + ready_event = self._ready_events.get(use.component_name) + if ready_event is None or not current_platform.is_cuda(): + return + torch.get_device_module().current_stream().wait_event(ready_event) + + def prefetch_for_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> bool: + if not current_platform.is_cuda(): + self.prepare_for_use(module, use, state) + return True + if _module_ready_on_local_device(module, dtype=use.target_dtype): + return True + if self._prefetch_stream is None: + self._prefetch_stream = torch.get_device_module().Stream( + device=get_local_torch_device() + ) + with torch.get_device_module().stream(self._prefetch_stream): + _module_to_local_device(module, dtype=use.target_dtype) + event = torch.get_device_module().Event() + event.record(self._prefetch_stream) + self._ready_events[use.component_name] = event + return True + + def enter(self, module: nn.Module) -> None: + param = next(module.parameters(), None) + if param is not None and param.device.type == "cpu": + _module_to_local_device(module) + + def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None: + param = next(module.parameters(), None) + if param is not None and param.device.type == "cuda": + module.to("cpu", non_blocking=True) + + def finish_use( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.wait_for_use(module, use, state) + self.exit(module) + self._ready_events.pop(use.component_name, None) + + def prepare_after_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.prefetch_for_use(module, use, state) + + def finish_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + *, + preferred: bool, + ) -> None: + if preferred and state.batch_is_warmup: + self.prepare_for_use(module, use, state) + self.wait_for_use(module, use, state) + return + if not preferred: + self.finish_use(module, use, state) + + +class LayerwiseOffloadStrategy(ComponentResidencyStrategy): + """A wrapper around LayerwiseOffloadManager to fit in a ComponentResidencyStrategy""" + + name = "layerwise" + + def enter(self, module: nn.Module) -> None: + if isinstance(module, OffloadableDiTMixin): + module.prepare_for_next_req() + + def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None: + if not isinstance(module, OffloadableDiTMixin): + return + for manager in module.layerwise_offload_managers: + manager.release_all() + + def prepare_after_request( + self, + module: nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.prepare_for_use(module, use, state) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a4b759f48..389b65087 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -37,6 +37,10 @@ from sglang.multimodal_gen.runtime.loader.weights_updater import ( WeightsUpdater, get_updatable_modules, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import ( + OffloadableDiTMixin, + iter_materialized_weights, +) from sglang.multimodal_gen.runtime.pipelines_core import ( ComposedPipelineBase, LoRAPipeline, @@ -47,10 +51,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch -from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( - OffloadableDiTMixin, - iter_materialized_weights, -) from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, globally_suppress_loggers, diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/layerwise_offload.py similarity index 99% rename from python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py rename to python/sglang/multimodal_gen/runtime/managers/layerwise_offload.py index 29d601845..b52610c5c 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/layerwise_offload.py @@ -199,6 +199,11 @@ class LayerwiseOffloadManager: self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer + # Keep non-layer parameters resident on GPU. Layer tensors have already + # been replaced by tiny device placeholders, so this does not reload the + # offloaded layer weights. + self.model.to(self.device) + # prefetch the first layer for warm-up self.prepare_for_next_req(non_blocking=False) diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 7ef2f6f45..1b55f9119 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -467,22 +467,20 @@ class Scheduler(SchedulerDisaggMixin): reqs = [item[1] for item in items] try: - processed_req = reqs[0] - if isinstance(processed_req, list) and processed_req: - is_warmup = processed_req[0].is_warmup + first_req = reqs[0] + if isinstance(first_req, list) and first_req: + is_warmup = first_req[0].is_warmup else: is_warmup = ( - processed_req.is_warmup - if isinstance(processed_req, Req) - else False + first_req.is_warmup if isinstance(first_req, Req) else False ) - handler = self.request_handlers.get(type(processed_req)) + handler = self.request_handlers.get(type(first_req)) if handler: output_batch = handler(reqs) else: output_batch = OutputBatch( - error=f"Unknown request type: {type(processed_req)}" + error=f"Unknown request type: {type(first_req)}" ) except Exception as e: logger.error( @@ -493,13 +491,11 @@ class Scheduler(SchedulerDisaggMixin): # 3. return results try: - if isinstance(processed_req, list) and processed_req: - is_warmup = processed_req[0].is_warmup + if isinstance(first_req, list) and first_req: + is_warmup = first_req[0].is_warmup else: is_warmup = ( - processed_req.is_warmup - if isinstance(processed_req, Req) - else False + first_req.is_warmup if isinstance(first_req, Req) else False ) if is_warmup: if output_batch.error is None: diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py index afc2e0e53..d4d4222a9 100644 --- a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py +++ b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py @@ -26,8 +26,8 @@ from sglang.multimodal_gen.runtime.layers.linear import ( from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 6d283f82d..75d79a4ae 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -13,7 +13,7 @@ from torch.nn.attention.flex_attention import ( flex_attention, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention # see https://github.com/pytorch/pytorch/issues/133254 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 19e0c9a70..295319027 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -33,8 +33,8 @@ from sglang.multimodal_gen.runtime.layers.linear import ( RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin def _rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 481296a2d..80fed17fb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -52,9 +52,9 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import ( CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index d087c6228..dfad883fb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -42,12 +42,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, apply_flashinfer_rope_qk_inplace, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 54b02861d..3cd445471 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -37,12 +37,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/helios.py b/python/sglang/multimodal_gen/runtime/models/dits/helios.py index ce62e521f..0187fff77 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/helios.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/helios.py @@ -46,8 +46,8 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import ( TimestepEmbedder, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py index 0bb8d2520..749cef4da 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py @@ -22,9 +22,9 @@ from sglang.multimodal_gen.runtime.layers.linear import ( RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 09a233ec9..9decdf49e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -37,13 +37,13 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import ( unpatchify, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.models.utils import modulate from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin class MMDoubleStreamBlock(nn.Module): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 31918203e..e3bc1c8fd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -31,9 +31,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py index f3deaf649..cbe8489fc 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py @@ -18,8 +18,8 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin # Reuse common functions and classes from mova_video_dit from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py index f6f520690..7e10767f5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py @@ -33,9 +33,9 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 2b51e4e97..827784f26 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -46,12 +46,12 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 868676264..56fb0c3da 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -8,8 +8,8 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbed from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index f5019ca57..eff9835b9 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -49,6 +49,7 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import ( TimestepEmbedder, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.models.utils import ( _use_aiter, @@ -58,7 +59,6 @@ from sglang.multimodal_gen.runtime.platforms import ( current_platform, ) from sglang.multimodal_gen.runtime.server_args import get_global_server_args -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.utils import add_prefix diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 832b45029..5a95f5952 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -41,9 +41,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger try: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index ad6f842a6..8e57ab99c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -16,6 +16,15 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp PipelineComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB +from sglang.multimodal_gen.runtime.managers.component_manager import ( + ComponentResidencyStrategy, + ComponentUse, + ResidencyState, +) +from sglang.multimodal_gen.runtime.managers.component_resident_strategies import ( + SnapshotModuleResidency, + SnapshotStrategy, +) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -321,44 +330,150 @@ class LTX2Pipeline(_BaseLTX2Pipeline): _add_ltx2_decoding_stage(self) -class LTX2TwoStageDeviceManager: - """ - Device residency manager for LTX-2.3 two-stage DiT switching. +class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy): + name = "ltx2_original" - Modes: - - resident: keep both DiTs on GPU; phase switch is pointer rebinding only. - - snapshot: keep CPU snapshots and prefetch the target DiT (DiT2 with pre-merged LoRA) with async H2D (similar to dit layerwise offload). - The DiT_1 will always be kept a replica in CPU. - - default snapshot behavior: allow stage1/stage2 overlap by prefetching - stage2 while stage1 is still running. - - snapshot low-VRAM behavior (`_snapshot_low_vram_mode=True`): evict - stage1 before stage2 prefetch and disable early overlap prefetch to - reduce peak VRAM, at the cost of higher phase-switch latency. - - default toggle: low-VRAM auto-enables on H100-like (<130 GiB) CUDA - GPUs, and stays disabled by default on higher-memory GPUs. It can be - overridden with `SGLANG_LTX2_SNAPSHOT_LOW_VRAM_MODE`. - - original: official two-stage semantics without premerged stage-2. + def __init__(self, manager: "LTX2TwoStageResidencyController") -> None: + self.manager = manager + + @property + def pipeline(self) -> "LTX2TwoStagePipeline": + return self.manager.pipeline + + @property + def server_args(self) -> ServerArgs: + return self.manager.server_args + + def _phase(self, use: ComponentUse) -> str: + if use.phase in ("stage1", "stage2"): + return use.phase + return "stage2" if use.component_name == "transformer_2" else "stage1" + + def initialize(self) -> None: + pass + + def prepare_for_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + phase = self._phase(use) + if phase != self.manager._active_phase: + self.enter_phase(phase) + + def wait_for_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.ensure_phase_ready(self._phase(use)) + + def finish_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + self.exit_phase(self._phase(use)) + + def prepare_after_request( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + phase = self._phase(use) + if phase != self.manager._active_phase: + self.enter_phase(phase) + + def enter_phase(self, phase: str) -> bool: + return False + + def exit_phase(self, phase: str | None, next_phase: str | None = None) -> None: + pass + + def ensure_phase_ready(self, phase: str | None) -> None: + """wait for the preparation to be ready""" + pass + + def _ensure_on_gpu(self, module_name: str) -> None: + module = self.pipeline.get_module(module_name) + if module is None: + return + param = next(module.parameters(), None) + if param is not None and param.device.type == "cpu": + module.to(get_local_torch_device(), non_blocking=True) + + @staticmethod + def _module_is_on_gpu(module: torch.nn.Module | None) -> bool: + return SnapshotModuleResidency.is_on_gpu(module) + + +class LTX2OriginalResidencyStrategy(LTX2TwoStageResidencyStrategy): + pass + + +class LTX2ResidentResidencyStrategy(LTX2TwoStageResidencyStrategy): + """A residency strategy for ltx two-stage pipeline with pre-merged lora, that keep both dits always resident""" + + name = "ltx2_resident" + + def initialize(self) -> None: + self._ensure_on_gpu("transformer") + self._ensure_on_gpu("transformer_2") + logger.info( + "Using resident LTX-2.3 two-stage transformers mode (both DiTs stay on GPU)" + ) + self.manager._active_phase = "stage1" + self.manager._sync_refinement_stage_transformer("stage1") + + def enter_phase(self, phase: str) -> bool: + self.manager._sync_refinement_stage_transformer(phase) + self.manager._active_phase = phase + return True + + +class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy): + """ + Snapshot mode keeps CPU snapshots and prefetches the target DiT with async H2D. (only with pre-merged lora enabled) + + The DiT_1 will always be kept a replica in CPU. + - default snapshot behavior: allow stage1/stage2 overlap by prefetching + stage2 while stage1 is still running. + - snapshot low-VRAM behavior (`_snapshot_low_vram_mode=True`): evict + stage1 before stage2 prefetch and disable early overlap prefetch to + reduce peak VRAM, at the cost of higher phase-switch latency. + - default toggle: low-VRAM auto-enables on H100-like (<130 GiB) CUDA + GPUs, and stays disabled by default on higher-memory GPUs. It can be + overridden with `SGLANG_LTX2_SNAPSHOT_LOW_VRAM_MODE`. """ - VALID_MODES = ("original", "snapshot", "resident") + name = "ltx2_snapshot" - def __init__(self, pipeline: "LTX2TwoStagePipeline", server_args: ServerArgs): - self.pipeline = pipeline - self.server_args = server_args - self.mode = self._resolve_mode(server_args) - self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {} - self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {} - self._active_phase: str | None = None - self._prefetch_stream: object | None = None - self._phase_ready_events: dict[str, object] = {} + def __init__(self, manager: "LTX2TwoStageResidencyController") -> None: + super().__init__(manager) + self._snapshot_strategy = SnapshotStrategy( + pin_cpu_memory=manager.server_args.pin_cpu_memory, + enable_async_prefetch=manager.server_args.dit_cpu_offload, + ) self._snapshot_low_vram_mode = self._resolve_snapshot_low_vram_mode() self._snapshot_release_empty_cache = get_bool_env_var( "SGLANG_LTX2_SNAPSHOT_RELEASE_EMPTY_CACHE", default="false", ) + @staticmethod + def _module_name_for_phase(phase: str | None) -> str | None: + if phase == "stage1": + return "transformer" + if phase == "stage2": + return "transformer_2" + return None + def _resolve_snapshot_low_vram_mode(self) -> bool: - if self.mode != "snapshot" or not current_platform.is_cuda(): + if not current_platform.is_cuda(): return False device_name = str(current_platform.get_device_name(0)).upper() device_total_memory_gb = ( @@ -384,243 +499,100 @@ class LTX2TwoStageDeviceManager: ) return enabled - @classmethod - def _resolve_mode(cls, server_args: ServerArgs) -> str: - mode = getattr(server_args, "ltx2_two_stage_device_mode", None) - if mode is None: - env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE") - mode = env_mode.lower() if env_mode else "snapshot" - if mode not in cls.VALID_MODES: - raise ValueError( - f"Invalid ltx2_two_stage_device_mode={mode!r}. " - f"Expected one of {cls.VALID_MODES}." - ) - return mode - - @property - def should_use_premerged(self) -> bool: - """Whether to keep a pre-merged stage-2 DiT for LTX-2.3 two-stage. - - We only enable this optimization for native LTX-2.3 two-stage and when - users did not explicitly provide a stage-1 LoRA path - """ - return ( - self.mode != "original" - and self.pipeline._should_merge_stage2_distilled_lora(self.server_args) - and getattr(self.pipeline, "_stage1_lora_path", None) is None - ) - def initialize(self) -> None: - if not self.should_use_premerged: - return - - self.pipeline._initialize_premerged_stage2_transformer(self.server_args) - if self.mode == "snapshot": - # Snapshot mode keeps both DiT CPU snapshots for cheap GPU release - # and re-hydrates stage-2 with async H2D when stage-1 finishes. - self._capture_module_cpu_snapshot("transformer") - self._capture_module_cpu_snapshot("transformer_2") - self._pin_stage1_transformer_if_beneficial() - elif self.mode == "resident": - self._ensure_on_gpu("transformer") - self._ensure_on_gpu("transformer_2") - logger.info( - "Using resident LTX-2.3 two-stage transformers mode (both DiTs stay on GPU)" - ) - self._active_phase = "stage1" - - self._sync_refinement_stage_transformer("stage1") - self._record_phase_ready_event("stage1") - - def switch_phase(self, phase: str) -> bool: - """Switch active two-stage DiT with minimal transfer/sync overhead.""" - if not self.should_use_premerged: - return False - if phase == self._active_phase: - return True - - if self.mode == "resident": - self._sync_refinement_stage_transformer(phase) - self._active_phase = phase - return True + # Snapshot mode keeps both DiT CPU snapshots for cheap GPU release + # and re-hydrates stage-2 with async H2D when stage-1 finishes. + self._capture_module_cpu_snapshot("transformer") + self._capture_module_cpu_snapshot("transformer_2") + self._pin_stage1_transformer_if_beneficial() + self.manager._sync_refinement_stage_transformer("stage1") + self._record_component_ready("transformer") + def enter_phase(self, phase: str) -> bool: if self.server_args.dit_cpu_offload: - target_name = "transformer_2" if phase == "stage2" else "transformer" - target_module = self.pipeline.get_module(target_name) - if self.mode == "snapshot" and self._snapshot_low_vram_mode: + target_module_name = self._module_name_for_phase(phase) + if target_module_name is None: + return False + target_module = self.pipeline.get_module(target_module_name) + if self._snapshot_low_vram_mode: # Trade a bit of phase-switch latency for lower peak VRAM: # evict stage-1 before stage-2 H2D. - if phase == "stage2" and phase not in self._phase_ready_events: - stage1_module = self.pipeline.get_module("transformer") - stage1_param = ( - next(stage1_module.parameters(), None) - if stage1_module is not None - else None - ) - if stage1_param is not None and stage1_param.device.type == "cuda": - self._release_module_to_cpu_snapshot("transformer") - if phase not in self._phase_ready_events: + if phase == "stage2" and not self._snapshot_strategy.is_ready( + target_module_name + ): + self._release_stage1_for_low_vram() + + # make sure the component is pre-fetched + if not self._snapshot_strategy.is_ready(target_module_name): if self._module_is_on_gpu(target_module): - self._record_phase_ready_event(phase) + self._record_component_ready(target_module_name) else: - self._schedule_phase_prefetch(phase, target_module) - - # Stage-2 is only consumed after stage-1 denoising + upsample. - # Kick off the H2D early in stage-1 to overlap transfer with compute. - if ( - phase == "stage1" - and "stage2" not in self._phase_ready_events - and not self._snapshot_low_vram_mode - ): - self._schedule_phase_prefetch( - "stage2", self.pipeline.get_module("transformer_2") - ) + self._snapshot_strategy.prefetch_component( + target_module_name, target_module + ) else: - self._record_phase_ready_event(phase) + component_name = self._module_name_for_phase(phase) + if component_name is not None: + self._record_component_ready(component_name) - self._sync_refinement_stage_transformer(phase) - self._active_phase = phase + self.manager._sync_refinement_stage_transformer(phase) + self.manager._active_phase = phase return True - def prefetch_stage2_after_stage1(self) -> None: - """Kick off stage-2 H2D right after stage-1 denoising to hide switch latency.""" + def prepare_after_request( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + phase = self._phase(use) + if phase != "stage1": + return + if self.server_args.dit_cpu_offload: + target_module = self.pipeline.get_module("transformer") + if self._module_is_on_gpu(target_module): + self._record_component_ready("transformer") + elif not self._snapshot_strategy.is_ready("transformer"): + self._snapshot_strategy.prefetch_component("transformer", target_module) + else: + self._record_component_ready("transformer") + self.manager._sync_refinement_stage_transformer("stage1") + self.manager._active_phase = "stage1" + + def finish_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> None: + phase = self._phase(use) + if self.server_args.dit_cpu_offload: + # release cuda storage + self._snapshot_strategy.release_component(use.component_name, module) if ( - not self.should_use_premerged - or self.mode != "snapshot" - or not self.server_args.dit_cpu_offload - ): - return - - if "stage2" in self._phase_ready_events: - return - if self._snapshot_low_vram_mode: - self._release_stage1_for_low_vram() - - self._schedule_phase_prefetch( - "stage2", self.pipeline.get_module("transformer_2") - ) - - def prepare_upsample_after_stage1(self) -> bool: - if ( - not self.should_use_premerged - or self.mode != "snapshot" - or not self.server_args.dit_cpu_offload - or not self._snapshot_low_vram_mode - ): - return False - if "stage2" in self._phase_ready_events: - return False - self._release_stage1_for_low_vram() - return True - - def ensure_phase_ready(self, phase: str | None) -> None: - if not self.should_use_premerged or phase not in ("stage1", "stage2"): - return - if self.mode == "resident": - return - ready_event = self._phase_ready_events.get(phase) - if ready_event is None or not current_platform.is_cuda(): - return - torch.get_device_module().current_stream().wait_event(ready_event) - - def release_premerged_transformers(self) -> None: - if not self.should_use_premerged or self.mode != "snapshot": - return - # Keep stage-1 resident across requests so the next request can start - # denoising immediately while stage-2 is prefetched in the background. - for module_name in ("transformer_2",): - module = self.pipeline.get_module(module_name) - param = next(module.parameters(), None) if module is not None else None - if param is not None and param.device.type == "cuda": - self._release_module_to_cpu_snapshot(module_name) - if ( - self._snapshot_release_empty_cache + phase == "stage2" + and self._snapshot_release_empty_cache and torch.get_device_module().is_available() ): torch.get_device_module().empty_cache() - self._record_phase_ready_event("stage1") - @staticmethod - def _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 ensure_phase_ready(self, phase: str | None) -> None: + component_name = self._module_name_for_phase(phase) + if component_name is None: + return + self._snapshot_strategy.wait_component_ready(component_name) def _capture_module_cpu_snapshot(self, module_name: str) -> None: - if module_name in self._cpu_param_snapshots: - return - module = self.pipeline.get_module(module_name) if module is None: raise ValueError(f"Module {module_name} is not available.") - - pin_memory = bool( - self.server_args.pin_cpu_memory and torch.get_device_module().is_available() - ) - self._cpu_param_snapshots[module_name] = { - name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory) - for name, param in module.named_parameters() - } - self._cpu_buffer_snapshots[module_name] = { - name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory) - for name, buffer in module.named_buffers() - } + self._snapshot_strategy.capture(module_name, module) def _release_module_to_cpu_snapshot(self, module_name: str) -> None: - """Replace module tensors with cached CPU snapshots to avoid D2H copies. - - This does not call `module.to("cpu")`. Instead, parameter and buffer storages - are rebound to pre-captured CPU tensors so CUDA storages can be released by - the allocator without an explicit D2H transfer. - """ module = self.pipeline.get_module(module_name) if module is None: return - - param_snapshots = self._cpu_param_snapshots.get(module_name) - buffer_snapshots = self._cpu_buffer_snapshots.get(module_name) - if param_snapshots is None or buffer_snapshots is None: - module.to("cpu") - return - - pin_memory = bool( - self.server_args.pin_cpu_memory and torch.get_device_module().is_available() - ) - for name, param in module.named_parameters(): - snapshot = param_snapshots.get(name) - if snapshot is None: - snapshot = self._clone_cpu_tensor_snapshot( - param.data, pin_memory=pin_memory - ) - param_snapshots[name] = snapshot - param.data = snapshot - - for name, buffer in module.named_buffers(): - snapshot = buffer_snapshots.get(name) - if snapshot is None: - snapshot = self._clone_cpu_tensor_snapshot( - buffer.data, pin_memory=pin_memory - ) - buffer_snapshots[name] = snapshot - # Preserve runtime-updated buffers (e.g., lazily built caches) when - # releasing back to CPU snapshots. - if buffer.device.type == "cuda": - snapshot.copy_(buffer.detach().to(device="cpu", dtype=snapshot.dtype)) - elif buffer.device.type == "cpu": - snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype)) - buffer.data = snapshot - - phase = "stage2" if module_name == "transformer_2" else "stage1" - self._phase_ready_events.pop(phase, None) + self._snapshot_strategy.release_component(module_name, module) def _release_stage1_for_low_vram(self) -> None: stage1_module = self.pipeline.get_module("transformer") @@ -632,66 +604,29 @@ class LTX2TwoStageDeviceManager: if stage1_param is not None and stage1_param.device.type == "cuda": self._release_module_to_cpu_snapshot("transformer") - def _ensure_on_gpu(self, module_name: str) -> None: - module = self.pipeline.get_module(module_name) - if module is None: - return - param = next(module.parameters(), None) - if param is not None and param.device.type == "cpu": - module.to(get_local_torch_device(), non_blocking=True) - - @staticmethod - def _module_is_on_gpu(module: torch.nn.Module | None) -> bool: - if module is None: - return False - param = next(module.parameters(), None) - return param is not None and param.device.type == "cuda" - - def _supports_async_phase_prefetch(self) -> bool: - return ( - self.mode == "snapshot" - and self.server_args.dit_cpu_offload - and current_platform.is_cuda() + def _record_component_ready(self, module_name: str) -> None: + self._snapshot_strategy.record_ready( + module_name, self.pipeline.get_module(module_name) ) - def _get_prefetch_stream(self): - if not self._supports_async_phase_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 _record_phase_ready_event(self, phase: str) -> None: - if not current_platform.is_cuda(): - self._phase_ready_events.pop(phase, None) - return - module_name = "transformer_2" if phase == "stage2" else "transformer" - module = self.pipeline.get_module(module_name) - if not self._module_is_on_gpu(module): - self._phase_ready_events.pop(phase, None) - return - event = torch.get_device_module().Event() - event.record(torch.get_device_module().current_stream()) - self._phase_ready_events[phase] = event - - def _schedule_phase_prefetch( - self, phase: str, module: torch.nn.Module | None - ) -> None: - if module is None: - self._phase_ready_events.pop(phase, None) - return - prefetch_stream = self._get_prefetch_stream() - if prefetch_stream is None: - module.to(get_local_torch_device(), non_blocking=True) - self._record_phase_ready_event(phase) - return - with torch.get_device_module().stream(prefetch_stream): - module.to(get_local_torch_device(), non_blocking=True) - event = torch.get_device_module().Event() - event.record(prefetch_stream) - self._phase_ready_events[phase] = event + def prefetch_for_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> bool: + if not self.server_args.dit_cpu_offload: + return True + phase = self._phase(use) + if phase == "stage2": + if self._snapshot_strategy.is_ready("transformer_2"): + return True + if self._snapshot_low_vram_mode and state.current_use is not None: + return False + if self._snapshot_low_vram_mode: + self._release_stage1_for_low_vram() + self._snapshot_strategy.prefetch_component(use.component_name, module) + return True def _pin_stage1_transformer_if_beneficial(self) -> None: """Optionally pin stage-1 DiT on GPU to remove first-stage cold H2D stall. @@ -717,7 +652,80 @@ class LTX2TwoStageDeviceManager: logger.info( "Pinned stage1 transformer on GPU for LTX-2.3 two-stage startup" ) - self._active_phase = "stage1" + self.manager._active_phase = "stage1" + + +class LTX2TwoStageResidencyController: + """ + LTX-2.3 two-stage residency controller. + It builds the selected LTX2 ComponentResidencyStrategy and keeps the + thin stage adapter methods that are specific to two-stage LoRA flow. + + Modes: + - resident: keep both DiTs on GPU; phase switch is pointer rebinding only. + - snapshot: keep CPU snapshots and prefetch the target DiT. + - original: official two-stage semantics without premerged stage-2. + """ + + VALID_MODES = ("original", "snapshot", "resident") + + def __init__(self, pipeline: "LTX2TwoStagePipeline", server_args: ServerArgs): + self.pipeline = pipeline + self.server_args = server_args + self.mode = self._resolve_mode(server_args) + self._active_phase: str | None = None + self._strategy = self._build_strategy() + + @classmethod + def _resolve_mode(cls, server_args: ServerArgs) -> str: + mode = server_args.ltx2_two_stage_device_mode + if mode is None: + env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE") + mode = env_mode.lower() if env_mode else "snapshot" + if mode not in cls.VALID_MODES: + raise ValueError( + f"Invalid ltx2_two_stage_device_mode={mode!r}. " + f"Expected one of {cls.VALID_MODES}." + ) + return mode + + def _build_strategy(self) -> LTX2TwoStageResidencyStrategy: + if self.mode == "snapshot": + return LTX2SnapshotResidencyStrategy(self) + if self.mode == "resident": + return LTX2ResidentResidencyStrategy(self) + return LTX2OriginalResidencyStrategy(self) + + @property + def strategy(self) -> ComponentResidencyStrategy: + return self._strategy + + @property + def should_use_premerged(self) -> bool: + """Whether to keep a pre-merged stage-2 DiT for LTX-2.3 two-stage. + + We only enable this optimization for native LTX-2.3 two-stage and when + users did not explicitly provide a stage-1 LoRA path + """ + return ( + self.mode != "original" + and self.pipeline._should_merge_stage2_distilled_lora(self.server_args) + and self.pipeline._stage1_lora_path is None + ) + + def initialize(self) -> None: + if not self.should_use_premerged: + return + self.pipeline._initialize_premerged_stage2_transformer(self.server_args) + self._strategy.initialize() + + def enter_phase(self, phase: str) -> bool: + """Switch active two-stage DiT with minimal transfer/sync overhead.""" + if not self.should_use_premerged: + return False + if phase == self._active_phase: + return True + return self._strategy.enter_phase(phase) def _sync_refinement_stage_transformer(self, phase: str) -> None: """Keep stage-2 refinement bound to the expected DiT for current phase.""" @@ -740,11 +748,18 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._device_manager = LTX2TwoStageDeviceManager(self, self.server_args) + self._ltx2_residency = LTX2TwoStageResidencyController(self, self.server_args) self._use_premerged_stage2_transformer = ( - self._device_manager.should_use_premerged + self._ltx2_residency.should_use_premerged ) - self._device_manager.initialize() + self._ltx2_residency.initialize() + if self._use_premerged_stage2_transformer: + self.component_residency_strategies["transformer"] = ( + self._ltx2_residency.strategy + ) + self.component_residency_strategies["transformer_2"] = ( + self._ltx2_residency.strategy + ) @staticmethod def _should_merge_stage2_distilled_lora(server_args: ServerArgs) -> bool: @@ -809,25 +824,8 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): merge_weights=True, ) - def release_premerged_transformers_to_cpu_snapshots(self) -> None: - """Release inactive premerged DiTs according to the selected device mode.""" - self._device_manager.release_premerged_transformers() - - def release_ltx2_phase_state(self, phase: str | None) -> None: - if phase == "stage2": - self.release_premerged_transformers_to_cpu_snapshots() - - def ensure_ltx2_phase_ready(self, phase: str | None) -> None: - self._device_manager.ensure_phase_ready(phase) - - def prefetch_ltx2_stage2_after_stage1(self) -> None: - self._device_manager.prefetch_stage2_after_stage1() - - def prepare_ltx2_upsample_after_stage1(self) -> bool: - return self._device_manager.prepare_upsample_after_stage1() - def should_skip_ltx2_lora_switch_stage(self) -> bool: - return self._use_premerged_stage2_transformer and self._device_manager.mode in ( + return self._use_premerged_stage2_transformer and self._ltx2_residency.mode in ( "snapshot", "resident", ) @@ -912,7 +910,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): if phase_signature == self._active_lora_signature: return - if self._device_manager.switch_phase( + if self._ltx2_residency.enter_phase( phase ) and self._can_short_circuit_lora_switch(phase, batch): self._active_lora_phase = phase diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index ffe9e3103..109b3f036 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -21,6 +21,11 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import ( from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, ) +from sglang.multimodal_gen.runtime.managers.component_manager import ( + ComponentResidencyManager, + ComponentResidencyStrategy, + get_global_component_residency_manager, +) from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) @@ -91,7 +96,9 @@ class ComposedPipelineBase(ABC): self.model_path: str = model_path self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} + self.component_residency_strategies: dict[str, ComponentResidencyStrategy] = {} self.executor = executor or self.build_executor(server_args=server_args) + self.component_residency_manager: ComponentResidencyManager | None = None if required_config_modules is not None: self._required_config_modules = required_config_modules @@ -738,6 +745,11 @@ class ComposedPipelineBase(ABC): main_process_only=True, ) + self.component_residency_manager = get_global_component_residency_manager( + self, server_args + ) + self.executor.component_residency_manager = self.component_residency_manager + return self.executor.execute_with_profiling(self.stages, batch, server_args) @torch.no_grad() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 2853ca4ad..3a4cf3419 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -55,7 +55,7 @@ class ParallelExecutor(PipelineExecutor): def _execute_stages( self, stages: List[PipelineStage], - payload: Any, + batch: Any, server_args: ServerArgs, run_stage: Callable[[PipelineStage, Any], Any], ) -> Any: @@ -66,30 +66,40 @@ class ParallelExecutor(PipelineExecutor): rank = get_world_rank() cfg_group = get_cfg_group() - # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY - for stage in stages: - paradigm = stage.parallelism_type + self.begin_component_residency_request(stages, batch, server_args) + try: + # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY + for stage_index, stage in enumerate(stages): + paradigm = stage.parallelism_type - if paradigm == StageParallelismType.MAIN_RANK_ONLY: - if rank == 0: - # Only main rank executes, others just wait - payload = run_stage(stage, payload) - torch.distributed.barrier() + if paradigm == StageParallelismType.MAIN_RANK_ONLY: + if rank == 0: + # Only main rank executes, others just wait + self.before_stage(stage, stage_index, batch, server_args) + batch = stage(batch, server_args) + self.after_stage(stage_index) + torch.distributed.barrier() - elif paradigm == StageParallelismType.CFG_PARALLEL: - obj_list = [payload] if rank == 0 else [] - broadcasted_list = broadcast_pyobj( - obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0 - ) - if rank != 0: - payload = broadcasted_list[0] - payload = run_stage(stage, payload) + elif paradigm == StageParallelismType.CFG_PARALLEL: + obj_list = [batch] if rank == 0 else [] + broadcasted_list = broadcast_pyobj( + obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0 + ) + if rank != 0: + batch = broadcasted_list[0] + self.before_stage(stage, stage_index, batch, server_args) + batch = stage(batch, server_args) + self.after_stage(stage_index) - torch.distributed.barrier() + torch.distributed.barrier() - elif paradigm == StageParallelismType.REPLICATED: - payload = run_stage(stage, payload) - return payload + elif paradigm == StageParallelismType.REPLICATED: + self.before_stage(stage, stage_index, batch, server_args) + batch = stage(batch, server_args) + self.after_stage(stage_index) + finally: + self.finish_component_residency_request() + return batch def execute( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index 88117093a..2d0ee385d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py @@ -45,6 +45,33 @@ class PipelineExecutor(ABC): def __init__(self, server_args): self.server_args = server_args + self.component_residency_manager = None + + def begin_component_residency_request( + self, + stages: List["PipelineStage"], + batch: Req, + server_args: ServerArgs, + ) -> None: + self.component_residency_manager.begin_request(stages, batch, server_args) + + def before_stage( + self, + stage: "PipelineStage", + stage_index: int, + batch: Req, + server_args: ServerArgs, + ) -> None: + stage.set_component_residency_manager(self.component_residency_manager) + self.component_residency_manager.before_stage( + stage, stage_index, batch, server_args + ) + + def after_stage(self, stage_index: int) -> None: + self.component_residency_manager.after_stage(stage_index) + + def finish_component_residency_request(self) -> None: + self.component_residency_manager.finish_request() def execute_with_profiling( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py index b1539863f..d1ed6fcf1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py @@ -29,11 +29,17 @@ class SyncExecutor(PipelineExecutor): run_stage: Callable[[PipelineStage, Any], Any], ) -> Any: """Execute all pipeline stages sequentially and step the profiler.""" - for stage in stages: - payload = run_stage(stage, payload) - profiler = SGLDiffusionProfiler.get_instance() - if profiler: - profiler.step_stage() + self.begin_component_residency_request(stages, payload, server_args) + try: + for stage_index, stage in enumerate(stages): + self.before_stage(stage, stage_index, payload, server_args) + payload = run_stage(stage, payload) + self.after_stage(stage_index) + profiler = SGLDiffusionProfiler.get_instance() + if profiler: + profiler.step_stage() + finally: + self.finish_component_residency_request() return payload def run_profile_all_stages( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index fb93a3f01..a8e13f222 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -169,7 +169,7 @@ class LoRAPipeline(ComposedPipelineBase): Yields: List of modules that had offload disabled. """ - from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( + from sglang.multimodal_gen.runtime.managers.layerwise_offload import ( OffloadableDiTMixin, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index dc5fe9bde..4183fbcf3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -69,8 +69,8 @@ class Req: # Primary encoder embeddings prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list) negative_prompt_embeds: list[torch.Tensor] | None = None - prompt_attention_mask: list[torch.Tensor] | None = None - negative_attention_mask: list[torch.Tensor] | None = None + prompt_attention_mask: list[torch.Tensor | None] | None = None + negative_attention_mask: list[torch.Tensor | None] | None = None clip_embedding_pos: list[torch.Tensor] | None = None clip_embedding_neg: list[torch.Tensor] | None = None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index c6f46f5da..a62cd2923 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -9,11 +9,15 @@ composed to create complete diffusion pipelines. """ from abc import ABC, abstractmethod +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import replace from enum import Enum, auto import torch from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.dedup import StageDedupMixin from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( @@ -53,6 +57,7 @@ class PipelineStage(StageDedupMixin, ABC): def __init__(self): self.server_args = get_global_server_args() + self._component_residency_manager = None def log_info(self, msg, *args): """Logs an informational message with the stage name as a prefix.""" @@ -105,6 +110,81 @@ class PipelineStage(StageDedupMixin, ABC): """ pass + def set_component_residency_manager(self, manager) -> None: + self._component_residency_manager = manager + + def _component_stage_name(self, stage_name: str | None = None) -> str: + return stage_name or self.__class__.__name__ + + def _active_component_stage_name(self) -> str: + manager = self._component_residency_manager + if manager is not None and manager.state.stage_name is not None: + return manager.state.stage_name + return self.__class__.__name__ + + def _finish_active_component_use(self) -> None: + if self._component_residency_manager is not None: + self._component_residency_manager.finish_active_use() + + @contextmanager + def _use_component( + self, + use: ComponentUse, + module=None, + ) -> Iterator[object | None]: + if self._component_residency_manager is None: + yield module + return + with self._component_residency_manager.use_component(use, module) as component: + yield component + + def _declared_component_use( + self, + *, + component_name: str, + phase: str | None = None, + target_dtype: torch.dtype | None = None, + ) -> ComponentUse: + manager = self._component_residency_manager + stage_name = self._active_component_stage_name() + server_args = manager.server_args if manager is not None else self.server_args + for use in self.component_uses(server_args, stage_name): + if use.component_name != component_name: + continue + if phase is not None and use.phase != phase: + continue + if target_dtype is not None: + return replace(use, target_dtype=target_dtype) + return use + raise ValueError( + f"{self.__class__.__name__} did not declare component use: " + f"{component_name}" + ) + + @contextmanager + def use_declared_component( + self, + *, + component_name: str, + module=None, + phase: str | None = None, + target_dtype: torch.dtype | None = None, + ) -> Iterator[object | None]: + """reference a component already declared in `component_uses`""" + use = self._declared_component_use( + component_name=component_name, + phase=phase, + target_dtype=target_dtype, + ) + with self._use_component(use, module) as component: + yield component + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + """Declares component uses of current stage for unified residency scheduling.""" + return [] + # Default role affinity: ENCODER. Override in subclasses for DENOISING/DECODER. @property def role_affinity(self) -> RoleType: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index b8e25636b..a6010a524 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -11,6 +11,7 @@ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import VAELoader +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( @@ -68,6 +69,20 @@ class DecodingStage(PipelineStage): self.pipeline = weakref.ref(pipeline) if pipeline else None self.component_name = component_name + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name, + self.component_name, + target_dtype=vae_dtype, + keep_ready_after_warmup=True, + ) + ] + @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: @@ -110,7 +125,13 @@ class DecodingStage(PipelineStage): return latents @torch.no_grad() - def decode(self, latents: torch.Tensor, server_args: ServerArgs) -> torch.Tensor: + def decode( + self, + latents: torch.Tensor, + server_args: ServerArgs, + *, + vae_dtype: torch.dtype, + ) -> torch.Tensor: """ Decode latent representations into pixel space using VAE. @@ -125,8 +146,6 @@ class DecodingStage(PipelineStage): Decoded video tensor with shape (batch, channels, frames, height, width), normalized to [0, 1] range and moved to CPU as float32 """ - vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - self.vae = self.vae.to(device=get_local_torch_device(), dtype=vae_dtype) latents = latents.to(get_local_torch_device()) vae_autocast_enabled = ( vae_dtype != torch.float32 @@ -175,22 +194,6 @@ class DecodingStage(PipelineStage): pipeline.add_module(self.component_name, self.vae) self.server_args.model_loaded[self.component_name] = True - def offload_model(self): - # Offload models if needed - self.maybe_free_model_hooks() - - if self.server_args.vae_cpu_offload: - self.vae.to("cpu", non_blocking=True) - - if torch.backends.mps.is_available(): - # Flush lazy MPS kernels before freeing weights to avoid hangs. - torch.mps.synchronize() - del self.vae - pipeline = self.pipeline() if self.pipeline else None - if pipeline is not None and self.component_name in pipeline.modules: - del pipeline.modules[self.component_name] - self.server_args.model_loaded[self.component_name] = False - @torch.no_grad() def forward( self, @@ -208,32 +211,42 @@ class DecodingStage(PipelineStage): # load vae if not already loaded (used for memory constrained devices) self.load_model() - frames = self.decode(batch.latents, server_args) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + with self.use_declared_component( + component_name=self.component_name, + module=self.vae, + ) as vae: + assert vae is not None + self.vae = vae - # decode trajectory latents if needed - if batch.return_trajectory_decoded: - assert ( - batch.trajectory_latents is not None - ), "batch should have trajectory latents" + frames = self.decode(batch.latents, server_args, vae_dtype=vae_dtype) - # 1. Batch trajectory decoding to improve GPU utilization - # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] - B, T, C, F, H, W = batch.trajectory_latents.shape - flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W) + # decode trajectory latents if needed + if batch.return_trajectory_decoded: + assert ( + batch.trajectory_latents is not None + ), "batch should have trajectory latents" - logger.info("decoding %s trajectory latents in batch", B * T) - # Use the optimized batch decode - all_decoded = self.decode(flat_latents, server_args) + # 1. Batch trajectory decoding to improve GPU utilization + # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] + B, T, C, F, H, W = batch.trajectory_latents.shape + flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W) - # 2. Reshape back - # Keep on GPU to allow faster vectorized post-processing - decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:]) + logger.info("decoding %s trajectory latents in batch", B * T) + # Use the optimized batch decode + all_decoded = self.decode( + flat_latents, server_args, vae_dtype=vae_dtype + ) - # Convert to list of tensors (per timestep) as expected by OutputBatch - # Each element in list is [B, channels, frames, H_out, W_out] - trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] - else: - trajectory_decoded = None + # 2. Reshape back + # Keep on GPU to allow faster vectorized post-processing + decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:]) + + # Convert to list of tensors (per timestep) as expected by OutputBatch + # Each element in list is [B, channels, frames, H_out, W_out] + trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] + else: + trajectory_decoded = None frames = server_args.pipeline_config.post_decoding(frames, server_args) @@ -248,8 +261,4 @@ class DecodingStage(PipelineStage): noise_pred=None, ) - # Keep VAE resident during warmup; the real request needs it next. - if not getattr(batch, "is_warmup", False): - self.offload_model() - return output_batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py index 362d81097..6f82300cc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py @@ -1,6 +1,7 @@ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage from sglang.multimodal_gen.runtime.platforms import current_platform @@ -25,6 +26,16 @@ class LTX2AVDecodingStage(DecodingStage): self.video_processor = VideoProcessor(vae_scale_factor=32) + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse(stage_name, "vae", target_dtype=torch.bfloat16), + ComponentUse(stage_name, "audio_vae"), + ComponentUse(stage_name, "vocoder"), + ] + @staticmethod def _ltx2_should_externally_denorm_video_latents(server_args: ServerArgs) -> bool: arch_config = server_args.pipeline_config.vae_config.arch_config @@ -33,45 +44,44 @@ class LTX2AVDecodingStage(DecodingStage): def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: self.load_model() - self.vae = self.vae.to(get_local_torch_device()) - self.vae.eval() - latents = batch.latents.to(get_local_torch_device()) - vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast original_dtype = vae_dtype - self.vae.to(torch.bfloat16) - latents = latents.to(torch.bfloat16) - if self._ltx2_should_externally_denorm_video_latents(server_args): - std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) - mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) - latents = latents * std + mean - latents = server_args.pipeline_config.preprocess_decoding( - latents, server_args, vae=self.vae - ) + with self.use_declared_component(component_name="vae", module=self.vae) as vae: + assert vae is not None + self.vae = vae + self.vae.eval() + latents = batch.latents.to(get_local_torch_device(), dtype=torch.bfloat16) + if self._ltx2_should_externally_denorm_video_latents(server_args): + std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) + mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) + latents = latents * std + mean + latents = server_args.pipeline_config.preprocess_decoding( + latents, server_args, vae=self.vae + ) - with torch.autocast( - device_type=current_platform.device_type, - dtype=vae_dtype, - enabled=vae_autocast_enabled, - ): - try: - if server_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - except Exception: - pass - decode_output = self.vae.decode(latents) - if isinstance(decode_output, tuple): - video = decode_output[0] - elif hasattr(decode_output, "sample"): - video = decode_output.sample - else: - video = decode_output + with torch.autocast( + device_type=current_platform.device_type, + dtype=vae_dtype, + enabled=vae_autocast_enabled, + ): + try: + if server_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + except Exception: + pass + decode_output = self.vae.decode(latents) + if isinstance(decode_output, tuple): + video = decode_output[0] + elif hasattr(decode_output, "sample"): + video = decode_output.sample + else: + video = decode_output - self.vae.to(original_dtype) + self.vae.to(original_dtype) video = self.video_processor.postprocess_video(video, output_type="np") output_batch = OutputBatch( @@ -90,49 +100,64 @@ class LTX2AVDecodingStage(DecodingStage): if audio_latents is not None: # Ensure device/dtype device = get_local_torch_device() - self.audio_vae = self.audio_vae.to(device) - self.vocoder = self.vocoder.to(device) - self.audio_vae.eval() - self.vocoder.eval() - try: - dtype = self.audio_vae.dtype - except AttributeError: - dtype = None - if dtype is None: + with self.use_declared_component( + component_name="audio_vae", + module=self.audio_vae, + ) as audio_vae: + assert audio_vae is not None + self.audio_vae = audio_vae + self.audio_vae.eval() try: - dtype = next(self.audio_vae.parameters()).dtype - except StopIteration: - dtype = torch.float32 - audio_latents = audio_latents.to(device, dtype=dtype) - try: - latents_std = self.audio_vae.latents_std - except AttributeError: - latents_std = None - if isinstance(latents_std, torch.Tensor) and torch.all(latents_std == 0): - logger.warning( - "audio_vae.latents_std is all zeros; audio denorm may be incorrect." - ) - try: - latents_mean = self.audio_vae.latents_mean - except AttributeError: - latents_mean = None - if isinstance(latents_mean, torch.Tensor) and isinstance( - latents_std, torch.Tensor - ): - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - if audio_latents.ndim == 4: - latents_mean = latents_mean.view( - 1, audio_latents.shape[1], 1, audio_latents.shape[3] + dtype = self.audio_vae.dtype + except AttributeError: + dtype = None + if dtype is None: + try: + dtype = next(self.audio_vae.parameters()).dtype + except StopIteration: + dtype = torch.float32 + audio_latents = audio_latents.to(device, dtype=dtype) + try: + latents_std = self.audio_vae.latents_std + except AttributeError: + latents_std = None + if isinstance(latents_std, torch.Tensor) and torch.all( + latents_std == 0 + ): + logger.warning( + "audio_vae.latents_std is all zeros; audio denorm may be incorrect." ) - latents_std = latents_std.view( - 1, audio_latents.shape[1], 1, audio_latents.shape[3] - ) - audio_latents = audio_latents * latents_std + latents_mean + try: + latents_mean = self.audio_vae.latents_mean + except AttributeError: + latents_mean = None + if isinstance(latents_mean, torch.Tensor) and isinstance( + latents_std, torch.Tensor + ): + latents_mean = latents_mean.to(device=device, dtype=dtype) + latents_std = latents_std.to(device=device, dtype=dtype) + if audio_latents.ndim == 4: + latents_mean = latents_mean.view( + 1, audio_latents.shape[1], 1, audio_latents.shape[3] + ) + latents_std = latents_std.view( + 1, audio_latents.shape[1], 1, audio_latents.shape[3] + ) + audio_latents = audio_latents * latents_std + latents_mean - with torch.no_grad(): - # Decode latents to spectrogram - spectrogram = self.audio_vae.decode(audio_latents, return_dict=False)[0] + with torch.no_grad(): + # Decode latents to spectrogram + spectrogram = self.audio_vae.decode( + audio_latents, return_dict=False + )[0] + + with self.use_declared_component( + component_name="vocoder", + module=self.vocoder, + ) as vocoder: + assert vocoder is not None + self.vocoder = vocoder + self.vocoder.eval() if hasattr(self.vocoder, "conv_in") and hasattr( self.vocoder.conv_in, "in_channels" ): @@ -143,7 +168,8 @@ class LTX2AVDecodingStage(DecodingStage): f"Vocoder expects channels*mel_bins={expected_in}, got {actual_in} from spectrogram shape {tuple(spectrogram.shape)}" ) # Decode spectrogram to waveform - waveform = self.vocoder(spectrogram) + with torch.no_grad(): + waveform = self.vocoder(spectrogram) output_batch.audio = waveform.cpu().float() try: pipeline_audio_cfg = server_args.pipeline_config.audio_vae_config @@ -170,5 +196,4 @@ class LTX2AVDecodingStage(DecodingStage): vocoder_sr or audio_vae_sr or pipeline_audio_sr ) - self.offload_model() return output_batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index e5b14784b..9dda74903 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -60,6 +60,7 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import ( from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import ( TransformerLoader, ) +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( @@ -87,11 +88,10 @@ from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import RolloutDenoisingMixin, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler -from sglang.multimodal_gen.utils import dict_to_3d_list +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, dict_to_3d_list from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) @@ -195,6 +195,49 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): self._cached_num_steps = None self._is_warmed_up = False + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + uses: list[ComponentUse] = [] + if self.vae is not None: + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + uses.append( + ComponentUse( + stage_name=stage_name, + component_name="vae", + target_dtype=vae_dtype, + ) + ) + for default_name, module in ( + ("transformer", self.transformer), + ("transformer_2", self.transformer_2), + ): + if module is None: + continue + component_name = self._component_name_for_stage_module(module, default_name) + uses.append( + ComponentUse( + stage_name=stage_name, + component_name=component_name, + phase=component_name, + preferred_ready_after_request=component_name == "transformer", + memory_intensive=True, + ) + ) + return uses + + def _component_name_for_stage_module( + self, module: nn.Module | None, default_name: str + ) -> str: + pipeline = self.pipeline() if self.pipeline else None + if pipeline is None or module is None: + return default_name + for name, candidate in pipeline.modules.items(): + if candidate is module: + return name + return default_name + def _maybe_enable_torch_compile(self, module: object) -> None: """ Compile a module with torch.compile, and enable inductor overlap tweak if available. @@ -572,13 +615,22 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): # TI2V specific preparations - before SP sharding if should_preprocess_for_wan_ti2v: - seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents( - self.vae, - latents, - target_dtype, - batch, - server_args, - ) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + with self.use_declared_component( + component_name="vae", + module=self.vae, + target_dtype=vae_dtype, + ) as vae: + assert vae is not None + self.vae = vae + seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents( + self.vae, + latents, + target_dtype, + vae_dtype, + batch, + server_args, + ) else: seq_len, z, reserved_frames_masks = ( None, @@ -765,7 +817,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): self, batch: Req ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]: """Return the prompt-embedding validator used by verify_input.""" - del batch return V.list_not_empty def _get_negative_prompt_embeds_validator( @@ -928,9 +979,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): ): self.save_sta_search_results(batch) - # Capture references before potential deletion on MPS - dits = list(filter(None, [self.transformer, self.transformer_2])) - # deallocate transformer if on mps pipeline = self.pipeline() if self.pipeline else None if torch.backends.mps.is_available() and not is_warmup: @@ -947,14 +995,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): torch.mps.current_allocated_memory(), ) - # reset offload managers with prefetching first layer for next forward - for dit in dits: - if isinstance(dit, OffloadableDiTMixin): - # release all DiT weights to avoid peak VRAM usage, which may increasing the latency for next req - # TODO: should be make this an option? - for manager in dit.layerwise_offload_managers: - manager.release_all() - def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): """Shard latents for Sequence Parallelism if applicable.""" if get_sp_world_size() <= 1: @@ -1026,35 +1066,32 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): if profiler: profiler.step_denoising_step() - def _manage_device_placement( + def _manage_dit_use_site( self, - model_to_use: nn.Module, - model_to_offload: nn.Module | None, - server_args: ServerArgs, - ): + current_model: nn.Module, + current_phase: str, + batch: Req, + ) -> None: """ - Manages the offload / load behavior of dit + manage dit's residency by reporting the active sequential use + + only applicable for dual-dit architecture like Wan + + Args: + current_model: the next active dit, transformer_1 or transformer_2 """ - if not server_args.dit_cpu_offload: - return + manager = self._component_residency_manager - # FSDP manages offloading internally - if server_args.use_fsdp_inference: - return - - # Offload the unused model if it's on CUDA - if ( - model_to_offload is not None - and next(model_to_offload.parameters()).device.type == "cuda" - ): - model_to_offload.to("cpu") - - # Load the model to use if it's on CPU - if ( - model_to_use is not None - and next(model_to_use.parameters()).device.type == "cpu" - ): - model_to_use.to(get_local_torch_device()) + component_name = manager.component_name_for_module(current_model, current_phase) + phase = str(batch.extra.get("ltx2_phase", current_phase)) + use = ComponentUse( + stage_name=self._active_component_stage_name(), + component_name=component_name, + phase=phase, + preferred_ready_after_request=component_name == "transformer", + memory_intensive=True, + ) + manager.begin_use(use) def _select_and_manage_model( self, @@ -1066,15 +1103,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): if boundary_timestep is None or t_int >= boundary_timestep: # High-noise stage current_model = self.transformer - model_to_offload = self.transformer_2 current_guidance_scale = batch.guidance_scale + current_phase = "transformer" else: # Low-noise stage current_model = self.transformer_2 - model_to_offload = self.transformer current_guidance_scale = batch.guidance_scale_2 + current_phase = "transformer_2" - self._manage_device_placement(current_model, model_to_offload, server_args) + self._manage_dit_use_site(current_model, current_phase, batch) assert current_model is not None, "The model for the current step is not set." return current_model, current_guidance_scale @@ -1195,6 +1232,8 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): (denoising_end_time - denoising_start_time) / len(ctx.timesteps), ) + self._finish_active_component_use() + # Rollout postprocessing must run BEFORE _finalize_denoising_loop so # the final scheduler.step output (ctx.latents) is still SP-sharded and # can be gathered uniformly alongside the per-step dit_trajectory via diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 0df3e14bb..1f22533da 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -2,6 +2,7 @@ import torch from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import is_ltx23_native_variant +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import ( clone_scheduler_runtime, ) @@ -10,7 +11,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import LTX2DenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -28,6 +28,20 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage): ) self.audio_vae = audio_vae + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name=stage_name, + component_name="transformer", + phase="stage1", + preferred_ready_after_request=True, + memory_intensive=True, + ) + ] + def _post_denoising_loop( self, batch: Req, @@ -83,24 +97,6 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage): batch.latents = latents batch.audio_latents = audio_latents - pipeline = self.pipeline() if self.pipeline else None - current_phase = ( - str(getattr(batch, "extra", {}).get("ltx2_phase", "")) - if hasattr(batch, "extra") - else "" - ) - release_phase_state = ( - getattr(pipeline, "release_ltx2_phase_state", None) - if pipeline is not None - else None - ) - if callable(release_phase_state): - release_phase_state(current_phase) - - if isinstance(self.transformer, OffloadableDiTMixin): - for manager in self.transformer.layerwise_offload_managers: - manager.release_all() - class LTX2RefinementStage(LTX2AVDenoisingStage): """Stage-2 refinement wrapper that re-noises distilled LTX latents once.""" @@ -125,6 +121,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): ) self.distilled_sigmas = torch.tensor(distilled_sigmas) + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + component_name = "transformer_2" + pipeline = self.pipeline() if self.pipeline else None + if pipeline is not None and "transformer_2" not in pipeline.modules: + component_name = "transformer" + return [ + ComponentUse( + stage_name=stage_name, + component_name=component_name, + phase="stage2", + memory_intensive=True, + ) + ] + @staticmethod def _randn_like_with_batch_generators( reference_tensor: torch.Tensor, batch: Req @@ -218,14 +231,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Run the distilled refinement schedule on top of the shared AV denoiser.""" batch.extra["ltx2_phase"] = "stage2" - pipeline = self.pipeline() if self.pipeline else None - ensure_phase_ready = ( - getattr(pipeline, "ensure_ltx2_phase_ready", None) - if pipeline is not None - else None - ) - if callable(ensure_phase_ready): - ensure_phase_ready("stage2") original_clean_latent_background = getattr( batch, "ltx2_ti2v_clean_latent_background", None ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 955fa7521..08957f1e0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -121,7 +121,7 @@ class DmdDenoisingStage(DenoisingStage): ) else: current_model = self.transformer - self._manage_device_placement(current_model, None, server_args) + self._manage_dit_use_site(current_model, "transformer", batch) # Expand latents for I2V noise_latents = latents.clone() latent_model_input = latents.to(target_dtype) @@ -232,45 +232,19 @@ class DmdDenoisingStage(DenoisingStage): if boundary_timestep is None or t_int >= boundary_timestep: # High-noise stage current_model = self.transformer - model_to_offload = self.transformer_2 current_guidance_scale = batch.guidance_scale + current_phase = "transformer" else: # Low-noise stage current_model = self.transformer_2 - model_to_offload = self.transformer current_guidance_scale = batch.guidance_scale_2 + current_phase = "transformer_2" - self._manage_device_placement(current_model, model_to_offload, server_args) + self._manage_dit_use_site(current_model, current_phase, batch) assert current_model is not None, "The model for the current step is not set." return current_model, current_guidance_scale - def _manage_device_placement( - self, - model_to_use: torch.nn.Module, - model_to_offload: torch.nn.Module | None, - server_args: ServerArgs, - ): - """ - Manages the offload / load behavior of dit - """ - if not server_args.dit_cpu_offload: - return - - # Offload the unused model if it's on CUDA - if ( - model_to_offload is not None - and next(model_to_offload.parameters()).device.type == "cuda" - ): - model_to_offload.to("cpu") - - # Load the model to use if it's on CPU - if ( - model_to_use is not None - and next(model_to_use.parameters()).device.type == "cpu" - ): - model_to_use.to(get_local_torch_device()) - def _handle_boundary_ratio( self, server_args, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py index 9ff56c787..8ac73ccb3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py @@ -8,6 +8,7 @@ Encoding stage for diffusion pipelines. import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage @@ -37,6 +38,19 @@ class EncodingStage(PipelineStage): super().__init__() self.vae: ParallelTiledVAE = vae + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name, + "vae", + target_dtype=vae_dtype, + ) + ] + @torch.no_grad() def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify encoding stage inputs.""" @@ -67,8 +81,6 @@ class EncodingStage(PipelineStage): """ assert batch.latents is not None and isinstance(batch.latents, torch.Tensor) - self.vae = self.vae.to(get_local_torch_device()) - # Setup VAE precision vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( @@ -81,27 +93,25 @@ class EncodingStage(PipelineStage): # Move to appropriate device and dtype latents = latents.to(get_local_torch_device()) - # Encode image to latents - with torch.autocast( - device_type=current_platform.device_type, - dtype=vae_dtype, - enabled=vae_autocast_enabled, - ): - if server_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - # if server_args.vae_sp: - # self.vae.enable_parallel() - if not vae_autocast_enabled: - latents = latents.to(vae_dtype) - latents = self.vae.encode(latents).mean + with self.use_declared_component(component_name="vae", module=self.vae) as vae: + assert vae is not None + self.vae = vae + + # Encode image to latents + with torch.autocast( + device_type=current_platform.device_type, + dtype=vae_dtype, + enabled=vae_autocast_enabled, + ): + if server_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + # if server_args.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + latents = latents.to(vae_dtype) + latents = self.vae.encode(latents).mean # Update batch with encoded latents batch.latents = latents - # Offload models if needed - self.maybe_free_model_hooks() - - if server_args.vae_cpu_offload: - self.vae.to("cpu") - return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 2791f7f2e..fe6187c0b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -22,6 +22,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( qwen_image_postprocess_text, ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.models.vision_utils import ( @@ -137,26 +138,16 @@ class ImageEncodingStage(PipelineStage): self.image_encoder = image_encoder self.text_encoder = text_encoder - def load_model(self): - if self.server_args.image_encoder_cpu_offload: - device = get_local_torch_device() - self.move_to_device(device) - - def offload_model(self): - if self.server_args.image_encoder_cpu_offload: - self.move_to_device("cpu") - - def move_to_device(self, device): - if self.server_args.use_fsdp_inference: - return - fields = [ - "image_processor", - "image_encoder", - ] - for field in fields: - processor = getattr(self, field, None) - if processor and hasattr(processor, "to"): - setattr(self, field, processor.to(device)) + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + uses = [] + if self.image_encoder is not None: + uses.append(ComponentUse(stage_name, "image_encoder")) + if self.text_encoder is not None: + uses.append(ComponentUse(stage_name, "text_encoder")) + return uses def encoding_qwen_image_edit(self, outputs, image_inputs): # encoder hidden state @@ -177,8 +168,6 @@ class ImageEncodingStage(PipelineStage): return batch cuda_device = get_local_torch_device() - self.load_model() - image_processor_kwargs = ( server_args.pipeline_config.prepare_image_processor_kwargs(batch) ) @@ -215,14 +204,20 @@ class ImageEncodingStage(PipelineStage): if self.image_encoder: # if an image encoder is provided - with set_forward_context(current_timestep=0, attn_metadata=None): - outputs = self.image_encoder( - **image_inputs, - **server_args.pipeline_config.image_encoder_extra_args, - ) - image_embeds = server_args.pipeline_config.postprocess_image( - outputs - ) + with self.use_declared_component( + component_name="image_encoder", + module=self.image_encoder, + ) as image_encoder: + assert image_encoder is not None + self.image_encoder = image_encoder + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = self.image_encoder( + **image_inputs, + **server_args.pipeline_config.image_encoder_extra_args, + ) + image_embeds = server_args.pipeline_config.postprocess_image( + outputs + ) batch.image_embeds.append(image_embeds) elif self.text_encoder: # if a text encoder is provided, e.g. Qwen-Image-Edit @@ -243,22 +238,28 @@ class ImageEncodingStage(PipelineStage): **neg_image_processor_kwargs, ).to(cuda_device) - with set_forward_context(current_timestep=0, attn_metadata=None): - outputs = self.text_encoder( - input_ids=image_inputs.input_ids, - attention_mask=image_inputs.attention_mask, - pixel_values=image_inputs.pixel_values, - image_grid_thw=image_inputs.image_grid_thw, - output_hidden_states=True, - ) - if batch.do_classifier_free_guidance: - neg_outputs = self.text_encoder( - input_ids=neg_image_inputs.input_ids, - attention_mask=neg_image_inputs.attention_mask, - pixel_values=neg_image_inputs.pixel_values, - image_grid_thw=neg_image_inputs.image_grid_thw, + with self.use_declared_component( + component_name="text_encoder", + module=self.text_encoder, + ) as text_encoder: + assert text_encoder is not None + self.text_encoder = text_encoder + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = self.text_encoder( + input_ids=image_inputs.input_ids, + attention_mask=image_inputs.attention_mask, + pixel_values=image_inputs.pixel_values, + image_grid_thw=image_inputs.image_grid_thw, output_hidden_states=True, ) + if batch.do_classifier_free_guidance: + neg_outputs = self.text_encoder( + input_ids=neg_image_inputs.input_ids, + attention_mask=neg_image_inputs.attention_mask, + pixel_values=neg_image_inputs.pixel_values, + image_grid_thw=neg_image_inputs.image_grid_thw, + output_hidden_states=True, + ) all_prompt_embeds.append( self.encoding_qwen_image_edit(outputs, image_inputs) @@ -273,8 +274,6 @@ class ImageEncodingStage(PipelineStage): if all_neg_prompt_embeds: batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0)) - self.offload_model() - return batch def build_dedup_fingerprint( @@ -328,20 +327,22 @@ class LTX2ImageEncodingStage(PipelineStage): self._condition_image_encoder = None self._condition_image_encoder_dir = None - # -- device management (mirrors ImageVAEEncodingStage) --------------- - - def load_model(self): - device = get_local_torch_device() - if self._condition_image_encoder is not None: - self._condition_image_encoder = self._condition_image_encoder.to(device) - else: - self.vae = self.vae.to(device) - - def offload_model(self): - if self.server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - if self._condition_image_encoder is not None: - self._condition_image_encoder = self._condition_image_encoder.to("cpu") + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + arch_config = server_args.pipeline_config.vae_config.arch_config + encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) + stage_name = self._component_stage_name(stage_name) + if encoder_subdir: + return [ + ComponentUse( + stage_name, + "condition_image_encoder", + ) + ] + if self.vae is None: + return [] + return [ComponentUse(stage_name, "vae")] # -- lazy condition encoder (LTX-2.3) -------------------------------- @@ -576,58 +577,61 @@ class LTX2ImageEncodingStage(PipelineStage): if len(batch.condition_image) == 1: batch.condition_image = batch.condition_image[0] - # 2. Load encoder(s) to device, cast to encode_dtype + # 2. Select encoder(s); residency manager moves it to device and dtype. use_condition_encoder = self._ensure_condition_image_encoder(server_args) - self.load_model() device = get_local_torch_device() encode_dtype = batch.latents.dtype - # Cast the active encoder to the latent precision (must match original - # behavior — running in a different dtype shifts the encoded latents). if use_condition_encoder: - self._condition_image_encoder = self._condition_image_encoder.to( - dtype=encode_dtype - ) + component_name = "condition_image_encoder" + encoder = self._condition_image_encoder else: - self.vae = self.vae.to(dtype=encode_dtype) + component_name = "vae" + encoder = self.vae packed_latents = [] - for conditioned_img in conditioned_imgs: - video_condition = self._pil_to_video_tensor( - conditioned_img, - width=int(batch.width), - height=int(batch.height), - device=device, - dtype=encode_dtype, - ) - - # 3. Encode + with self.use_declared_component( + component_name=component_name, + module=encoder, + target_dtype=encode_dtype, + ) as active_encoder: + assert active_encoder is not None if use_condition_encoder: - latent = self._condition_encode(video_condition, server_args).to( - dtype=encode_dtype - ) + self._condition_image_encoder = active_encoder else: - latent = self._vae_encode(video_condition, server_args, batch.generator) + self.vae = active_encoder - packed = server_args.pipeline_config.maybe_pack_latents( - latent, latent.shape[0], batch - ) - if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): - raise ValueError("Expected packed image latents [B, S0, D].") - if int(packed.shape[1]) != expected_tokens: - raise ValueError( - f"LTX-2 conditioning token count mismatch: " - f"{packed.shape[1]=} {expected_tokens=}." + for conditioned_img in conditioned_imgs: + video_condition = self._pil_to_video_tensor( + conditioned_img, + width=int(batch.width), + height=int(batch.height), + device=device, + dtype=encode_dtype, ) - packed_latents.append(packed) - # Restore VAE to its config dtype (shared with decoding stage) - if not use_condition_encoder: - original_dtype = PRECISION_TO_TYPE[ - server_args.pipeline_config.vae_precision - ] - self.vae = self.vae.to(dtype=original_dtype) + # 3. Encode + if use_condition_encoder: + latent = self._condition_encode(video_condition, server_args).to( + dtype=encode_dtype + ) + else: + latent = self._vae_encode( + video_condition, server_args, batch.generator + ) + + packed = server_args.pipeline_config.maybe_pack_latents( + latent, latent.shape[0], batch + ) + if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): + raise ValueError("Expected packed image latents [B, S0, D].") + if int(packed.shape[1]) != expected_tokens: + raise ValueError( + f"LTX-2 conditioning token count mismatch: " + f"{packed.shape[1]=} {expected_tokens=}." + ) + packed_latents.append(packed) batch.image_latent = ( packed_latents[0] if len(packed_latents) == 1 else packed_latents @@ -643,7 +647,6 @@ class LTX2ImageEncodingStage(PipelineStage): batch.height, ) - self.offload_model() return batch def build_dedup_fingerprint( @@ -688,12 +691,18 @@ class ImageVAEEncodingStage(PipelineStage): super().__init__() self.vae: ParallelTiledVAE = vae - def load_model(self): - self.vae = self.vae.to(get_local_torch_device()) - - def offload_model(self): - if self.server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name, + "vae", + target_dtype=vae_dtype, + ) + ] def forward( self, @@ -707,7 +716,6 @@ class ImageVAEEncodingStage(PipelineStage): if batch.condition_image is None: return batch - self.load_model() num_frames = batch.num_frames images = ( @@ -721,110 +729,115 @@ class ImageVAEEncodingStage(PipelineStage): server_args.pipeline_config, "prepare_condition_image_latent_ids", None ) condition_latents = [] if callable(prepare_condition_image_latent_ids) else None - for image in images: - image = self.preprocess( - image, - ).to(get_local_torch_device(), dtype=torch.float32) + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32 + ) and not server_args.disable_autocast - # (B, C, H, W) -> (B, C, 1, H, W) - image = image.unsqueeze(2) + with self.use_declared_component(component_name="vae", module=self.vae) as vae: + assert vae is not None + self.vae = vae - if num_frames == 1: - video_condition = image - else: - video_condition = torch.cat( - [ - image, - image.new_zeros( - image.shape[0], - image.shape[1], - num_frames - 1, - image.shape[3], - image.shape[4], - ), - ], - dim=2, + for image in images: + image = self.preprocess( + image, + ).to(get_local_torch_device(), dtype=torch.float32) + + # (B, C, H, W) -> (B, C, 1, H, W) + image = image.unsqueeze(2) + + if num_frames == 1: + video_condition = image + else: + video_condition = torch.cat( + [ + image, + image.new_zeros( + image.shape[0], + image.shape[1], + num_frames - 1, + image.shape[3], + image.shape[4], + ), + ], + dim=2, + ) + video_condition = video_condition.to( + device=get_local_torch_device(), dtype=torch.float32 ) - video_condition = video_condition.to( - device=get_local_torch_device(), dtype=torch.float32 - ) - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - vae_autocast_enabled = ( - vae_dtype != torch.float32 - ) and not server_args.disable_autocast + # Encode Image + with torch.autocast( + device_type=current_platform.device_type, + dtype=vae_dtype, + enabled=vae_autocast_enabled, + ): + if server_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + # if server_args.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + video_condition = video_condition.to(vae_dtype) + latent_dist: DiagonalGaussianDistribution = self.vae.encode( + video_condition + ) + # for auto_encoder from diffusers + if isinstance(latent_dist, AutoencoderKLOutput): + latent_dist = latent_dist.latent_dist - # Encode Image - with torch.autocast( - device_type=current_platform.device_type, - dtype=vae_dtype, - enabled=vae_autocast_enabled, - ): - if server_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - # if server_args.vae_sp: - # self.vae.enable_parallel() - if not vae_autocast_enabled: - video_condition = video_condition.to(vae_dtype) - latent_dist: DiagonalGaussianDistribution = self.vae.encode( - video_condition + generator = batch.generator + if generator is None: + raise ValueError("Generator must be provided") + + sample_mode = ( + server_args.pipeline_config.vae_config.encode_sample_mode() ) - # for auto_encoder from diffusers - if isinstance(latent_dist, AutoencoderKLOutput): - latent_dist = latent_dist.latent_dist - generator = batch.generator - if generator is None: - raise ValueError("Generator must be provided") - - sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode() - - latent_condition = self.retrieve_latents( - latent_dist, generator, sample_mode=sample_mode - ) - latent_condition = server_args.pipeline_config.postprocess_vae_encode( - latent_condition, self.vae - ) - normalized_latent_condition = ( - server_args.pipeline_config.normalize_vae_encode( + latent_condition = self.retrieve_latents( + latent_dist, generator, sample_mode=sample_mode + ) + latent_condition = server_args.pipeline_config.postprocess_vae_encode( latent_condition, self.vae ) - ) - if normalized_latent_condition is None: - scaling_factor, shift_factor = ( - server_args.pipeline_config.get_decode_scale_and_shift( - device=latent_condition.device, - dtype=latent_condition.dtype, - vae=self.vae, + normalized_latent_condition = ( + server_args.pipeline_config.normalize_vae_encode( + latent_condition, self.vae ) ) + if normalized_latent_condition is None: + scaling_factor, shift_factor = ( + server_args.pipeline_config.get_decode_scale_and_shift( + device=latent_condition.device, + dtype=latent_condition.dtype, + vae=self.vae, + ) + ) - # apply shift & scale if needed - if isinstance(shift_factor, torch.Tensor): - shift_factor = shift_factor.to(latent_condition.device) + # apply shift & scale if needed + if isinstance(shift_factor, torch.Tensor): + shift_factor = shift_factor.to(latent_condition.device) - if isinstance(scaling_factor, torch.Tensor): - scaling_factor = scaling_factor.to(latent_condition.device) + if isinstance(scaling_factor, torch.Tensor): + scaling_factor = scaling_factor.to(latent_condition.device) - latent_condition -= shift_factor - latent_condition = latent_condition * scaling_factor - else: - latent_condition = normalized_latent_condition + latent_condition -= shift_factor + latent_condition = latent_condition * scaling_factor + else: + latent_condition = normalized_latent_condition - if condition_latents is not None: - condition_latents.append(latent_condition) + if condition_latents is not None: + condition_latents.append(latent_condition) - image_latent = server_args.pipeline_config.postprocess_image_latent( - latent_condition, batch - ) - all_image_latents.append(image_latent) + image_latent = server_args.pipeline_config.postprocess_image_latent( + latent_condition, batch + ) + all_image_latents.append(image_latent) batch.image_latent = torch.cat(all_image_latents, dim=1) if condition_latents is not None: prepare_condition_image_latent_ids(condition_latents, batch) - self.offload_model() return batch def build_dedup_fingerprint( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index 91c0b2274..b20996e5c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -1160,22 +1160,12 @@ class LTX2DenoisingStage(DenoisingStage): self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs ) -> None: """Reset the mirrored audio scheduler before the shared loop begins.""" - if ctx.stage in ("stage1", "stage2"): + if is_ltx2_two_stage_pipeline_name( + server_args.pipeline_class_name + ) and ctx.stage in ("stage1", "stage2"): pipeline = self.pipeline() if self.pipeline else None - switch_lora_phase = ( - getattr(pipeline, "switch_lora_phase", None) - if pipeline is not None - else None - ) - if callable(switch_lora_phase): - switch_lora_phase(ctx.stage, batch=batch) - ensure_phase_ready = ( - getattr(pipeline, "ensure_ltx2_phase_ready", None) - if pipeline is not None - else None - ) - if callable(ensure_phase_ready): - ensure_phase_ready(ctx.stage) + if pipeline is not None: + pipeline.switch_lora_phase(ctx.stage, batch=batch) super()._before_denoising_loop(ctx, batch, server_args) if ctx.audio_scheduler is None: raise ValueError("LTX-2 audio scheduler was not prepared.") @@ -1194,7 +1184,6 @@ class LTX2DenoisingStage(DenoisingStage): ): """Preserve the legacy LTX-2 attention-metadata contract.""" # Legacy LTX-2 paths used the plain attention-metadata builder call here. - del ctx, t_int, timesteps_cpu return self._build_attn_metadata(step_index, batch, server_args) def _run_denoising_step( @@ -1996,7 +1985,6 @@ class LTX2DenoisingStage(DenoisingStage): def _get_prompt_embeds_validator(self, batch: Req): """Allow either tensor or list prompt embeddings for LTX-2 prompts.""" - del batch return lambda x: V.is_tensor(x) or V.list_not_empty(x) def _get_negative_prompt_embeds_validator(self, batch: Req): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py index 899a18937..5c2cebdab 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py @@ -14,6 +14,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -47,11 +48,20 @@ class HeliosDecodingStage(DecodingStage): # Load VAE if needed self.load_model() + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] # Decode each chunk separately and concatenate in pixel space video_chunks = [] - for chunk_latents in latent_chunks: - chunk_video = self.decode(chunk_latents, server_args) - video_chunks.append(chunk_video) + with self.use_declared_component( + component_name=self.component_name, + module=self.vae, + ) as vae: + assert vae is not None + self.vae = vae + for chunk_latents in latent_chunks: + chunk_video = self.decode( + chunk_latents, server_args, vae_dtype=vae_dtype + ) + video_chunks.append(chunk_video) frames = torch.cat(video_chunks, dim=2) frames = server_args.pipeline_config.post_decoding(frames, server_args) @@ -64,5 +74,4 @@ class HeliosDecodingStage(DecodingStage): metrics=batch.metrics, ) - self.offload_model() return output_batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py index e208f3643..84c70bbe0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -13,7 +13,7 @@ import numpy as np import torch import torch.nn.functional as F -from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import ( get_or_create_request_scheduler, @@ -107,6 +107,20 @@ class HeliosChunkedDenoisingStage(PipelineStage): def parallelism_type(self): return StageParallelismType.REPLICATED + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name=stage_name, + component_name="transformer", + phase="transformer", + preferred_ready_after_request=True, + memory_intensive=True, + ) + ] + def _denoise_one_chunk( self, latents, @@ -483,10 +497,15 @@ class HeliosChunkedDenoisingStage(PipelineStage): is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk gamma = pipeline_config.gamma - # Move transformer to GPU if CPU-offloaded - if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: - if next(self.transformer.parameters()).device.type == "cpu": - self.transformer.to(get_local_torch_device()) + transformer_use = ComponentUse( + self.__class__.__name__, + "transformer", + phase="transformer", + preferred_ready_after_request=True, + memory_intensive=True, + ) + manager = self._component_residency_manager + manager.begin_use(transformer_use, module=self.transformer) # Get encoder outputs (prompt_embeds is a list of tensors, one per encoder) prompt_embeds = batch.prompt_embeds @@ -718,10 +737,6 @@ class HeliosChunkedDenoisingStage(PipelineStage): history_latents = torch.cat([history_latents, latents], dim=2) chunk_latents_list.append(latents) - # Move transformer back to CPU after denoising - if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: - if next(self.transformer.parameters()).device.type != "cpu": - self.transformer.to("cpu") torch.cuda.empty_cache() # Store per-chunk latents for chunk-by-chunk VAE decode (matches diffusers behavior). diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py index a33242f08..9ab2d6f86 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py @@ -46,6 +46,7 @@ from sglang.multimodal_gen.runtime.models.dits.mova_video_dit import ( # Create aliases for backward compatibility video_sinusoidal_embedding_1d = sinusoidal_embedding_1d audio_sinusoidal_embedding_1d = sinusoidal_embedding_1d +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, @@ -62,7 +63,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args -from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler @@ -159,6 +159,32 @@ class MOVADenoisingStage(PipelineStage): self._cached_num_steps = None self._torch_compiled = False + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + uses = [ + ComponentUse(stage_name, "audio_dit"), + ComponentUse(stage_name, "dual_tower_bridge"), + ComponentUse( + stage_name, + "video_dit", + phase="video_dit", + preferred_ready_after_request=True, + memory_intensive=True, + ), + ] + if self.video_dit_2 is not None: + uses.append( + ComponentUse( + stage_name, + "video_dit_2", + phase="video_dit_2", + memory_intensive=True, + ) + ) + return uses + @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: @@ -329,27 +355,6 @@ class MOVADenoisingStage(PipelineStage): ) -> object | None: return None - def _manage_device_placement( - self, - model_to_use: nn.Module | None, - model_to_offload: nn.Module | None, - server_args: ServerArgs, - ): - if not server_args.dit_cpu_offload: - return - - if ( - model_to_offload is not None - and next(model_to_offload.parameters()).device.type == "cuda" - ): - model_to_offload.to("cpu") - - if ( - model_to_use is not None - and next(model_to_use.parameters()).device.type == "cpu" - ): - model_to_use.to(get_local_torch_device()) - def _select_visual_dit( self, timestep: float, @@ -358,24 +363,52 @@ class MOVADenoisingStage(PipelineStage): scheduler, ): if boundary_ratio is None or self.video_dit_2 is None: - self._manage_device_placement(self.video_dit, None, server_args) + self._manage_video_dit_use(self.video_dit, "video_dit") return self.video_dit boundary_timestep = boundary_ratio * scheduler.num_train_timesteps if timestep >= boundary_timestep: current_model = self.video_dit - model_to_offload = self.video_dit_2 + current_name = "video_dit" else: current_model = self.video_dit_2 - model_to_offload = self.video_dit + current_name = "video_dit_2" - self._manage_device_placement(current_model, model_to_offload, server_args) + self._manage_video_dit_use(current_model, current_name) return current_model + def _manage_video_dit_use( + self, current_model: nn.Module, default_name: str + ) -> bool: + manager = self._component_residency_manager + if manager is None: + return False + + component_name = manager.component_name_for_module(current_model, default_name) + use = ComponentUse( + stage_name=self._active_component_stage_name(), + component_name=component_name, + phase=component_name, + preferred_ready_after_request=component_name == "video_dit", + memory_intensive=True, + ) + manager.begin_use(use, module=current_model) + return True + def _ensure_shared_models_on_device(self, server_args: ServerArgs): """Ensure shared denoising modules are on the active device when cpu offload is enabled.""" - self._manage_device_placement(self.audio_dit, None, server_args) - self._manage_device_placement(self.dual_tower_bridge, None, server_args) + manager = self._component_residency_manager + if manager is None: + return + stage_name = self._active_component_stage_name() + manager.ensure_ready( + ComponentUse(stage_name, "audio_dit"), + module=self.audio_dit, + ) + manager.ensure_ready( + ComponentUse(stage_name, "dual_tower_bridge"), + module=self.dual_tower_bridge, + ) def _apply_guidance_rescale( self, @@ -598,9 +631,7 @@ class MOVADenoisingStage(PipelineStage): if not is_warmup and hasattr(self, "step_profile"): self.step_profile() - for dit in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]): - if isinstance(dit, OffloadableDiTMixin): - dit.prepare_for_next_req() + self._finish_active_component_use() return batch @@ -911,6 +942,16 @@ class MOVADecodingStage(PipelineStage): self.video_vae = video_vae self.audio_vae = audio_vae + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + return [ + ComponentUse(stage_name, "video_vae", target_dtype=vae_dtype), + ComponentUse(stage_name, "audio_vae"), + ] + @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: @@ -919,36 +960,45 @@ class MOVADecodingStage(PipelineStage): @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: - self.video_vae = self.video_vae.to(get_local_torch_device()) - self.audio_vae = self.audio_vae.to(get_local_torch_device()) - - video_latents = server_args.pipeline_config.denormalize_video_latents( - batch.latents, self.video_vae - ) - vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast - with torch.autocast( - device_type=current_platform.device_type, - dtype=vae_dtype, - enabled=vae_autocast_enabled, - ): - if server_args.pipeline_config.vae_tiling: - self.video_vae.enable_tiling() - if not vae_autocast_enabled: - video_latents = video_latents.to(vae_dtype) - decode_output = self.video_vae.decode(video_latents) - video = _ensure_tensor_decode_output(decode_output) + with self.use_declared_component( + component_name="video_vae", + module=self.video_vae, + ) as video_vae: + assert video_vae is not None + self.video_vae = video_vae + video_latents = server_args.pipeline_config.denormalize_video_latents( + batch.latents, self.video_vae + ) + + with torch.autocast( + device_type=current_platform.device_type, + dtype=vae_dtype, + enabled=vae_autocast_enabled, + ): + if server_args.pipeline_config.vae_tiling: + self.video_vae.enable_tiling() + if not vae_autocast_enabled: + video_latents = video_latents.to(vae_dtype) + decode_output = self.video_vae.decode(video_latents) + video = _ensure_tensor_decode_output(decode_output) video = (video / 2 + 0.5).clamp(0, 1) - with torch.autocast( - device_type=current_platform.device_type, dtype=torch.float32 - ): - audio = self.audio_vae.decode(batch.audio_latents) + with self.use_declared_component( + component_name="audio_vae", + module=self.audio_vae, + ) as audio_vae: + assert audio_vae is not None + self.audio_vae = audio_vae + with torch.autocast( + device_type=current_platform.device_type, dtype=torch.float32 + ): + audio = self.audio_vae.decode(batch.audio_latents) output_batch = OutputBatch( output=video, audio=audio, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py index 57c89359d..5bd25f31e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py @@ -8,6 +8,7 @@ from diffusers.image_processor import VaeImageProcessor from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vision_utils import load_image from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -117,13 +118,9 @@ class QwenImageLayeredBeforeDenoisingStage(PipelineStage): self.vae = vae.to(torch.bfloat16) from transformers import Qwen2_5_VLForConditionalGeneration - self.text_encoder = ( - Qwen2_5_VLForConditionalGeneration.from_pretrained( - model_path, subfolder="text_encoder" - ) - .to(get_local_torch_device()) - .to(torch.bfloat16) - ) + self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + model_path, subfolder="text_encoder" + ).to(torch.bfloat16) self.tokenizer = tokenizer self.processor = processor self.transformer = transformer @@ -158,6 +155,17 @@ generalizations\n - Describe all visible information in the image, while do not the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n""" self.default_sample_size = 128 + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name, "qwen_layered_text_encoder", target_dtype=torch.bfloat16 + ), + ComponentUse(stage_name, "vae", target_dtype=torch.bfloat16), + ] + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -300,21 +308,23 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - self.vae = self.vae.to(get_local_torch_device()) - if isinstance(generator, list): - image_latents = [ - retrieve_latents( - self.vae.encode(image[i : i + 1]), - generator=generator[i], - sample_mode="argmax", + with self.use_declared_component(component_name="vae", module=self.vae) as vae: + assert vae is not None + self.vae = vae + if isinstance(generator, list): + image_latents = [ + retrieve_latents( + self.vae.encode(image[i : i + 1]), + generator=generator[i], + sample_mode="argmax", + ) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents( + self.vae.encode(image), generator=generator, sample_mode="argmax" ) - for i in range(image.shape[0]) - ] - image_latents = torch.cat(image_latents, dim=0) - else: - image_latents = retrieve_latents( - self.vae.encode(image), generator=generator, sample_mode="argmax" - ) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.latent_channels, 1, 1, 1) @@ -326,7 +336,6 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as .to(image_latents.device, image_latents.dtype) ) image_latents = (image_latents - latents_mean) / latents_std - self.vae.to("cpu") return image_latents def prepare_latents( @@ -447,20 +456,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as image = image.to(dtype=torch.bfloat16) prompt = batch.prompt - if not prompt or prompt.isspace(): - prompt = self.get_image_caption( - prompt_image, use_en_prompt=use_en_prompt, device=device + with self.use_declared_component( + component_name="qwen_layered_text_encoder", + module=self.text_encoder, + ) as text_encoder: + assert text_encoder is not None + self.text_encoder = text_encoder + if not prompt or prompt.isspace(): + prompt = self.get_image_caption( + prompt_image, use_en_prompt=use_en_prompt, device=device + ) + + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + device=device, ) - prompt_embeds, prompt_embeds_mask = self.encode_prompt( - prompt=prompt, - device=device, - ) - - negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( - prompt=batch.negative_prompt, - device=device, - ) + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + prompt=batch.negative_prompt, + device=device, + ) num_channels_latents = self.transformer.config.in_channels // 4 latents, image_latents = self.prepare_latents( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py index b26d00fa6..709eeaf47 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py @@ -33,6 +33,7 @@ def prepare_wan_ti2v_latents( vae: object, latents: torch.Tensor, target_dtype: torch.dtype, + vae_dtype: torch.dtype, batch: Req, server_args: ServerArgs, ) -> tuple[int, torch.Tensor, list[torch.Tensor]]: @@ -43,10 +44,10 @@ def prepare_wan_ti2v_latents( assert batch.image_latent is None, "TI2V task should not have image latents" assert vae is not None, "VAE is not provided for TI2V task" - vae = vae.to(batch.condition_image.device) - z = vae.encode(batch.condition_image).mean.float() - if getattr(vae, "device", None) != "cpu" and server_args.vae_cpu_offload: - vae = vae.to("cpu") + condition_image = batch.condition_image.to( + device=get_local_torch_device(), dtype=vae_dtype + ) + z = vae.encode(condition_image).mean.float() if hasattr(vae, "shift_factor") and vae.shift_factor is not None: if isinstance(vae.shift_factor, torch.Tensor): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index 565d95243..a8cb97ce6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -15,6 +15,7 @@ import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage @@ -68,6 +69,19 @@ class TextEncodingStage(PipelineStage): self.tokenizers = tokenizers self.text_encoders = text_encoders + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name=stage_name, + component_name="text_encoder" if i == 0 else f"text_encoder_{i + 1}", + preferred_ready_after_request=i == 0, + ) + for i in range(len(self.text_encoders)) + ] + @torch.no_grad() def forward( self, @@ -166,6 +180,21 @@ class TextEncodingStage(PipelineStage): return tok_kwargs + def _manage_text_encoder_use(self, encoder_index: int) -> None: + manager = self._component_residency_manager + if manager is None: + return + component_name = ( + "text_encoder" + if encoder_index == 0 + else f"text_encoder_{encoder_index + 1}" + ) + use = self._declared_component_use(component_name=component_name) + # TODO: Keep this begin-only interval until manager supports explicit + # declared-use interval grouping. Wrapping each encoder call separately + # can offload between positive and negative prompt encoding. + manager.before_use(use) + def _forward_text_encoder(self, text_encoder, encoder_forward_kwargs): if not getattr(text_encoder, "uses_sglang_forward_context", True): return text_encoder(**encoder_forward_kwargs) @@ -251,7 +280,7 @@ class TextEncodingStage(PipelineStage): embeds_list: list[torch.Tensor] = [] pooled_embeds_list: list[torch.Tensor] = [] - attn_masks_list: list[torch.Tensor] = [] + attn_masks_list: list[torch.Tensor | None] = [] preprocess_funcs = server_args.pipeline_config.preprocess_text_funcs postprocess_funcs = server_args.pipeline_config.postprocess_text_funcs @@ -308,6 +337,7 @@ class TextEncodingStage(PipelineStage): encoder_forward_kwargs["attention_mask"] = attention_mask if "use_cache" in inspect.signature(text_encoder.forward).parameters: encoder_forward_kwargs["use_cache"] = False + self._manage_text_encoder_use(i) outputs: BaseEncoderOutput = self._forward_text_encoder( text_encoder, encoder_forward_kwargs ) @@ -345,11 +375,19 @@ class TextEncodingStage(PipelineStage): if postprocessed_attention_mask is not None else None ) - elif attention_mask is not None: + elif attention_mask is not None and list(attention_mask.shape) == list( + prompt_embeds.shape[:2] + ): mask_to_store = attention_mask.to(device=target_device) else: mask_to_store = torch.ones( - input_ids.shape[:2], device=target_device + prompt_embeds.shape[:2], + device=target_device, + dtype=( + attention_mask.dtype + if attention_mask is not None + else torch.long + ), ) attn_masks_list.append(mask_to_store) @@ -379,13 +417,23 @@ class TextEncodingStage(PipelineStage): ) stacked_embeds = torch.stack(embeds_list, dim=0) if return_attention_mask: - base_mask_shape = list(attn_masks_list[0].shape) - for m in attn_masks_list[1:]: + stackable_masks = [ + ( + mask + if mask is not None + else torch.ones( + embed.shape[:2], device=embed.device, dtype=torch.long + ) + ) + for embed, mask in zip(embeds_list, attn_masks_list, strict=True) + ] + base_mask_shape = list(stackable_masks[0].shape) + for m in stackable_masks[1:]: if list(m.shape) != base_mask_shape: raise ValueError( - f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}" + f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in stackable_masks]}" ) - stacked_masks = torch.stack(attn_masks_list, dim=0) + stacked_masks = torch.stack(stackable_masks, dim=0) return stacked_embeds, stacked_masks return stacked_embeds diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py index d8f5fb24c..b88a16e06 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py @@ -1,6 +1,7 @@ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -44,18 +45,10 @@ class LTX2LoRASwitchStage(PipelineStage): self.phase = phase def forward(self, batch: Req, server_args: ServerArgs) -> Req: - switch_fn = getattr(self.pipeline, "switch_lora_phase", None) - should_skip_switch_stage = getattr( - self.pipeline, "should_skip_ltx2_lora_switch_stage", None - ) - if callable(should_skip_switch_stage) and should_skip_switch_stage(): + if self.pipeline.should_skip_ltx2_lora_switch_stage(): batch.extra["ltx2_phase"] = self.phase return batch - if not callable(switch_fn): - raise ValueError( - "LTX2LoRASwitchStage requires pipeline.switch_lora_phase()" - ) - switch_fn(self.phase, batch=batch) + self.pipeline.switch_lora_phase(self.phase, batch=batch) batch.extra["ltx2_phase"] = self.phase return batch @@ -63,13 +56,31 @@ class LTX2LoRASwitchStage(PipelineStage): class LTX2UpsampleStage(PipelineStage): """Upsample Stage-1 video latents and prepare Stage-2 inputs.""" - def __init__(self, spatial_upsampler, vae, audio_vae=None, pipeline=None): + def __init__( + self, + spatial_upsampler, + vae, + audio_vae=None, + pipeline=None, + ): super().__init__() self.spatial_upsampler = spatial_upsampler self.vae = vae self.audio_vae = audio_vae self.pipeline = pipeline + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + uses = [ + ComponentUse(stage_name, "spatial_upsampler"), + ComponentUse(stage_name, "vae"), + ] + if self.audio_vae is not None: + uses.append(ComponentUse(stage_name, "audio_vae")) + return uses + def _upsample_video_latents( self, latents: torch.Tensor, server_args: ServerArgs, device: torch.device ) -> torch.Tensor: @@ -118,25 +129,8 @@ class LTX2UpsampleStage(PipelineStage): ) def forward(self, batch: Req, server_args: ServerArgs) -> Req: - delay_stage2_prefetch = False - if self.pipeline is not None: - prepare_upsample = getattr( - self.pipeline, "prepare_ltx2_upsample_after_stage1", None - ) - if callable(prepare_upsample): - delay_stage2_prefetch = prepare_upsample() - prefetch_stage2 = ( - getattr(self.pipeline, "prefetch_ltx2_stage2_after_stage1", None) - if self.pipeline is not None - else None - ) - if callable(prefetch_stage2) and not delay_stage2_prefetch: - prefetch_stage2() - device = get_local_torch_device() latents = self._upsample_video_latents(batch.latents, server_args, device) - if callable(prefetch_stage2) and delay_stage2_prefetch: - prefetch_stage2() logger.info("Upsampled video latents: %s", list(latents.shape)) self._restore_full_resolution(batch) batch.image_latent = None diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 73b256683..c8a98fd0a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -9,10 +9,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( _ModelOptFp8OffloadAdapter, ) -from sglang.multimodal_gen.runtime.utils import ( +from sglang.multimodal_gen.runtime.managers import ( layerwise_offload as layerwise_offload_mod, ) -from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( +from sglang.multimodal_gen.runtime.managers.layerwise_offload import ( LayerwiseOffloadManager, )