[diffusion]: Fix ZImage SP sharding for caption and latent (#17301)

Co-authored-by: rhyshen <rhyshen@tencent.com>
Co-authored-by: florianzhao <florianzhao@tencent.com>

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

@@ -1,4 +1,5 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import math
 from dataclasses import dataclass, field
 from typing import Callable
 
@@ -14,7 +15,13 @@ from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
-    PipelineConfig,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_parallel_rank,
+    get_sp_world_size,
 )
 
 
@@ -37,10 +44,8 @@ class TransformersModelConfig(EncoderConfig):
 
 @dataclass
 class ZImagePipelineConfig(ImagePipelineConfig):
-
     should_use_guidance: bool = False
     task_type: ModelTaskType = ModelTaskType.T2I
-
     dit_config: DiTConfig = field(default_factory=ZImageDitConfig)
     vae_config: VAEConfig = field(default_factory=FluxVAEConfig)
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
@@ -50,11 +55,14 @@ class ZImagePipelineConfig(ImagePipelineConfig):
     preprocess_text_funcs: tuple[Callable, ...] = field(
         default_factory=lambda: (zimage_preprocess_text,)
     )
-
     postprocess_text_funcs: tuple[Callable, ...] = field(
         default_factory=lambda: (zimage_postprocess_text,)
     )
 
+    SEQ_LEN_MULTIPLE: int = 32
+    PATCH_SIZE: int = 2
+    F_PATCH_SIZE: int = 1
+
     def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
         # flatten to 1-d list
         inputs = tokenizer.apply_chat_template(
@@ -70,19 +78,141 @@ class ZImagePipelineConfig(ImagePipelineConfig):
         )
         return inputs
 
+    @staticmethod
+    def _ceil_to_multiple(x: int, m: int) -> int:
+        if m <= 0:
+            return x
+        return int(math.ceil(x / m) * m)
+
+    def _build_zimage_sp_plan(self, batch) -> dict:
+        """Build a minimal SP plan on batch for zimage (spatial sharding + cap sharding)."""
+        sp_size = get_sp_world_size()
+        rank = get_sp_parallel_rank()
+
+        raw_latent_shape = getattr(batch, "raw_latent_shape", None)
+        if raw_latent_shape is not None and len(raw_latent_shape) >= 5:
+            H = int(raw_latent_shape[3])
+            W = int(raw_latent_shape[4])
+        else:
+            H = int(
+                batch.height // self.vae_config.arch_config.spatial_compression_ratio
+            )
+            W = int(
+                batch.width // self.vae_config.arch_config.spatial_compression_ratio
+            )
+
+        # Rule: shard along the larger spatial dimension (W/H), implemented via optional H/W transpose.
+        # Choose the larger of H and W for sharding, so H_eff = max(H, W).
+        swap_hw = W > H
+        H_eff = W if swap_hw else H
+        W_eff = H if swap_hw else W
+
+        # ZImage uses PATCH_SIZE=2 for spatial patchify; shard in token space and convert back to latent rows.
+        H_tok = H_eff // self.PATCH_SIZE
+        W_tok = W_eff // self.PATCH_SIZE
+        H_tok_pad = self._ceil_to_multiple(H_tok, sp_size)
+        H_tok_local = H_tok_pad // sp_size
+        h0_tok = rank * H_tok_local
+
+        # Cap/text sharding: avoid duplicating cap tokens across ranks.
+        cap_len = (
+            int(batch.prompt_embeds[0].size(0))
+            if getattr(batch, "prompt_embeds", None)
+            else 0
+        )
+        cap_total = self._ceil_to_multiple(cap_len, self.SEQ_LEN_MULTIPLE * sp_size)
+        cap_local = cap_total // sp_size
+        cap_start = rank * cap_local
+
+        plan = {
+            "sp_size": sp_size,
+            "rank": rank,
+            "swap_hw": swap_hw,
+            "H": H,
+            "W": W,
+            "H_eff": H_eff,
+            "W_eff": W_eff,
+            "H_tok": H_tok,
+            "W_tok": W_tok,
+            "H_tok_pad": H_tok_pad,
+            "H_tok_local": H_tok_local,
+            "h0_tok": h0_tok,
+            "cap_total": cap_total,
+            "cap_local": cap_local,
+            "cap_start": cap_start,
+        }
+        batch._zimage_sp_plan = plan
+        return plan
+
+    def _get_zimage_sp_plan(self, batch) -> dict:
+        plan = getattr(batch, "_zimage_sp_plan", None)
+        sp_size = get_sp_world_size()
+        if plan is None or plan.get("sp_size") != sp_size:
+            plan = self._build_zimage_sp_plan(batch)
+        return plan
+
+    def _shard_cap(self, cap: torch.Tensor, plan: dict) -> torch.Tensor:
+        """cap: [L, D] -> [cap_local, D], padded by repeating last token."""
+        if plan["sp_size"] <= 1:
+            return cap
+        # print(f"cap shape: {cap.shape}")  # [L, 2560] for zimage-turbo
+        L = cap.size(0)
+        cap_total = plan["cap_total"]
+        if cap_total > L:
+            cap = torch.cat([cap, cap[-1:].repeat(cap_total - L, 1)], dim=0)
+        start = plan["cap_start"]
+        local = plan["cap_local"]
+        return cap[start : start + local]
+
+    def get_pos_prompt_embeds(self, batch):
+        # Keep ZImage model signature: encoder_hidden_states is List[Tensor]
+        if get_sp_world_size() <= 1:
+            return batch.prompt_embeds
+        plan = self._get_zimage_sp_plan(batch)
+        return [self._shard_cap(batch.prompt_embeds[0], plan)]
+
     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)
