[diffusion] refactor: streamline denoising stages (#22633)

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

@@ -151,7 +151,6 @@ def maybe_unpad_latents(latents, batch):
     return latents
 
 
-# config for a single pipeline
 @dataclass
 class PipelineConfig:
     """The base configuration class for a generation pipeline."""
@@ -360,10 +359,28 @@ class PipelineConfig:
     def preprocess_decoding(self, latents, server_args=None, vae=None):
         return latents
 
+    @staticmethod
+    def _gather_sp_tensor(tensor: torch.Tensor, *, dim: int) -> torch.Tensor:
+        """All-gather an SP-sharded tensor along the specified logical dimension."""
+        return sequence_model_parallel_all_gather(tensor.contiguous(), dim=dim)
+
+    @staticmethod
+    def _trim_sp_gather_padding(
+        tensor: torch.Tensor, *, orig_len: int | None, dim: int
+    ) -> torch.Tensor:
+        """Trim padding introduced before SP sharding back to the original length."""
+        if orig_len is None:
+            return tensor
+        orig_len = int(orig_len)
+        if orig_len <= 0 or tensor.shape[dim] <= orig_len:
+            return tensor
+        slices = [slice(None)] * tensor.ndim
+        slices[dim] = slice(orig_len)
+        return tensor[tuple(slices)]
+
     def gather_latents_for_sp(self, latents, batch=None):
         # For video latents [B, C, T_local, H, W], gather along time dim=2
-        latents = sequence_model_parallel_all_gather(latents, dim=2)
+        return self._gather_sp_tensor(latents, dim=2)
-        return latents
 
     def can_shard_audio_latents_for_sp(self, audio_latents) -> bool:
         """Return whether this pipeline uses packed audio latents that can be SP-sharded."""
@@ -405,9 +422,9 @@ class PipelineConfig:
         noise_pred = self.gather_latents_for_sp(noise_pred)
         raw_latent_shape = getattr(batch, "raw_latent_shape", None)
         if raw_latent_shape is not None and noise_pred.dim() == 3:
-            orig_s = raw_latent_shape[1]
+            noise_pred = self._trim_sp_gather_padding(
-            if noise_pred.shape[1] > orig_s:
+                noise_pred, orig_len=raw_latent_shape[1], dim=1
-                noise_pred = noise_pred[:, :orig_s, :]
+            )
         return noise_pred
 
     def preprocess_vae_image(self, batch, vae_image_processor):
@@ -850,8 +867,7 @@ class ImagePipelineConfig(PipelineConfig):
 
     def gather_latents_for_sp(self, latents, batch=None):
         # For image latents [B, S_local, D], gather along sequence dim=1
-        latents = sequence_model_parallel_all_gather(latents, dim=1)
+        return self._gather_sp_tensor(latents, dim=1)
-        return latents
 
     def _unpad_and_unpack_latents(self, latents, batch):
         vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
@@ -908,7 +924,7 @@ class SpatialImagePipelineConfig(ImagePipelineConfig):
         if latents.dim() != 4:
             return super().gather_latents_for_sp(latents, batch=batch)
         # Gather along dim=2 (H') to match shard_latents_for_sp
-        return sequence_model_parallel_all_gather(latents, dim=2)
+        return self._gather_sp_tensor(latents, dim=2)
 
 
 @dataclass

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

@@ -19,7 +19,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
 from sglang.multimodal_gen.runtime.distributed import (
     get_sp_parallel_rank,
     get_sp_world_size,
-    sequence_model_parallel_all_gather,
 )
 
 
@@ -383,7 +382,7 @@ class LTX2PipelineConfig(PipelineConfig):
         if get_sp_world_size() <= 1:
             return latents
         if isinstance(latents, torch.Tensor) and latents.ndim == 3:
-            return sequence_model_parallel_all_gather(latents.contiguous(), dim=1)
+            return self._gather_sp_tensor(latents, dim=1)
         return super().gather_latents_for_sp(latents, batch=batch)
 
     def shard_audio_latents_for_sp(self, batch, audio_latents):
@@ -432,18 +431,21 @@ class LTX2PipelineConfig(PipelineConfig):
         )
 
     def gather_audio_latents_for_sp(self, audio_latents, batch):
+        """Gather packed audio latents after SP and trim any pad-only tail tokens."""
         if get_sp_world_size() <= 1:
             return audio_latents
         if not (isinstance(audio_latents, torch.Tensor) and audio_latents.ndim == 3):
             return audio_latents
 
