From 36dc32ff810d337de67f427d828a0573174cd4de Mon Sep 17 00:00:00 2001 From: JohnJan Date: Wed, 27 May 2026 20:43:02 +0800 Subject: [PATCH] [Feature] support SP for FLUX.2-dev (#3244) --- docs/user_guide/diffusion_features.md | 2 +- .../test_flux_2_dev_expansion.py | 28 +- .../models/flux2/flux2_transformer.py | 291 +++++++++++++++--- .../diffusion/models/flux2/pipeline_flux2.py | 4 +- 4 files changed, 277 insertions(+), 48 deletions(-) diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 2a296a059..52c24ab9e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -113,7 +113,7 @@ The following tables show which models support each feature: | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index 953cb448a..aaf118112 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -3,6 +3,8 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode. Coverage: - CPU offload +- Ulysses sequence parallelism +- Ring sequence parallelism This test verifies that FLUX.2-dev can be launched with CPU offload enabled, accepts text-to-image requests through the OpenAI-compatible API, and returns @@ -28,7 +30,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) def _get_flux_2_dev_feature_cases(model: str): - """Return FLUX.2-dev diffusion feature cases for CPU offload.""" + """Return FLUX.2-dev diffusion feature cases for CPU offload and SP.""" return [ pytest.param( @@ -53,6 +55,30 @@ def _get_flux_2_dev_feature_cases(model: str): id="parallel_cfg_2", marks=PARALLEL_FEATURE_MARKS, ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + "--ulysses-degree", + "2", + ], + ), + id="ulysses_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + "--ring-degree", + "2", + ], + ), + id="ring_2", + marks=PARALLEL_FEATURE_MARKS, + ), ] diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index ef2cdd025..0c6896c13 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -25,7 +25,13 @@ 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.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -91,6 +97,7 @@ class Flux2FeedForward(nn.Module): class Flux2Attention(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, query_dim: int, heads: int = 8, dim_head: int = 64, @@ -106,6 +113,7 @@ class Flux2Attention(nn.Module): prefix: str = "", ): super().__init__() + self.parallel_config = parallel_config self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.query_dim = query_dim @@ -206,34 +214,82 @@ class Flux2Attention(nn.Module): 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) - - if image_rotary_emb is not None: - cos, sin = image_rotary_emb - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) - - attn_metadata = None - if attention_mask is not None: - if attention_mask.dim() == 3: - attention_mask = attention_mask.unsqueeze(1) - attn_metadata = AttentionMetadata(attn_mask=attention_mask) - - hidden_states = self.attn(query, key, value, attn_metadata) - hidden_states = hidden_states.flatten(2, 3).to(query.dtype) - - if encoder_hidden_states is not None: - context_len = encoder_hidden_states.shape[1] - encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( - [context_len, hidden_states.shape[1] - context_len], - dim=1, + sp_size = self.parallel_config.sequence_parallel_size + use_sp_joint_attention = ( + sp_size is not None + and sp_size > 1 + and is_forward_context_available() + and not get_forward_context().split_text_embed_in_sp ) - # Contiguous for FP8 quantization in RowParallelLinear - encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) + + if use_sp_joint_attention and image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + txt_len = encoder_query.shape[1] + txt_cos, img_cos = cos[:txt_len], cos[txt_len:] + txt_sin, img_sin = sin[:txt_len], sin[txt_len:] + + query = self.rope(query, img_cos, img_sin) + key = self.rope(key, img_cos, img_sin) + encoder_query = self.rope(encoder_query, txt_cos, txt_sin) + encoder_key = self.rope(encoder_key, txt_cos, txt_sin) + + attn_metadata = AttentionMetadata( + joint_query=encoder_query, + joint_key=encoder_key, + joint_value=encoder_value, + joint_strategy="front", + ) + hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None) + encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None) + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask + if encoder_hidden_states_mask is not None: + attn_metadata.joint_attn_mask = encoder_hidden_states_mask + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + txt_len = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [txt_len, hidden_states.shape[1] - txt_len], + dim=1, + ) + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) + else: + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + context_len = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [context_len, hidden_states.shape[1] - context_len], + dim=1, + ) + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) + else: + query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) hidden_states = self.to_out[0](hidden_states.contiguous()) hidden_states = self.to_out[1](hidden_states) @@ -250,6 +306,7 @@ class Flux2ParallelSelfAttention(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, query_dim: int, heads: int = 8, dim_head: int = 64, @@ -265,6 +322,7 @@ class Flux2ParallelSelfAttention(nn.Module): prefix: str = "", ): super().__init__() + self.parallel_config = parallel_config self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.query_dim = query_dim @@ -327,20 +385,59 @@ class Flux2ParallelSelfAttention(nn.Module): query = self.norm_q(query) key = self.norm_k(key) - if image_rotary_emb is not None: + sp_size = self.parallel_config.sequence_parallel_size + text_seq_len = kwargs.get("text_seq_len", None) + use_sp_single_stream = ( + sp_size is not None + and sp_size > 1 + and text_seq_len is not None + and is_forward_context_available() + and not get_forward_context().split_text_embed_in_sp + ) + + if use_sp_single_stream and image_rotary_emb is not None: cos, sin = image_rotary_emb cos = cos.to(query.dtype) sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:] + txt_sin, img_sin = sin[:text_seq_len], sin[text_seq_len:] - attn_metadata = None - if attention_mask is not None: - if attention_mask.dim() == 3: - attention_mask = attention_mask.unsqueeze(1) - attn_metadata = AttentionMetadata(attn_mask=attention_mask) + 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] - attn_output = self.attn(query, key, value, attn_metadata) + img_query = self.rope(img_query, img_cos, img_sin) + img_key = self.rope(img_key, img_cos, img_sin) + text_query = self.rope(text_query, txt_cos, txt_sin) + text_key = self.rope(text_key, txt_cos, txt_sin) + + attn_metadata = AttentionMetadata( + joint_query=text_query, + joint_key=text_key, + joint_value=text_value, + joint_strategy="front", + ) + hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None) + encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None) + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask + if encoder_hidden_states_mask is not None: + attn_metadata.joint_attn_mask = encoder_hidden_states_mask + + attn_output = self.attn(img_query, img_key, img_value, attn_metadata) + else: + query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + attn_output = self.attn(query, key, value, attn_metadata) attn_output = attn_output.flatten(2, 3).to(query.dtype) mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) @@ -352,6 +449,7 @@ class Flux2ParallelSelfAttention(nn.Module): class Flux2SingleTransformerBlock(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, dim: int, num_attention_heads: int, attention_head_dim: int, @@ -364,6 +462,7 @@ class Flux2SingleTransformerBlock(nn.Module): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2ParallelSelfAttention( + parallel_config=parallel_config, query_dim=dim, dim_head=attention_head_dim, heads=num_attention_heads, @@ -387,6 +486,13 @@ class Flux2SingleTransformerBlock(nn.Module): split_hidden_states: bool = False, text_seq_len: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for Flux2SingleTransformerBlock with SP support. + + In SP mode: image hidden_states is chunked (B, img_len/SP, D), + text encoder_hidden_states is full (B, txt_len, D). + The block concatenates them for joint attention. + """ if encoder_hidden_states is not None: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -400,6 +506,7 @@ class Flux2SingleTransformerBlock(nn.Module): attn_output = self.attn( hidden_states=norm_hidden_states, image_rotary_emb=image_rotary_emb, + text_seq_len=text_seq_len, **joint_attention_kwargs, ) @@ -416,6 +523,7 @@ class Flux2SingleTransformerBlock(nn.Module): class Flux2TransformerBlock(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, dim: int, num_attention_heads: int, attention_head_dim: int, @@ -432,6 +540,7 @@ class Flux2TransformerBlock(nn.Module): self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2Attention( + parallel_config=parallel_config, query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, @@ -543,6 +652,47 @@ class Flux2PosEmbed(nn.Module): return freqs_cos, freqs_sin +class Flux2RopePrepare(nn.Module): + """Prepares RoPE embeddings for sequence parallel. + + This module encapsulates the RoPE computation for Flux.2-dev. + For dual-stream attention, text components (outputs 0, 1) are replicated + across SP ranks, while image components (outputs 2, 3) are sharded. + + NOTE: The hidden_states projection is handled separately in forward() + so that _sp_plan can shard it at the root level. + """ + + def __init__(self, pos_embed: Flux2PosEmbed): + super().__init__() + self.pos_embed = pos_embed + + def forward( + self, + img_ids: torch.Tensor, + txt_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Compute RoPE embeddings for text and image sequences. + + Args: + img_ids: Image position IDs (img_seq_len, n_axes) + txt_ids: Text position IDs (txt_seq_len, n_axes) + + Returns: + Tuple of cosine / sine components for text & image + in the order: (txt_cos, txt_sin, img_cos, img_sin) + + NOTE: careful about output orders if this is refactored in the + future; we need to match the _sp_plan indices, since text + components (0 & 1) need to be replicated across SP ranks, + while image components (2 & 3) must be sharded. + """ + img_freqs_cos, img_freqs_sin = self.pos_embed(img_ids) + txt_freqs_cos, txt_freqs_sin = self.pos_embed(txt_ids) + return txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin + + class Flux2TimestepGuidanceEmbeddings(nn.Module): def __init__( self, @@ -601,9 +751,22 @@ class Flux2Modulation(nn.Module): class Flux2Transformer2DModel(nn.Module): """ The Transformer model introduced in Flux 2. + + Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + _sp_plan = { + "": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), + }, + "rope_prepare": { + 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + }, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + """SP plan: shard hidden_states at root level, shard img_freqs at rope_prepare, gather output at proj_out.""" @staticmethod def _is_transformer_block(name: str, module) -> bool: @@ -628,6 +791,7 @@ class Flux2Transformer2DModel(nn.Module): eps: float = 1e-6, guidance_embeds: bool = True, quant_config: "QuantizationConfig | None" = None, + od_config: OmniDiffusionConfig | None = None, ): super().__init__() self.guidance_embeds = guidance_embeds @@ -652,7 +816,13 @@ class Flux2Transformer2DModel(nn.Module): guidance_embeds=guidance_embeds, ) + if od_config is not None: + self.parallel_config = od_config.parallel_config + else: + self.parallel_config = DiffusionParallelConfig() + self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope) + self.rope_prepare = Flux2RopePrepare(self.pos_embed) self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, @@ -670,6 +840,7 @@ class Flux2Transformer2DModel(nn.Module): self.transformer_blocks = nn.ModuleList( [ Flux2TransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -686,6 +857,7 @@ class Flux2Transformer2DModel(nn.Module): self.single_transformer_blocks = nn.ModuleList( [ Flux2SingleTransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -722,6 +894,9 @@ class Flux2Transformer2DModel(nn.Module): joint_attention_kwargs = joint_attention_kwargs or {} num_txt_tokens = encoder_hidden_states.shape[1] + sp_size = self.parallel_config.sequence_parallel_size + if sp_size and sp_size > 1 and is_forward_context_available(): + get_forward_context().split_text_embed_in_sp = False timestep = timestep.to(hidden_states.dtype) * 1000 if guidance is not None: @@ -742,18 +917,43 @@ class Flux2Transformer2DModel(nn.Module): txt_ids = txt_ids[0] if is_torch_npu_available(): - freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu()) - image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu()) - freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu()) - text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu()) + txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids.cpu(), txt_ids.cpu()) + txt_freqs_cos = txt_freqs_cos.npu() + txt_freqs_sin = txt_freqs_sin.npu() + img_freqs_cos = img_freqs_cos.npu() + img_freqs_sin = img_freqs_sin.npu() else: - image_rotary_emb = self.pos_embed(img_ids) - text_rotary_emb = self.pos_embed(txt_ids) + txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids, txt_ids) concat_rotary_emb = ( - torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), - torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), + torch.cat([txt_freqs_cos, img_freqs_cos], dim=0), + torch.cat([txt_freqs_sin, img_freqs_sin], dim=0), ) + # Create separate masks for image and text portions for Ulysses SP joint attention + hidden_states_mask = None + encoder_hidden_states_mask = None + if is_forward_context_available(): + ctx = get_forward_context() + else: + ctx = None + if ctx is not None and ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + batch_size = hidden_states.shape[0] + img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + hidden_states_mask = torch.ones( + batch_size, + img_padded_seq_len, + dtype=torch.bool, + device=hidden_states.device, + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + if hidden_states_mask.all(): + hidden_states_mask = None + + if hidden_states_mask is not None: + joint_attention_kwargs["hidden_states_mask"] = hidden_states_mask + if encoder_hidden_states_mask is not None: + joint_attention_kwargs["encoder_hidden_states_mask"] = encoder_hidden_states_mask + for index_block, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, @@ -773,6 +973,7 @@ class Flux2Transformer2DModel(nn.Module): temb_mod_params=single_stream_mod, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, + text_seq_len=num_txt_tokens, ) hidden_states = hidden_states[:, num_txt_tokens:, ...] diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 8700309e5..948f922d4 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -398,7 +398,9 @@ class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarM self._execution_device ) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel) - self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs) + self.transformer = Flux2Transformer2DModel( + quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs + ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)