[diffusion] Skip automatic Wan/MOVA DiT layerwise offload on high-end GPUs (#21248)

python/sglang/multimodal_gen/runtime/server_args.py

@@ -36,6 +36,10 @@ from sglang.multimodal_gen.runtime.utils.common import (
     is_valid_ipv6_address,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    CYAN,
+    GREEN,
+    RED,
+    RESET,
     _sanitize_for_logging,
     configure_logger,
     init_logger,
@@ -49,6 +53,18 @@ from sglang.multimodal_gen.utils import (
 
 logger = init_logger(__name__)
 
+# Derived from single-H200 benchmarking (~140.4 GiB total) at the maximum
+# supported 720p workloads with dit_layerwise_offload=False and
+# num_inference_steps=1:
+# - Wan-AI/Wan2.2-T2V-A14B-Diffusers, 1280x720, 81 frames:
+#   peak_reserved=108076 MB (~105.5 GiB), peak_allocated=97665 MB (~95.4 GiB)
+# - OpenMOSS-Team/MOVA-720p, 1280x720, 193 frames:
+#   peak_reserved=130264 MB (~127.2 GiB), peak_allocated=108819 MB (~106.3 GiB)
+# Also, on H200, enabling dit_layerwise_offload regressed latency noticeably on
+# our validated Wan/MOVA workloads, so use a 130 GiB cutoff to keep H200-class
+# GPUs on the faster no-offload default while preserving some headroom.
+WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
+
 
 class Backend(str, Enum):
     """
@@ -434,15 +450,33 @@ class ServerArgs:
         if not envs.SGLANG_CACHE_DIT_ENABLED:
             pipeline_name_lower = self.pipeline_config.__class__.__name__.lower()
             if (
-                ("wan" in pipeline_name_lower or "mova" in pipeline_name_lower)
+                "wan" in pipeline_name_lower or "mova" in pipeline_name_lower
-                and self.dit_layerwise_offload is None
+            ) and self.dit_layerwise_offload is None:
-                and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
+                auto_enable_layerwise_offload = (
-            ):
+                    current_platform.enable_dit_layerwise_offload_for_wan_by_default()
-                logger.info(
-                    f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} "
-                    "for low memory and performance balance"
                 )
-                self.dit_layerwise_offload = True
+                if auto_enable_layerwise_offload and current_platform.is_cuda():
+                    device_total_memory_gb = (
+                        current_platform.get_device_total_memory() / BYTES_PER_GB
+                    )
+                    if (
+                        device_total_memory_gb
+                        >= WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB
+                    ):
+                        logger.info(
+                            "Skipping automatic dit_layerwise_offload for %s on a high-memory CUDA GPU (e.g. H200/B200/B300-class, %.2f GiB total)",
+                            self.pipeline_config.__class__.__name__,
+                            device_total_memory_gb,
+                        )
+                        auto_enable_layerwise_offload = False
+                        self.dit_layerwise_offload = False
+
+                if auto_enable_layerwise_offload:
+                    logger.info(
+                        f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} "
+                        "for low memory and performance balance"
+                    )
+                    self.dit_layerwise_offload = True
 
     def _adjust_autocast(self):
         if self.disable_autocast is None:
@@ -1057,6 +1091,18 @@ class ServerArgs:
                     "Please disable either --dit-layerwise-offload or SGLANG_CACHE_DIT_ENABLED."
                 )
 
+            logger.warning(
+                "dit_layerwise_offload is enabled: %slower GPU memory usage%s, but %smay reduce throughput or increase latency%s. "
+                "%sIf you are using multi-GPU deployment and already have enough memory headroom, prefer keeping dit_layerwise_offload disabled.%s "
+                "Please tune this based on your memory headroom and performance target.",
+                GREEN,
+                RESET,
+                RED,
+                RESET,
+                CYAN,
+                RESET,
+            )
+
     def _validate_parallelism(self):
         if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0:
             raise ValueError(