diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index aee206d52..c8b255064 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 03a8b5cf0..b93ea9fe0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py index 709eeaf47..4556bf0f6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_ti2v_helpers.py b/python/sglang/multimodal_gen/test/unit/test_wan_ti2v_helpers.py new file mode 100644 index 000000000..a150005f9 --- /dev/null +++ b/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()