From badcd0289623ca18bc56e154bbdc961b789758f1 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 7 Jan 2026 10:22:08 +0800 Subject: [PATCH] [diffusion] chore: automatically enable dit_layerwise_offload for Wan (#16499) --- python/sglang/multimodal_gen/runtime/server_args.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 9095995a9..95b45bff7 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -17,6 +17,7 @@ from dataclasses import field from enum import Enum from typing import Any, Optional +from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig, STA_Mode from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, @@ -923,10 +924,19 @@ class ServerArgs: def check_server_args(self) -> None: """Validate inference arguments for consistency""" + # layerwise offload if current_platform.is_mps(): self.use_fsdp_inference = False self.dit_layerwise_offload = False + if not envs.SGLANG_CACHE_DIT_ENABLED: + # TODO: need a better way to tell this + if "wan" in self.pipeline_config.__class__.__name__.lower(): + logger.info( + "Automatically enable dit_layerwise_offload for Wan for best performance" + ) + self.dit_layerwise_offload = True + if self.dit_layerwise_offload: if self.use_fsdp_inference: logger.warning( @@ -938,7 +948,7 @@ class ServerArgs: "dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload." ) self.dit_cpu_offload = False - if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true": + if envs.SGLANG_CACHE_DIT_ENABLED: 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, "