[Diffusion] Keep FastHunyuan VAE resident on high-memory GPUs (#28773)

python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py

@@ -20,6 +20,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     PipelineConfig,
     TextConditioningOutput,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 
 PROMPT_TEMPLATE_ENCODE_VIDEO = (
     "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
@@ -158,3 +161,9 @@ class FastHunyuanConfig(HunyuanConfig):
 
     # No need to re-specify guidance_scale or embedded_cfg_scale as they
     # already have the desired values from HunyuanConfig
+
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_disable_component_offload_min_available_memory_gb=150,
+            auto_disable_component_offload_components=("vae",),
+        )

python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py

@@ -6,7 +6,7 @@ ModelDeploymentConfig provides model-specific config on how to deploy a model op
 from dataclasses import dataclass
 from typing import Literal
 
-OffloadComponentName = Literal["dit", "text_encoder", "image_encoder"]
+OffloadComponentName = Literal["dit", "text_encoder", "image_encoder", "vae"]
 
 
 @dataclass(frozen=True)

python/sglang/multimodal_gen/runtime/server_args_auto_tune.py

@@ -107,6 +107,20 @@ class ServerArgsAutoTuner:
             components = (
                 self._deployment_config().auto_disable_component_offload_components
             )
+            if (
+                args.layerwise_offload_components is not None
+                and not args.is_arg_explicitly_set("layerwise_offload_components")
+            ):
+                layerwise_components = [
+                    component_name
+                    for component_name in args.layerwise_offload_components
+                    if component_name not in components
+                ]
+                if layerwise_components != args.layerwise_offload_components:
+                    args.layerwise_offload_components = layerwise_components or None
+                    changed.append(
+                        f"layerwise_offload_components={args.layerwise_offload_components}"
+                    )
             if (
                 args.dit_cpu_offload
                 and "dit" in components