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

This commit is contained in:
Xiaoyu Zhang
2026-03-25 18:45:30 +08:00
committed by GitHub
parent 3d2a61cbf6
commit e4ad10520b
@@ -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)
and self.dit_layerwise_offload is None
and 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"
"wan" in pipeline_name_lower or "mova" in pipeline_name_lower
) and self.dit_layerwise_offload is None:
auto_enable_layerwise_offload = (
current_platform.enable_dit_layerwise_offload_for_wan_by_default()
)
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(