[diffusion] refactor: simplify DmdDenoisingStage (#14269)

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -54,7 +54,10 @@ def _build_sampling_params_from_request(
     background: Optional[str],
     image_path: Optional[str] = None,
 ) -> SamplingParams:
-    width, height = _parse_size(size)
+    if size is None:
+        width, height = None, None
+    else:
+        width, height = _parse_size(size)
     ext = _choose_ext(output_format, background)
     server_args = get_global_server_args()
     # Build user params
@@ -149,7 +152,7 @@ async def edits(
     model: Optional[str] = Form(None),
     n: Optional[int] = Form(1),
     response_format: Optional[str] = Form(None),
-    size: Optional[str] = Form("1024x1024"),
+    size: Optional[str] = Form(None),
     output_format: Optional[str] = Form(None),
     background: Optional[str] = Form("auto"),
     user: Optional[str] = Form(None),

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -47,7 +47,10 @@ router = APIRouter(prefix="/v1/videos", tags=["videos"])
 def _build_sampling_params_from_request(
     request_id: str, request: VideoGenerationsRequest
 ) -> SamplingParams:
-    width, height = _parse_size(request.size or "720x1280")
+    if request.size is None:
+        width, height = None, None
+    else:
+        width, height = _parse_size(request.size)
     seconds = request.seconds if request.seconds is not None else 4
     # Prefer user-provided fps/num_frames from request; fallback to defaults
     fps_default = 24

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

@@ -20,6 +20,7 @@ from einops import rearrange
 from tqdm.auto import tqdm
 
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode
+from sglang.multimodal_gen.configs.pipeline_configs.wan import Wan2_2_TI2V_5B_Config
 from sglang.multimodal_gen.runtime.distributed import (
     cfg_model_parallel_all_reduce,
     get_local_torch_device,
@@ -184,6 +185,118 @@ 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 _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs):
         """
         Prepare all necessary invariant variables for the denoising loop.
@@ -264,144 +377,38 @@ class DenoisingStage(PipelineStage):
         else:
             boundary_timestep = None
 
-        # TI2V specific preparations - BEFORE SP sharding
+        # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config
-        z, z_sp, reserved_frames_masks, reserved_frames_mask_sp, seq_len = (
+        should_preprocess_for_wan_ti2v = (
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
-        # FIXME: should probably move to latent preparation stage, to handle with offload
-        if (
             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
-            # 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):
+        # TI2V specific preparations - before SP sharding
-                z = z * self.vae.scaling_factor.to(z.device, z.dtype)
+        if should_preprocess_for_wan_ti2v:
-            else:
+            seq_len, z, reserved_frames_masks = self._preprocess_latents_for_ti2v(
-                z = z * self.vae.scaling_factor
+                latents, target_dtype, batch, server_args
-            # 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 = (
+        else:
-                server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            seq_len, z, reserved_frames_masks = (
-            )
+                None,
-            patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size
+                None,
-            seq_len = (
+                None,
-                ((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()
             )
 
-        # Handle sequence parallelism AFTER TI2V processing
+        # Handle sequence parallelism after TI2V processing
         self._preprocess_sp_latents(batch, server_args)
         latents = batch.latents
 
         # Shard z and reserved_frames_mask for TI2V if SP is enabled
-        if (
+        if should_preprocess_for_wan_ti2v:
-            server_args.pipeline_config.task_type == ModelTaskType.TI2V
+            reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v(
-            and batch.condition_image is not None
+                z, reserved_frames_masks, batch
-            and get_sp_world_size() > 1
+            )
-        ):
-            sp_world_size = get_sp_world_size()
-            rank_in_sp_group = get_sp_parallel_rank()
-
-            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
         else:
-            # TI2V not enabled or SP not enabled
+            reserved_frames_mask_sp, z_sp = (
-            z_sp = z
-            reserved_frames_mask_sp = (
                 reserved_frames_masks[0] if reserved_frames_masks is not None else None
-            )  # Extract tensor
+            ), z
 
         guidance = self.get_or_build_guidance(
             # TODO: replace with raw_latent_shape?
@@ -682,15 +689,18 @@ class DenoisingStage(PipelineStage):
         server_args: ServerArgs,
         t_device,
         target_dtype,
-        seq_len,
+        seq_len: int | None,
         reserved_frames_mask,
     ):
         bsz = batch.raw_latent_shape[0]
-        # expand timestep
+        should_preprocess_for_wan_ti2v = (
-        if (
             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.
             # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0))
             # is applied consistently *before* it's used by any rank.
@@ -732,10 +742,12 @@ class DenoisingStage(PipelineStage):
         """
         For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep
         """
-        if (
+        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.
             # 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.

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

@@ -3,20 +3,8 @@
 import time
 
 import torch
-from einops import rearrange
 
-from sglang.multimodal_gen.runtime.distributed import (
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
-    get_local_torch_device,
-    get_sp_parallel_rank,
-    get_sp_world_size,
-    sequence_model_parallel_all_gather,
-)
-from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import (
-    SlidingTileAttentionBackend,
-)
-from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import (
-    VideoSparseAttentionBackend,
-)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler,
@@ -24,10 +12,6 @@ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler
 from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
-from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
-    st_attn_available,
-    vsa_available,
-)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
@@ -36,7 +20,6 @@ from sglang.multimodal_gen.utils import dict_to_3d_list
 logger = init_logger(__name__)
 
 
-# TODO: use base methods of DenoisingStage
 class DmdDenoisingStage(DenoisingStage):
     """
     Denoising stage for DMD.
@@ -46,6 +29,29 @@ class DmdDenoisingStage(DenoisingStage):
         super().__init__(transformer, scheduler)
         self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
 
+    def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
+        # 1. to shard latents (B, C, T, H, W) along dim 2
+        super()._preprocess_sp_latents(batch, server_args)
+
+        # 2. DMD expects (B, T, C, H, W) for the main latents in the loop
+        if batch.latents is not None:
+            batch.latents = batch.latents.permute(0, 2, 1, 3, 4)
+
+        # Note: batch.image_latent is kept as (B, C, T, H, W) here
+
+    def _postprocess_sp_latents(
+        self,
+        batch: Req,
+        latents: torch.Tensor,
+        trajectory_tensor: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # 1. convert back from DMD's (B, T, C, H, W) to standard (B, C, T, H, W)
+        # this is because base gather_latents_for_sp expects dim=2 for T
+        latents = latents.permute(0, 2, 1, 3, 4)
+
+        # 2. use base method to gather
+        return super()._postprocess_sp_latents(batch, latents, trajectory_tensor)
+
     def forward(
         self,
         batch: Req,
@@ -53,38 +59,25 @@ class DmdDenoisingStage(DenoisingStage):
     ) -> Req:
         """
         Run the denoising loop.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The batch with denoised latents.
         """
-        # Setup precision and autocast settings
+        prepared_vars = self._prepare_denoising_loop(batch, server_args)
-        # TODO(will): make the precision configurable for inference
-        # target_dtype = PRECISION_TO_TYPE[server_args.precision]
-        target_dtype = torch.bfloat16
-        autocast_enabled = (
-            target_dtype != torch.float32
-        ) and not server_args.disable_autocast
 
-        # Get timesteps and calculate warmup steps
+        target_dtype = prepared_vars["target_dtype"]
-        timesteps = batch.timesteps
+        autocast_enabled = prepared_vars["autocast_enabled"]
+        num_warmup_steps = prepared_vars["num_warmup_steps"]
+        latents = prepared_vars["latents"]
+        video_raw_latent_shape = latents.shape
 
-        # TODO(will): remove this once we add input/output validation for stages
+        timesteps = torch.tensor(
-        if timesteps is None:
+            server_args.pipeline_config.dmd_denoising_steps,
-            raise ValueError("Timesteps must be provided")
+            dtype=torch.long,
-        num_inference_steps = batch.num_inference_steps
+            device=get_local_torch_device(),
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        )
 
-        # Prepare image latents and embeddings for I2V generation
+        # prepare image_kwargs
         image_embeds = batch.image_embeds
         if len(image_embeds) > 0:
-            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [img.to(target_dtype) for img in image_embeds]
-            image_embeds = [
-                image_embed.to(target_dtype) for image_embed in image_embeds
-            ]
 
         image_kwargs = self.prepare_extra_func_kwargs(
             self.transformer.forward,
@@ -94,56 +87,12 @@ class DmdDenoisingStage(DenoisingStage):
             },
         )
 
-        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+        pos_cond_kwargs = prepared_vars["pos_cond_kwargs"]
-            self.transformer.forward,
+        prompt_embeds = prepared_vars["prompt_embeds"]
-            {
-                "encoder_hidden_states_2": batch.clip_embedding_pos,
-                "encoder_attention_mask": batch.prompt_attention_mask,
-            },
-        )
 
-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, server_args)
-
-        # Get latents and embeddings
-        assert batch.latents is not None, "latents must be provided"
-        latents = batch.latents
-        latents = latents.permute(0, 2, 1, 3, 4)
-
-        video_raw_latent_shape = latents.shape
-        prompt_embeds = batch.prompt_embeds
-        assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan"
-        timesteps = torch.tensor(
-            server_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long,
-            device=get_local_torch_device(),
-        )
-
-        # Handle sequence parallelism if enabled
-        sp_world_size, rank_in_sp_group = (
-            get_sp_world_size(),
-            get_sp_parallel_rank(),
-        )
-        sp_group = sp_world_size > 1
-        if sp_group:
-            latents = rearrange(
-                latents, "b (n t) c h w -> b n t c h w", n=sp_world_size
-            ).contiguous()
-            latents = latents[:, rank_in_sp_group, :, :, :, :]
-            if batch.image_latent is not None:
-                image_latent = rearrange(
-                    batch.image_latent,
-                    "b c (n t) h w -> b c n t h w",
-                    n=sp_world_size,
-                ).contiguous()
-
-                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                batch.image_latent = image_latent
-
-        # Run denoising loop
         denoising_loop_start_time = time.time()
         self.start_profile(batch=batch)
