diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 68f00a667..b851c27c8 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/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),