[diffusion] parallel: pad tokens for video models under sp (#14833)

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

@@ -102,6 +102,14 @@ def shard_rotary_emb_for_sp(emb):
         return 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]
+    if latents.shape[1] > target_tokens:
+        latents = latents[:, :target_tokens, :]
+    return latents
+
+
 # config for a single pipeline
 @dataclass
 class PipelineConfig:
@@ -310,6 +318,7 @@ class PipelineConfig:
         return batch.negative_prompt_embeds
 
     def post_denoising_loop(self, latents, batch):
+        latents = maybe_unpad_latents(latents, batch)
         return latents
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
@@ -655,10 +664,7 @@ class ImagePipelineConfig(PipelineConfig):
         height = 2 * (int(batch.height) // (vae_scale_factor * 2))
         width = 2 * (int(batch.width) // (vae_scale_factor * 2))
 
-        # If SP padding was applied, remove extra tokens before reshaping
-        target_tokens = (height // 2) * (width // 2)
-        if latents.shape[1] > target_tokens:
-            latents = latents[:, :target_tokens, :]
+        latents = maybe_unpad_latents(latents, batch)
 
         latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
         latents = latents.permute(0, 3, 1, 4, 2, 5)

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -228,7 +228,7 @@ class SamplingParams:
             logger.debug(f"Setting num_frames to 1 because this is an image-gen model")
             self.num_frames = 1
             self.data_type = DataType.IMAGE
-        else:
+        elif self.adjust_frames:
             # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.
             # If we apply it after, adjust_num_frames might modify the frame count
             # and break the divisibility constraint (alignment) required by num_gpus.
@@ -536,8 +536,8 @@ class SamplingParams:
             default=SamplingParams.adjust_frames,
             help=(
                 "Enable/disable adjusting num_frames to evenly split latent frames across GPUs "
-                "and satisfy model temporal constraints. Default: true. "
-                "Examples: --adjust-frames, --adjust-frames true, --adjust-frames false."
+                "and satisfy model temporal constraints. If disabled, tokens might be padded for SP."
+                "Default: true. Examples: --adjust-frames, --adjust-frames true, --adjust-frames false."
             ),
         )
         return parser