[diffusion] fix: fix ZImage SP sharding for 5D latents and unpad frames (#16418)

Signed-off-by: Chi <chixie.mcisaac@gmail.com>

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

@@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
+    PipelineConfig,
 )
 
 
@@ -69,8 +70,24 @@ class ZImagePipelineConfig(ImagePipelineConfig):
         )
         return inputs
 
+    def shard_latents_for_sp(self, batch, latents):
+        if latents.dim() == 5:
+            return PipelineConfig.shard_latents_for_sp(self, batch, latents)
+        return super().shard_latents_for_sp(batch, latents)
+
+    def gather_latents_for_sp(self, latents):
+        if latents.dim() == 5:
+            return PipelineConfig.gather_latents_for_sp(self, latents)
+        return super().gather_latents_for_sp(latents)
+
     def post_denoising_loop(self, latents, batch):
         bs, channels, num_frames, height, width = latents.shape
+        raw_latent_shape = getattr(batch, "raw_latent_shape", None)
+        if raw_latent_shape is not None and num_frames > raw_latent_shape[2]:
+            latents = latents[:, :, : raw_latent_shape[2], :, :]
+            num_frames = raw_latent_shape[2]
+        if num_frames != 1:
+            return latents[:, :, 0, :, :]
         return latents.view(bs, channels, height, width)
 
     def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -529,6 +529,31 @@
             "expected_avg_denoise_ms": 94.15,
             "expected_median_denoise_ms": 102.03
         },
+        "zimage_image_t2i_2_gpus": {
+            "stages_ms": {
+                "InputValidationStage": 0.08,
+                "TextEncodingStage": 118.66,
+                "ConditioningStage": 0.01,
+                "TimestepPreparationStage": 1.5,
+                "LatentPreparationStage": 0.12,
+                "DenoisingStage": 1304.07,
+                "DecodingStage": 37.83
+            },
+            "denoise_step_ms": {
+                "0": 49.76,
+                "1": 155.22,
+                "2": 155.98,
+                "3": 156.16,
+                "4": 157.04,
+                "5": 156.54,
+                "6": 156.29,
+                "7": 157.36,
+                "8": 156.05
+            },
+            "expected_e2e_ms": 1464.87,
+            "expected_avg_denoise_ms": 144.49,
+            "expected_median_denoise_ms": 156.16
+        },
         "zimage_image_t2i_warmup": {
             "stages_ms": {
                 "InputValidationStage": 0.02,

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -578,6 +578,16 @@ TWO_GPU_CASES_B = [
         ),
         T2I_sampling_params,
     ),
+    DiffusionTestCase(
+        "zimage_image_t2i_2_gpus",
+        DiffusionServerArgs(
+            model_path="Tongyi-MAI/Z-Image-Turbo",
+            modality="image",
+            num_gpus=2,
+            ulysses_degree=2,
+        ),
+        T2I_sampling_params,
+    ),
     DiffusionTestCase(
         "flux_image_t2i_2_gpus",
         DiffusionServerArgs(