From b098b1ae241fc4b79c88fa45afae37214da1c794 Mon Sep 17 00:00:00 2001
From: Mick <mickjagger19@icloud.com>
Date: Mon, 15 Dec 2025 20:25:43 +0800
Subject: [PATCH] [diffusion] fix: fix video model sp when resolution is not
 specified (#15047)

Co-authored-by: Brain97 <Brain97@users.noreply.github.com>
---
 .../configs/pipeline_configs/base.py          | 30 ++++++++++++-------
 .../pipelines_core/stages/input_validation.py | 14 +++++----
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
index 816f1765d..17698b75b 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
@@ -104,7 +104,8 @@ def shard_rotary_emb_for_sp(emb):
 
 def maybe_unpad_latents(latents, batch):
     # If SP padding was applied, remove extra tokens before reshaping
-    target_tokens = batch.raw_latent_shape[-1] * batch.raw_latent_shape[-2]
+    width, height = batch.raw_latent_shape[-1], batch.raw_latent_shape[-2]
+    target_tokens = width * height
     if latents.shape[1] > target_tokens:
         latents = latents[:, :target_tokens, :]
     return latents
@@ -309,13 +310,23 @@ class PipelineConfig:
         if latents.dim() != 5:
             return latents, False
         time_dim = latents.shape[2]
-        if time_dim > 0 and time_dim % sp_world_size == 0:
-            sharded_tensor = rearrange(
-                latents, "b c (n t) h w -> b c n t h w", n=sp_world_size
-            ).contiguous()
-            sharded_tensor = sharded_tensor[:, :, rank_in_sp_group, :, :, :]
-            return sharded_tensor, True
-        return latents, False
+
+        # Pad to next multiple of SP degree if needed
+        if time_dim > 0 and time_dim % sp_world_size != 0:
+            pad_len = sp_world_size - (time_dim % sp_world_size)
+            pad = torch.zeros(
+                (*latents.shape[:2], pad_len, *latents.shape[3:]),
+                dtype=latents.dtype,
+                device=latents.device,
+            )
+            latents = torch.cat([latents, pad], dim=2)
+
+        assert latents.shape[2] % sp_world_size == 0
+        sharded_tensor = rearrange(
+            latents, "b c (n t) h w -> b c n t h w", n=sp_world_size
+        ).contiguous()
+        sharded_tensor = sharded_tensor[:, :, rank_in_sp_group, :, :, :]
+        return sharded_tensor, True
 
     def get_pos_prompt_embeds(self, batch):
         return batch.prompt_embeds
@@ -636,6 +647,7 @@ class ImagePipelineConfig(PipelineConfig):
         return sigmas
 
     def shard_latents_for_sp(self, batch, latents):
+        # latents: [B, H * W, C]
         sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank()
         seq_len = latents.shape[1]
 
@@ -648,8 +660,6 @@ class ImagePipelineConfig(PipelineConfig):
                 device=latents.device,
             )
             latents = torch.cat([latents, pad], dim=1)
-            # Record padding length for later unpad
-            batch.sp_seq_pad = int(getattr(batch, "sp_seq_pad", 0)) + pad_len
 
         sharded_tensor = rearrange(
             latents, "b (n s) d -> b n s d", n=sp_world_size
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
index 8254297d9..6b17eb923 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
@@ -293,12 +293,16 @@ class InputValidationStage(PipelineStage):
         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}."
-            )
+        # Validate height and width
+        def check_size(value: int, name: str):
+            if value % (8 * server_args.num_gpus) != 0:
+                raise ValueError(
+                    f"{name} must be divisible by (8 x num_gpus) but {value} % (8 * {server_args.num_gpus}) != 0."
+                )
+
+        check_size(batch.height, "Height")
+        check_size(batch.width, "Width")
         result.add_check("seeds", batch.seeds, V.list_not_empty)
         result.add_check("generator", batch.generator, V.generator_or_list_generators)
         return result