[diffusion] feat: allow --dit-cpu-offload with --dit-layerwise-offload (#26925)

Co-authored-by: Yiqi Yang <yiqi.yang@kiwiar.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/server_args.py

@@ -893,21 +893,18 @@ class ServerArgs(DisaggArgsMixin):
 
         if explicitly_set_component_names is not None:
             self.layerwise_offload_components = explicitly_set_component_names
-            self._disable_cpu_offload_for_layerwise_components(
+            self._disable_non_dit_cpu_offload_for_layerwise_components(
                 explicitly_set_component_names
             )
             return
 
-    def _disable_cpu_offload_for_layerwise_components(
+    def _disable_non_dit_cpu_offload_for_layerwise_components(
         self, component_names: list[str]
     ) -> None:
-        # Layerwise offload owns H2D/D2H for selected component weights.
+        # non-DiT layerwise offload replaces the corresponding component-level CPU offload
         flag_names = cpu_offload_flags_for_layerwise_components(component_names)
         disabled_flag_names: list[str] = []
 
-        if "dit_cpu_offload" in flag_names and self.dit_cpu_offload is not False:
-            self.dit_cpu_offload = False
-            disabled_flag_names.append("dit_cpu_offload")
         if (
             "text_encoder_cpu_offload" in flag_names
             and self.text_encoder_cpu_offload is not False
@@ -1254,7 +1251,9 @@ class ServerArgs(DisaggArgsMixin):
             default=ServerArgs.dit_layerwise_offload,
             help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. "
             "It selects only the DiT layerwise group. Cannot be used together with cache-dit "
-            "(SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.",
+            "(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. May be combined with "
+            "--dit-cpu-offload, in which case DiT weights stay on host memory and only the "
+            "layers needed for the current step are brought on-device (lowest peak GPU memory).",
         )
         parser.add_argument(
             "--layerwise-offload-components",
@@ -1836,20 +1835,14 @@ class ServerArgs(DisaggArgsMixin):
             if self.dit_offload_prefetch_size < 0.0:
                 raise ValueError("dit_offload_prefetch_size must be non-negative")
 
-            should_disable_dit_cpu_offload = self.is_dit_layerwise_offload_selected
+            is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected
-            if self.use_fsdp_inference and should_disable_dit_cpu_offload:
+            if self.use_fsdp_inference and is_dit_layerwise_offload_selected:
                 logger.warning(
                     "layerwise offload is selected for DiT components, automatically disabling use_fsdp_inference."
                 )
                 self.use_fsdp_inference = False
 
-            if should_disable_dit_cpu_offload and self.dit_cpu_offload is not False:
+            if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected:
-                logger.warning(
-                    "layerwise offload is selected for DiT components, automatically disabling dit_cpu_offload."
-                )
-                self.dit_cpu_offload = False
-
-            if envs.SGLANG_CACHE_DIT_ENABLED and should_disable_dit_cpu_offload:
                 raise ValueError(
                     "DiT layerwise offload cannot be enabled together with cache-dit. "
                     "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, "

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

@@ -559,7 +559,7 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertFalse(args.text_encoder_cpu_offload)
         self.assertEqual(args.layerwise_offload_components, ["image_encoder", "vae"])
 
-    def test_layerwise_components_disable_matching_cpu_offloads(self):
+    def test_layerwise_components_disable_matching_non_dit_cpu_offloads(self):
         args = self._from_dict_with_task_type(
             ModelTaskType.T2V,
             memory_gb=16,
@@ -580,11 +580,34 @@ class TestOffloadDefaults(unittest.TestCase):
         args._adjust_layerwise_offload_components()
 
         self.assertTrue(args.layerwise_offload_components)
-        self.assertFalse(args.dit_cpu_offload)
+        # dit_cpu_offload is complementary to DiT layerwise offload (keeps
+        # weights off-device during load), so it must be preserved here.
+        self.assertTrue(args.dit_cpu_offload)
         self.assertFalse(args.text_encoder_cpu_offload)
         self.assertFalse(args.image_encoder_cpu_offload)
         self.assertFalse(args.vae_cpu_offload)
 
+    def test_dit_layerwise_offload_preserves_dit_cpu_offload(self):
+        """Combining --dit-cpu-offload with --dit-layerwise-offload must keep both on.
+
+        dit_cpu_offload controls initial residency (host memory), while
+        dit_layerwise_offload only swaps layers on/off device at inference.
+        Force-disabling dit_cpu_offload here would push the full DiT to GPU at
+        load time and OOM low-VRAM cards.
+        """
+        args = self._from_dict_with_task_type(
+            ModelTaskType.T2I,
+            memory_gb=32,
+            kwargs={
+                "dit_cpu_offload": True,
+                "dit_layerwise_offload": True,
+            },
+        )
+
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertTrue(args.dit_layerwise_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()
@@ -704,7 +727,10 @@ class TestOffloadDefaults(unittest.TestCase):
 
                 self.assertTrue(args.layerwise_offload_components)
                 self.assertFalse(args.use_fsdp_inference)
-                self.assertFalse(args.dit_cpu_offload)
+                # dit_cpu_offload is complementary to DiT layerwise offload:
+                # layerwise only moves layers on/off device at runtime, while
+                # dit_cpu_offload keeps the initial weights on host memory.
+                self.assertTrue(args.dit_cpu_offload)
                 self.assertFalse(args.text_encoder_cpu_offload)
                 self.assertFalse(args.image_encoder_cpu_offload)
                 self.assertEqual(args.dit_offload_prefetch_size, 2)
@@ -744,7 +770,7 @@ class TestOffloadDefaults(unittest.TestCase):
 
         self.assertTrue(args.layerwise_offload_components)
         self.assertFalse(args.use_fsdp_inference)
-        self.assertFalse(args.dit_cpu_offload)
+        self.assertTrue(args.dit_cpu_offload)
         self.assertFalse(args.text_encoder_cpu_offload)
         self.assertFalse(args.image_encoder_cpu_offload)
         self.assertEqual(
@@ -839,7 +865,9 @@ class TestOffloadDefaults(unittest.TestCase):
             },
         )
 
-        self.assertFalse(args.dit_cpu_offload)
+        # dit_cpu_offload defaults to True from _adjust_offload and is now
+        # preserved alongside DiT layerwise offload (the two are complementary).
+        self.assertTrue(args.dit_cpu_offload)
         self.assertEqual(args.layerwise_offload_components, ["dit"])
 
     def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self):
@@ -875,7 +903,7 @@ class TestOffloadDefaults(unittest.TestCase):
         )
 
         self.assertFalse(args.use_fsdp_inference)
-        self.assertFalse(args.dit_cpu_offload)
+        self.assertTrue(args.dit_cpu_offload)
         self.assertTrue(args.layerwise_offload_components)
         self.assertTrue(args.text_encoder_cpu_offload)
         self.assertTrue(args.image_encoder_cpu_offload)
@@ -1124,7 +1152,7 @@ class TestOffloadDefaults(unittest.TestCase):
 
         self.assertFalse(args.use_fsdp_inference)
         self.assertTrue(args.layerwise_offload_components)
-        self.assertFalse(args.dit_cpu_offload)
+        self.assertTrue(args.dit_cpu_offload)
         self.assertFalse(args.text_encoder_cpu_offload)
         self.assertFalse(args.image_encoder_cpu_offload)
         self.assertEqual(