[diffusion] feat: add an arg for controlling the number of prefetched layers in layerwise-offload (#17693)

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

@@ -742,7 +742,10 @@ class DenoisingStage(PipelineStage):
         # reset offload managers with prefetching first layer for next forward
         for dit in filter(None, [self.transformer, self.transformer_2]):
             if isinstance(dit, OffloadableDiTMixin):
-                dit.prepare_for_next_denoise()
+                # 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."""

python/sglang/multimodal_gen/runtime/server_args.py

@@ -8,6 +8,7 @@ import argparse
 import dataclasses
 import inspect
 import json
+import math
 import os
 import random
 import sys
@@ -285,6 +286,7 @@ class ServerArgs:
     # CPU offload parameters
     dit_cpu_offload: bool | None = None
     dit_layerwise_offload: bool | None = None
+    dit_offload_prefetch_size: float = 0.0
     text_encoder_cpu_offload: bool | None = None
     image_encoder_cpu_offload: bool | None = None
     vae_cpu_offload: bool | None = None
@@ -617,6 +619,12 @@ class ServerArgs:
             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(
+            "--dit-offload-prefetch-size",
+            type=float,
+            default=ServerArgs.dit_offload_prefetch_size,
+            help="The size of prefetch for dit-layerwise-offload. If the value is between 0.0 and 1.0, it is treated as a ratio of the total number of layers. If the value is >= 1, it is treated as the absolute number of layers. 0.0 means prefetch 1 layer (lowest memory). Values above 0.5 might have peak memory close to no offload but worse performance.",
+        )
         parser.add_argument(
             "--use-fsdp-inference",
             action=StoreBoolean,
@@ -949,6 +957,21 @@ class ServerArgs:
             self.use_fsdp_inference = False
             self.dit_layerwise_offload = False
 
+        if self.dit_offload_prefetch_size > 1 and (
+            isinstance(self.dit_offload_prefetch_size, float)
+            and not self.dit_offload_prefetch_size.is_integer()
+        ):
+            self.dit_offload_prefetch_size = int(
+                math.floor(self.dit_offload_prefetch_size)
+            )
+            logger.info(
+                f"Invalid --dit-offload-prefetch-size value passed, truncated to: {self.dit_offload_prefetch_size}"
+            )
+        if 0.5 <= self.dit_offload_prefetch_size < 1.0:
+            logger.info(
+                f"We do not recommend --dit-offload-prefetch-size to be between 0.5 and 1.0"
+            )
+
         if not envs.SGLANG_CACHE_DIT_ENABLED:
             # TODO: need a better way to tell this
             if (
@@ -962,6 +985,9 @@ class ServerArgs:
                 self.dit_layerwise_offload = True
 
         if self.dit_layerwise_offload:
+            assert (
+                self.dit_offload_prefetch_size >= 0.0
+            ), "dit_offload_prefetch_size must be non-negative"
             if self.use_fsdp_inference:
                 logger.warning(
                     "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference."

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

@@ -33,12 +33,13 @@ class LayerwiseOffloadManager:
         num_layers: int,
         enabled: bool,
         pin_cpu_memory: bool = True,
+        prefetch_size: int = 1,
     ) -> None:
         self.model = model
         self.layers_attr_str = layers_attr_str
         self.num_layers = num_layers
         self.pin_cpu_memory = pin_cpu_memory
-
+        self.prefetch_size = min(max(1, prefetch_size), self.num_layers)
         self.enabled = bool(enabled and torch.cuda.is_available())
         if not self.enabled:
             return
@@ -57,6 +58,8 @@ class LayerwiseOffloadManager:
         self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
         # layer indices that are already in gpu
         self._gpu_layers: Set[int] = set()
+        # layer_idx -> torch.cuda.Event for fine-grained sync, to make sure the weight is resident in pre-hook
+        self._prefetch_events: Dict[int, torch.cuda.Event] = {}
 
         self._named_parameters: Dict[str, torch.nn.Parameter] = {}
         self._named_buffers: Dict[str, torch.Tensor] = {}
@@ -127,13 +130,19 @@ class LayerwiseOffloadManager:
                 self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer
 
         # prefetch the first layer for warm-up
-        self.prepare_for_next_denoise(non_blocking=False)
+        self.prepare_for_next_req(non_blocking=False)
 
         self.register_forward_hooks()
-        logger.info("LayerwiseOffloadManager initialized")
+        logger.info(
+            f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, total num layers: {self.num_layers}"
+        )
 
-    def prepare_for_next_denoise(self, non_blocking=True):
+    def prepare_for_next_req(self, non_blocking=True):
-        self.prefetch_layer(0, non_blocking=non_blocking)
+        """
+        Prepare for the next round of denoising loop with prefetching the necessary layers
+        """
+        for i in range(self.prefetch_size):
+            self.prefetch_layer(i, non_blocking=non_blocking)
         if not non_blocking and self.copy_stream is not None:
             torch.cuda.current_stream().wait_stream(self.copy_stream)
 
@@ -147,6 +156,9 @@ class LayerwiseOffloadManager:
 
     @torch.compiler.disable
     def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
+        """
+        idempotent
+        """
         if not self.enabled or self.device is None or self.copy_stream is None:
             return
         if layer_idx < 0 or layer_idx >= self.num_layers:
@@ -167,6 +179,11 @@ class LayerwiseOffloadManager:
                 gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking)
                 gpu_buffers[dtype] = gpu_buffer
 
+        # record the prefetch event of this layer
+        event = torch.cuda.Event()
+        event.record(self.copy_stream)
+        self._prefetch_events[layer_idx] = event
+
         # restore model's weights by their metadata using gpu buffer
         for name, meta in self._weight_metadata[layer_idx].items():
             dtype = meta["dtype"]
@@ -182,8 +199,16 @@ class LayerwiseOffloadManager:
 
     @torch.compiler.disable
     def release_layer(self, layer_idx: int) -> None:
+        """
+        lightweight release layer weights
+        Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched
+        """
         if not self.enabled or self.device is None:
             return
+
+        # clear prefetch event, since it's useless and needs to be reset
+        self._prefetch_events.pop(layer_idx, None)
+
         if layer_idx <= 0:
             return
 
@@ -259,14 +284,24 @@ class LayerwiseOffloadManager:
 
         def make_pre_hook(i):
             def hook(module, input):
-                self.prefetch_layer(i + 1, non_blocking=True)
+                # wait only for the current layer if it's being prefetched
+                if i == 0:
+                    self.prepare_for_next_req(non_blocking=False)
+                if i in self._prefetch_events:
+                    torch.cuda.current_stream().wait_event(self._prefetch_events[i])
+
+                # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed
+                if i % self.prefetch_size == 0:
+                    for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size):
+                        layer_to_prefetch = j % self.num_layers
+                        self.prefetch_layer(layer_to_prefetch, non_blocking=True)
 
             return hook
 
         def make_post_hook(i):
             def hook(module, input, output):
-                if self.copy_stream is not None:
+                # previous, we wait here, until the copy stream for next layer is finished,
-                    torch.cuda.current_stream().wait_stream(self.copy_stream)
+                # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer
                 self.release_layer(i)
 
             return hook
@@ -292,7 +327,7 @@ class OffloadableDiTMixin:
 
     # the list of names of a DiT's layers/blocks
     layer_names: List[str]
-    layerwise_offload_managers: list[LayerwiseOffloadManager] | None = None
+    layerwise_offload_managers: list[LayerwiseOffloadManager] = []
 
     def configure_layerwise_offload(self, server_args: ServerArgs):
         self.layerwise_offload_managers = []
@@ -303,12 +338,20 @@ class OffloadableDiTMixin:
                 continue
 
             num_layers = len(module_list)
+            if server_args.dit_offload_prefetch_size < 1.0:
+                prefetch_size = 1 + int(
+                    round(server_args.dit_offload_prefetch_size * (num_layers - 1))
+                )
+            else:
+                prefetch_size = int(server_args.dit_offload_prefetch_size)
+
             manager = LayerwiseOffloadManager(
                 model=self,
                 layers_attr_str=layer_name,
                 num_layers=num_layers,
                 enabled=True,
                 pin_cpu_memory=server_args.pin_cpu_memory,
+                prefetch_size=prefetch_size,
             )
             self.layerwise_offload_managers.append(manager)
 
@@ -316,11 +359,11 @@ class OffloadableDiTMixin:
             f"Enabled layerwise offload for {self.__class__.__name__} on modules: {self.layer_names}"
         )
 
