From 82f7afb881064250739747113227186d553464e8 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 13 Aug 2026 23:20:55 +0800 Subject: [PATCH] [Diffusion] Make auto residency decisions component-scoped (#34615) --- .../runtime/server_args/auto_tune.py | 35 ++++++-- .../test/unit/test_server_args.py | 82 ++++++++++++++++++- 2 files changed, 109 insertions(+), 8 deletions(-) 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 62e2a82d1..2ea98df0f 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -11,10 +11,12 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config impo ModelDeploymentConfig, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + LAYERWISE_OFFLOAD_ALL_COMPONENTS, LAYERWISE_OFFLOAD_DIT_GROUP, LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP, LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP, LAYERWISE_OFFLOAD_VAE_GROUP, + normalize_layerwise_offload_components, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -122,13 +124,31 @@ class ServerArgsAutoTuner: def maybe_adjust_auto_component_residency_after_offload(self) -> None: args = self.server_args - if ( - args.performance_mode != "auto" - or self._explicit_memory_policy - or current_platform.is_cpu() - ): + if args.performance_mode != "auto" or current_platform.is_cpu(): return + # Explicitness is component-scoped below. For example, explicitly + # disabling DiT layerwise offload must not freeze an unrelated, + # implicit ``dit_cpu_offload=True`` default on a high-memory GPU. + # Each mutation below already preserves its own explicit CLI flag. + + explicit_cpu_components = args.is_arg_explicitly_set("cpu_offload_components") + explicit_layerwise_components = ( + normalize_layerwise_offload_components(args.layerwise_offload_components) + if args.is_arg_explicitly_set("layerwise_offload_components") + else None + ) + explicit_dit_layerwise = bool( + args.is_arg_explicitly_set("dit_layerwise_offload") + and args.dit_layerwise_offload + ) or bool( + explicit_layerwise_components + and ( + LAYERWISE_OFFLOAD_DIT_GROUP in explicit_layerwise_components + or LAYERWISE_OFFLOAD_ALL_COMPONENTS in explicit_layerwise_components + ) + ) + min_available_gb = self._get_min_available_device_memory_gb() deployment_config = self._deployment_config() disable_threshold_gb = self._resolve_keep_resident_min_available_gb( @@ -161,6 +181,8 @@ class ServerArgsAutoTuner: args.dit_cpu_offload and "dit" in components and not args.is_arg_explicitly_set("dit_cpu_offload") + and not explicit_cpu_components + and not explicit_dit_layerwise ): args.dit_cpu_offload = False changed.append("dit_cpu_offload=False") @@ -168,6 +190,7 @@ class ServerArgsAutoTuner: args.text_encoder_cpu_offload and LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in components and not args.is_arg_explicitly_set("text_encoder_cpu_offload") + and not explicit_cpu_components ): args.text_encoder_cpu_offload = False changed.append("text_encoder_cpu_offload=False") @@ -175,6 +198,7 @@ class ServerArgsAutoTuner: args.image_encoder_cpu_offload and LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in components and not args.is_arg_explicitly_set("image_encoder_cpu_offload") + and not explicit_cpu_components ): args.image_encoder_cpu_offload = False changed.append("image_encoder_cpu_offload=False") @@ -182,6 +206,7 @@ class ServerArgsAutoTuner: args.vae_cpu_offload and LAYERWISE_OFFLOAD_VAE_GROUP in components and not args.is_arg_explicitly_set("vae_cpu_offload") + and not explicit_cpu_components ): args.vae_cpu_offload = False changed.append("vae_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 1b61d2562..ac67027e1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -998,6 +998,80 @@ class TestOffloadDefaults(unittest.TestCase): self.assertTrue(args.dit_layerwise_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) + def test_explicit_layerwise_false_keeps_independent_auto_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "dit_layerwise_offload": False, + }, + ) + + self.assertFalse(args.dit_layerwise_offload) + self.assertFalse(args.dit_cpu_offload) + + def test_explicit_dit_cpu_offload_is_preserved_by_auto_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "dit_layerwise_offload": False, + "dit_cpu_offload": True, + }, + ) + + self.assertFalse(args.dit_layerwise_offload) + self.assertTrue(args.dit_cpu_offload) + + def test_explicit_layerwise_true_preserves_initial_dit_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "dit_layerwise_offload": True, + }, + ) + + self.assertTrue(args.dit_layerwise_offload) + self.assertTrue(args.dit_cpu_offload) + + def test_explicit_vae_cpu_offload_is_preserved_by_auto_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "dit_layerwise_offload": False, + "vae_cpu_offload": True, + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.vae_cpu_offload) + + def test_explicit_cpu_offload_components_are_preserved_by_auto_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "cpu_offload_components": ["dit", "vae"], + }, + ) + + self.assertTrue(args.dit_cpu_offload) + self.assertTrue(args.vae_cpu_offload) + + def test_explicit_dit_layerwise_component_preserves_initial_residency(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "performance_mode": "auto", + "layerwise_offload_components": ["dit"], + }, + ) + + self.assertTrue(args.dit_cpu_offload) + self.assertEqual(args.layerwise_offload_components, ["dit"]) + def test_pipeline_configs_declare_auto_tune_hints(self): qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config() wan_deployment = WanT2V480PConfig().get_model_deployment_config() @@ -1724,10 +1798,12 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - self.assertTrue(args.dit_cpu_offload) + # Explicit FSDP selection must not freeze unrelated, implicit DiT + # residency decisions on a high-memory GPU. + self.assertFalse(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) - # explicit use_fsdp_inference skips the residency pass, but the layerwise - # filter still drops vae (kept resident); encoders stay offloaded + # The layerwise filter still drops VAE (kept resident); encoders stay + # offloaded. self.assertEqual( args.layerwise_offload_components, ["text_encoder", "image_encoder"],