anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] support SP for FLUX.2-dev (#3244)

Browse Source
This commit is contained in:
JohnJan authored and GitHub committed 2026-05-27 20:43:02 +08:00
1 parent 9abf9545df
commit 36dc32ff81
4 files changed
+277 -48
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion_features.md
+1 -1
View File
@@ -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** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
tests/e2e/online_serving/test_flux_2_dev_expansion.py
+27 -1
View File
@@ -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,
),
]
vllm_omni/diffusion/models/flux2/flux2_transformer.py
+246 -45
View File
@@ -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:, ...]
vllm_omni/diffusion/models/flux2/pipeline_flux2.py
+3 -1
View File
@@ -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)