+
         with self.progress_bar(total=len(timesteps)) as progress_bar:
             for i, t in enumerate(timesteps):
                 # Skip if interrupted
@@ -171,16 +120,11 @@ class DmdDenoisingStage(DenoisingStage):
 
                     # Prepare inputs for transformer
                     t_expand = t.repeat(latent_model_input.shape[0])
-                    guidance_expand = (
+
-                        torch.tensor(
+                    guidance_expand = self.get_or_build_guidance(
-                            [server_args.pipeline_config.embedded_cfg_scale]
+                        latent_model_input.shape[0],
-                            * latent_model_input.shape[0],
+                        target_dtype,
-                            dtype=torch.float32,
+                        get_local_torch_device(),
-                            device=get_local_torch_device(),
-                        ).to(target_dtype)
-                        * 1000.0
-                        if server_args.pipeline_config.embedded_cfg_scale is not None
-                        else None
                     )
 
                     # Predict noise residual
@@ -189,41 +133,13 @@ class DmdDenoisingStage(DenoisingStage):
                         dtype=target_dtype,
                         enabled=autocast_enabled,
                     ):
-                        if (
+                        attn_metadata = self._build_attn_metadata(i, batch, server_args)
-                            vsa_available
-                            and self.attn_backend == VideoSparseAttentionBackend
-                        ):
-                            self.attn_metadata_builder_cls = (
-                                self.attn_backend.get_builder_cls()
-                            )
-
-                            if self.attn_metadata_builder_cls is not None:
-                                self.attn_metadata_builder = (
-                                    self.attn_metadata_builder_cls()
-                                )
-                                # TODO(will): clean this up
-                                attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                    current_timestep=i,  # type: ignore
-                                    raw_latent_shape=batch.raw_latent_shape[2:5],  # type: ignore
-                                    patch_size=server_args.pipeline_config.dit_config.patch_size,  # type: ignore
-                                    STA_param=batch.STA_param,  # type: ignore
-                                    VSA_sparsity=server_args.VSA_sparsity,  # type: ignore
-                                    device=get_local_torch_device(),  # type: ignore
-                                )  # type: ignore
-                                assert (
-                                    attn_metadata is not None
-                                ), "attn_metadata cannot be None"
-                            else:
-                                attn_metadata = None
-                        else:
-                            attn_metadata = None
 
                         batch.is_cfg_negative = False
                         with set_forward_context(
                             current_timestep=i,
                             attn_metadata=attn_metadata,
                             forward_batch=batch,
-                            # server_args=server_args
                         ):
                             # Run transformer
                             pred_noise = self.transformer(
@@ -251,13 +167,6 @@ class DmdDenoisingStage(DenoisingStage):
                                 dtype=pred_video.dtype,
                                 generator=batch.generator[0],
                             ).to(self.device)
