[Diffusion] Make auto residency decisions component-scoped (#34615)

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 (
+        if args.performance_mode != "auto" or current_platform.is_cpu():
-            args.performance_mode != "auto"
-            or self._explicit_memory_policy
-            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")

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
+        # The layerwise filter still drops VAE (kept resident); encoders stay
-        # filter still drops vae (kept resident); encoders stay offloaded
+        # offloaded.
         self.assertEqual(
             args.layerwise_offload_components,
             ["text_encoder", "image_encoder"],