Skip to content

[Feat] support SP for FLUX.2-klein #1250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
RuixiangMa wants to merge 13 commits into
base: main
Choose a base branch
from
+283 −56
Open

[Feat] support SP for FLUX.2-klein #1250

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
e8fb739
[Feat] support SP for FLUX.2-klein
RuixiangMa Feb 6, 2026
f2daa52
upd
RuixiangMa Feb 7, 2026
225ec8a
upd
RuixiangMa Feb 7, 2026
c57c866
upd
RuixiangMa Feb 8, 2026
b33ce10
upd
RuixiangMa Feb 9, 2026
7166799
Merge branch 'main' into spforflux2klein
RuixiangMa Mar 3, 2026
780ed85
upd
RuixiangMa Mar 3, 2026
8a7e687
upd
RuixiangMa Mar 9, 2026
63c4658
upd
RuixiangMa Mar 9, 2026
219d4f7
Merge branch 'main' into spforflux2klein
RuixiangMa Mar 10, 2026
876d492
upd
RuixiangMa Mar 10, 2026
ec5f9f2
Merge branch 'main' into spforflux2klein
RuixiangMa Mar 11, 2026
6a1e5f2
upd
RuixiangMa Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/user_guide/diffusion/parallelism_acceleration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ The following table shows which models are currently supported by parallelism me
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |

!!! note "TP Limitations for Diffusion Models"
Expand Down
59 changes: 55 additions & 4 deletions vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
)
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.normalization import AdaLayerNormContinuous
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
Expand All @@ -39,8 +40,16 @@

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.distributed.parallel_state import (
get_sp_group,
)
from vllm_omni.diffusion.distributed.sp_sharding import sp_shard_with_padding
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.diffusion.layers.rope import RotaryEmbedding

logger = init_logger(__name__)


class Flux2SwiGLU(nn.Module):
"""SwiGLU activation used by Flux2."""
Expand Down Expand Up @@ -334,6 +343,12 @@ def forward(


class Flux2SingleTransformerBlock(nn.Module):
"""
Single-stream Transformer block for Flux 2 with SP (Sequence Parallelism) support.

SP handling is delegated to Flux2Attention via the forward context.
"""

def __init__(
self,
dim: int,
Expand Down Expand Up @@ -367,6 +382,13 @@ def forward(
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)
Expand Down Expand Up @@ -556,6 +578,8 @@ def forward(self, temb: torch.Tensor) -> tuple[tuple[torch.Tensor, torch.Tensor,
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"]
Expand All @@ -580,6 +604,7 @@ def __init__(
rope_theta: int = 2000,
eps: float = 1e-6,
guidance_embeds: bool = True,
od_config: OmniDiffusionConfig = None,
):
super().__init__()
self.out_channels = out_channels or in_channels
Expand All @@ -601,6 +626,13 @@ def __init__(
guidance_embeds=guidance_embeds,
)

if od_config is not None:
self.parallel_config = od_config.parallel_config
else:
from vllm_omni.diffusion.data import DiffusionParallelConfig

self.parallel_config = DiffusionParallelConfig()

self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
in_channels=timestep_guidance_channels,
Expand Down Expand Up @@ -672,6 +704,12 @@ def forward(

num_txt_tokens = encoder_hidden_states.shape[1]

sp_size = self.parallel_config.sequence_parallel_size
sp_pad_size = 0
if sp_size > 1:
hidden_states, sp_pad_size = sp_shard_with_padding(hidden_states, dim=1)
Copy link
Collaborator

@wtomin wtomin Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am a little concerned about the padding behavior here. It looks like flux2_klein does not support hidden_states_mask (like qwen_image), therefore, the padded tokens will participate in the attention computation, which is wrong.

If using Non-intrusive _sp_plan, sp_padding_size will be automatically set in get_forward_context(), however, hidden_states_mask is required to exclude the padded tokens.

# vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
        hidden_states_mask = None  # default
        if self.parallel_config is not None and self.parallel_config.sequence_parallel_size > 1:
            ctx = get_forward_context()
            if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
                # Create mask for the full (padded) sequence
                # valid positions = True, padding positions = False
                batch_size = hidden_states.shape[0]
                padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
                hidden_states_mask = torch.ones(
                    batch_size,
                    padded_seq_len,
                    dtype=torch.bool,
                    device=hidden_states.device,
                )
                hidden_states_mask[:, ctx.sp_original_seq_len :] = False

Copy link
Collaborator

@wtomin wtomin Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@RuixiangMa I would like to invite you to join the discussion here #1324.

Copy link
Collaborator

@wtomin wtomin Feb 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have two suggestions:

  1. refer to [Feature] Support Wan2.2 output with irregular shapes #1279 , you may use auto_pad=True and use hidden_states_mask to exclude the padded tokens.
  2. You may set SequenceParallelInput(auto_pad=False), this will raise an error when seq_len not divisible by sp_size, which I think might be rare for Flux.2-klein. We can take care of rare cases in the future, after we make our discussion clear in [RFC]: Ulysses-SP Constraints Solution #1324.

get_forward_context().split_text_embed_in_sp = False

timestep = timestep.to(hidden_states.dtype) * 1000
if guidance is not None:
guidance = guidance.to(hidden_states.dtype) * 1000
Expand All @@ -690,11 +728,19 @@ def forward(
if txt_ids.ndim == 3:
txt_ids = txt_ids[0]

image_rotary_emb = self.pos_embed(img_ids)
text_rotary_emb = self.pos_embed(txt_ids)
img_freqs_cos, img_freqs_sin = self.pos_embed(img_ids)
txt_freqs_cos, txt_freqs_sin = self.pos_embed(txt_ids)

if sp_size > 1:
img_freqs_cos, _ = sp_shard_with_padding(img_freqs_cos, dim=0)
img_freqs_sin, _ = sp_shard_with_padding(img_freqs_sin, dim=0)
if get_forward_context().split_text_embed_in_sp:
txt_freqs_cos, _ = sp_shard_with_padding(txt_freqs_cos, dim=0)
txt_freqs_sin, _ = sp_shard_with_padding(txt_freqs_sin, dim=0)

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),
)

for block in self.transformer_blocks:
Expand Down Expand Up @@ -722,6 +768,11 @@ def forward(
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)

if self.parallel_config.sequence_parallel_size > 1:
output = get_sp_group().all_gather(output, dim=1)
if sp_pad_size > 0:
output = output[:, :-sp_pad_size, ...]

if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
Expand Down
2 changes: 1 addition & 1 deletion vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ def __init__(
).to(self._execution_device)

transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
self.transformer = Flux2Transformer2DModel(**transformer_kwargs)
self.transformer = Flux2Transformer2DModel(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)
Expand Down