diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py b/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py index fcd0ebf8d..4e02c43c7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py +++ b/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): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index 5826c2789..b0e013f99 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/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 --- diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 350c27b45..099762774 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 470516511..5bd23f436 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/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)