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

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: