minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-05-01 11:10:41 +08:00
committed by GitHub
parent 108bfd8b6a
commit 9d84268705
46 changed files with 2477 additions and 973 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/weights_updater.py
+1 -1
View File
@@ -53,9 +53,9 @@ from sglang.multimodal_gen.runtime.loader.utils import (
from sglang.multimodal_gen.runtime.loader.weight_utils import (
safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/managers/component_manager.py
+673
View File
@@ -0,0 +1,673 @@
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from functools import lru_cache
from typing import Mapping, MutableMapping, Protocol, Sequence, TypeVar
import torch
import torch.nn as nn
from sglang.multimodal_gen.runtime.managers.component_resident_strategies import (
ComponentResidencyStrategy,
LayerwiseOffloadStrategy,
ResidentStrategy,
VanillaD2HStrategy,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
_T = TypeVar("_T")
@dataclass(slots=True)
class ComponentUse:
"""Describes one stage/use-site access to a pipeline component."""
stage_name: str
# Pipeline module key: transformer / video_dit / text_encoder / ...
component_name: str
# Model-specific phase for sequential components, e.g. stage1 or stage2.
# TODO: Replace this with ordered timeline identity. In an all-sequential
# pipeline, use-site identity should come from the declared ComponentUse
# order instead of a per-use `phase` field.
phase: str | None = None
# Whether the manager may prepare this component for the next request.
preferred_ready_after_request: bool = False
# Whether cross-stage prefetch may prepare this use before the use-site.
allow_prefetch: bool = True
# Whether this use is expensive enough that earlier timeline prefetch matters.
# TODO: Replace this boolean hint with a budget-aware lookahead planner:
# estimate memory/load cost and reuse distance, keep small and early-request
# components resident within budget, prefetch as soon as VRAM slack appears,
# and release completed components only when the budget requires it.
memory_intensive: bool = False
# Optional module dtype required by this use-site.
target_dtype: torch.dtype | None = None
# Some components are intentionally kept ready between warmup and the first
# real request to avoid measuring a cold H2D in the user-visible request.
keep_ready_after_warmup: bool = False
@dataclass(slots=True)
class ResidencyState:
"""
Necessary internal runtime info of ComponentResidencyManager
"""
stages: Sequence["ComponentResidencyStage"] = ()
stage_index: int = -1
stage_name: str | None = None
next_stage_name: str | None = None
current_use: ComponentUse | None = None
# the ComponentUses of the preceding stages
future_uses: tuple[ComponentUse, ...] = ()
batch_is_warmup: bool = False
manager_mode: str = "static"
trace_enabled: bool = False
class ResidencyBatch(Protocol):
is_warmup: bool
class ComponentResidencyStage(Protocol):
def component_uses(
self, server_args: ServerArgs, stage_name: str | None = None
) -> list[ComponentUse]: ...
class ComponentResidencyPipeline(Protocol):
modules: Mapping[str, object]
_stage_name_mapping: Mapping[str, ComponentResidencyStage]
component_residency_strategies: MutableMapping[str, "ComponentResidencyStrategy"]
def build_dit_residency_strategy(
module: nn.Module,
server_args: ServerArgs,
) -> ComponentResidencyStrategy:
if (
isinstance(module, OffloadableDiTMixin)
and module.layerwise_offload_managers
and any(manager.enabled for manager in module.layerwise_offload_managers)
):
# only if dit_layerwise_offload is enabled
return LayerwiseOffloadStrategy()
if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
# handles offload by vanalla D2H
return VanillaD2HStrategy()
return ResidentStrategy()
def is_fsdp_managed_module(module: nn.Module) -> bool:
return module.__class__.__name__.startswith("FSDP")
def build_component_residency_strategy(
component_name: str,
module: nn.Module,
server_args: ServerArgs,
) -> ComponentResidencyStrategy:
if component_name in {
"transformer",
"transformer_2",
"video_dit",
"video_dit_2",
"audio_dit",
"dual_tower_bridge",
}:
return build_dit_residency_strategy(module, server_args)
if component_name.startswith("text_encoder") or component_name.endswith(
"text_encoder"
):
if (
server_args.text_encoder_cpu_offload
and not server_args.use_fsdp_inference
and not is_fsdp_managed_module(module)
):
return VanillaD2HStrategy()
return ResidentStrategy()
if component_name == "image_encoder":
if server_args.image_encoder_cpu_offload and not server_args.use_fsdp_inference:
return VanillaD2HStrategy()
return ResidentStrategy()
if component_name in {
"vae",
"video_vae",
"audio_vae",
"vocoder",
"spatial_upsampler",
"condition_image_encoder",
}:
if server_args.vae_cpu_offload and not server_args.use_fsdp_inference:
return VanillaD2HStrategy()
return ResidentStrategy()
return ResidentStrategy()
class ComponentResidencyManager:
"""Executor-owned component lifecycle coordinator. Provide hooks for a PipelineExecutor
Hooks are called around executor progress:
before request: collect a flat ordered ComponentUse timeline.
before stage: update current/next stage context only.
begin use: finish previous active use, prepare current use, wait until ready.
end use: finish or keep current use, then prefetch the next heavy timeline use.
finish request: finish active use and schedule preferred next-request prefetch.
The manager instance is global and rebound to the active pipeline before request execution.
This manager is designed only for sequential execution order for now
"""
def __init__(
self, pipeline: ComponentResidencyPipeline, server_args: ServerArgs
) -> None:
self.pipeline = pipeline
self.server_args = server_args
self.state = ResidencyState(trace_enabled=False)
self._stage_names_by_id: dict[int, str] = {}
self._stage_uses_by_index: list[tuple[ComponentUse, ...]] = []
self._ordered_uses: tuple[ComponentUse, ...] = ()
self._current_use_index: int = -1
self._active_use: ComponentUse | None = None
self._active_use_module: nn.Module | None = None
self._prefetched_use_keys: set[tuple[str, str, str | None]] = set()
self._custom_strategies: dict[str, ComponentResidencyStrategy] = dict(
pipeline.component_residency_strategies
)
self._uses_seen: dict[str, ComponentUse] = {}
@property
def enabled(self) -> bool:
return True
def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None:
custom_strategies = dict(pipeline.component_residency_strategies)
if pipeline is not self.pipeline:
self.strategy_for.cache_clear()
self._should_keep_single_dit.cache_clear()
self._active_use = None
self._active_use_module = None
self._uses_seen.clear()
self._prefetched_use_keys.clear()
elif custom_strategies != self._custom_strategies:
self.strategy_for.cache_clear()
self.pipeline = pipeline
self._custom_strategies = custom_strategies
self._stage_names_by_id = {
id(stage): name for name, stage in pipeline._stage_name_mapping.items()
}
def refresh_server_args(self, server_args: ServerArgs) -> None:
if server_args is not self.server_args:
self.strategy_for.cache_clear()
self.server_args = server_args
def register_strategy(
self, component_name: str, strategy: ComponentResidencyStrategy
) -> None:
self.pipeline.component_residency_strategies[component_name] = strategy
self._custom_strategies[component_name] = strategy
self.strategy_for.cache_clear()
def begin_request(
self,
stages: Sequence[ComponentResidencyStage],
batch: ResidencyBatch,
server_args: ServerArgs,
) -> None:
"""A hook called before processing an actual request"""
self.refresh_server_args(server_args)
self.state = ResidencyState(
stages=stages, batch_is_warmup=batch.is_warmup, trace_enabled=False
)
self._active_use = None
self._active_use_module = None
self._current_use_index = -1
self._prefetched_use_keys.clear()
self._uses_seen.clear()
if self.enabled:
self._stage_uses_by_index = [
tuple(stage.component_uses(server_args, self.stage_name(stage)))
for stage in stages
]
self._ordered_uses = tuple(
use for uses in self._stage_uses_by_index for use in uses
)
else:
self._stage_uses_by_index = []
self._ordered_uses = ()
self._trace(
"request_start",
detail=f"stages={len(stages)} uses={len(self._ordered_uses)}",
)
def before_stage(
self,
stage: ComponentResidencyStage,
stage_index: int,
batch: ResidencyBatch,
server_args: ServerArgs,
) -> None:
"""called after stage starts"""
if not self.enabled:
return
# update state before entering the stage
self.state.stage_index = stage_index
self.state.stage_name = self.stage_name(stage)
self.state.next_stage_name = self._next_stage_name(stage_index)
self._trace("stage_enter", detail=f"index={stage_index}")
def after_stage(self, stage_index: int) -> None:
"""called after stage exits"""
if not self.enabled:
return
self._trace("stage_exit", detail=f"index={stage_index}")
def before_use(self, use: ComponentUse) -> None:
"""component use-site starts"""
if not self.enabled:
return
self.begin_use(use)
def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
"""Begin one sequential component use interval. this is idempotent
1. Finish the previous active use if this is a different timeline use.
2. Prepare the current component.
3. Wait until the current component is ready, then prefetch the next heavy use.
"""
if self._active_use is not None and self._same_use(self._active_use, use):
return
if self._active_use is not None:
# finish previous active use
self._finish_use(
self._active_use,
module=self._active_use_module,
keep_on_warmup=self._active_use.keep_ready_after_warmup,
)
self._active_use = None
self._active_use_module = None
self.state.current_use = None
self._mark_current_use(use)
self._prepare_forward_use(use, module=module)
self._active_use = use
self._active_use_module = module
self._prefetch_next_memory_intensive_use()
def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
"""End one sequential component use interval.
1. Finish or keep the current component.
2. Clear it as the active use.
3. Prefetch the next memory-intensive use without waiting.
"""
if self._active_use is None or not self._same_use(self._active_use, use):
return
self._finish_use(
self._active_use,
module=self._active_use_module or module,
keep_on_warmup=self._active_use.keep_ready_after_warmup,
)
self._active_use = None
self._active_use_module = None
self.state.current_use = None
self._prefetch_next_memory_intensive_use()
@contextmanager
def use_component(
self, use: ComponentUse, module: nn.Module | None = None
) -> Iterator[nn.Module | None]:
self.begin_use(use, module=module)
try:
yield module if module is not None else self.get_module(use.component_name)
finally:
self.end_use(use, module=module)
def call_component(
self,
use: ComponentUse,
module: Callable[..., _T],
*args,
**kwargs,
) -> _T:
with self.use_component(use):
return module(*args, **kwargs)
def prefetch_use(self, use: ComponentUse) -> None:
"""Prepare a future use without blocking the current use."""
if not self.enabled:
return
self._prefetch_use(use)
def ensure_ready(self, use: ComponentUse, module: nn.Module | None = None) -> None:
"""Prepare a shared component and wait without making it the active use."""
if not self.enabled:
return
self._prepare_forward_use(use, module=module)
def prefetch_checkpoint(self, anchor: ComponentUse | None = None) -> None:
"""Give the manager a timeline overlap point.
1. Locate the anchor or current use in the ordered timeline.
2. Find the next prefetchable memory-intensive use.
3. Prepare it opportunistically without waiting.
"""
if not self.enabled:
return
if anchor is not None:
self._mark_current_use(anchor)
self._prefetch_next_memory_intensive_use()
def finish_active_use(self, *, prefetch_next: bool = True) -> None:
"""Finish the currently active sequential use, if any."""
if self._active_use is None:
return
active_use = self._active_use
self._finish_use(
active_use,
module=self._active_use_module,
keep_on_warmup=active_use.keep_ready_after_warmup,
)
self._active_use = None
self._active_use_module = None
self.state.current_use = None
if prefetch_next:
self._prefetch_next_memory_intensive_use()
def _prepare_forward_use(
self, use: ComponentUse, module: nn.Module | None = None
) -> None:
"""Prepare a component that is about to run and wait until it is ready."""
module = module or self.get_module(use.component_name)
if module is None:
self._trace("skip_missing", use)
return
strategy = self.strategy_for(use.component_name, module)
self._uses_seen[use.component_name] = use
self.state.current_use = use
self._trace("prepare", use, strategy, module)
strategy.prepare_for_use(module, use, self.state)
self._trace("wait", use, strategy, module)
strategy.wait_for_use(module, use, self.state)
def _prefetch_use(self, use: ComponentUse) -> None:
"""Prepare a future component opportunistically without waiting.
This is called when the component is memory-intensive so it may takes a long time to prefetch.
manager will perform the prefetch at some checkpoints, if necessary
"""
if not use.allow_prefetch:
return
module = self.get_module(use.component_name)
if module is None:
self._trace("skip_missing", use)
return
strategy = self.strategy_for(use.component_name, module)
if isinstance(strategy, VanillaD2HStrategy) and self._active_use is not None:
# Avoid making two vanilla-offloaded heavy components resident before
# a budget-aware planner can prove the overlap is safe.
self._trace("prefetch_skip_active_vanilla", use, strategy, module)
return
self._uses_seen[use.component_name] = use
self._trace("prefetch", use, strategy, module)
if strategy.prefetch_for_use(module, use, self.state):
self._prefetched_use_keys.add(self._use_key(use))
def after_use(self, use: ComponentUse) -> None:
if not self.enabled:
return
self.end_use(use)
def _finish_use(
self,
use: ComponentUse,
*,
module: nn.Module | None = None,
keep_on_warmup: bool,
) -> None:
"""finish a specific use by keeping them resident or call finish_use hook"""
module = module or self.get_module(use.component_name)
if module is None:
self._trace("skip_missing", use)
return
should_keep = (
keep_on_warmup and self.state.batch_is_warmup
) or self._should_keep_after_use(use)
if should_keep:
self._trace(
"keep",
use,
self.strategy_for(use.component_name, module),
module,
)
return
strategy = self.strategy_for(use.component_name, module)
self._trace("finish", use, strategy, module)
strategy.finish_use(module, use, self.state)
def finish_request(self) -> None:
if not self.enabled and not self._uses_seen and self._active_use is None:
return
# 1. Close the currently active sequential use.
self.finish_active_use(prefetch_next=False)
# 2. Pick components that should be ready for the next request.
preferred_uses = self._preferred_request_end_uses()
# 3. Finish everything else, or prepare preferred uses for request tail.
for component_name, use in list(self._uses_seen.items()):
module = self.get_module(component_name)
if module is None:
continue
if self.state.batch_is_warmup and use.keep_ready_after_warmup:
self._trace(
"request_keep_warmup",
use,
self.strategy_for(component_name, module),
module,
)
continue
preferred = component_name in preferred_uses
if not preferred and self._should_keep_single_dit(component_name):
self._trace(
"keep",
use,
self.strategy_for(component_name, module),
module,
detail="single_dit",
)
continue
strategy = self.strategy_for(component_name, module)
if preferred and not self.state.batch_is_warmup:
self._trace("request_prefetch", use, strategy, module)
strategy.prepare_after_request(module, use, self.state)
else:
action = "request_resident" if preferred else "request_finish"
self._trace(action, use, strategy, module)
strategy.finish_request(module, use, self.state, preferred=preferred)
self._trace("request_end")
def stage_name(self, stage: ComponentResidencyStage) -> str:
return self._stage_names_by_id.get(id(stage), stage.__class__.__name__)
def component_name_for_module(self, module: nn.Module | None, default: str) -> str:
if module is None:
return default
for name, candidate in self.pipeline.modules.items():
if candidate is module:
return name
return default
def get_module(self, component_name: str) -> nn.Module | None:
module = self.pipeline.modules.get(component_name)
return module if isinstance(module, nn.Module) else None
@lru_cache(maxsize=None)
def strategy_for(
self, component_name: str, module: nn.Module
) -> ComponentResidencyStrategy:
"""Return the pre-registered strategy for a specific component"""
custom_strategy = self._custom_strategies.get(component_name)
if custom_strategy is not None:
return custom_strategy
return build_component_residency_strategy(
component_name, module, self.server_args
)
def _stage_uses(self, stage_index: int) -> tuple[ComponentUse, ...]:
"""Returns the ComponentUse(s) of a specific stage"""
if stage_index < 0 or stage_index >= len(self._stage_uses_by_index):
return ()
return self._stage_uses_by_index[stage_index]
def _next_stage_name(self, stage_index: int) -> str | None:
next_index = stage_index + 1
if next_index < 0 or next_index >= len(self.state.stages):
return None
return self.stage_name(self.state.stages[next_index])
def _mark_current_use(self, use: ComponentUse) -> None:
index = self._locate_use_index(use)
if index is None:
self._current_use_index = len(self._ordered_uses)
self.state.future_uses = ()
return
self._current_use_index = index
self.state.future_uses = self._ordered_uses[index + 1 :]
def _locate_use_index(self, use: ComponentUse) -> int | None:
for index in range(self._current_use_index + 1, len(self._ordered_uses)):
if self._same_use(self._ordered_uses[index], use):
return index
for index, candidate in enumerate(self._ordered_uses):
if self._same_use(candidate, use):
return index
return None
def _prefetch_next_memory_intensive_use(self) -> None:
for use in self._ordered_uses[self._current_use_index + 1 :]:
if not use.memory_intensive:
continue
if self._use_key(use) in self._prefetched_use_keys:
return
self.prefetch_use(use)
return
def _should_keep_after_use(self, use: ComponentUse) -> bool:
future_component_names = {
future.component_name for future in self.state.future_uses
}
if use.component_name in future_component_names:
return True
if self._should_keep_single_dit(use.component_name):
return True
return False
@lru_cache(maxsize=None)
def _should_keep_single_dit(self, component_name: str) -> bool:
modules = self.pipeline.modules
return (component_name == "transformer" and "transformer_2" not in modules) or (
component_name == "video_dit" and "video_dit_2" not in modules
)
def _preferred_request_end_use(self) -> ComponentUse | None:
"""Returns a ComponentUse preferred to be resident after a request finishes, to prepare for next request"""
for uses in self._stage_uses_by_index:
for use in uses:
if use.preferred_ready_after_request:
return use
for uses in self._stage_uses_by_index:
if uses:
return uses[0]
return None
def _preferred_request_end_uses(self) -> dict[str, ComponentUse]:
preferred_uses: dict[str, ComponentUse] = {}
for uses in self._stage_uses_by_index:
for use in uses:
if use.preferred_ready_after_request:
preferred_uses[use.component_name] = use
for use in self._uses_seen.values():
if use.preferred_ready_after_request:
preferred_uses[use.component_name] = use
if preferred_uses:
return preferred_uses
preferred_use = self._preferred_request_end_use()
if preferred_use is None:
return {}
return {preferred_use.component_name: preferred_use}
@staticmethod
def _same_use(lhs: ComponentUse, rhs: ComponentUse) -> bool:
return lhs.component_name == rhs.component_name and lhs.phase == rhs.phase
@staticmethod
def _use_key(use: ComponentUse) -> tuple[str, str, str | None]:
return (use.stage_name, use.component_name, use.phase)
def _trace(
self,
action: str,
use: ComponentUse | None = None,
strategy: ComponentResidencyStrategy | None = None,
module: nn.Module | None = None,
*,
component_name: str | None = None,
detail: str = "",
) -> None:
if not self.state.trace_enabled:
return
if use is not None:
component_name = use.component_name
device = self._module_device(module)
logger.info(
"[component_residency] action=%s stage=%s next_stage=%s component=%s "
"strategy=%s phase=%s device=%s warmup=%s mode=%s %s",
action,
self.state.stage_name,
self.state.next_stage_name,
component_name,
strategy.name if strategy is not None else None,
use.phase if use is not None else None,
device,
self.state.batch_is_warmup,
self.state.manager_mode,
detail,
)
def _module_device(self, module: nn.Module | None) -> str | None:
if module is None:
return None
param = next(module.parameters(), None)
if param is not None:
return param.device.type
buffer = next(module.buffers(), None)
return buffer.device.type if buffer is not None else None
_GLOBAL_COMPONENT_RESIDENCY_MANAGER: ComponentResidencyManager | None = None
def get_global_component_residency_manager(
pipeline: ComponentResidencyPipeline,
server_args: ServerArgs,
) -> ComponentResidencyManager:
global _GLOBAL_COMPONENT_RESIDENCY_MANAGER
if _GLOBAL_COMPONENT_RESIDENCY_MANAGER is None:
_GLOBAL_COMPONENT_RESIDENCY_MANAGER = ComponentResidencyManager(
pipeline, server_args
)
else:
_GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_server_args(server_args)
_GLOBAL_COMPONENT_RESIDENCY_MANAGER.refresh_pipeline(pipeline)
return _GLOBAL_COMPONENT_RESIDENCY_MANAGER
python/sglang/multimodal_gen/runtime/managers/component_resident_strategies.py
+502
View File
@@ -0,0 +1,502 @@
"""
Basic Component Resident Strategy Utilities for defining usage of components, to let ComponentResidencyManager to coordinate
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
import torch.nn as nn
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
if TYPE_CHECKING:
from sglang.multimodal_gen.runtime.managers.component_manager import (
ComponentUse,
ResidencyState,
)
logger = init_logger(__name__)
def _module_to_local_device(
module: nn.Module, *, dtype: torch.dtype | None = None
) -> None:
device = get_local_torch_device()
tensor = _module_reference_tensor(module)
if tensor is not None and tensor.device == device:
if dtype is None or tensor.dtype == dtype:
return
if dtype is None:
module.to(device, non_blocking=True)
else:
module.to(device, dtype=dtype, non_blocking=True)
def _module_reference_tensor(module: nn.Module) -> torch.Tensor | None:
tensor = next(module.parameters(), None)
if tensor is None:
tensor = next(module.buffers(), None)
return tensor
def _module_ready_on_local_device(
module: nn.Module, *, dtype: torch.dtype | None = None
) -> bool:
tensor = _module_reference_tensor(module)
if tensor is None:
return True
if tensor.device != get_local_torch_device():
return False
return dtype is None or tensor.dtype == dtype
class ComponentResidencyStrategy:
"""Baseclass for describing how a component should be treated (regarding where its weights locates)
e.g., a LayerwiseOffloadStrategy would override:
enter: to prefetch some layers before DiT is used, and
exits: to release GPU weight snapshot after DiT is used
to achieve desired behavior
"""
name = "resident"
def prepare_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
"""hook called"""
self.enter(module)
def wait_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
"""Wait for the preparation to be ready, only applicable for async device syncs"""
pass
def finish_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
"""Finish a specific component use"""
self.exit(module)
def prepare_after_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
"""Called after a request is finished, to prepare for the upcoming request"""
pass
def finish_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
*,
preferred: bool,
) -> None:
if preferred:
self.prepare_for_use(module, use, state)
self.wait_for_use(module, use, state)
else:
self.finish_use(module, use, state)
def prefetch_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> bool:
self.prepare_for_use(module, use, state)
return True
def enter(self, module: nn.Module) -> None:
pass
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
pass
class ResidentStrategy(ComponentResidencyStrategy):
name = "resident"
def prepare_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
if use.target_dtype is not None:
_module_to_local_device(module, dtype=use.target_dtype)
class SnapshotModuleResidency:
"""Reusable snapshot-based module residency primitive.
This helper only knows how to:
- keep CPU parameter/buffer snapshots,
- prefetch a module (H2D) to the local device on a CUDA side stream
- release a module by rebinding tensors to those snapshots,
- track and wait for readiness events.
It deliberately does not know about pipeline stages, phases, or model-specific
ordering. Strategy subclasses decide when each primitive is called.
"""
def __init__(self, *, pin_cpu_memory: bool, enable_async_prefetch: bool) -> None:
self.pin_cpu_memory = pin_cpu_memory
self.enable_async_prefetch = enable_async_prefetch
self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {}
self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {}
self._prefetch_stream: object | None = None
self._ready_events: dict[str, object] = {}
@staticmethod
def is_on_gpu(module: nn.Module | None) -> bool:
if module is None:
return False
param = next(module.parameters(), None)
return param is not None and param.device.type == "cuda"
def is_ready(self, component_name: str) -> bool:
return component_name in self._ready_events
def wait_ready(self, component_name: str) -> None:
"""wait for the (H2D) stream to be ready"""
ready_event = self._ready_events.get(component_name)
if ready_event is None or not current_platform.is_cuda():
return
torch.get_device_module().current_stream().wait_event(ready_event)
def record_ready(self, component_name: str, module: nn.Module | None) -> None:
if not current_platform.is_cuda():
self._ready_events.pop(component_name, None)
return
if not self.is_on_gpu(module):
self._ready_events.pop(component_name, None)
return
event = torch.get_device_module().Event()
event.record(torch.get_device_module().current_stream())
self._ready_events[component_name] = event
@staticmethod
def _clone_cpu_tensor_snapshot(
tensor: torch.Tensor, *, pin_memory: bool
) -> torch.Tensor:
snapshot = tensor.detach()
if snapshot.device.type == "cpu":
if pin_memory and not snapshot.is_pinned():
return snapshot.pin_memory()
return snapshot
cpu_tensor = snapshot.to("cpu")
if pin_memory:
return cpu_tensor.pin_memory()
return cpu_tensor
def _should_pin_memory(self) -> bool:
return bool(self.pin_cpu_memory and torch.get_device_module().is_available())
def capture(self, component_name: str, module: nn.Module) -> None:
"""Capture a CPU snapshot for a component"""
if component_name in self._cpu_param_snapshots:
return
pin_memory = self._should_pin_memory()
self._cpu_param_snapshots[component_name] = {
name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory)
for name, param in module.named_parameters()
}
self._cpu_buffer_snapshots[component_name] = {
name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory)
for name, buffer in module.named_buffers()
}
def release_to_snapshot(
self,
component_name: str,
module: nn.Module,
*,
copy_runtime_buffers: bool = False,
) -> None:
"""Release CUDA storages by rebinding tensors to cached CPU snapshots.
This does not call `module.to("cpu")`. Instead, parameter and buffer
storages are rebound to pre-captured CPU tensors so CUDA storages can be
released by the allocator without an explicit D2H transfer.
"""
param_snapshots = self._cpu_param_snapshots.get(component_name)
buffer_snapshots = self._cpu_buffer_snapshots.get(component_name)
if param_snapshots is None or buffer_snapshots is None:
module.to("cpu")
self._ready_events.pop(component_name, None)
return
pin_memory = self._should_pin_memory()
for name, param in module.named_parameters():
snapshot = param_snapshots.get(name)
if snapshot is None:
snapshot = self._clone_cpu_tensor_snapshot(
param.data, pin_memory=pin_memory
)
param_snapshots[name] = snapshot
param.data = snapshot
for name, buffer in module.named_buffers():
snapshot = buffer_snapshots.get(name)
if snapshot is None:
snapshot = self._clone_cpu_tensor_snapshot(
buffer.data, pin_memory=pin_memory
)
buffer_snapshots[name] = snapshot
if copy_runtime_buffers:
# Preserve runtime-updated buffers (e.g., lazily built caches) when
# releasing back to CPU snapshots.
if buffer.device.type == "cuda":
snapshot.copy_(
buffer.detach().to(device="cpu", dtype=snapshot.dtype)
)
elif buffer.device.type == "cpu":
snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype))
buffer.data = snapshot
self._ready_events.pop(component_name, None)
def _supports_async_prefetch(self) -> bool:
return self.enable_async_prefetch and current_platform.is_cuda()
def _get_prefetch_stream(self):
"""returns a stream is async-prefetch is enabled"""
if not self._supports_async_prefetch():
return None
if self._prefetch_stream is None:
self._prefetch_stream = torch.get_device_module().Stream(
device=get_local_torch_device()
)
return self._prefetch_stream
def prefetch_to_device(self, component_name: str, module: nn.Module | None) -> None:
if module is None:
self._ready_events.pop(component_name, None)
return
prefetch_stream = self._get_prefetch_stream()
if prefetch_stream is None:
# if the async prefetching is disabled
module.to(get_local_torch_device(), non_blocking=True)
self.record_ready(component_name, module)
return
with torch.get_device_module().stream(prefetch_stream):
module.to(get_local_torch_device(), non_blocking=True)
event = torch.get_device_module().Event()
event.record(prefetch_stream)
self._ready_events[component_name] = event
class SnapshotStrategy(ComponentResidencyStrategy):
"""Snapshot residency: async H2D before use and light snapshot release after use."""
name = "snapshot"
def __init__(
self,
*,
pin_cpu_memory: bool,
enable_async_prefetch: bool,
copy_runtime_buffers_on_release: bool = False,
) -> None:
self._snapshot_residency = SnapshotModuleResidency(
pin_cpu_memory=pin_cpu_memory,
enable_async_prefetch=enable_async_prefetch,
)
self._copy_runtime_buffers_on_release = copy_runtime_buffers_on_release
def capture(self, component_name: str, module: nn.Module) -> None:
self._snapshot_residency.capture(component_name, module)
def is_ready(self, component_name: str) -> bool:
return self._snapshot_residency.is_ready(component_name)
def record_ready(self, component_name: str, module: nn.Module | None) -> None:
self._snapshot_residency.record_ready(component_name, module)
def prefetch_component(self, component_name: str, module: nn.Module | None) -> None:
if SnapshotModuleResidency.is_on_gpu(module):
self._snapshot_residency.record_ready(component_name, module)
return
self._snapshot_residency.prefetch_to_device(component_name, module)
def wait_component_ready(self, component_name: str) -> None:
self._snapshot_residency.wait_ready(component_name)
def release_component(self, component_name: str, module: nn.Module) -> None:
self._snapshot_residency.release_to_snapshot(
component_name,
module,
copy_runtime_buffers=self._copy_runtime_buffers_on_release,
)
def prepare_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.prefetch_component(use.component_name, module)
def wait_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.wait_component_ready(use.component_name)
def finish_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.release_component(use.component_name, module)
def prepare_after_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.prepare_for_use(module, use, state)
class VanillaD2HStrategy(ComponentResidencyStrategy):
"""A strategy that performs native torch D2H and H2D for a component"""
name = "vanilla"
def __init__(self) -> None:
self._prefetch_stream: object | None = None
self._ready_events: dict[str, object] = {}
def prepare_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
_module_to_local_device(module, dtype=use.target_dtype)
def wait_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
ready_event = self._ready_events.get(use.component_name)
if ready_event is None or not current_platform.is_cuda():
return
torch.get_device_module().current_stream().wait_event(ready_event)
def prefetch_for_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> bool:
if not current_platform.is_cuda():
self.prepare_for_use(module, use, state)
return True
if _module_ready_on_local_device(module, dtype=use.target_dtype):
return True
if self._prefetch_stream is None:
self._prefetch_stream = torch.get_device_module().Stream(
device=get_local_torch_device()
)
with torch.get_device_module().stream(self._prefetch_stream):
_module_to_local_device(module, dtype=use.target_dtype)
event = torch.get_device_module().Event()
event.record(self._prefetch_stream)
self._ready_events[use.component_name] = event
return True
def enter(self, module: nn.Module) -> None:
param = next(module.parameters(), None)
if param is not None and param.device.type == "cpu":
_module_to_local_device(module)
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
param = next(module.parameters(), None)
if param is not None and param.device.type == "cuda":
module.to("cpu", non_blocking=True)
def finish_use(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.wait_for_use(module, use, state)
self.exit(module)
self._ready_events.pop(use.component_name, None)
def prepare_after_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.prefetch_for_use(module, use, state)
def finish_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
*,
preferred: bool,
) -> None:
if preferred and state.batch_is_warmup:
self.prepare_for_use(module, use, state)
self.wait_for_use(module, use, state)
return
if not preferred:
self.finish_use(module, use, state)
class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
"""A wrapper around LayerwiseOffloadManager to fit in a ComponentResidencyStrategy"""
name = "layerwise"
def enter(self, module: nn.Module) -> None:
if isinstance(module, OffloadableDiTMixin):
module.prepare_for_next_req()
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
if not isinstance(module, OffloadableDiTMixin):
return
for manager in module.layerwise_offload_managers:
manager.release_all()
def prepare_after_request(
self,
module: nn.Module,
use: ComponentUse,
state: ResidencyState,
) -> None:
self.prepare_for_use(module, use, state)
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+4 -4
View File
@@ -37,6 +37,10 @@ from sglang.multimodal_gen.runtime.loader.weights_updater import (
WeightsUpdater,
get_updatable_modules,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import (
OffloadableDiTMixin,
iter_materialized_weights,
)
from sglang.multimodal_gen.runtime.pipelines_core import (
ComposedPipelineBase,
LoRAPipeline,
@@ -47,10 +51,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
OffloadableDiTMixin,
iter_materialized_weights,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import (
configure_logger,
globally_suppress_loggers,
python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py → python/sglang/multimodal_gen/runtime/managers/layerwise_offload.py
+5
View File
@@ -199,6 +199,11 @@ class LayerwiseOffloadManager:
self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer
# Keep non-layer parameters resident on GPU. Layer tensors have already
# been replaced by tiny device placeholders, so this does not reload the
# offloaded layer weights.
self.model.to(self.device)
# prefetch the first layer for warm-up
self.prepare_for_next_req(non_blocking=False)
python/sglang/multimodal_gen/runtime/managers/scheduler.py
+9 -13
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py
+1 -1
View File
@@ -26,8 +26,8 @@ from sglang.multimodal_gen.runtime.layers.linear import (
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
+1 -1
View File
@@ -13,7 +13,7 @@ from torch.nn.attention.flex_attention import (
flex_attention,
)
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
# see https://github.com/pytorch/pytorch/issues/133254
python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py
+1 -1
View File
@@ -33,8 +33,8 @@ from sglang.multimodal_gen.runtime.layers.linear import (
RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
def _rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+1 -1
View File
@@ -52,9 +52,9 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
CombinedTimestepGuidanceTextProjEmbeddings,
CombinedTimestepTextProjEmbeddings,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+1 -1
View File
@@ -42,12 +42,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
)
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+1 -1
View File
@@ -37,12 +37,12 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
)
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/helios.py
+1 -1
View File
@@ -46,8 +46,8 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
TimestepEmbedder,
)
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
+1 -1
View File
@@ -22,9 +22,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+1 -1
View File
@@ -37,13 +37,13 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
unpatchify,
)
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.models.utils import modulate
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
)
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
class MMDoubleStreamBlock(nn.Module):
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+1 -1
View File
@@ -31,9 +31,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py
+1 -1
View File
@@ -18,8 +18,8 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
# Reuse common functions and classes from mova_video_dit
from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d
python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py
+1 -1
View File
@@ -33,9 +33,9 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+1 -1
View File
@@ -46,12 +46,12 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
)
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
python/sglang/multimodal_gen/runtime/models/dits/sana.py
+1 -1
View File
@@ -8,8 +8,8 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbed
from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+1 -1
View File
@@ -49,6 +49,7 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import (
TimestepEmbedder,
)
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.models.utils import (
_use_aiter,
@@ -58,7 +59,6 @@ from sglang.multimodal_gen.runtime.platforms import (
current_platform,
)
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.utils import add_prefix
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+1 -1
View File
@@ -41,9 +41,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
_apply_rotary_emb,
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
try:
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+317 -319
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+12
View File
@@ -21,6 +21,11 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import (
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PipelineComponentLoader,
)
from sglang.multimodal_gen.runtime.managers.component_manager import (
ComponentResidencyManager,
ComponentResidencyStrategy,
get_global_component_residency_manager,
)
from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
PipelineExecutor,
)
@@ -91,7 +96,9 @@ class ComposedPipelineBase(ABC):
self.model_path: str = model_path
self._stages: list[PipelineStage] = []
self._stage_name_mapping: dict[str, PipelineStage] = {}
self.component_residency_strategies: dict[str, ComponentResidencyStrategy] = {}
self.executor = executor or self.build_executor(server_args=server_args)
self.component_residency_manager: ComponentResidencyManager | None = None
if required_config_modules is not None:
self._required_config_modules = required_config_modules
@@ -738,6 +745,11 @@ class ComposedPipelineBase(ABC):
main_process_only=True,
)
self.component_residency_manager = get_global_component_residency_manager(
self, server_args
)
self.executor.component_residency_manager = self.component_residency_manager
return self.executor.execute_with_profiling(self.stages, batch, server_args)
@torch.no_grad()
python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py
+31 -21
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py
+27
View File
@@ -45,6 +45,33 @@ class PipelineExecutor(ABC):
def __init__(self, server_args):
self.server_args = server_args
self.component_residency_manager = None
def begin_component_residency_request(
self,
stages: List["PipelineStage"],
batch: Req,
server_args: ServerArgs,
) -> None:
self.component_residency_manager.begin_request(stages, batch, server_args)
def before_stage(
self,
stage: "PipelineStage",
stage_index: int,
batch: Req,
server_args: ServerArgs,
) -> None:
stage.set_component_residency_manager(self.component_residency_manager)
self.component_residency_manager.before_stage(
stage, stage_index, batch, server_args
)
def after_stage(self, stage_index: int) -> None:
self.component_residency_manager.after_stage(stage_index)
def finish_component_residency_request(self) -> None:
self.component_residency_manager.finish_request()
def execute_with_profiling(
self,
python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py
+11 -5
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
+1 -1
View File
@@ -169,7 +169,7 @@ class LoRAPipeline(ComposedPipelineBase):
Yields:
List of modules that had offload disabled.
"""
from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
from sglang.multimodal_gen.runtime.managers.layerwise_offload import (
OffloadableDiTMixin,
)
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+2 -2
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
+80
View File
@@ -9,11 +9,15 @@ composed to create complete diffusion pipelines.
"""
from abc import ABC, abstractmethod
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import replace
from enum import Enum, auto
import torch
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.dedup import StageDedupMixin
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
@@ -53,6 +57,7 @@ class PipelineStage(StageDedupMixin, ABC):
def __init__(self):
self.server_args = get_global_server_args()
self._component_residency_manager = None
def log_info(self, msg, *args):
"""Logs an informational message with the stage name as a prefix."""
@@ -105,6 +110,81 @@ class PipelineStage(StageDedupMixin, ABC):
"""
pass
def set_component_residency_manager(self, manager) -> None:
self._component_residency_manager = manager
def _component_stage_name(self, stage_name: str | None = None) -> str:
return stage_name or self.__class__.__name__
def _active_component_stage_name(self) -> str:
manager = self._component_residency_manager
if manager is not None and manager.state.stage_name is not None:
return manager.state.stage_name
return self.__class__.__name__
def _finish_active_component_use(self) -> None:
if self._component_residency_manager is not None:
self._component_residency_manager.finish_active_use()
@contextmanager
def _use_component(
self,
use: ComponentUse,
module=None,
) -> Iterator[object | None]:
if self._component_residency_manager is None:
yield module
return
with self._component_residency_manager.use_component(use, module) as component:
yield component
def _declared_component_use(
self,
*,
component_name: str,
phase: str | None = None,
target_dtype: torch.dtype | None = None,
) -> ComponentUse:
manager = self._component_residency_manager
stage_name = self._active_component_stage_name()
server_args = manager.server_args if manager is not None else self.server_args
for use in self.component_uses(server_args, stage_name):
if use.component_name != component_name:
continue
if phase is not None and use.phase != phase:
continue
if target_dtype is not None:
return replace(use, target_dtype=target_dtype)
return use
raise ValueError(
f"{self.__class__.__name__} did not declare component use: "
f"{component_name}"
)
@contextmanager
def use_declared_component(
self,
*,
component_name: str,
module=None,
phase: str | None = None,
target_dtype: torch.dtype | None = None,
) -> Iterator[object | None]:
"""reference a component already declared in `component_uses`"""
use = self._declared_component_use(
component_name=component_name,
phase=phase,
target_dtype=target_dtype,
)
with self._use_component(use, module) as component:
yield component
def component_uses(
self, server_args: ServerArgs, stage_name: str | None = None
) -> list[ComponentUse]:
"""Declares component uses of current stage for unified residency scheduling."""
return []
# Default role affinity: ENCODER. Override in subclasses for DENOISING/DECODER.
@property
def role_affinity(self) -> RoleType:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
+53 -44
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py
+98 -73
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+88 -49
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+32 -27
View File
@@ -2,6 +2,7 @@ import torch
from diffusers.utils.torch_utils import randn_tensor
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import is_ltx23_native_variant
from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse
from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
clone_scheduler_runtime,
)
@@ -10,7 +11,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import
LTX2DenoisingStage,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
@@ -28,6 +28,20 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
)
self.audio_vae = audio_vae
def component_uses(
self, server_args: ServerArgs, stage_name: str | None = None
) -> list[ComponentUse]:
stage_name = self._component_stage_name(stage_name)
return [
ComponentUse(
stage_name=stage_name,
component_name="transformer",
phase="stage1",
preferred_ready_after_request=True,
memory_intensive=True,
)
]
def _post_denoising_loop(
self,
batch: Req,
@@ -83,24 +97,6 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
batch.latents = latents
batch.audio_latents = audio_latents
pipeline = self.pipeline() if self.pipeline else None
current_phase = (
str(getattr(batch, "extra", {}).get("ltx2_phase", ""))
if hasattr(batch, "extra")
else ""
)
release_phase_state = (
getattr(pipeline, "release_ltx2_phase_state", None)
if pipeline is not None
else None
)
if callable(release_phase_state):
release_phase_state(current_phase)
if isinstance(self.transformer, OffloadableDiTMixin):
for manager in self.transformer.layerwise_offload_managers:
manager.release_all()
class LTX2RefinementStage(LTX2AVDenoisingStage):
"""Stage-2 refinement wrapper that re-noises distilled LTX latents once."""
@@ -125,6 +121,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
)
self.distilled_sigmas = torch.tensor(distilled_sigmas)
def component_uses(
self, server_args: ServerArgs, stage_name: str | None = None
) -> list[ComponentUse]:
stage_name = self._component_stage_name(stage_name)
component_name = "transformer_2"
pipeline = self.pipeline() if self.pipeline else None
if pipeline is not None and "transformer_2" not in pipeline.modules:
component_name = "transformer"
return [
ComponentUse(
stage_name=stage_name,
component_name=component_name,
phase="stage2",
memory_intensive=True,
)
]
@staticmethod
def _randn_like_with_batch_generators(
reference_tensor: torch.Tensor, batch: Req
@@ -218,14 +231,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
"""Run the distilled refinement schedule on top of the shared AV denoiser."""
batch.extra["ltx2_phase"] = "stage2"
pipeline = self.pipeline() if self.pipeline else None
ensure_phase_ready = (
getattr(pipeline, "ensure_ltx2_phase_ready", None)
if pipeline is not None
else None
)
if callable(ensure_phase_ready):
ensure_phase_ready("stage2")
original_clean_latent_background = getattr(
batch, "ltx2_ti2v_clean_latent_background", None
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py
+4 -30
View File
@@ -121,7 +121,7 @@ class DmdDenoisingStage(DenoisingStage):
)
else:
current_model = self.transformer
self._manage_device_placement(current_model, None, server_args)
self._manage_dit_use_site(current_model, "transformer", batch)
# Expand latents for I2V
noise_latents = latents.clone()
latent_model_input = latents.to(target_dtype)
@@ -232,45 +232,19 @@ class DmdDenoisingStage(DenoisingStage):
if boundary_timestep is None or t_int >= boundary_timestep:
# High-noise stage
current_model = self.transformer
model_to_offload = self.transformer_2
current_guidance_scale = batch.guidance_scale
current_phase = "transformer"
else:
# Low-noise stage
current_model = self.transformer_2
model_to_offload = self.transformer
current_guidance_scale = batch.guidance_scale_2
current_phase = "transformer_2"
self._manage_device_placement(current_model, model_to_offload, server_args)
self._manage_dit_use_site(current_model, current_phase, batch)
assert current_model is not None, "The model for the current step is not set."
return current_model, current_guidance_scale
def _manage_device_placement(
self,
model_to_use: torch.nn.Module,
model_to_offload: torch.nn.Module | None,
server_args: ServerArgs,
):
"""
Manages the offload / load behavior of dit
"""
if not server_args.dit_cpu_offload:
return
# Offload the unused model if it's on CUDA
if (
model_to_offload is not None
and next(model_to_offload.parameters()).device.type == "cuda"
):
model_to_offload.to("cpu")
# Load the model to use if it's on CPU
if (
model_to_use is not None
and next(model_to_use.parameters()).device.type == "cpu"
):
model_to_use.to(get_local_torch_device())
def _handle_boundary_ratio(
self,
server_args,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py
+31 -21
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
+202 -189
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+5 -17
View File
@@ -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):
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py
+13 -4
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py
+24 -9
View File
@@ -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).
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py
+103 -53
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py
+49 -34
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py
+5 -4
View File
@@ -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):
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py
+55 -7
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py
+22 -28
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+2 -2
View File
@@ -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,
)