-                            if sp_group:
-                                noise = rearrange(
-                                    noise,
-                                    "b (n t) c h w -> b n t c h w",
-                                    n=sp_world_size,
-                                ).contiguous()
-                                noise = noise[:, rank_in_sp_group, :, :, :, :]
                             latents = self.scheduler.add_noise(
                                 pred_video.flatten(0, 1),
                                 noise.flatten(0, 1),
@@ -284,11 +193,12 @@ class DmdDenoisingStage(DenoisingStage):
                 (denoising_loop_end_time - denoising_loop_start_time) / len(timesteps),
             )
 
-        # Gather results if using sequence parallelism
+        self._post_denoising_loop(
-        if sp_group:
+            batch=batch,
-            latents = sequence_model_parallel_all_gather(latents, dim=1)
+            latents=latents,
-        latents = latents.permute(0, 2, 1, 3, 4)
+            trajectory_latents=[],
-        # Update batch with final latents
+            trajectory_timesteps=[],
-        batch.latents = latents
+            server_args=server_args,
+        )
 
         return batch

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

@@ -93,8 +93,8 @@ class InputValidationStage(PipelineStage):
 
             # adjust output image size
             calculated_width, calculated_height = calculated_size
-            width = calculated_width if batch.width_not_provided else batch.width
+            width = batch.width or calculated_width
-            height = calculated_height if batch.height_not_provided else batch.height
+            height = batch.height or calculated_height
             multiple_of = (
                 server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2
             )
