[diffusion] fix: fix wan ti2v sp timestep padding (#27876)

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

@@ -182,6 +182,9 @@ def shard_rotary_emb_for_sp(emb):
 def maybe_unpad_latents(latents, batch):
     # If SP padding was applied, remove extra tokens before reshaping
     raw_shape = batch.raw_latent_shape
+    if len(raw_shape) == 5 and latents.dim() == 5:
+        return latents[:, :, : raw_shape[2], :, :]
+
     if len(raw_shape) == 3:
         # Sequence format [B, S, D]: use seq_len directly
         target_tokens = raw_shape[1]

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

@@ -1322,6 +1322,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 target_dtype,
                 seq_len,
                 reserved_frames_mask,
+                server_args.pipeline_config.dit_config.arch_config.patch_size,
             )
         else:
             timestep = t_device.repeat(bsz)

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

@@ -3,6 +3,7 @@
 import math
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
@@ -110,7 +111,14 @@ def prepare_wan_ti2v_sp_inputs(
         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:
+            if time_dim > 0:
+                pad_len = (sp_world_size - time_dim % sp_world_size) % sp_world_size
+                if pad_len:
+                    reserved_frames_mask = F.pad(
+                        reserved_frames_mask,
+                        (0, 0, 0, 0, 0, pad_len),
+                        value=1.0,
+                    )
                 reserved_frames_mask_sp_tensor = rearrange(
                     reserved_frames_mask,
                     "c (n t) h w -> c n t h w",
@@ -138,15 +146,14 @@ def expand_wan_ti2v_timestep(
     target_dtype: torch.dtype,
     seq_len: int,
     reserved_frames_mask: torch.Tensor | None,
+    patch_size: tuple[int, int, int],
 ) -> 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()
+    local_seq_len = _get_wan_ti2v_local_seq_len(batch, seq_len, patch_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()
@@ -161,6 +168,17 @@ def expand_wan_ti2v_timestep(
     return t_device.repeat(batch_size, local_seq_len)
 
 
+def _get_wan_ti2v_local_seq_len(
+    batch: Req, fallback_seq_len: int, patch_size: tuple[int, int, int]
+) -> int:
+    latents = batch.latents
+    if latents is None or latents.ndim != 5:
+        return fallback_seq_len
+
+    _, _, local_t, latent_h, latent_w = latents.shape
+    return int(local_t * latent_h * latent_w // (patch_size[1] * patch_size[2]))
+
+
 def blend_wan_ti2v_latents(
     latents: torch.Tensor,
     reserved_frames_mask: torch.Tensor | None,

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

@@ -0,0 +1,84 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import maybe_unpad_latents
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v import (
+    expand_wan_ti2v_timestep,
+    prepare_wan_ti2v_sp_inputs,
+)
+
+
+class TestWanTI2VHelpers(unittest.TestCase):
+    def test_sp_mask_is_padded_before_sharding(self):
+        mask = torch.ones(1, 21, 4, 4)
+        mask[:, 0] = 0
+        batch = SimpleNamespace(did_sp_shard_latents=True)
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v.get_sp_world_size",
+                return_value=2,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v.get_sp_parallel_rank",
+                return_value=0,
+            ),
+        ):
+            mask_rank0, _ = prepare_wan_ti2v_sp_inputs(None, [mask], batch)
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v.get_sp_world_size",
+                return_value=2,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v.get_sp_parallel_rank",
+                return_value=1,
+            ),
+        ):
+            mask_rank1, _ = prepare_wan_ti2v_sp_inputs(None, [mask], batch)
+
+        self.assertEqual(mask_rank0.shape, (1, 11, 4, 4))
+        self.assertEqual(mask_rank1.shape, (1, 11, 4, 4))
+        self.assertTrue(torch.all(mask_rank0[:, 0] == 0))
+        self.assertTrue(torch.all(mask_rank1 == 1))
+
+    def test_expanded_timestep_uses_local_latent_shape(self):
+        mask = torch.ones(1, 11, 4, 4)
+        mask[:, 0] = 0
+        batch = SimpleNamespace(
+            raw_latent_shape=(1, 16, 21, 4, 4),
+            latents=torch.empty(1, 16, 11, 4, 4),
+        )
+
+        with patch(
+            "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v.get_sp_parallel_rank",
+            return_value=0,
+        ):
+            timestep = expand_wan_ti2v_timestep(
+                batch,
+                torch.tensor(1000.0),
+                torch.float32,
+                seq_len=84,
+                reserved_frames_mask=mask,
+                patch_size=(1, 2, 2),
+            )
+
+        self.assertEqual(timestep.shape, (1, 44))
+        self.assertTrue(torch.all(timestep[:, :4] == 0))
+        self.assertTrue(torch.all(timestep[:, 4:] == 1000))
+
+    def test_video_latent_sp_padding_is_trimmed_on_time_dim(self):
+        batch = SimpleNamespace(raw_latent_shape=(1, 16, 21, 4, 4))
+        latents = torch.empty(1, 16, 22, 4, 4)
+
+        trimmed = maybe_unpad_latents(latents, batch)
+
+        self.assertEqual(trimmed.shape, (1, 16, 21, 4, 4))
+
+
+if __name__ == "__main__":
+    unittest.main()