From 8c7f704c23ea73a42d89e74c882b3762365183b4 Mon Sep 17 00:00:00 2001 From: XU Mingshi <91017482+mxuax@users.noreply.github.com> Date: Mon, 12 Jan 2026 16:22:50 +0800 Subject: [PATCH] [Diffusion][Feature] Implement SP support in LongCatImageTransformer (#721) Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Signed-off-by: mxuax --- .../acceleration/parallelism_acceleration.md | 4 +- .../longcat_image_transformer.py | 270 ++++++++++++++++-- 2 files changed, 242 insertions(+), 32 deletions(-) diff --git a/docs/user_guide/acceleration/parallelism_acceleration.md b/docs/user_guide/acceleration/parallelism_acceleration.md index ac4c8c5fa..dfacd2183 100644 --- a/docs/user_guide/acceleration/parallelism_acceleration.md +++ b/docs/user_guide/acceleration/parallelism_acceleration.md @@ -20,8 +20,8 @@ The following table shows which models are currently supported by parallelism me | Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | |-------|------------------|------------|---------|--------------|--------------------------| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ | ❌ | ❌ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ | ❌ | ❌ | +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ | | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | | **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ | diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index cf587ffe0..d2a55046d 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -6,7 +6,7 @@ from typing import Any import torch import torch.nn as nn -from diffusers.models.embeddings import TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle from vllm.logger import init_logger @@ -15,9 +15,15 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, +) +from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.utils.platform_utils import is_npu logger = init_logger(__name__) @@ -98,7 +104,6 @@ class LongCatImageAttention(nn.Module): self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) - self.rope = RotaryEmbedding(is_neox_style=False) self.attn = Attention( num_heads=heads, head_size=self.head_dim, @@ -106,6 +111,55 @@ class LongCatImageAttention(nn.Module): causal=False, ) + def _sp_attention_with_rope( + self, + img_query: torch.Tensor, + img_key: torch.Tensor, + img_value: torch.Tensor, + text_query: torch.Tensor, + text_key: torch.Tensor, + text_value: torch.Tensor, + text_seq_len: int, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, + ) -> torch.Tensor: + """ + Apply RoPE separately to text and image Q/K, then run SP attention with joint tensors. + + This is the common SP attention pattern used by both dual-stream (added_kv_proj_dim) + and single-stream (no added_kv_proj_dim) blocks. + + Args: + img_query/key/value: Image Q/K/V tensors (chunked in SP mode) + text_query/key/value: Text Q/K/V tensors (full, not chunked) + text_seq_len: Length of text sequence for splitting RoPE + image_rotary_emb: (freqs_cos, freqs_sin) containing [txt_pos, img_pos] + + Returns: + Attention output with shape (B, txt_len + img_len/SP, H, D) + """ + if image_rotary_emb is not None: + freqs_cos, freqs_sin = image_rotary_emb + txt_rotary_emb = (freqs_cos[:text_seq_len], freqs_sin[:text_seq_len]) + img_rotary_emb_split = (freqs_cos[text_seq_len:], freqs_sin[text_seq_len:]) + # Apply RoPE to image Q/K + img_query = apply_rotary_emb(img_query, img_rotary_emb_split, sequence_dim=1) + img_key = apply_rotary_emb(img_key, img_rotary_emb_split, sequence_dim=1) + # Apply RoPE to text Q/K + text_query = apply_rotary_emb(text_query, txt_rotary_emb, sequence_dim=1) + text_key = apply_rotary_emb(text_key, txt_rotary_emb, sequence_dim=1) + + return self.attn( + img_query, + img_key, + img_value, + AttentionMetadata( + joint_query=text_query, + joint_key=text_key, + joint_value=text_value, + joint_strategy="front", + ), + ) + def forward( self, hidden_states: torch.Tensor, @@ -113,6 +167,22 @@ class LongCatImageAttention(nn.Module): image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: + """ + Forward pass with SP-aware joint attention. + + Input shapes (in SP mode): + - hidden_states: (B, img_seq_len // SP, D) - image hidden states (chunked) + - encoder_hidden_states: (B, txt_seq_len, D) - text hidden states (full) + + SP Mode (sequence_parallel_size > 1): + - Image Q/K/V: processed with AllToAll or Ring communication + - Text Q/K/V: passed as joint tensors, broadcasted to all ranks + - Output: attention over (text + image) with proper SP handling + + Non-SP Mode (sequence_parallel_size = 1): + - Standard concatenation of text + image Q/K/V + - Regular attention over the full sequence + """ qkv, _ = self.to_qkv(hidden_states) query, key, value = qkv.chunk(3, dim=-1) @@ -132,29 +202,85 @@ class LongCatImageAttention(nn.Module): encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) + # Apply RMSNorm to text Q/K encoder_query = self.norm_added_q(encoder_query) encoder_key = self.norm_added_k(encoder_key) - query = torch.cat([encoder_query, query], dim=1) - key = torch.cat([encoder_key, key], dim=1) - value = torch.cat([encoder_value, value], dim=1) + # Check if SP is enabled from forward context (set by LongCatImageTransformer2DModel) + forward_ctx = get_forward_context() + sp_size = forward_ctx.sequence_parallel_size + use_sp_joint_attention = sp_size > 1 and not forward_ctx.split_text_embed_in_sp - if image_rotary_emb is not None: - cos, sin = image_rotary_emb # [S, D/2] - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + if use_sp_joint_attention: + # SP Mode: Use common helper for RoPE + joint attention + hidden_states = self._sp_attention_with_rope( + img_query=query, + img_key=key, + img_value=value, + text_query=encoder_query, + text_key=encoder_key, + text_value=encoder_value, + text_seq_len=encoder_query.shape[1], + image_rotary_emb=image_rotary_emb, + ) + else: + # Non-SP Mode: Concat first, then apply RoPE to full sequence + joint_query = torch.cat([encoder_query, query], dim=1) + joint_key = torch.cat([encoder_key, key], dim=1) + joint_value = torch.cat([encoder_value, value], dim=1) - hidden_states = self.attn( - query, - key, - value, - ) + if image_rotary_emb is not None: + # Apply RoPE to full (text + image) sequence + joint_query = apply_rotary_emb(joint_query, image_rotary_emb, sequence_dim=1) + joint_key = apply_rotary_emb(joint_key, image_rotary_emb, sequence_dim=1) + + hidden_states = self.attn( + joint_query, + joint_key, + joint_value, + ) + else: + # No added_kv_proj_dim: single stream attention (e.g., from SingleTransformerBlock) + # hidden_states is the combined (text + image) sequence + # In SP mode, image part is chunked: (B, txt_len + img_len/SP, D) + + # Check if SP is enabled and we have text_seq_len info + forward_ctx = get_forward_context() + sp_size = forward_ctx.sequence_parallel_size + text_seq_len = kwargs.get("text_seq_len", None) + use_sp_single_stream = sp_size > 1 and not forward_ctx.split_text_embed_in_sp and text_seq_len is not None + + if use_sp_single_stream: + # SP Mode for single-stream block: + # Split QKV into text and image parts, then use common helper + hidden_states = self._sp_attention_with_rope( + img_query=query[:, text_seq_len:], + img_key=key[:, text_seq_len:], + img_value=value[:, text_seq_len:], + text_query=query[:, :text_seq_len], + text_key=key[:, :text_seq_len], + text_value=value[:, :text_seq_len], + text_seq_len=text_seq_len, + image_rotary_emb=image_rotary_emb, + ) + else: + # Non-SP Mode: standard path + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + hidden_states = self.attn( + query, + key, + value, + ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) if encoder_hidden_states is not None: + # Split output back into text and image portions + # In SP mode: seq_len = txt_seq_len + img_seq_len // SP + # In non-SP mode: seq_len = txt_seq_len + img_seq_len encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 ) @@ -268,15 +394,16 @@ class LongCatImagePosEmbed(nn.Module): is_npu = ids.device.type == "npu" freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 for i in range(n_axes): - freqs_cis = get_1d_rotary_pos_embed( + cos, sin = get_1d_rotary_pos_embed( self.axes_dim[i], pos[:, i], theta=self.theta, - use_real=False, + repeat_interleave_real=True, + use_real=True, freqs_dtype=freqs_dtype, ) - cos_out.append(freqs_cis.real) - sin_out.append(freqs_cis.imag) + cos_out.append(cos) + sin_out.append(sin) freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) return freqs_cos, freqs_sin @@ -297,6 +424,13 @@ class LongCatImageTimestepEmbeddings(nn.Module): class LongCatImageSingleTransformerBlock(nn.Module): + """ + Single-stream Transformer block for LongCat with SP (Sequence Parallelism) support. + + SP handling is delegated to LongCatImageAttention via the text_seq_len parameter. + This keeps the block logic clean and centralizes SP logic in the attention layer. + """ + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) @@ -306,6 +440,7 @@ class LongCatImageSingleTransformerBlock(nn.Module): self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) + # SP handling is delegated to LongCatImageAttention via text_seq_len kwarg self.attn = LongCatImageAttention( query_dim=dim, dim_head=attention_head_dim, @@ -324,16 +459,28 @@ class LongCatImageSingleTransformerBlock(nn.Module): image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, joint_attention_kwargs: dict[str, Any] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - text_seq_len = encoder_hidden_states.shape[1] - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + """ + Forward pass for SingleTransformerBlock with SP support. - residual = hidden_states - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + SP handling is delegated to LongCatImageAttention.forward via text_seq_len kwarg. + This keeps the block logic clean and centralizes SP logic in the attention layer. + """ + text_seq_len = encoder_hidden_states.shape[1] + + # Concatenate text and image + # In SP mode: image is chunked (B, img_len/SP, D), text is full (B, txt_len, D) + combined = torch.cat([encoder_hidden_states, hidden_states], dim=1) + residual = combined + norm_hidden_states, gate = self.norm(combined, emb=temb) mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + + # Delegate SP handling to LongCatImageAttention by passing text_seq_len + # LongCatImageAttention will detect SP mode and handle text/image splitting internally joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( hidden_states=norm_hidden_states, image_rotary_emb=image_rotary_emb, + text_seq_len=text_seq_len, # Pass text_seq_len for SP mode handling **joint_attention_kwargs, ) @@ -341,6 +488,7 @@ class LongCatImageSingleTransformerBlock(nn.Module): gate = gate.unsqueeze(1) hidden_states = gate * self.proj_out(hidden_states) hidden_states = residual + hidden_states + if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) @@ -351,12 +499,9 @@ class LongCatImageSingleTransformerBlock(nn.Module): class LongCatImageTransformer2DModel(nn.Module): """ The Transformer model introduced in Flux. - """ - _repeated_blocks = [ - "LongCatImageTransformerBlock", - "LongCatImageSingleTransformerBlock", - ] + Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. + """ def __init__( self, @@ -376,6 +521,9 @@ class LongCatImageTransformer2DModel(nn.Module): self.inner_dim = num_attention_heads * attention_head_dim self.pooled_projection_dim = pooled_projection_dim + # Store parallel config for SP support + self.parallel_config = od_config.parallel_config + self.pos_embed = LongCatImagePosEmbed(theta=10000, axes_dim=axes_dims_rope) self.time_embed = LongCatImageTimestepEmbeddings(embedding_dim=self.inner_dim) @@ -423,6 +571,31 @@ class LongCatImageTransformer2DModel(nn.Module): guidance: torch.Tensor = None, return_dict: bool = True, ) -> torch.FloatTensor | Transformer2DModelOutput: + # Before: hidden_states shape = (B, img_seq_len, in_channels) + # After: hidden_states shape = (B, img_seq_len // SP, in_channels) + sp_size = self.parallel_config.sequence_parallel_size + # Store SP size in forward context for sub-modules to access + get_forward_context().sequence_parallel_size = sp_size + if sp_size > 1: + sp_world_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + original_shape = hidden_states.shape + hidden_states = torch.chunk(hidden_states, sp_world_size, dim=1)[sp_rank] + # LongCat uses dual-stream (text + image) with joint attention + # Text embeddings should be replicated across SP ranks for correctness + get_forward_context().split_text_embed_in_sp = False + # Debug log (only first forward) + if not hasattr(self, "_sp_forward_logged"): + self._sp_forward_logged = True + logger.info( + f"[LongCat Transformer] SP enabled: sp_size={sp_size}, world_size={sp_world_size}, " + f"rank={sp_rank}, original_shape={original_shape}, chunked_shape={hidden_states.shape}" + ) + else: + if not hasattr(self, "_sp_forward_logged"): + self._sp_forward_logged = True + logger.info(f"[LongCat Transformer] SP disabled: sp_size={sp_size}") + hidden_states = self.x_embedder(hidden_states) timestep = timestep.to(hidden_states.dtype) * 1000 @@ -438,6 +611,39 @@ class LongCatImageTransformer2DModel(nn.Module): else: image_rotary_emb = self.pos_embed(ids) + # SP: Chunk RoPE embeddings along sequence dimension + if self.parallel_config.sequence_parallel_size > 1: + sp_world_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + freqs_cos, freqs_sin = image_rotary_emb + txt_len = txt_ids.shape[0] + + # Split RoPE into text and image portions + # txt_freqs: (txt_seq_len, head_dim) - keep full for all ranks + # img_freqs: (img_seq_len, head_dim) -> (img_seq_len // SP, head_dim) + txt_freqs_cos = freqs_cos[:txt_len] + txt_freqs_sin = freqs_sin[:txt_len] + img_freqs_cos = freqs_cos[txt_len:] + img_freqs_sin = freqs_sin[txt_len:] + + # Chunk image RoPE for each SP rank + # img_freqs_cos: (img_seq_len // SP, head_dim) + # img_freqs_sin: (img_seq_len // SP, head_dim) + img_freqs_cos = torch.chunk(img_freqs_cos, sp_world_size, dim=0)[sp_rank] + img_freqs_sin = torch.chunk(img_freqs_sin, sp_world_size, dim=0)[sp_rank] + + # Optionally chunk text RoPE if split_text_embed_in_sp is True + if get_forward_context().split_text_embed_in_sp: + txt_freqs_cos = torch.chunk(txt_freqs_cos, sp_world_size, dim=0)[sp_rank] + txt_freqs_sin = torch.chunk(txt_freqs_sin, sp_world_size, dim=0)[sp_rank] + + # Reconstruct image_rotary_emb with chunked values + # Final shape: (txt_seq_len + img_seq_len // SP, head_dim) + image_rotary_emb = ( + torch.cat([txt_freqs_cos, img_freqs_cos], dim=0), + torch.cat([txt_freqs_sin, img_freqs_sin], dim=0), + ) + for index_block, block in enumerate(self.transformer_blocks): if torch.is_grad_enabled() and self.gradient_checkpointing and self.use_checkpoint[index_block]: encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( @@ -475,6 +681,10 @@ class LongCatImageTransformer2DModel(nn.Module): hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) + # SP: All-gather output to reconstruct full sequence + if self.parallel_config.sequence_parallel_size > 1: + output = get_sp_group().all_gather(output, dim=1) + if not return_dict: return (output,)