[Diffusion] Support layerwise offload for mova (#18272)

This commit is contained in:
Xiaoyu Zhang
2026-02-05 13:16:07 +08:00
committed by GitHub
parent c910829708
commit dff3ba202a
@@ -642,7 +642,7 @@ class ServerArgs:
"--dit-layerwise-offload", "--dit-layerwise-offload",
action=StoreBoolean, action=StoreBoolean,
default=ServerArgs.dit_layerwise_offload, default=ServerArgs.dit_layerwise_offload,
help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan). " help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan, MOVA). "
"Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.", "Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.",
) )
parser.add_argument( parser.add_argument(
@@ -998,13 +998,14 @@ class ServerArgs:
if not envs.SGLANG_CACHE_DIT_ENABLED: if not envs.SGLANG_CACHE_DIT_ENABLED:
# TODO: need a better way to tell this # TODO: need a better way to tell this
pipeline_name_lower = self.pipeline_config.__class__.__name__.lower()
if ( if (
"wan" in self.pipeline_config.__class__.__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() and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
): ):
logger.info( logger.info(
"Automatically enable dit_layerwise_offload for Wan for best performance" f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} for low memory and performance balance"
) )
self.dit_layerwise_offload = True self.dit_layerwise_offload = True