-    def prepare_for_next_denoise(self):
+    def prepare_for_next_req(self):
         if self.layerwise_offload_managers is None:
             return
         for manager in self.layerwise_offload_managers:
-            manager.prepare_for_next_denoise(non_blocking=True)
+            manager.prepare_for_next_req(non_blocking=True)
 
     def disable_offload(self) -> None:
         """Disable layerwise offload: load all layers to GPU and remove hooks."""
@@ -338,7 +381,5 @@ class OffloadableDiTMixin:
         for manager in self.layerwise_offload_managers:
             if manager.enabled:
                 manager.sync_all_layers_to_cpu()
-                for layer_idx in list(manager._gpu_layers):
+                manager.release_all()
-                    if layer_idx > 0:
-                        manager.release_layer(layer_idx)
                 manager.register_forward_hooks()

python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py

@@ -53,19 +53,22 @@ def _openai_client(port: int) -> OpenAI:
 
 
 def _build_server_extra_args(case: DiffusionTestCase) -> str:
+    server_args = case.server_args
     a = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
-    a += f" --num-gpus {case.server_args.num_gpus}"
+    a += f" --num-gpus {server_args.num_gpus}"
-    if case.server_args.tp_size is not None:
+    if server_args.tp_size is not None:
-        a += f" --tp-size {case.server_args.tp_size}"
+        a += f" --tp-size {server_args.tp_size}"
-    if case.server_args.ulysses_degree is not None:
+    if server_args.ulysses_degree is not None:
-        a += f" --ulysses-degree {case.server_args.ulysses_degree}"
+        a += f" --ulysses-degree {server_args.ulysses_degree}"
-    if case.server_args.dit_layerwise_offload:
+    if server_args.dit_layerwise_offload:
         a += " --dit-layerwise-offload true"