-        audio_latents = sequence_model_parallel_all_gather(
+        audio_latents = self._gather_sp_tensor(
-            audio_latents.contiguous(), dim=1
+            audio_latents,
+            dim=1,
+        )
+        return self._trim_sp_gather_padding(
+            audio_latents,
+            orig_len=getattr(batch, "sp_audio_orig_num_frames", None),
+            dim=1,
         )
-        orig_num_frames = int(batch.sp_audio_orig_num_frames)
-        if orig_num_frames > 0:
-            audio_latents = audio_latents[:, :orig_num_frames, :]
-        return audio_latents
 
     def prepare_video_rope_coords_for_sp(
         self,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py

@@ -56,6 +56,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation impo
 from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import (
     LTX2AVLatentPreparationStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
+    LTX2DenoisingStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import (
     LTX2TextConnectorStage,
 )
@@ -80,6 +83,7 @@ __all__ = [
     "LTX2AVLatentPreparationStage",
     "DenoisingStage",
     "DmdDenoisingStage",
+    "LTX2DenoisingStage",
     "LTX2AVDenoisingStage",
     "CausalDMDDenoisingStage",
     "EncodingStage",

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -10,13 +10,13 @@ import math
 import os
 import time
 import weakref
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field, fields
 from functools import lru_cache
 from typing import Any
 
 import torch
 import torch.nn as nn
-from einops import rearrange
 from tqdm.auto import tqdm
 
 from sglang.multimodal_gen import envs
@@ -25,9 +25,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import (
     Flux2PipelineConfig,
     FluxPipelineConfig,
 )
-from sglang.multimodal_gen.configs.pipeline_configs.wan import (
-    Wan2_2_TI2V_5B_Config,
-)
 from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
 from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
     CacheDitConfig,
@@ -41,7 +38,6 @@ from sglang.multimodal_gen.runtime.distributed import (
     cfg_model_parallel_all_reduce,
     get_local_torch_device,
     get_sp_group,
-    get_sp_parallel_rank,
     get_sp_world_size,
     get_tp_group,
     get_world_group,
@@ -68,6 +64,13 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
     StageParallelismType,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v import (
+    blend_wan_ti2v_latents,
+    expand_wan_ti2v_timestep,
+    prepare_wan_ti2v_latents,
+    prepare_wan_ti2v_sp_inputs,
+    should_apply_wan_ti2v,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     StageValidators as V,
 )
@@ -89,12 +92,60 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
-from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like
+from sglang.multimodal_gen.utils import dict_to_3d_list
 from sglang.srt.utils.common import get_compiler_backend
 
 logger = init_logger(__name__)
 
 
+@dataclass(slots=True)
+class DenoisingContext:
+    """Loop-scoped state shared across the denoising skeleton and its hooks."""
+
+    extra_step_kwargs: dict[str, Any]
+    target_dtype: torch.dtype
+    autocast_enabled: bool
+    timesteps: torch.Tensor
+    num_inference_steps: int
+    num_warmup_steps: int
+    image_kwargs: dict[str, Any]
+    pos_cond_kwargs: dict[str, Any]
+    neg_cond_kwargs: dict[str, Any]
+    latents: torch.Tensor
+    boundary_timestep: float | None
+    z: torch.Tensor | None
+    reserved_frames_mask: torch.Tensor | None
+    seq_len: int | None
+    guidance: torch.Tensor
+    is_warmup: bool
+    trajectory_timesteps: list[torch.Tensor] = field(default_factory=list)
+    trajectory_latents: list[torch.Tensor] = field(default_factory=list)
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def __getitem__(self, key: str) -> Any:
+        return getattr(self, key)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return getattr(self, key, default)
+
+    def to_kwargs(self) -> dict[str, Any]:
+        """Return a shallow field mapping for derived context construction."""
+        return {item.name: getattr(self, item.name) for item in fields(self)}
+
+
+@dataclass(slots=True)
+class DenoisingStepState:
+    """Per-step hot-path state computed once and reused within a denoising step."""
+
+    step_index: int
+    t_host: torch.Tensor
+    t_device: torch.Tensor
+    t_int: int
+    current_model: Any
+    current_guidance_scale: Any
+    attn_metadata: Any | None
+
+
 class DenoisingStage(PipelineStage):
     """
     Stage for running the denoising loop in diffusion pipelines.
@@ -431,118 +482,6 @@ class DenoisingStage(PipelineStage):
         # return StageParallelismType.CFG_PARALLEL if get_global_server_args().enable_cfg_parallel else StageParallelismType.REPLICATED
         return StageParallelismType.REPLICATED
 
-    def _preprocess_latents_for_ti2v(
-        self, latents, target_dtype, batch, server_args: ServerArgs
-    ):
-        # FIXME: should probably move to latent preparation stage, to handle with offload
-        # Wan2.2 TI2V directly replaces the first frame of the latent with
-        # the image latent instead of appending along the channel dim
-        assert batch.image_latent is None, "TI2V task should not have image latents"
-        assert self.vae is not None, "VAE is not provided for TI2V task"
-        self.vae = self.vae.to(batch.condition_image.device)
-        z = self.vae.encode(batch.condition_image).mean.float()
-        if self.vae.device != "cpu" and server_args.vae_cpu_offload:
-            self.vae = self.vae.to("cpu")
-        if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
-            if isinstance(self.vae.shift_factor, torch.Tensor):
-                z -= self.vae.shift_factor.to(z.device, z.dtype)
-            else:
-                z -= self.vae.shift_factor
-
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            z = z * self.vae.scaling_factor.to(z.device, z.dtype)
-        else:
-            z = z * self.vae.scaling_factor
-        # z: [B, C, 1, H, W]
-        latent_model_input = latents.to(target_dtype)
-        # Keep as [B, C, T, H, W] for proper broadcasting
-        assert latent_model_input.ndim == 5
-
-        # Create mask with proper shape [B, C, T, H, W]
-        latent_for_mask = latent_model_input.squeeze(0)  # [C, T, H, W]
-        _, reserved_frames_masks = masks_like([latent_for_mask], zero=True)
-        reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0)  # [1, C, T, H, W]
-
-        # replace GLOBAL first frame with image - proper broadcasting
-        # z: [B, C, 1, H, W], reserved_frames_mask: [1, C, T, H, W]
-        # Both will broadcast correctly
-        latents = (
-            1.0 - reserved_frames_mask
-        ) * z + reserved_frames_mask * latent_model_input
-        assert latents.ndim == 5
-        latents = latents.to(get_local_torch_device())
-        batch.latents = latents
-
-        F = batch.num_frames
-        temporal_scale = (
-            server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
-        )
-        spatial_scale = (
-            server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
-        )
-        patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size
-        seq_len = (
-            ((F - 1) // temporal_scale + 1)
-            * (batch.height // spatial_scale)
-            * (batch.width // spatial_scale)
-            // (patch_size[1] * patch_size[2])
-        )
-        seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size()
-        return seq_len, z, reserved_frames_masks
-
-    def _postprocess_latents_for_ti2v(self, z, reserved_frames_masks, batch):
-        rank_in_sp_group = get_sp_parallel_rank()
-        sp_world_size = get_sp_world_size()
-
-        if getattr(batch, "did_sp_shard_latents", False):
-            # Shard z (image latent) along time dimension
-            # z shape: [1, C, 1, H, W] - only first frame
-            # Only rank 0 has the first frame after sharding
-            if z.shape[2] == 1:
-                # z is single frame, only rank 0 needs it
-                if rank_in_sp_group == 0:
-                    z_sp = z
-                else:
-                    # Other ranks don't have the first frame
-                    z_sp = None
-            else:
-                # Should not happen for TI2V
-                z_sp = z
-
-            # Shard reserved_frames_mask along time dimension to match sharded latents
-            # reserved_frames_mask is a list from masks_like, extract reserved_frames_mask[0] first
-            # reserved_frames_mask[0] shape: [C, T, H, W]
-            # All ranks need their portion of reserved_frames_mask for timestep calculation
-            if reserved_frames_masks is not None:
-                reserved_frames_mask = reserved_frames_masks[
-                    0
-                ]  # Extract tensor from list
-                time_dim = reserved_frames_mask.shape[1]  # [C, T, H, W]
-                if time_dim > 0 and time_dim % sp_world_size == 0:
-                    reserved_frames_mask_sp_tensor = rearrange(
-                        reserved_frames_mask,
-                        "c (n t) h w -> c n t h w",
-                        n=sp_world_size,
-                    ).contiguous()
-                    reserved_frames_mask_sp_tensor = reserved_frames_mask_sp_tensor[
-                        :, rank_in_sp_group, :, :, :
-                    ]
-                    reserved_frames_mask_sp = (
-                        reserved_frames_mask_sp_tensor  # Store as tensor, not list
-                    )
-                else:
-                    reserved_frames_mask_sp = reserved_frames_mask
-            else:
-                reserved_frames_mask_sp = None
-        else:
-            # SP not enabled or latents not sharded
-            z_sp = z
-            reserved_frames_mask_sp = (
-                reserved_frames_masks[0] if reserved_frames_masks is not None else None
-            )  # Extract tensor
-
-        return reserved_frames_mask_sp, z_sp
-
     def _handle_boundary_ratio(
         self,
         server_args,
@@ -572,7 +511,7 @@ class DenoisingStage(PipelineStage):
         Prepare all necessary invariant variables for the denoising loop.
 
         Returns:
-            A dictionary containing all the prepared variables for the denoising loop.
+            A context object containing the invariant state for the denoising loop.
         """
         assert self.transformer is not None
         pipeline = self.pipeline() if self.pipeline else None
@@ -642,17 +581,16 @@ class DenoisingStage(PipelineStage):
             assert neg_prompt_embeds is not None
             # Removed Tensor truthiness assert to avoid GPU sync
 
-        # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config
+        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
-        should_preprocess_for_wan_ti2v = (
-            server_args.pipeline_config.task_type == ModelTaskType.TI2V
-            and batch.condition_image is not None
-            and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config
-        )
 
         # TI2V specific preparations - before SP sharding
         if should_preprocess_for_wan_ti2v:
-            seq_len, z, reserved_frames_masks = self._preprocess_latents_for_ti2v(
+            seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents(
-                latents, target_dtype, batch, server_args
+                self.vae,
+                latents,
+                target_dtype,
+                batch,
+                server_args,
             )
         else:
             seq_len, z, reserved_frames_masks = (
@@ -667,7 +605,7 @@ class DenoisingStage(PipelineStage):
 
         # Shard z and reserved_frames_mask for TI2V if SP is enabled
         if should_preprocess_for_wan_ti2v:
-            reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v(
+            reserved_frames_mask_sp, z_sp = prepare_wan_ti2v_sp_inputs(
                 z, reserved_frames_masks, batch
             )
         else:
@@ -732,26 +670,199 @@ class DenoisingStage(PipelineStage):
         else:
             neg_cond_kwargs = {}
 
-        return {
+        return DenoisingContext(
-            "extra_step_kwargs": extra_step_kwargs,
+            extra_step_kwargs=extra_step_kwargs,
-            "target_dtype": target_dtype,
+            target_dtype=target_dtype,
-            "autocast_enabled": autocast_enabled,
+            autocast_enabled=autocast_enabled,
-            "timesteps": timesteps,
+            timesteps=timesteps,
-            "num_inference_steps": num_inference_steps,
+            num_inference_steps=num_inference_steps,
-            "num_warmup_steps": num_warmup_steps,
+            num_warmup_steps=num_warmup_steps,
-            "image_kwargs": image_kwargs,
+            image_kwargs=image_kwargs,
-            "pos_cond_kwargs": pos_cond_kwargs,
+            pos_cond_kwargs=pos_cond_kwargs,
-            "neg_cond_kwargs": neg_cond_kwargs,
+            neg_cond_kwargs=neg_cond_kwargs,
-            "latents": latents,
+            latents=latents,
-            "prompt_embeds": prompt_embeds,
+            boundary_timestep=boundary_timestep,
-            "neg_prompt_embeds": neg_prompt_embeds,
+            z=z_sp,
-            "boundary_timestep": boundary_timestep,
+            reserved_frames_mask=reserved_frames_mask_sp,
-            "z": z_sp,  # Use SP-sharded version
+            seq_len=seq_len,
-            # ndim == 5
+            guidance=guidance,
-            "reserved_frames_mask": reserved_frames_mask_sp,  # Use SP-sharded version
+            is_warmup=batch.is_warmup,
-            "seq_len": seq_len,
+        )
-            "guidance": guidance,
+
-        }
+    def _before_denoising_loop(
+        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
+    ) -> None:
+        """Prepare scheduler state before entering the shared denoising loop."""
+        self.scheduler.set_begin_index(0)
+
+    def _prepare_step_state(
+        self,
+        ctx: DenoisingContext,
+        batch: Req,
+        server_args: ServerArgs,
+        step_index: int,
+        t_host: torch.Tensor,
+        timesteps_cpu: torch.Tensor,
+    ) -> DenoisingStepState:
+        """Build the per-step state shared by the loop and model-specific hooks."""
+        t_int = int(t_host.item())
+        t_device = ctx.timesteps[step_index]
+        current_model, current_guidance_scale = self._select_and_manage_model(
+            t_int=t_int,
+            boundary_timestep=ctx.boundary_timestep,
+            server_args=server_args,
+            batch=batch,
+        )
+        attn_metadata = self._prepare_step_attn_metadata(
+            ctx=ctx,
+            batch=batch,
+            server_args=server_args,
+            step_index=step_index,
+            t_int=t_int,
+            timesteps_cpu=timesteps_cpu,
+        )
+        return DenoisingStepState(
+            step_index=step_index,
+            t_host=t_host,
+            t_device=t_device,
+            t_int=t_int,
+            current_model=current_model,
+            current_guidance_scale=current_guidance_scale,
+            attn_metadata=attn_metadata,
+        )
+
+    def _prepare_step_attn_metadata(
+        self,
+        ctx: DenoisingContext,
+        batch: Req,
+        server_args: ServerArgs,
+        step_index: int,
+        t_int: int,
+        timesteps_cpu: torch.Tensor,
+    ) -> Any | None:
+        """Build attention metadata for the current denoising step."""
+        # Keep attention metadata preparation overridable so model-specific stages
+        # can preserve their original semantics without duplicating step state setup.
+        return self._build_attn_metadata(
+            step_index,
+            batch,
+            server_args,
+            timestep_value=t_int,
+            timesteps=timesteps_cpu,
+        )
+
+    def _get_prompt_embeds_validator(
+        self, batch: Req
+    ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
+        """Return the prompt-embedding validator used by verify_input."""
+        del batch
+        return V.list_not_empty
+
+    def _get_negative_prompt_embeds_validator(
+        self, batch: Req
+    ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
+        """Return the negative-prompt validator used by verify_input."""
+        return lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)
+
+    def _run_denoising_step(
+        self,
+        ctx: DenoisingContext,
+        step: DenoisingStepState,
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> None:
+        """Run one scheduler-backed denoising step in the shared base path.
+
+        Model-specific stages should override this instead of the whole loop whenever possible to achieve better performance
+        """
+        # 1. Prepare latent inputs in the model's compute dtype.
+        latent_model_input = ctx.latents.to(ctx.target_dtype)
+        if batch.image_latent is not None:
+            assert (
+                not server_args.pipeline_config.task_type == ModelTaskType.TI2V
+            ), "image latents should not be provided for TI2V task"
+            latent_model_input = torch.cat(
+                [latent_model_input, batch.image_latent], dim=1
+            ).to(ctx.target_dtype)
+
+        # 2. Expand the timestep to the shape expected by the current model.
+        timestep = self.expand_timestep_before_forward(
+            batch,
+            server_args,
+            step.t_device,
+            ctx.target_dtype,
+            ctx.seq_len,
+            ctx.reserved_frames_mask,
+        )
+
+        # 3. Apply scheduler-side input scaling before the model forward.
+        latent_model_input = self.scheduler.scale_model_input(
+            latent_model_input, step.t_device
+        )
+
+        # 4. Run the model prediction path, including CFG when enabled.
+        noise_pred = self._predict_noise_with_cfg(
+            current_model=step.current_model,
+            latent_model_input=latent_model_input,
+            timestep=timestep,
+            batch=batch,
+            timestep_index=step.step_index,
+            attn_metadata=step.attn_metadata,
+            target_dtype=ctx.target_dtype,
+            current_guidance_scale=step.current_guidance_scale,
+            image_kwargs=ctx.image_kwargs,
+            pos_cond_kwargs=ctx.pos_cond_kwargs,
+            neg_cond_kwargs=ctx.neg_cond_kwargs,
+            server_args=server_args,
+            guidance=ctx.guidance,
+            latents=ctx.latents,
+        )
+        if server_args.comfyui_mode:
+            batch.noise_pred = noise_pred
+
+        # 5. Advance the scheduler state with the predicted noise.
+        ctx.latents = self.scheduler.step(
+            model_output=noise_pred,
+            timestep=step.t_device,
+            sample=ctx.latents,
+            **ctx.extra_step_kwargs,
+            return_dict=False,
+        )[0]
+
+        # 6. Re-apply any model-specific latent constraints after the update.
+        ctx.latents = self.post_forward_for_ti2v_task(
+            batch,
+            server_args,
+            ctx.reserved_frames_mask,
+            ctx.latents,
+            ctx.z,
+        )
+
+    def _record_trajectory(
+        self,
+        ctx: DenoisingContext,
+        step: DenoisingStepState,
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> None:
+        """Append the current step to the returned latent trajectory, if requested."""
+        if not batch.return_trajectory_latents:
+            return
+        ctx.trajectory_timesteps.append(step.t_host)
+        ctx.trajectory_latents.append(ctx.latents)
+
+    def _finalize_denoising_loop(
+        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
+    ) -> None:
+        """Finalize the shared loop by handing state to post-denoising processing."""
+        self._post_denoising_loop(
+            batch=batch,
+            latents=ctx.latents,
+            trajectory_latents=ctx.trajectory_latents,
+            trajectory_timesteps=ctx.trajectory_timesteps,
+            server_args=server_args,
+            is_warmup=ctx.is_warmup,
+        )
 
     def _post_denoising_loop(
         self,
@@ -972,45 +1083,18 @@ class DenoisingStage(PipelineStage):
         reserved_frames_mask,
     ):
         bsz = batch.raw_latent_shape[0]
-        should_preprocess_for_wan_ti2v = (
+        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
-            server_args.pipeline_config.task_type == ModelTaskType.TI2V
-            and batch.condition_image is not None
-            and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config
-        )
 
         # expand timestep
         if should_preprocess_for_wan_ti2v:
-            # Explicitly cast t_device to the target float type at the beginning.
+            assert seq_len is not None, "Wan TI2V requires a token sequence length."
-            # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0))
+            timestep = expand_wan_ti2v_timestep(
-            # is applied consistently *before* it's used by any rank.
+                batch,
-            t_device_rounded = t_device.to(target_dtype)
+                t_device,
-
+                target_dtype,
-            local_seq_len = seq_len
+                seq_len,
-            if get_sp_world_size() > 1 and getattr(
+                reserved_frames_mask,
-                batch, "did_sp_shard_latents", False
-            ):
-                local_seq_len = seq_len // get_sp_world_size()
-
-            if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None:
-                # Rank 0 has the first frame, create a special timestep tensor
-                # NOTE: The spatial downsampling in the next line is suspicious but kept
-                # to match original model's potential training configuration.
-                temp_ts = (
-                    reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded
-                ).flatten()
-
-                # Pad to full local sequence length
-                temp_ts = torch.cat(
-                    [
-                        temp_ts,
-                        temp_ts.new_ones(local_seq_len - temp_ts.size(0))
-                        * t_device_rounded,
-                    ]
             )
-                timestep = temp_ts.unsqueeze(0).repeat(bsz, 1)
-            else:
-                # Other ranks get a uniform timestep tensor of the correct shape [B, local_seq_len]
-                timestep = t_device.repeat(bsz, local_seq_len)
         else:
             timestep = t_device.repeat(bsz)
         return timestep
@@ -1018,27 +1102,10 @@ class DenoisingStage(PipelineStage):
     def post_forward_for_ti2v_task(
         self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z
     ):
-        """
+        """Re-apply Wan TI2V first-frame conditioning after each denoising step."""
-        For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep
+        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
-        """
-        should_preprocess_for_wan_ti2v = (
-            server_args.pipeline_config.task_type == ModelTaskType.TI2V
-            and batch.condition_image is not None
-            and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config
-        )
         if should_preprocess_for_wan_ti2v:
-            # Apply TI2V mask blending with SP-aware z and reserved_frames_mask.
+            latents = blend_wan_ti2v_latents(latents, reserved_frames_mask, z)
-            # This ensures the first frame is always the condition image after each step.
-            # This is only applied on rank 0, where z is not None.
-            if z is not None and reserved_frames_mask is not None:
-                # z: [1, C, 1, H, W]
-                # latents: [1, C, T_local, H, W]
-                # reserved_frames_mask: [C, T_local, H, W]
-                # Unsqueeze mask to [1, C, T_local, H, W] for broadcasting.
-                # z will broadcast along the time dimension.
-                latents = (
-                    1.0 - reserved_frames_mask.unsqueeze(0)
-                ) * z + reserved_frames_mask.unsqueeze(0) * latents
 
         return latents
 
@@ -1051,159 +1118,56 @@ class DenoisingStage(PipelineStage):
         """
         Run the denoising loop.
         """
-        # Prepare variables for the denoising loop
+        ctx = self._prepare_denoising_loop(batch, server_args)
-
-        prepared_vars = self._prepare_denoising_loop(batch, server_args)
-        extra_step_kwargs = prepared_vars["extra_step_kwargs"]
-        target_dtype = prepared_vars["target_dtype"]
-        autocast_enabled = prepared_vars["autocast_enabled"]
-        timesteps = prepared_vars["timesteps"]
-        num_inference_steps = prepared_vars["num_inference_steps"]
-        num_warmup_steps = prepared_vars["num_warmup_steps"]
-        image_kwargs = prepared_vars["image_kwargs"]
-        pos_cond_kwargs = prepared_vars["pos_cond_kwargs"]
-        neg_cond_kwargs = prepared_vars["neg_cond_kwargs"]
-        latents = prepared_vars["latents"]
-        boundary_timestep = prepared_vars["boundary_timestep"]
-        z = prepared_vars["z"]
-        reserved_frames_mask = prepared_vars["reserved_frames_mask"]
-        seq_len = prepared_vars["seq_len"]
-        guidance = prepared_vars["guidance"]
-
-        # Initialize lists for ODE trajectory
-        trajectory_timesteps: list[torch.Tensor] = []
-        trajectory_latents: list[torch.Tensor] = []
-
-        # Run denoising loop
         denoising_start_time = time.time()
-
+        self._before_denoising_loop(ctx, batch, server_args)
         # to avoid device-sync caused by timestep comparison
-        is_warmup = batch.is_warmup
+        timesteps_cpu = ctx.timesteps.cpu()
-        self.scheduler.set_begin_index(0)
-        timesteps_cpu = timesteps.cpu()
         num_timesteps = timesteps_cpu.shape[0]
         with torch.autocast(
             device_type=current_platform.device_type,
-            dtype=target_dtype,
+            dtype=ctx.target_dtype,
-            enabled=autocast_enabled,
+            enabled=ctx.autocast_enabled,
         ):
-            with self.progress_bar(total=num_inference_steps) as progress_bar:
+            with self.progress_bar(total=ctx.num_inference_steps) as progress_bar:
-                for i, t_host in enumerate(timesteps_cpu):
+                for step_index, t_host in enumerate(timesteps_cpu):
                     with StageProfiler(
-                        f"denoising_step_{i}",
+                        f"denoising_step_{step_index}",
                         logger=logger,
                         metrics=batch.metrics,
                         perf_dump_path_provided=batch.perf_dump_path is not None,
                         record_as_step=True,
                     ):
-                        t_int = int(t_host.item())
+                        step = self._prepare_step_state(
-                        t_device = timesteps[i]
+                            ctx,
-                        current_model, current_guidance_scale = (
-                            self._select_and_manage_model(
-                                t_int=t_int,
-                                boundary_timestep=boundary_timestep,
-                                server_args=server_args,
-                                batch=batch,
-                            )
-                        )
-
-                        # Expand latents for I2V
-                        latent_model_input = latents.to(target_dtype)
-                        if batch.image_latent is not None:
-                            assert (
-                                not server_args.pipeline_config.task_type
-                                == ModelTaskType.TI2V
-                            ), "image latents should not be provided for TI2V task"
-                            latent_model_input = torch.cat(
-                                [latent_model_input, batch.image_latent], dim=1
-                            ).to(target_dtype)
-
-                        timestep = self.expand_timestep_before_forward(
                             batch,
                             server_args,
-                            t_device,
+                            step_index,
-                            target_dtype,
+                            t_host,
-                            seq_len,
+                            timesteps_cpu,
-                            reserved_frames_mask,
                         )
+                        self._run_denoising_step(ctx, step, batch, server_args)
+                        self._record_trajectory(ctx, step, batch, server_args)
 
-                        latent_model_input = self.scheduler.scale_model_input(
+                        if step_index == num_timesteps - 1 or (
-                            latent_model_input, t_device
+                            (step_index + 1) > ctx.num_warmup_steps
-                        )
+                            and (step_index + 1) % self.scheduler.order == 0
-
-                        # Predict noise residual
-                        attn_metadata = self._build_attn_metadata(
-                            i,
-                            batch,
-                            server_args,
-                            timestep_value=t_int,
-                            timesteps=timesteps_cpu,
-                        )
-                        noise_pred = self._predict_noise_with_cfg(
-                            current_model=current_model,
-                            latent_model_input=latent_model_input,
-                            timestep=timestep,
-                            batch=batch,
-                            timestep_index=i,
-                            attn_metadata=attn_metadata,
-                            target_dtype=target_dtype,
-                            current_guidance_scale=current_guidance_scale,
-                            image_kwargs=image_kwargs,
-                            pos_cond_kwargs=pos_cond_kwargs,
-                            neg_cond_kwargs=neg_cond_kwargs,
-                            server_args=server_args,
-                            guidance=guidance,
-                            latents=latents,
-                        )
-                        # Save noise_pred to batch for external access (e.g., ComfyUI)
-                        if server_args.comfyui_mode:
-                            batch.noise_pred = noise_pred
-
-                        # Compute the previous noisy sample
-                        latents = self.scheduler.step(
-                            model_output=noise_pred,
-                            timestep=t_device,
-                            sample=latents,
-                            **extra_step_kwargs,
-                            return_dict=False,
-                        )[0]
-
-                        latents = self.post_forward_for_ti2v_task(
-                            batch, server_args, reserved_frames_mask, latents, z
-                        )
-
-                        # save trajectory latents if needed
-                        if batch.return_trajectory_latents:
-                            trajectory_timesteps.append(t_host)
-                            trajectory_latents.append(latents)
-
-                        # Update progress bar
-                        if i == num_timesteps - 1 or (
-                            (i + 1) > num_warmup_steps
-                            and (i + 1) % self.scheduler.order == 0
                             and progress_bar is not None
                         ):
                             progress_bar.update()
 
-                        if not is_warmup:
+                        if not ctx.is_warmup:
                             self.step_profile()
 
         denoising_end_time = time.time()
 
-        if num_timesteps > 0 and not is_warmup:
+        if num_timesteps > 0 and not ctx.is_warmup:
             self.log_info(
                 "average time per step: %.4f seconds",
-                (denoising_end_time - denoising_start_time) / len(timesteps),
+                (denoising_end_time - denoising_start_time) / len(ctx.timesteps),
             )
 
-        self._post_denoising_loop(
+        self._finalize_denoising_loop(ctx, batch, server_args)
-            batch=batch,
-            latents=latents,
-            trajectory_latents=trajectory_latents,
-            trajectory_timesteps=trajectory_timesteps,
-            server_args=server_args,
-            is_warmup=is_warmup,
-        )
         return batch
 
     # TODO: this will extends the preparation stage, should let subclass/passed-in variables decide which to prepare
@@ -1846,7 +1810,11 @@ class DenoisingStage(PipelineStage):
         result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
         # disable temporarily for image-generation models
         # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check(
+            "prompt_embeds",
+            batch.prompt_embeds,
+            self._get_prompt_embeds_validator(batch),
+        )
         result.add_check("image_embeds", batch.image_embeds, V.is_list)
         # result.add_check(
         #     "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)
@@ -1865,7 +1833,7 @@ class DenoisingStage(PipelineStage):
         result.add_check(
             "negative_prompt_embeds",
             batch.negative_prompt_embeds,
-            lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x),
+            self._get_negative_prompt_embeds_validator(batch),
         )
         return result
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py

@@ -0,0 +1 @@
+"""Model-specific helpers and stages for diffusion pipeline components."""

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py

@@ -0,0 +1,174 @@
+"""WAN TI2V-specific helpers shared by the generic denoising stage."""
+
+import math
+
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
+from sglang.multimodal_gen.configs.pipeline_configs.wan import (
+    Wan2_2_TI2V_5B_Config,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.utils import masks_like
+
+
+def should_apply_wan_ti2v(batch: Req, server_args: ServerArgs) -> bool:
+    """Return whether the request should use the Wan2.2 TI2V latent path."""
+
+    return bool(
+        server_args.pipeline_config.task_type == ModelTaskType.TI2V
+        and batch.condition_image is not None
+        and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config
+    )
+
+
+def prepare_wan_ti2v_latents(
+    vae: object,
+    latents: torch.Tensor,
+    target_dtype: torch.dtype,
+    batch: Req,
+    server_args: ServerArgs,
+) -> tuple[int, torch.Tensor, list[torch.Tensor]]:
+    """Encode the conditioning image and splice it into Wan TI2V latents."""
+
+    # Wan2.2 TI2V directly replaces the first frame of the latent with
+    # the image latent instead of appending along the channel dim.
+    assert batch.image_latent is None, "TI2V task should not have image latents"
+    assert vae is not None, "VAE is not provided for TI2V task"
+
+    vae = vae.to(batch.condition_image.device)
+    z = vae.encode(batch.condition_image).mean.float()
+    if getattr(vae, "device", None) != "cpu" and server_args.vae_cpu_offload:
+        vae = vae.to("cpu")
+
+    if hasattr(vae, "shift_factor") and vae.shift_factor is not None:
+        if isinstance(vae.shift_factor, torch.Tensor):
+            z -= vae.shift_factor.to(z.device, z.dtype)
+        else:
+            z -= vae.shift_factor
+
+    if isinstance(vae.scaling_factor, torch.Tensor):
+        z = z * vae.scaling_factor.to(z.device, z.dtype)
+    else:
+        z = z * vae.scaling_factor
+
+    latent_model_input = latents.to(target_dtype)
+    assert latent_model_input.ndim == 5
+
+    latent_for_mask = latent_model_input.squeeze(0)
+    _, reserved_frames_masks = masks_like([latent_for_mask], zero=True)
+    reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0)
+
+    latents = (
+        1.0 - reserved_frames_mask
+    ) * z + reserved_frames_mask * latent_model_input
+    assert latents.ndim == 5
+    batch.latents = latents.to(get_local_torch_device())
+
+    num_frames = batch.num_frames
+    temporal_scale = (
+        server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
+    )
+    spatial_scale = (
+        server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+    )
+    patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size
+    seq_len = (
+        ((num_frames - 1) // temporal_scale + 1)
+        * (batch.height // spatial_scale)
+        * (batch.width // spatial_scale)
+        // (patch_size[1] * patch_size[2])
+    )
+    seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size()
+    return seq_len, z, reserved_frames_masks
+
+
+def prepare_wan_ti2v_sp_inputs(
+    z: torch.Tensor | None,
+    reserved_frames_masks: list[torch.Tensor] | None,
+    batch: Req,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """Shard Wan TI2V image-conditioning state to match SP-sharded video latents."""
+
+    rank_in_sp_group = get_sp_parallel_rank()
+    sp_world_size = get_sp_world_size()
+
+    if getattr(batch, "did_sp_shard_latents", False):
+        if z is not None and z.shape[2] == 1:
+            z_sp = z if rank_in_sp_group == 0 else None
+        else:
+            z_sp = z
+
+        if reserved_frames_masks is not None:
+            reserved_frames_mask = reserved_frames_masks[0]
+            time_dim = reserved_frames_mask.shape[1]
+            if time_dim > 0 and time_dim % sp_world_size == 0:
+                reserved_frames_mask_sp_tensor = rearrange(
+                    reserved_frames_mask,
+                    "c (n t) h w -> c n t h w",
+                    n=sp_world_size,
+                ).contiguous()
+                reserved_frames_mask_sp = reserved_frames_mask_sp_tensor[
+                    :, rank_in_sp_group, :, :, :
+                ]
+            else:
+                reserved_frames_mask_sp = reserved_frames_mask
+        else:
+            reserved_frames_mask_sp = None
+    else:
+        z_sp = z
+        reserved_frames_mask_sp = (
+            reserved_frames_masks[0] if reserved_frames_masks is not None else None
+        )
+
+    return reserved_frames_mask_sp, z_sp
+
+
+def expand_wan_ti2v_timestep(
+    batch: Req,
+    t_device: torch.Tensor,
+    target_dtype: torch.dtype,
+    seq_len: int,
+    reserved_frames_mask: torch.Tensor | None,
+) -> torch.Tensor:
+    """Expand the timestep tensor for Wan TI2V's first-frame masking semantics."""
+
+    batch_size = batch.raw_latent_shape[0]
+    t_device_rounded = t_device.to(target_dtype)
+
+    local_seq_len = seq_len
+    if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False):
+        local_seq_len = seq_len // get_sp_world_size()
+
+    if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None:
+        temp_ts = (reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded).flatten()
+        temp_ts = torch.cat(
+            [
+                temp_ts,
+                temp_ts.new_ones(local_seq_len - temp_ts.size(0)) * t_device_rounded,
+            ]
+        )
+        return temp_ts.unsqueeze(0).repeat(batch_size, 1)
+
+    return t_device.repeat(batch_size, local_seq_len)
+
+
+def blend_wan_ti2v_latents(
+    latents: torch.Tensor,
+    reserved_frames_mask: torch.Tensor | None,
+    z: torch.Tensor | None,
+) -> torch.Tensor:
+    """Restore Wan TI2V's conditioned first frame after each denoising step."""
+
+    if z is None or reserved_frames_mask is None:
+        return latents
+    return (
+        1.0 - reserved_frames_mask.unsqueeze(0)
+    ) * z + reserved_frames_mask.unsqueeze(0) * latents

scripts/ci/utils/diffusion/comparison_configs.json

@@ -135,6 +135,23 @@
                 }
             }
         },
+        {
+            "id": "ltx2.3_twostage_t2v_2gpus",
+            "model": "Lightricks/LTX-2.3",
+            "task": "text-to-video",
+            "prompt": "A cat and a dog baking a cake together in a kitchen.",
+            "width": 768,
+            "height": 512,
+            "num_frames": 121,
+            "seed": 42,
+            "num_gpus": 2,
+            "frameworks": {
+                "sglang": {
+                    "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --pipeline-class-name LTX2TwoStagePipeline",
+                    "extra_env": {}
+                }
+            }
+        },
         {
             "id": "wan22_i2v_a14b_720p",
             "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",