[diffusion] Auto residency: per-pipeline containment fixes and CI baselines (4/4)

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

@@ -14,9 +14,9 @@ class ModelDeploymentConfig:
     dit_layerwise_offload_modes: tuple[Literal["auto", "memory"], ...] = ()
     auto_dit_offload_prefetch_size: float | None = None
     keep_resident_min_available_gb: float | None = None
-    # Per-model resident defaults. Auto mode additionally keeps an image DiT
-    # resident above the image workload memory threshold; video DiT placement
-    # stays with the model's FSDP/layerwise policy.
+    # Fallback residency for offline/no-warmup deployments. Eligible server
+    # warmup paths start load-safe and replace these coarse thresholds with a
+    # measured per-phase placement plan before reporting ready.
     keep_resident_components: tuple[OffloadComponentName, ...] = ("vae",)
     fsdp_auto_min_available_memory_gb: float | None = None
     fsdp_auto_requires_cfg: bool = True

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

@@ -165,11 +165,15 @@ def __post_init__(self):
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         return ModelDeploymentConfig(
             dit_layerwise_offload_modes=("memory",),
+            # The two-stage path is not numerically invariant when its text
+            # encoders or VAE use layerwise offload. Keep them resident when
+            # there is enough headroom, while preserving the low-memory
+            # layerwise path on consumer GPUs.
+            keep_resident_min_available_gb=70,
+            keep_resident_components=("text_encoder", "vae"),
             # 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,
-            keep_resident_min_available_gb=120,
-            keep_resident_components=("dit", "vae"),
         )
 
     # --- Latent shape ---

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -231,10 +231,12 @@ def set_lora_weights(
                           If False, append to existing list (for multi-LoRA support).
             output_offset: Optional constant output term paired with this adapter
         """
-        lora_A_param = torch.nn.Parameter(
-            A
-        )  # share storage with weights in the pipeline
-        lora_B_param = torch.nn.Parameter(B)
+        # This runtime never trains adapters. In particular, weights loaded
+        # during a warmup may be inference tensors; leaving Parameter's
+        # requires_grad default enabled makes a later dynamic-LoRA request try
+        # to read their unavailable version counters while sharding views.
+        lora_A_param = torch.nn.Parameter(A, requires_grad=False)
+        lora_B_param = torch.nn.Parameter(B, requires_grad=False)
         output_offset_param = (
             torch.nn.Parameter(output_offset, requires_grad=False)
             if output_offset is not None

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -22,6 +22,7 @@
     ComponentResidencyStrategy,
     ComponentUse,
     ResidencyState,
+    build_component_residency_strategy,
 )
 from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler,
@@ -456,6 +457,12 @@ class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy):
 
     def __init__(self, manager: "LTX2TwoStageResidencyController") -> None:
         self.manager = manager
+        self._placement_strategies: dict[
+            tuple[str, int, str], ComponentResidencyStrategy
+        ] = {}
+
+    def supports_auto_residency(self) -> bool:
+        return True
 
     @property
     def pipeline(self) -> "LTX2TwoStagePipeline":
@@ -473,6 +480,19 @@ def _phase(self, use: ComponentUse) -> str:
     def initialize(self) -> None:
         pass
 
+    def _placement_strategy(
+        self, module: torch.nn.Module, use: ComponentUse
+    ) -> ComponentResidencyStrategy:
+        mode = self.server_args.residency_mode(use.component_name)
+        key = (use.component_name, id(module), mode)
+        strategy = self._placement_strategies.get(key)
+        if strategy is None:
+            strategy = build_component_residency_strategy(
+                use.component_name, module, self.server_args
+            )
+            self._placement_strategies[key] = strategy
+        return strategy
+
     def prepare_for_use(
         self,
         module: torch.nn.Module,
@@ -482,13 +502,28 @@ def prepare_for_use(
         phase = self._phase(use)
         if phase != self.manager._active_phase:
             self.enter_phase(phase)
+        self._placement_strategy(module, use).prepare_for_use(module, use, state)
+
+    def prefetch_for_use(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+    ) -> bool:
+        phase = self._phase(use)
+        if phase != self.manager._active_phase:
+            self.enter_phase(phase)
+        return self._placement_strategy(module, use).prefetch_for_use(
+            module, use, state
+        )
 
     def wait_for_use(
         self,
         module: torch.nn.Module,
         use: ComponentUse,
         state: ResidencyState,
     ) -> None:
+        self._placement_strategy(module, use).wait_for_use(module, use, state)
         self.ensure_phase_ready(self._phase(use))
 
     def finish_use(
@@ -497,6 +532,7 @@ def finish_use(
         use: ComponentUse,
         state: ResidencyState,
     ) -> None:
+        self._placement_strategy(module, use).finish_use(module, use, state)
         self.exit_phase(self._phase(use))
 
     def finish_request(
@@ -507,8 +543,11 @@ def finish_request(
         *,
         preferred: bool,
     ) -> None:
+        self._placement_strategy(module, use).finish_request(
+            module, use, state, preferred=preferred
+        )
         if not preferred:
-            self.finish_use(module, use, state)
+            self.exit_phase(self._phase(use))
             return
         phase = self._phase(use)
         if phase != self.manager._active_phase:

python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py

@@ -293,26 +293,48 @@ def _temporarily_disable_offload(
             yield []
             return
 
-        # clear device cache to free up unused memory
-        if torch.get_device_module().is_available():
-            torch.get_device_module().synchronize()
-            torch.get_device_module().empty_cache()
+        offloaded_module_names = [
+            module_name
+            for module_name in module_names
+            if (
+                (module := self.modules.get(module_name)) is not None
+                and is_layerwise_offloaded_module(module)
+            )
+        ]
+        # Record every target, including coarse component offload. The current
+        # coarse path can already materialize the full component here, and the
+        # planner must preserve that phase if it later chooses layerwise mode.
+        residency_manager = getattr(self, "component_residency_manager", None)
+        residency_transition = (
+            residency_manager.full_weight_transition(module_names)
+            if residency_manager is not None
+            else nullcontext()
+        )
 
-        offload_disabled_modules = []
-        for module_name in module_names:
-            module = self.modules.get(module_name)
-            if isinstance(module, torch.nn.Module):
-                restore_weight_snapshot(module)
-            if module is not None and is_layerwise_offloaded_module(module):
+        with residency_transition:
+            # clear device cache to free up unused memory
+            if torch.get_device_module().is_available():
+                torch.get_device_module().synchronize()
+                torch.get_device_module().empty_cache()
+
+            # snapshot-offloaded components merge into their restored weights
+            for module_name in module_names:
+                module = self.modules.get(module_name)
+                if isinstance(module, torch.nn.Module):
+                    restore_weight_snapshot(module)
+
+            offload_disabled_modules = []
+            for module_name in offloaded_module_names:
+                module = self.modules[module_name]
                 module.disable_offload()
                 offload_disabled_modules.append(module)
 
-        try:
-            yield offload_disabled_modules
-        finally:
-            # Re-enable layerwise offload: sync weights to CPU and restore hooks
-            for module in offload_disabled_modules:
-                module.enable_offload()
+            try:
+                yield offload_disabled_modules
+            finally:
+                # Re-enable layerwise offload: sync weights to CPU and restore hooks
+                for module in offload_disabled_modules:
+                    module.enable_offload()
 
     def _needs_lora_weight_update_context(
         self,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py

@@ -246,7 +246,7 @@ def _encode_delight_prompts(
         return torch.cat([positive, negative, negative])
 
     @torch.no_grad()
-    def _run_delight(self, image: Image.Image) -> Image.Image:
+    def _run_delight(self, image: Image.Image, batch: Req) -> Image.Image:
         assert self.delight_transformer is not None
         assert self.delight_vae is not None
         assert self.delight_scheduler is not None
@@ -278,7 +278,14 @@ def _run_delight(self, image: Image.Image) -> Image.Image:
             )
 
         scheduler = self.delight_scheduler
-        scheduler.set_timesteps(self.config.delight_num_inference_steps, device=device)
+        target_steps = self.config.delight_num_inference_steps
+        measured_steps = (
+            min(target_steps, max(1, int(batch.num_inference_steps)))
+            if batch.is_warmup
+            else target_steps
+        )
+        batch.record_stage_iterations(measured_steps, target_steps)
+        scheduler.set_timesteps(measured_steps, device=device)
         generator = torch.Generator(device="cpu").manual_seed(42)
         latent_channels = self.delight_transformer.config.out_channels
         latents = randn_tensor(
@@ -350,11 +357,11 @@ def _run_delight(self, image: Image.Image) -> Image.Image:
             (composited.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
         )
 
-    def _prepare_reference_image(self, image_path: str) -> Image.Image:
+    def _prepare_reference_image(self, image_path: str, batch: Req) -> Image.Image:
         image = self._load_input_image(image_path)
         if not self.config.delight_enable:
             return image
-        return self._run_delight(image)
+        return self._run_delight(image, batch)
 
     def _render_multiview(
         self, mesh: Any
@@ -404,7 +411,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
 
         with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
             mesh_future = executor.submit(self._unwrap_mesh, mesh)
-            image_future = executor.submit(self._prepare_reference_image, image_path)
+            image_future = executor.submit(
+                self._prepare_reference_image, image_path, batch
+            )
             paint_mesh = mesh_future.result()
             delighted_image = image_future.result()
 
@@ -542,15 +551,27 @@ def _camera_index(azimuth: float, elevation: float) -> int:
             base, divisor = 12, 1
         return base + base_index // divisor
 
-    def _timesteps(self, device: torch.device) -> torch.Tensor:
+    def _timesteps(
+        self, device: torch.device, batch: Req | None = None
+    ) -> torch.Tensor:
+        target_steps = (
+            10
+            if self.config.paint_turbo_mode
+            else (self.config.paint_num_inference_steps)
+        )
+        measured_steps = (
+            min(target_steps, max(1, int(batch.num_inference_steps)))
+            if batch is not None and batch.is_warmup
+            else target_steps
+        )
+        if batch is not None:
+            batch.record_stage_iterations(measured_steps, target_steps)
         if not self.config.paint_turbo_mode:
-            self.scheduler.set_timesteps(
-                self.config.paint_num_inference_steps, device=device
-            )
+            self.scheduler.set_timesteps(measured_steps, device=device)
             return self.scheduler.timesteps
 
         self.scheduler.set_timesteps(
-            num_inference_steps=10,
+            num_inference_steps=measured_steps,
             original_inference_steps=30,
             device=device,
         )
@@ -639,7 +660,7 @@ def _prepare_denoising_inputs(self, batch: Req) -> PaintDenoisingInputs:
         if position_attention_mask is not None:
             model_kwargs["position_attn_mask"] = position_attention_mask
 
-        timesteps = self._timesteps(device)
+        timesteps = self._timesteps(device, batch)
         latent_channels = self.transformer.config.in_channels
         latent_size = render_size // self.vae_scale_factor
         latents = randn_tensor(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py

@@ -50,7 +50,10 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         if self.pipeline.should_skip_ltx2_lora_switch_stage():
             batch.extra["ltx2_phase"] = self.phase
             return batch
-        self.pipeline.switch_lora_phase(self.phase, batch=batch)
+        # The installed adapter survives this stage. Give its tensors version
+        # counters for later offload stages, which run outside inference mode.
+        with torch.inference_mode(False), torch.no_grad():
+            self.pipeline.switch_lora_phase(self.phase, batch=batch)
         batch.extra["ltx2_phase"] = self.phase
         return batch
 

python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json

@@ -5,6 +5,12 @@
             "ssim_threshold": 0.86,
             "psnr_threshold": 19.5,
             "mean_abs_diff_threshold": 8.5
+        },
+        "wan2_1_t2v_1.3b": {
+            "clip_threshold": 0.95,
+            "ssim_threshold": 0.85,
+            "psnr_threshold": 25.0,
+            "mean_abs_diff_threshold": 8.0
         }
     }
 }

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

@@ -2,6 +2,7 @@
 
 import unittest
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import torch
 from diffusers import AutoencoderKL as DiffusersAutoencoderKL
@@ -222,6 +223,24 @@ def test_uses_standard_lcm_schedule_without_custom_timesteps(self):
         )
         self.assertFalse(stage.scheduler.custom_timesteps)
 
+    def test_warmup_caps_paint_steps_and_records_full_target(self):
+        stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage)
+        stage.config = Hunyuan3D2PipelineConfig(paint_turbo_mode=True)
+        stage.scheduler = LCMScheduler(
+            num_train_timesteps=1000,
+            original_inference_steps=50,
+        )
+        batch = SimpleNamespace(
+            is_warmup=True,
+            num_inference_steps=4,
+            record_stage_iterations=Mock(),
+        )
+
+        timesteps = stage._timesteps(torch.device("cpu"), batch)
+
+        self.assertEqual(len(timesteps), 4)
+        batch.record_stage_iterations.assert_called_once_with(4, 10)
+
 
 if __name__ == "__main__":
     unittest.main()