[diffusion] fix: fix model-driven dit layerwise offload auto policy (#34401)

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

@@ -58,7 +58,7 @@ class LingBotVideoMoEPipelineConfig(PipelineConfig):
         self.vae_config.load_decoder = True
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
-        return ModelDeploymentConfig(auto_dit_layerwise_offload=True)
+        return ModelDeploymentConfig(dit_layerwise_offload_modes=("memory",))
 
     def get_pos_prompt_embeds(self, batch):
         return batch.prompt_embeds[0]

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

@@ -14,6 +14,9 @@ import torch
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.models import DiTConfig
 from sglang.multimodal_gen.configs.models.dits import LingBotWorldVideoConfig
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.wan import Wan2_2_I2V_A14B_Config
 from sglang.multimodal_gen.runtime.realtime.session import (
     BaseRealtimeState,
@@ -284,6 +287,9 @@ class LingBotWorldI2VConfig(Wan2_2_I2V_A14B_Config):
         default_factory=lambda: (lingbot_prompt_clean,)
     )
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(dit_layerwise_offload_modes=("memory",))
+
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
         kwargs = super().prepare_pos_cond_kwargs(batch, device, rotary_emb, dtype)
         if batch.c2ws_plucker_emb is not None:

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

@@ -11,9 +11,8 @@ OffloadComponentName = Literal["dit", "text_encoder", "image_encoder", "vae"]
 
 @dataclass(frozen=True)
 class ModelDeploymentConfig:
-    auto_dit_layerwise_offload: bool = False
-    # if the available memory is bigger than this value, keep dit resident instead of apply layerwise-offload
-    auto_dit_layerwise_offload_high_memory_disable_gb: float | None = None
+    dit_layerwise_offload_modes: tuple[Literal["auto", "memory"], ...] = ()
+    auto_dit_offload_prefetch_size: float | None = None
     keep_resident_min_available_gb: float | None = None
     # only vae -- it is tiny so keeping it resident barely shifts memory; large
     # encoders stay offloaded and dit placement stays with the FSDP/dit-layerwise

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

@@ -58,8 +58,9 @@ class MOVAPipelineConfig(PipelineConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
-            auto_dit_layerwise_offload_high_memory_disable_gb=130,
+            dit_layerwise_offload_modes=("auto", "memory"),
+            keep_resident_min_available_gb=130,
+            keep_resident_components=("dit", "vae"),
         )
 
     def _center_crop_and_resize(

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

@@ -161,7 +161,7 @@ class SanaWMPipelineConfig(PipelineConfig):
     # --- Deployment ---
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
             # 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,
@@ -328,7 +328,7 @@ class SanaWMRealtimeConfig(SanaWMPipelineConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
             keep_resident_min_available_gb=120,
             keep_resident_components=("dit",),
             auto_enable_cfg_parallel=False,

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

@@ -98,7 +98,7 @@ class WanT2V480PConfig(PipelineConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
         )
 
     def expand_conditioning_to_sample_batch(self, batch):
@@ -138,7 +138,7 @@ class TurboWanT2V1_3B480PConfig(TurboWanT2V480PConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
             keep_resident_min_available_gb=60,
             keep_resident_components=(
                 "text_encoder",
@@ -184,7 +184,7 @@ class WanI2V480PConfig(WanT2V480PConfig, WanI2VCommonConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
         )
 
 
@@ -225,7 +225,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
-            auto_dit_layerwise_offload=True,
+            dit_layerwise_offload_modes=("memory",),
             keep_resident_min_available_gb=60,
             keep_resident_components=(
                 "text_encoder",
@@ -277,6 +277,12 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
         self.dit_config.boundary_ratio = self.boundary_ratio
         self.dit_config.torch_compile_mode = "default"
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            dit_layerwise_offload_modes=("auto", "memory"),
+            auto_dit_offload_prefetch_size=2,
+        )
+
 
 @dataclass
 class Wan2_2_I2V_A14B_Config(WanI2V720PConfig):
@@ -288,6 +294,12 @@ class Wan2_2_I2V_A14B_Config(WanI2V720PConfig):
         super().__post_init__()
         self.dit_config.boundary_ratio = self.boundary_ratio
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            dit_layerwise_offload_modes=("auto", "memory"),
+            auto_dit_offload_prefetch_size=2,
+        )
+
 
 # =============================================
 # ============= Causal Self-Forcing =============

python/sglang/multimodal_gen/runtime/platforms/cpu.py

@@ -118,6 +118,6 @@ class CpuPlatform(Platform):
         return "sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator.CpuCommunicator"
 
     @classmethod
-    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
-        """Whether to enable DIT layerwise offload by default on the current platform."""
+    def enable_dit_layerwise_offload_by_default(cls) -> bool:
+        """Whether automatic DiT layerwise offload is enabled on this platform."""
         return False

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -413,8 +413,8 @@ class Platform:
         return CpuArchEnum.UNSPECIFIED
 
     @classmethod
-    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
-        """Whether to enable DIT layerwise offload by default on the current platform."""
+    def enable_dit_layerwise_offload_by_default(cls) -> bool:
+        """Whether automatic DiT layerwise offload is enabled on this platform."""
         return True
 
     @classmethod

python/sglang/multimodal_gen/runtime/platforms/npu.py

@@ -189,6 +189,6 @@ class NPUPlatformBase(Platform):
         return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
 
     @classmethod
-    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
-        """The performance of the layerwise_offload feature depends on the device's memory size and the memory size occupied by the model. Use --dit-layerwise-offload True if it suitable for your case."""
+    def enable_dit_layerwise_offload_by_default(cls) -> bool:
+        """Whether automatic DiT layerwise offload is enabled on this platform."""
         return False

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -399,6 +399,6 @@ class RocmPlatform(Platform):
         return patched
 
     @classmethod
-    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
-        """ROCm performs better without DIT layerwise offload on Wan."""
+    def enable_dit_layerwise_offload_by_default(cls) -> bool:
+        """Whether automatic DiT layerwise offload is enabled on this platform."""
         return False

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

@@ -426,7 +426,7 @@ class ServerArgsAutoTuner:
         components = self._filter_high_memory_resident_components(components)
         if self._should_auto_enable_dit_layerwise_offload():
             components.insert(0, LAYERWISE_OFFLOAD_DIT_GROUP)
-            self._set_default_wan_dit_offload_prefetch_size()
+            self._set_default_dit_offload_prefetch_size()
         return components
 
     def _filter_high_memory_resident_components(
@@ -465,52 +465,30 @@ class ServerArgsAutoTuner:
 
     def _should_auto_enable_dit_layerwise_offload(self) -> bool:
         args = self.server_args
-
-        # only for wan for now
-        if not self._is_wan_pipeline_config():
-            return False
-        if not self._deployment_config().auto_dit_layerwise_offload:
+        deployment_config = self._deployment_config()
+        if args.performance_mode not in deployment_config.dit_layerwise_offload_modes:
             return False
 
         if (
             args.pipeline_config.dmd_denoising_steps is not None
-            or not current_platform.enable_dit_layerwise_offload_for_wan_by_default()
+            or not current_platform.enable_dit_layerwise_offload_by_default()
             or envs.SGLANG_CACHE_DIT_ENABLED
             or args.use_fsdp_inference
             or args.is_arg_explicitly_set("dit_cpu_offload")
         ):
             return False
 
-        # memory mode is memory-first: keep the broad Wan DiT layerwise policy
-        # unless a guard above says it conflicts with another placement path
-        if args.performance_mode == "memory":
-            return True
+        return True
 
-        # auto mode is performance-first: profiling only showed clear wins for
-        # Wan2.2 A14B, where coarse DiT CPU offload creates large step spikes
-        return (
-            args.performance_mode == "auto" and self._is_wan2_2_a14b_pipeline_config()
-        )
-
-    def _is_wan2_2_a14b_pipeline_config(self) -> bool:
-        config_name = self.server_args.pipeline_config.__class__.__name__
-        return config_name.startswith("Wan2_2_") and "A14B" in config_name
-
-    def _set_default_wan_dit_offload_prefetch_size(self) -> None:
+    def _set_default_dit_offload_prefetch_size(self) -> None:
         args = self.server_args
+        prefetch_size = self._deployment_config().auto_dit_offload_prefetch_size
         if (
             args.performance_mode == "auto"
-            and self._is_wan2_2_a14b_pipeline_config()
+            and prefetch_size is not None
             and not args.is_arg_explicitly_set("dit_offload_prefetch_size")
         ):
-            # p2 was the fastest stable default in the Wan2.2 A14B sweep
-            args.dit_offload_prefetch_size = 2
-
-    def _is_wan_pipeline_config(self) -> bool:
-        return any(
-            cls.__module__.endswith(".wan")
-            for cls in self.server_args.pipeline_config.__class__.mro()
-        )
+            args.dit_offload_prefetch_size = prefetch_size
 
     def _auto_uses_dit_offload(self) -> bool:
         args = self.server_args

python/sglang/multimodal_gen/test/server/perf_baselines/h100.json

@@ -1229,12 +1229,12 @@
         },
         "wan2_2_ti2v_5b": {
             "stages_ms": {
-                "InputValidationStage": 380.0,
+                "InputValidationStage": 706.07,
                 "TextEncodingStage": 328.44,
                 "LatentPreparationStage": 0.13,
                 "TimestepPreparationStage": 2.32,
                 "DenoisingStage": 15171.58,
-                "DecodingStage": 1527.14
+                "DecodingStage": 1100.28
             },
             "denoise_step_ms": {
                 "0": 183.83,

python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py

@@ -135,9 +135,11 @@ class TestSanaWMPipelineConfig(unittest.TestCase):
         self.assertIs(kwargs["camera_conditions"], camera_conditions)
         self.assertIs(kwargs["chunk_plucker"], chunk_plucker)
 
-    def test_get_model_deployment_config_enables_dit_layerwise_offload(self) -> None:
+    def test_deployment_config_enables_memory_mode_dit_layerwise_offload(
+        self,
+    ) -> None:
         deployment = self.config.get_model_deployment_config()
-        self.assertTrue(deployment.auto_dit_layerwise_offload)
+        self.assertEqual(deployment.dit_layerwise_offload_modes, ("memory",))
         self.assertEqual(deployment.fsdp_auto_min_available_memory_gb, 60)
 
     def test_text_encoder_padding_matches_cfg_concat_contract(self) -> None:

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

@@ -100,7 +100,7 @@ def _mock_cuda_platform(
             side_effect=get_available_gpu_memory,
         ),
         patch(
-            "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_for_wan_by_default",
+            "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_by_default",
             return_value=True,
         ),
     ):
@@ -765,7 +765,7 @@ class TestOffloadDefaults(unittest.TestCase):
                 return_value=True,
             ),
             patch(
-                "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_for_wan_by_default",
+                "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_by_default",
                 return_value=True,
             ),
             patch(
@@ -911,17 +911,21 @@ class TestOffloadDefaults(unittest.TestCase):
         sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config()
 
         self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
-        self.assertFalse(qwen_deployment.auto_dit_layerwise_offload)
+        self.assertEqual(qwen_deployment.dit_layerwise_offload_modes, ())
 
         self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
-        self.assertTrue(wan_deployment.auto_dit_layerwise_offload)
+        self.assertEqual(wan_deployment.dit_layerwise_offload_modes, ("memory",))
 
         self.assertIsNone(mova_deployment.fsdp_auto_min_available_memory_gb)
-        self.assertTrue(mova_deployment.auto_dit_layerwise_offload)
+        self.assertEqual(
+            mova_deployment.dit_layerwise_offload_modes, ("auto", "memory")
+        )
+        self.assertEqual(mova_deployment.keep_resident_min_available_gb, 130)
+        self.assertEqual(mova_deployment.keep_resident_components, ("dit", "vae"))
 
         self.assertEqual(zimage_deployment.fsdp_auto_min_available_memory_gb, 40)
         self.assertTrue(zimage_deployment.fsdp_auto_requires_cfg)
-        self.assertFalse(zimage_deployment.auto_dit_layerwise_offload)
+        self.assertEqual(zimage_deployment.dit_layerwise_offload_modes, ())
 
         self.assertEqual(ltx_deployment.keep_resident_min_available_gb, 70)
         self.assertEqual(ltx_deployment.keep_resident_components, ("dit",))
@@ -939,7 +943,7 @@ class TestOffloadDefaults(unittest.TestCase):
         )
 
         self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60)
-        self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload)
+        self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",))
 
         # fasthunyuan no longer pins 150gb -- falls back to the global video default
         fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config()
@@ -1210,7 +1214,7 @@ class TestOffloadDefaults(unittest.TestCase):
             ["text_encoder", "image_encoder", "vae"],
         )
 
-    def test_auto_mova_layerwise_offload_does_not_implicitly_add_dit(self):
+    def test_auto_mova_layerwise_offload_adds_dit_below_memory_threshold(self):
         args = self._from_dict_with_pipeline_config(
             MOVAPipelineConfig(),
             kwargs={
@@ -1222,7 +1226,37 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertTrue(args.dit_cpu_offload)
         self.assertEqual(
             args.layerwise_offload_components,
-            ["text_encoder", "image_encoder", "vae"],
+            ["dit", "text_encoder", "image_encoder", "vae"],
+        )
+
+    def test_auto_mova_keeps_dit_resident_at_memory_threshold(self):
+        args = self._from_dict_with_pipeline_config(
+            MOVAPipelineConfig(),
+            memory_gb=140,
+            kwargs={
+                "model_path": "OpenMOSS-Team/MOVA-360p",
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertEqual(
+            args.layerwise_offload_components,
+            ["text_encoder", "image_encoder"],
+        )
+
+    def test_memory_sana_wm_layerwise_offload_adds_dit(self):
+        args = self._from_dict_with_pipeline_config(
+            SanaWMPipelineConfig(),
+            kwargs={
+                "model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
+                "performance_mode": "memory",
+            },
+        )
+
+        self.assertEqual(
+            args.layerwise_offload_components,
+            ["dit", "text_encoder", "image_encoder", "vae"],
         )
 
     def test_auto_fastwan_layerwise_offload_does_not_implicitly_add_dit(self):