diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_video_moe.py b/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_video_moe.py index 08c90dc44..4f7bfd400 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_video_moe.py +++ b/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] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py b/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py index fc31a6226..2585be683 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/lingbot_world.py +++ b/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: diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index c73d9243b..ca37a4765 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py b/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py index 851dfb8c5..ac5f69009 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py +++ b/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( 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 691f13685..5826c2789 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/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, diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index 893851e69..f3fc5d087 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/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 ============= diff --git a/python/sglang/multimodal_gen/runtime/platforms/cpu.py b/python/sglang/multimodal_gen/runtime/platforms/cpu.py index e6f46b115..bc493f3d4 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cpu.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index a0be9394e..a62478f98 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/npu.py b/python/sglang/multimodal_gen/runtime/platforms/npu.py index 0d765b66d..27034bb69 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/npu.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 92e02164d..2dac954fb 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/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 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 d2228463d..0e6a993d3 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index c0c97da53..0d5b6e490 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py index 39f847002..c6d145515 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py +++ b/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: 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 da595e168..06ba57208 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/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):