@@ -182,16 +182,6 @@ class InputValidationStage(PipelineStage):
                 "`negative_prompt_embeds` must be provided"
             )
 
-        # Validate height and width
-        if batch.height is None or batch.width is None:
-            raise ValueError(
-                "Height and width must be provided. Please set `height` and `width`."
-            )
-        if batch.height % 8 != 0 or batch.width % 8 != 0:
-            raise ValueError(
-                f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}."
-            )
-
         # Validate number of inference steps
         if batch.num_inference_steps <= 0:
             raise ValueError(
@@ -234,8 +224,7 @@ class InputValidationStage(PipelineStage):
             lambda _: V.string_or_list_strings(batch.prompt)
             or V.list_not_empty(batch.prompt_embeds),
         )
-        result.add_check("height", batch.height, V.positive_int)
+
-        result.add_check("width", batch.width, V.positive_int)
         result.add_check(
             "num_inference_steps", batch.num_inference_steps, V.positive_int
         )
@@ -249,6 +238,13 @@ class InputValidationStage(PipelineStage):
     def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
         """Verify input validation stage outputs."""
         result = VerificationResult()
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        # Validate height and width
+        if batch.height % 8 != 0 or batch.width % 8 != 0:
+            raise ValueError(
+                f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}."
+            )
         result.add_check("seeds", batch.seeds, V.list_not_empty)
         result.add_check("generator", batch.generator, V.generator_or_list_generators)
         return result

python/sglang/multimodal_gen/test/slack_utils.py

@@ -9,6 +9,8 @@ from datetime import datetime
 from urllib.parse import urlparse
 from urllib.request import urlopen
 
+from sglang.multimodal_gen.runtime.utils.perf_logger import get_git_commit_hash
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -47,7 +49,12 @@ except Exception as e:
 
 def _get_status_message(run_id, current_case_id, thread_messages=None):
     date_str = datetime.now().strftime("%d/%m")
-    base_header = f"*🧵 for nightly test of {date_str}*\n*GitHub Run ID:* {run_id}\n*Total Tasks:* {len(ALL_CASES)}"
+    base_header = f""""*🧵 for nightly test of {date_str}*
+*Git Revision:* {get_git_commit_hash()}
+*GitHub Run ID:* {run_id}
+*Total Tasks:* {len(ALL_CASES)}
+
+"""
 
     if not ALL_CASES:
         return base_header