[diffusion] auto-keep video DiT resident on high-memory GPUs (#38001)

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

@@ -13,6 +13,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -95,6 +98,13 @@ class HeliosT2VConfig(PipelineConfig):
         self.vae_config.load_encoder = False
         self.vae_config.load_decoder = True
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            dit_layerwise_offload_modes=("memory",),
+            keep_resident_min_available_gb=120,
+            keep_resident_components=("dit", "vae"),
+        )
+
 
 @dataclass
 class HeliosMidConfig(HeliosT2VConfig):

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

@@ -165,6 +165,8 @@ class SanaWMPipelineConfig(PipelineConfig):
             # Conservative auto-FSDP gate for the 720p world-model path. Users
             # can still force FSDP explicitly on smaller cards.
             fsdp_auto_min_available_memory_gb=60,
+            keep_resident_min_available_gb=120,
+            keep_resident_components=("dit", "vae"),
         )
 
     # --- Latent shape ---

python/sglang/multimodal_gen/runtime/server_args/auto_tune.py

@@ -175,6 +175,8 @@ class ServerArgsAutoTuner:
                 args.dit_cpu_offload
                 and "dit" in components
                 and args.explicit_residency_mode("transformer") is None
+                and not args.is_arg_explicitly_set("dit_cpu_offload")
+                and not args.is_arg_explicitly_set("dit_layerwise_offload")
                 and not explicit_dit_layerwise
             ):
                 args.dit_cpu_offload = False

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -17,6 +17,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     PipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+from sglang.multimodal_gen.configs.pipeline_configs.helios import (
+    HeliosDistilledConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import FastHunyuanConfig
 from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
     LingBotWorldCausalDMDConfig,
@@ -1595,6 +1598,13 @@ class TestOffloadDefaults(unittest.TestCase):
 
         self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60)
         self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",))
+        self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 120)
+        self.assertEqual(sana_wm_deployment.keep_resident_components, ("dit", "vae"))
+
+        helios_deployment = HeliosDistilledConfig().get_model_deployment_config()
+        self.assertEqual(helios_deployment.keep_resident_min_available_gb, 120)
+        self.assertEqual(helios_deployment.keep_resident_components, ("dit", "vae"))
+        self.assertEqual(helios_deployment.dit_layerwise_offload_modes, ("memory",))
 
         fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config()
         self.assertEqual(fast_hunyuan_deployment.keep_resident_min_available_gb, 60)