chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] perf: improve black-forest-labs/FLUX.2-dev (#14040)

Browse Source
This commit is contained in:
Mick
2025-11-27 14:49:52 +08:00
committed by GitHub
parent 077ca70ee4
commit 6edffc6391
8 changed files with 280 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
python/sglang/multimodal_gen/configs/models/encoders/mistral.py
View File

@@ -1,67 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models.encoders.base import (
TextEncoderArchConfig,
TextEncoderConfig,
)
def _is_transformer_layer(n: str, m) -> bool:
return "layers" in n and str.isdigit(n.split(".")[-1])
def _is_embeddings(n: str, m) -> bool:
return n.endswith("embed_tokens")
def _is_final_norm(n: str, m) -> bool:
return n.endswith("norm")
@dataclass
class Mistral3ArchConfig(TextEncoderArchConfig):
vocab_size: int = 32000
hidden_size: int = 4096
intermediate_size: int = 11008
num_hidden_layers: int = 32
num_attention_heads: int = 32
num_key_value_heads: int | None = None
hidden_act: str = "silu"
max_position_embeddings: int = 2048
initializer_range: float = 0.02
rms_norm_eps: float = 1e-6
use_cache: bool = True
pad_token_id: int = 0
bos_token_id: int = 1
eos_token_id: int = 2
pretraining_tp: int = 1
tie_word_embeddings: bool = False
rope_theta: float = 10000.0
rope_scaling: float | None = None
attention_bias: bool = False
attention_dropout: float = 0.0
mlp_bias: bool = False
head_dim: int | None = None
hidden_state_skip_layer: int = 2
text_len: int = 256
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0), # type: ignore
(".gate_up_proj", ".up_proj", 1), # type: ignore
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
)
@dataclass
class Mistral3Config(TextEncoderConfig):
arch_config: TextEncoderArchConfig = field(default_factory=Mistral3ArchConfig)
prefix: str = "mistral3"

10
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
View File

