diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 5c0989045..afa353b87 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -343,9 +343,9 @@ class PipelineConfig: def shard_latents_for_sp(self, batch, latents): # general logic for video models - if batch.enable_sequence_shard: - return latents, False sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() + if batch.enable_sequence_shard and sp_world_size > 1: + return latents, False if latents.dim() != 5: return latents, False time_dim = latents.shape[2] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index ba0074cb5..059d5b582 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -968,6 +968,28 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): # batch_size, 6, inner_dim timestep_proj = timestep_proj.unflatten(1, (6, -1)) + if sequence_shard_enabled and ts_seq_len is not None: + if seq_shard_pad > 0: + pad = torch.zeros( + ( + batch_size, + seq_shard_pad, + timestep_proj.shape[2], + timestep_proj.shape[3], + ), + dtype=timestep_proj.dtype, + device=timestep_proj.device, + ) + timestep_proj = torch.cat([timestep_proj, pad], dim=1) + timestep_proj = timestep_proj.view( + batch_size, + self.sp_size, + local_seq_len, + timestep_proj.shape[2], + timestep_proj.shape[3], + ) + timestep_proj = timestep_proj[:, sp_rank, :, :, :] + if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1