+        sp_size = get_sp_world_size()
+        if sp_size <= 1 or latents.dim() != 5:
+            return latents, False
+
+        plan = self._get_zimage_sp_plan(batch)
+
+        # Layout: [B, C, T, H, W]. Always shard on dim=3 by optionally swapping H/W.
+        if plan["swap_hw"]:
+            latents = latents.transpose(3, 4).contiguous()
+
+        # Pad on effective-H so that H_tok is divisible by sp.
+        H_eff = latents.size(3)
+
+        H_tok = H_eff // self.PATCH_SIZE
+        pad_tok = plan["H_tok_pad"] - H_tok
+        pad_lat = pad_tok * self.PATCH_SIZE
+        if pad_lat > 0:
+            pad = latents[:, :, :, -1:, :].repeat(1, 1, 1, pad_lat, 1)
+            latents = torch.cat([latents, pad], dim=3)
+        h0 = plan["h0_tok"] * self.PATCH_SIZE
+        h1 = (plan["h0_tok"] + plan["H_tok_local"]) * self.PATCH_SIZE
+        latents = latents[:, :, :, h0:h1, :]
+
+        batch._zimage_sp_swap_hw = plan["swap_hw"]
+        return latents, True
 
     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)
+        # Gather on effective-H dim=3 (matches shard_latents_for_sp); swap-back is handled in post_denoising_loop.
+        if get_sp_world_size() <= 1 or latents.dim() != 5:
+            return latents
+        return sequence_model_parallel_all_gather(latents, dim=3)
 
     def post_denoising_loop(self, latents, batch):
-        bs, channels, num_frames, height, width = latents.shape
+        # Restore swapped H/W and crop padded spatial dims before final reshape.
+        if latents.dim() == 5 and getattr(batch, "_zimage_sp_swap_hw", False):
+            latents = latents.transpose(3, 4).contiguous()
         raw_latent_shape = getattr(batch, "raw_latent_shape", None)
+        if raw_latent_shape is not None and latents.dim() == 5:
+            latents = latents[:, :, :, : raw_latent_shape[3], : raw_latent_shape[4]]
+
+        bs, channels, num_frames, height, width = latents.shape
         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]
@@ -102,12 +232,41 @@ class ZImagePipelineConfig(ImagePipelineConfig):
             grids = torch.meshgrid(axes, indexing="ij")
             return torch.stack(grids, dim=-1)
 
-        PATCH_SIZE = 2
-        F_PATCH_SIZE = 1
-        SEQ_MULTI_OF = 32
+        sp_size = get_sp_world_size()
+        if sp_size > 1:
+            # SP path: build local-only freqs_cis matching local cap/x.
+            plan = self._get_zimage_sp_plan(batch)
+
+            # cap (local)
+            cap_pos_ids = create_coordinate_grid(
+                size=(plan["cap_local"], 1, 1),
+                start=(1 + plan["cap_start"], 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            cap_freqs_cis = rotary_emb(cap_pos_ids)
+
+            # image (local, effective H-shard). Use cap_total for a stable offset across ranks/passes.
+            F_tokens = 1
+            H_tokens_local = plan["H_tok_local"]
+            W_tokens = plan["W_tok"]
+            img_pos_ids = create_coordinate_grid(
+                size=(F_tokens, H_tokens_local, W_tokens),
+                start=(plan["cap_total"] + 1, plan["h0_tok"], 0),
+                device=device,
+            ).flatten(0, 2)
+            img_pad_len = (-img_pos_ids.shape[0]) % self.SEQ_LEN_MULTIPLE
+            if img_pad_len:
+                pad_ids = create_coordinate_grid(
+                    size=(1, 1, 1), start=(0, 0, 0), device=device
+                ).flatten(0, 2)
+                img_pos_ids = torch.cat(
+                    [img_pos_ids, pad_ids.repeat(img_pad_len, 1)], dim=0
+                )
+            x_freqs_cis = rotary_emb(img_pos_ids)
+            return (cap_freqs_cis, x_freqs_cis)
 
         cap_ori_len = prompt_embeds.size(0)
-        cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
+        cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE
         cap_padded_pos_ids = create_coordinate_grid(
             size=(cap_ori_len + cap_padding_len, 1, 1),
             start=(1, 0, 0),
@@ -119,11 +278,11 @@ class ZImagePipelineConfig(ImagePipelineConfig):
         H = height // self.vae_config.arch_config.spatial_compression_ratio
         W = width // self.vae_config.arch_config.spatial_compression_ratio
 
-        pH, pW = PATCH_SIZE, PATCH_SIZE
-        pF = F_PATCH_SIZE
+        pH, pW = self.PATCH_SIZE, self.PATCH_SIZE
+        pF = self.F_PATCH_SIZE
         F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
         image_ori_len = F_tokens * H_tokens * W_tokens
-        image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
+        image_padding_len = (-image_ori_len) % self.SEQ_LEN_MULTIPLE
 
         image_ori_pos_ids = create_coordinate_grid(
             size=(F_tokens, H_tokens, W_tokens),