@@ -14,9 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import (
TextEncoderConfig,
)
from sglang.multimodal_gen.configs.models.encoders.base import TextEncoderArchConfig
from sglang.multimodal_gen.configs.models.encoders.mistral import (
Mistral3Config,
_is_embeddings,
from sglang.multimodal_gen.configs.models.encoders.qwen_image import (
_is_transformer_layer,
)
from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig, FluxVAEConfig
@@ -346,7 +344,7 @@ class Flux2MistralTextArchConfig(TextEncoderArchConfig):
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_transformer_layer, _is_embeddings]
default_factory=lambda: [_is_transformer_layer]
)
def __post_init__(self):
@@ -386,7 +384,6 @@ def format_text_input(prompts: List[str], system_message: str = None):
def flux_2_preprocess_text(prompt: str):
print(f"{prompt=}")
system_message = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation."
return format_text_input([prompt], system_message=system_message)
@@ -405,9 +402,6 @@ class Flux2PipelineConfig(FluxPipelineConfig):
task_type: ModelTaskType = ModelTaskType.I2I
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (Mistral3Config(),)
)
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
text_encoder_configs: tuple[EncoderConfig, ...] = field(

2
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

@@ -70,7 +70,7 @@ class RMSNorm(CustomOp):
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
shape = x.shape
x = x.view(-1, shape[-1])
x = x.reshape(-1, shape[-1])
if residual is not None:
residual_shape = residual.shape
residual = residual.view(-1, shape[-1])

1
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -148,7 +148,6 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
)
self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
# Scaled dot product attention
self.attn = USPAttention(
num_heads=num_heads,
head_size=self.head_dim,

410
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
View File

@@ -12,21 +12,24 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from typing import Any, Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.attention import AttentionModuleMixin
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.embeddings import (
TimestepEmbedding,
Timesteps,
get_1d_rotary_pos_embed,
)
from diffusers.models.normalization import AdaLayerNormContinuous
from diffusers.models.transformers.transformer_flux2 import Flux2PosEmbed
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
@@ -110,42 +113,99 @@ class Flux2FeedForward(nn.Module):
return x
class Flux2AttnProcessor:
_attention_backend = None
_parallel_config = None
class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version."
)
def __call__(
def __init__(
self,
query_dim: int,
num_heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
added_kv_proj_dim: Optional[int] = None,
added_proj_bias: Optional[bool] = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
elementwise_affine: bool = True,
):
super().__init__()
self.head_dim = dim_head
self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads
self.query_dim = query_dim
self.out_dim = out_dim if out_dim is not None else query_dim
self.heads = out_dim // dim_head if out_dim is not None else num_heads
self.use_bias = bias
self.dropout = dropout
self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias
self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
# QK Norm
self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps)
self.to_out = torch.nn.ModuleList([])
self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
self.to_out.append(torch.nn.Dropout(dropout))
if added_kv_proj_dim is not None:
self.norm_added_q = RMSNorm(dim_head, eps=eps)
self.norm_added_k = RMSNorm(dim_head, eps=eps)
self.add_q_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_k_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_v_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
self.attn = USPAttention(
num_heads=num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
causal=False,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.SAGE_ATTN,
},
)
def forward(
self,
attn: "Flux2Attention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
attention_mask: Optional[torch.Tensor] = None,
encoder_hidden_states: Optional[torch.Tensor] = None,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value = (
_get_qkv_projections(attn, hidden_states, encoder_hidden_states)
_get_qkv_projections(self, hidden_states, encoder_hidden_states)
)
query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))
query = query.unflatten(-1, (self.heads, -1))
key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
query = self.norm_q(query)
key = self.norm_k(key)
if attn.added_kv_proj_dim is not None:
encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
if self.added_kv_proj_dim is not None:
encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
encoder_query = attn.norm_added_q(encoder_query)
encoder_key = attn.norm_added_k(encoder_key)
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)
@@ -160,14 +220,8 @@ class Flux2AttnProcessor:
key, cos, sin, is_neox_style=False, interleaved=True
)
hidden_states = dispatch_attention_fn(
query,
key,
value,
attn_mask=attention_mask,
backend=self._attention_backend,
parallel_config=self._parallel_config,
)
hidden_states = self.attn(query, key, value)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
@@ -179,10 +233,10 @@ class Flux2AttnProcessor:
],
dim=1,
)
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
hidden_states = self.to_out[0](hidden_states)
hidden_states = self.to_out[1](hidden_states)
if encoder_hidden_states is not None:
return hidden_states, encoder_hidden_states
@@ -190,165 +244,6 @@ class Flux2AttnProcessor:
return hidden_states
class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
_default_processor_cls = Flux2AttnProcessor
_available_processors = [Flux2AttnProcessor]
def __init__(
self,
query_dim: int,
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
added_kv_proj_dim: Optional[int] = None,
added_proj_bias: Optional[bool] = True,
out_bias: bool = True,
eps: float = 1e-5,
out_dim: int = None,
elementwise_affine: bool = True,
processor=None,
):
super().__init__()
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
self.out_dim = out_dim if out_dim is not None else query_dim
self.heads = out_dim // dim_head if out_dim is not None else heads
self.use_bias = bias
self.dropout = dropout
self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias
self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
# QK Norm
self.norm_q = torch.nn.RMSNorm(
dim_head, eps=eps, elementwise_affine=elementwise_affine
)
self.norm_k = torch.nn.RMSNorm(
dim_head, eps=eps, elementwise_affine=elementwise_affine
)
self.to_out = torch.nn.ModuleList([])
self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
self.to_out.append(torch.nn.Dropout(dropout))
if added_kv_proj_dim is not None:
self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
self.add_q_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_k_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_v_proj = torch.nn.Linear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
if processor is None:
processor = self._default_processor_cls()
self.set_processor(processor)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(
inspect.signature(self.processor.__call__).parameters.keys()
)
unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters]
if len(unused_kwargs) > 0:
logger.warning(
f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
)
kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
return self.processor(
self,
hidden_states,
encoder_hidden_states,
attention_mask,
freqs_cis,
**kwargs,
)
class Flux2ParallelSelfAttnProcessor:
_attention_backend = None
_parallel_config = None
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version."
)
def __call__(
self,
attn: "Flux2ParallelSelfAttention",
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> torch.Tensor:
# Parallel in (QKV + MLP in) projection
hidden_states = attn.to_qkv_mlp_proj(hidden_states)
qkv, mlp_hidden_states = torch.split(
hidden_states,
[3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor],
dim=-1,
)
# Handle the attention logic
query, key, value = qkv.chunk(3, dim=-1)
query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))
query = attn.norm_q(query)
key = attn.norm_k(key)
if freqs_cis is not None:
cos, sin = freqs_cis
query = _apply_rotary_emb(
query, cos, sin, is_neox_style=False, interleaved=True
)
key = _apply_rotary_emb(
key, cos, sin, is_neox_style=False, interleaved=True
)
hidden_states = dispatch_attention_fn(
query,
key,
value,
attn_mask=attention_mask,
backend=self._attention_backend,
parallel_config=self._parallel_config,
)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
# Handle the feedforward (FF) logic
mlp_hidden_states = attn.mlp_act_fn(mlp_hidden_states)
# Concatenate and parallel output projection
hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
hidden_states = attn.to_out(hidden_states)
return hidden_states
class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
"""
Flux 2 parallel self-attention for the Flux 2 single-stream transformer blocks.
@@ -358,15 +253,13 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block.
"""
_default_processor_cls = Flux2ParallelSelfAttnProcessor
_available_processors = [Flux2ParallelSelfAttnProcessor]
# Does not support QKV fusion as the QKV projections are always fused
_supports_qkv_fusion = False
def __init__(
self,
query_dim: int,
heads: int = 8,
num_heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
@@ -376,15 +269,14 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
elementwise_affine: bool = True,
mlp_ratio: float = 4.0,
mlp_mult_factor: int = 2,
processor=None,
):
super().__init__()
self.head_dim = dim_head
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads
self.query_dim = query_dim
self.out_dim = out_dim if out_dim is not None else query_dim
self.heads = out_dim // dim_head if out_dim is not None else heads
self.heads = out_dim // dim_head if out_dim is not None else num_heads
self.use_bias = bias
self.dropout = dropout
@@ -402,21 +294,26 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
self.mlp_act_fn = Flux2SwiGLU()
# QK Norm
self.norm_q = torch.nn.RMSNorm(
dim_head, eps=eps, elementwise_affine=elementwise_affine
)
self.norm_k = torch.nn.RMSNorm(
dim_head, eps=eps, elementwise_affine=elementwise_affine
)
self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps)
# Fused attention output projection + MLP output projection
self.to_out = torch.nn.Linear(
self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias
)
if processor is None:
processor = self._default_processor_cls()
self.set_processor(processor)
self.attn = USPAttention(
num_heads=num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
causal=False,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.SAGE_ATTN,
},
)
def forward(
self,
@@ -425,16 +322,44 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs,
) -> torch.Tensor:
attn_parameters = set(
inspect.signature(self.processor.__call__).parameters.keys()
# Parallel in (QKV + MLP in) projection
hidden_states = self.to_qkv_mlp_proj(hidden_states)
qkv, mlp_hidden_states = torch.split(
hidden_states,
[3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor],
dim=-1,
)
unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters]
if len(unused_kwargs) > 0:
logger.warning(
f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
# Handle the attention logic
query, key, value = qkv.chunk(3, dim=-1)
query = query.unflatten(-1, (self.heads, -1))
key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1))
query = self.norm_q(query)
key = self.norm_k(key)
if freqs_cis is not None:
cos, sin = freqs_cis
query = _apply_rotary_emb(
query, cos, sin, is_neox_style=False, interleaved=True
)
kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
return self.processor(self, hidden_states, attention_mask, freqs_cis, **kwargs)
key = _apply_rotary_emb(
key, cos, sin, is_neox_style=False, interleaved=True
)
hidden_states = self.attn(query, key, value)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
# Handle the feedforward (FF) logic
mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
# Concatenate and parallel output projection
hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
hidden_states = self.to_out(hidden_states)
return hidden_states
class Flux2SingleTransformerBlock(nn.Module):
@@ -457,14 +382,13 @@ class Flux2SingleTransformerBlock(nn.Module):
self.attn = Flux2ParallelSelfAttention(
query_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
num_heads=num_attention_heads,
out_dim=dim,
bias=bias,
out_bias=bias,
eps=eps,
mlp_ratio=mlp_ratio,
mlp_mult_factor=2,
processor=Flux2ParallelSelfAttnProcessor(),
)
def forward(
@@ -529,13 +453,12 @@ class Flux2TransformerBlock(nn.Module):
query_dim=dim,
added_kv_proj_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
num_heads=num_attention_heads,
out_dim=dim,
bias=bias,
added_proj_bias=bias,
out_bias=bias,
eps=eps,
processor=Flux2AttnProcessor(),
)
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -674,6 +597,38 @@ class Flux2Modulation(nn.Module):
)
class Flux2PosEmbed(nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
def forward(self, ids: torch.Tensor) -> torch.Tensor:
# Expected ids shape: [S, len(self.axes_dim)]
cos_out = []
sin_out = []
pos = ids.float()
is_mps = ids.device.type == "mps"
is_npu = ids.device.type == "npu"
freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
# Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1]
for i in range(len(self.axes_dim)):
cos, sin = get_1d_rotary_pos_embed(
self.axes_dim[i],
pos[..., i],
theta=self.theta,
repeat_interleave_real=True,
use_real=True,
freqs_dtype=freqs_dtype,
)
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
class Flux2Transformer2DModel(CachableDiT):
"""
The Transformer model introduced in Flux 2.
@@ -701,7 +656,6 @@ class Flux2Transformer2DModel(CachableDiT):
self.inner_dim = num_attention_heads * attention_head_dim
# 1. Sinusoidal positional embedding for RoPE on image and text tokens
# self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
# 2. Combined timestep + guidance embedding

17
python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py
View File

@@ -18,9 +18,9 @@ from typing import Iterable, Optional, Union
import torch
from torch import nn
from transformers import Cache, DynamicCache, LlavaConfig, Mistral3Config, MistralConfig
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import create_causal_mask
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.models.mistral3.modeling_mistral3 import (
Mistral3CausalLMOutputWithPast,
Mistral3ModelOutputWithPast,
@@ -32,7 +32,9 @@ from transformers.models.mistral.modeling_mistral import (
apply_rotary_pos_emb,
)
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
@@ -89,6 +91,17 @@ class MistralAttention(nn.Module):
self.is_causal = True
self.num_heads = config.num_attention_heads
self.num_key_value_heads = config.num_key_value_heads
self.attn = USPAttention(
num_heads=self.num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
causal=False,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
},
)
def forward(
self,
@@ -118,7 +131,7 @@ class MistralAttention(nn.Module):
key_states, value_states, self.layer_idx, cache_kwargs
)
attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"]
attention_interface = sdpa_attention_forward
attn_output, attn_weights = attention_interface(
self,
query_states,

67
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -228,6 +228,73 @@
"expected_avg_denoise_ms": 165.83,
"expected_median_denoise_ms": 169.33
},
"flux_2_image_t2i": {
"stages_ms": {
"InputValidationStage": 0.05,
"TextEncodingStage": 530.93,
"ImageVAEEncodingStage": 0.0,
"ConditioningStage": 0.02,
"LatentPreparationStage": 12.71,
"TimestepPreparationStage": 2.91,
"DenoisingStage": 26403.1,
"DecodingStage": 286.85
},
"denoise_step_ms": {
"0": 511.3,
"1": 132.57,
"2": 541.19,
"3": 518.93,
"4": 541.2,
"5": 520.28,
"6": 532.47,
"7": 525.68,
"8": 538.25,
"9": 525.84,
"10": 526.13,
"11": 525.67,
"12": 524.63,
"13": 530.57,
"14": 530.46,
"15": 529.94,
"16": 532.47,
"17": 527.88,
"18": 527.7,
"19": 525.08,
"20": 525.72,
"21": 529.3,
"22": 522.59,
"23": 529.75,
"24": 523.46,
"25": 528.72,
"26": 526.92,
"27": 528.62,
"28": 522.77,
"29": 528.35,
"30": 528.05,
"31": 528.89,
"32": 525.34,
"33": 530.36,
"34": 529.19,
"35": 526.92,
"36": 528.16,
"37": 525.03,
"38": 527.33,
"39": 527.96,
"40": 527.81,
"41": 524.79,
"42": 528.46,
"43": 532.49,
"44": 526.95,
"45": 533.14,
"46": 529.32,
"47": 528.51,
"48": 532.14,
"49": 529.29
},
"expected_e2e_ms": 27648.69,
"expected_avg_denoise_ms": 520.09,
"expected_median_denoise_ms": 528.0
},
"flux_image_t2i_2_gpus": {
"stages_ms": {
"InputValidationStage": 0.03,

26
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -242,19 +242,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
output_size="1024x1024",
),
),
# DiffusionTestCase(
# "flux_2_image_t2i",
# DiffusionServerArgs(
# model_path="black-forest-labs/FLUX.2-dev",
# modality="image",
# warmup_text=1,
# warmup_edit=0,
# ),
# DiffusionSamplingParams(
# prompt="A futuristic cityscape at sunset with flying cars",
# output_size="1024x1024",
# ),
# ),
DiffusionTestCase(
"flux_2_image_t2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
modality="image",
warmup_text=1,
warmup_edit=0,
),
DiffusionSamplingParams(
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
),
),
# === Text and Image to Image (TI2I) ===
# TODO: Timeout with Torch2.9. Add back when it can pass CI
# DiffusionTestCase(