[diffusion] perf: support zero-cost weight offload and overlap with compute for wan-series (#15511)

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -32,6 +32,7 @@ If you only need to use the distributed environment without model parallelism,
 """
 import contextlib
 import os
+import time
 import weakref
 from collections import namedtuple
 from collections.abc import Callable
@@ -67,8 +68,6 @@ _DP: Optional[GroupCoordinator] = None
 _DIT: Optional[GroupCoordinator] = None
 _VAE: Optional[GroupCoordinator] = None
 
-logger = init_logger(__name__)
-
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
 
@@ -433,6 +432,9 @@ def initialize_model_parallel(
         ring_group=PROCESS_GROUP.RING_PG,
     )
 
+    if ulysses_degree > 1:
+        _warmup_ulysses_communication()
+
     global _TP
     assert _TP is None, "Tensor parallel group is already initialized"
     _TP = init_parallel_group_coordinator(
@@ -945,6 +947,49 @@ def get_ring_parallel_rank():
     return get_sp_group().ring_rank
 
 
+def _warmup_ulysses_communication():
+    """
+    Warmup NCCL communication for Ulysses all-to-all to avoid first-step latency.
+
+    This function performs a dummy all-to-all operation to initialize NCCL communication
+    channels, which can take several seconds on the first call.
+    """
+    logger.info("Warming up Ulysses all-to-all communication...")
+
+    try:
+        import torch.distributed._functional_collectives as ft_c
+
+        ulysses_pg = get_sp_group().ulysses_group
+        if ulysses_pg is None:
+            logger.warning("Ulysses group not initialized, skipping warmup")
+            return
+
+        warmup_start = time.time()
+
+        device = torch.device(f"cuda:{get_world_group().local_rank}")
+        dummy_tensor = torch.zeros(1024, device=device, dtype=torch.float32)
+
+        output = ft_c.all_to_all_single(
+            dummy_tensor,
+            output_split_sizes=None,
+            input_split_sizes=None,
+            group=ulysses_pg,
+        )
+
+        if isinstance(output, ft_c.AsyncCollectiveTensor):
+            output = output.wait()
+
+        torch.cuda.synchronize()
+
+        warmup_time = (time.time() - warmup_start) * 1000
+        logger.info(f"Ulysses communication warmup completed in {warmup_time:.2f}ms")
+
+    except Exception as e:
+        logger.warning(
+            f"Ulysses communication warmup failed: {e}. Continuing without warmup."
+        )
+
+
 # PP
 def get_pp_group() -> PipelineGroupCoordinator:
     assert _PP is not None, "pipeline model parallel group is not initialized"

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

@@ -43,6 +43,9 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
     get_hf_config,
 )
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+    LayerwiseOffloadManager,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
@@ -709,6 +712,23 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
+        if server_args.dit_layerwise_offload and hasattr(model, "blocks"):
+            try:
+                num_layers = len(getattr(model, "blocks"))
+            except Exception:
+                num_layers = None
+
+            if isinstance(num_layers, int) and num_layers > 0:
+                mgr = LayerwiseOffloadManager(
+                    model,
+                    module_list_attr="blocks",
+                    num_layers=num_layers,
+                    enabled=True,
+                    pin_cpu_memory=server_args.pin_cpu_memory,
+                    auto_initialize=True,
+                )
+                setattr(model, "_layerwise_offload_manager", mgr)
+
         return model
 
 

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

@@ -792,10 +792,25 @@ class WanTransformer3DModel(CachableDiT):
             if enable_teacache:
                 original_hidden_states = hidden_states.clone()
 
-            for block in self.blocks:
+            offload_mgr = getattr(self, "_layerwise_offload_manager", None)
-                hidden_states = block(
+            if offload_mgr is not None and getattr(offload_mgr, "enabled", False):
-                    hidden_states, encoder_hidden_states, timestep_proj, freqs_cis
+                for i, block in enumerate(self.blocks):
-                )
+                    with offload_mgr.layer_scope(
+                        prefetch_layer_idx=i + 1,
+                        release_layer_idx=i,
+                        non_blocking=True,
+                    ):
+                        hidden_states = block(
+                            hidden_states,
+                            encoder_hidden_states,
+                            timestep_proj,
+                            freqs_cis,
+                        )
+            else:
+                for block in self.blocks:
+                    hidden_states = block(
+                        hidden_states, encoder_hidden_states, timestep_proj, freqs_cis
+                    )
             # if teacache is enabled, we need to cache the original hidden states
             if enable_teacache:
                 self.maybe_cache_states(hidden_states, original_hidden_states)

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

@@ -728,6 +728,17 @@ class DenoisingStage(PipelineStage):
             latents, batch
         )
 
+        offload_mgr = getattr(self.transformer, "_layerwise_offload_manager", None)
+        if offload_mgr is not None and getattr(offload_mgr, "enabled", False):
+            offload_mgr.release_all()
+
+        if self.transformer_2 is not None:
+            offload_mgr_2 = getattr(
+                self.transformer_2, "_layerwise_offload_manager", None
+            )
+            if offload_mgr_2 is not None and getattr(offload_mgr_2, "enabled", False):
+                offload_mgr_2.release_all()
+
         # Save STA mask search results if needed
         if (
             st_attn_available

python/sglang/multimodal_gen/runtime/server_args.py

@@ -229,6 +229,7 @@ class ServerArgs:
     # CPU offload parameters
     dit_cpu_offload: bool = True
     use_fsdp_inference: bool = False
+    dit_layerwise_offload: bool = False
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
     vae_cpu_offload: bool = True
@@ -507,6 +508,13 @@ class ServerArgs:
             action=StoreBoolean,
             help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
         )
+        parser.add_argument(
+            "--dit-layerwise-offload",
+            action=StoreBoolean,
+            default=ServerArgs.dit_layerwise_offload,
+            help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan). "
+            "Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.",
+        )
         parser.add_argument(
             "--use-fsdp-inference",
             action=StoreBoolean,
@@ -844,6 +852,26 @@ class ServerArgs:
         """Validate inference arguments for consistency"""
         if current_platform.is_mps():
             self.use_fsdp_inference = False
+            self.dit_layerwise_offload = False
+
+        if self.dit_layerwise_offload:
+            if self.use_fsdp_inference:
+                logger.warning(
+                    "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference."
+                )
+                self.use_fsdp_inference = False
+            if self.dit_cpu_offload:
+                logger.warning(
+                    "dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload."
+                )
+                self.dit_cpu_offload = False
+            if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true":
+                raise ValueError(
+                    "dit_layerwise_offload cannot be enabled together with cache-dit. "
+                    "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, "
+                    "causing shape mismatch errors. "
+                    "Please disable either --dit-layerwise-offload or SGLANG_CACHE_DIT_ENABLED."
+                )
 
         # autocast
         if self.disable_autocast is None:

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

@@ -0,0 +1,197 @@
+import re
+from contextlib import contextmanager
+from typing import Dict, Set
+
+import torch
+
+
+# Adapted from skywork AI Infra diffusion optimize
+class LayerwiseOffloadManager:
+    """A lightweight layerwise CPU offload manager.
+
+    This utility offloads per-layer parameters/buffers from GPU to CPU, and
+    supports async H2D prefetch using a dedicated CUDA stream.
+
+    Typical usage:
+    - Construct the manager with the target model and the list-like module
+      attribute that represents transformer blocks (e.g. ``blocks``).
+    - Call :meth:`initialize` once to offload weights and prefetch layer 0.
+    - During forward, call :meth:`prefetch_layer` for the next layer and
+      :meth:`release_layer` for the finished layer.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        *,
+        module_list_attr: str,
+        num_layers: int,
+        enabled: bool,
+        pin_cpu_memory: bool = True,
+        auto_initialize: bool = False,
+    ) -> None:
+        self.model = model
+        self.module_list_attr = module_list_attr
+        self.num_layers = num_layers
+        self.pin_cpu_memory = pin_cpu_memory
+
+        self.enabled = bool(enabled and torch.cuda.is_available())
+        self.device = (
+            torch.device("cuda", torch.cuda.current_device()) if self.enabled else None
+        )
+        self.copy_stream = torch.cuda.Stream() if self.enabled else None
+
+        self._layer_name_re = re.compile(
+            rf"(^|\.){re.escape(module_list_attr)}\.(\d+)(\.|$)"
+        )
+
+        self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
+        self._cpu_dtypes: Dict[int, Dict[str, torch.dtype]] = {}
+        self._gpu_layers: Dict[int, Set[str]] = {}
+
+        self._named_parameters: Dict[str, torch.nn.Parameter] = {}
+        self._named_buffers: Dict[str, torch.Tensor] = {}
+
+        if auto_initialize:
+            self.initialize()
+
+    def _match_layer_idx(self, name: str) -> int | None:
+        m = self._layer_name_re.search(name)
+        if not m:
+            return None
+        try:
+            return int(m.group(2))
+        except Exception:
+            return None
+
+    def _offload_tensor(self, name: str, tensor: torch.Tensor, layer_idx: int) -> None:
+        if layer_idx not in self._cpu_weights:
+            self._cpu_weights[layer_idx] = {}
+            self._cpu_dtypes[layer_idx] = {}
+
+        cpu_weight = tensor.detach().to("cpu")
+        if self.pin_cpu_memory:
+            cpu_weight = cpu_weight.pin_memory()
+        self._cpu_weights[layer_idx][name] = cpu_weight
+        self._cpu_dtypes[layer_idx][name] = tensor.dtype
+
+        if self.device is not None:
+            tensor.data = torch.empty((1,), device=self.device, dtype=tensor.dtype)
+
+    @torch.compiler.disable
+    def initialize(self) -> None:
+        if not self.enabled:
+            return
+
+        self._named_parameters = dict(self.model.named_parameters())
+        self._named_buffers = dict(self.model.named_buffers())
+
+        for name, param in self._named_parameters.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            self._offload_tensor(name, param, layer_idx)
+
+        for name, buf in self._named_buffers.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            self._offload_tensor(name, buf, layer_idx)
+
+        self.prefetch_layer(0, non_blocking=False)
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+    @torch.compiler.disable
+    def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
+        if not self.enabled or self.device is None or self.copy_stream is None:
+            return
+        if layer_idx < 0 or layer_idx >= self.num_layers:
+            return
+        if layer_idx in self._gpu_layers:
+            return
+        if layer_idx not in self._cpu_weights:
+            return
+
+        self.copy_stream.wait_stream(torch.cuda.current_stream())
+
+        param_names: Set[str] = set()
+
+        for name, cpu_weight in self._cpu_weights[layer_idx].items():
+            if name in self._named_parameters:
+                target = self._named_parameters[name]
+            else:
+                target = self._named_buffers[name]
+
+            gpu_weight = torch.empty(
+                cpu_weight.shape,
+                dtype=self._cpu_dtypes[layer_idx][name],
+                device=self.device,
+            )
+            with torch.cuda.stream(self.copy_stream):
+                gpu_weight.copy_(cpu_weight, non_blocking=non_blocking)
+            target.data = gpu_weight
+            param_names.add(name)
+
+        self._gpu_layers[layer_idx] = param_names
+
+    @contextmanager
+    def layer_scope(
+        self,
+        *,
+        prefetch_layer_idx: int | None,
+        release_layer_idx: int | None,
+        non_blocking: bool = True,
+    ):
+        """A helper context manager to improve readability at call sites.
+
+        It optionally prefetches ``prefetch_layer_idx`` before entering the
+        context, and waits for the copy stream then releases
+        ``release_layer_idx`` on exit.
+        """
+        if self.enabled and prefetch_layer_idx is not None:
+            self.prefetch_layer(prefetch_layer_idx, non_blocking=non_blocking)
+        try:
+            yield
+        finally:
+            if self.enabled and self.copy_stream is not None:
+                torch.cuda.current_stream().wait_stream(self.copy_stream)
+            if self.enabled and release_layer_idx is not None:
+                self.release_layer(release_layer_idx)
+
+    @torch.compiler.disable
+    def release_layer(self, layer_idx: int) -> None:
+        if not self.enabled or self.device is None:
+            return
+        if layer_idx <= 0:
+            return
+
+        param_names = self._gpu_layers.pop(layer_idx, None)
+        if not param_names:
+            return
+
+        for name in param_names:
+            if name in self._named_parameters:
+                target = self._named_parameters[name]
+            else:
+                target = self._named_buffers[name]
+            target.data = torch.empty((1,), device=self.device, dtype=target.dtype)
+
+    @torch.compiler.disable
+    def release_all(self) -> None:
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        layer_indices = list(self._gpu_layers.keys())
+        for layer_idx in layer_indices:
+            param_names = self._gpu_layers.pop(layer_idx, None)
+            if not param_names:
+                continue
+            for name in param_names:
+                if name in self._named_parameters:
+                    target = self._named_parameters[name]
+                else:
+                    target = self._named_buffers[name]
+                target.data = torch.empty((1,), device=self.device, dtype=target.dtype)