-    if case.server_args.ring_degree is not None:
+    if server_args.dit_offload_prefetch_size:
-        a += f" --ring-degree {case.server_args.ring_degree}"
+        a += f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}"
-    if case.server_args.lora_path:
+    if server_args.ring_degree is not None:
-        a += f" --lora-path {case.server_args.lora_path}"
+        a += f" --ring-degree {server_args.ring_degree}"
-    if case.server_args.enable_warmup:
+    if server_args.lora_path:
+        a += f" --lora-path {server_args.lora_path}"
+    if server_args.enable_warmup:
         a += " --enable-warmup"
     return a
 

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -76,6 +76,11 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     if server_args.dit_layerwise_offload:
         extra_args += f" --dit-layerwise-offload true"
 
+    if server_args.dit_offload_prefetch_size:
+        extra_args += (
+            f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}"
+        )
+
     if server_args.text_encoder_cpu_offload:
         extra_args += f" --text-encoder-cpu-offload"
 

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -164,6 +164,7 @@ class DiffusionServerArgs:
     enable_warmup: bool = False
 
     dit_layerwise_offload: bool = False
+    dit_offload_prefetch_size: int | float | None = None
     enable_cache_dit: bool = False
     text_encoder_cpu_offload: bool = False
 
@@ -369,6 +370,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
             model_path="black-forest-labs/FLUX.2-dev",
             modality="image",
             dit_layerwise_offload=True,
+            dit_offload_prefetch_size=2,
         ),
         T2I_sampling_params,
     ),