From f066036c8bac1970209551e1acfafc545af4bbd6 Mon Sep 17 00:00:00 2001
From: Chi McIsaac <153383231+qimcis@users.noreply.github.com>
Date: Wed, 7 Jan 2026 00:24:04 -0500
Subject: [PATCH] [diffusion] fix: fix ZImage SP sharding for 5D latents and
 unpad frames (#16418)

Signed-off-by: Chi <chixie.mcisaac@gmail.com>
---
 .../configs/pipeline_configs/zimage.py        | 17 +++++++++++++
 .../test/server/perf_baselines.json           | 25 +++++++++++++++++++
 .../test/server/testcase_configs.py           | 10 ++++++++
 3 files changed, 52 insertions(+)

diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
index b9e86b6ad..68f00a667 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
+++ b/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):
diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json
index 80e6e88cd..a14903ac7 100644
--- a/python/sglang/multimodal_gen/test/server/perf_baselines.json
+++ b/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,
diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py
index c0bc33d17..083f2f947 100644
--- a/python/sglang/multimodal_gen/test/server/testcase_configs.py
+++ b/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(