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

python/sglang/multimodal_gen/configs/models/encoders/mistral.py

@@ -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"

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -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 (
+from sglang.multimodal_gen.configs.models.encoders.qwen_image import (
-    Mistral3Config,
-    _is_embeddings,
     _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(

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -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])

python/sglang/multimodal_gen/runtime/models/dits/flux.py

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

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -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 (
-from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+    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:
+class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
-    _attention_backend = None
-    _parallel_config = None
 
-    def __init__(self):
+    def __init__(
-        if not hasattr(F, "scaled_dot_product_attention"):
+        self,
-            raise ImportError(
+        query_dim: int,
-                f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version."
+        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 __call__(
+    def forward(
         self,
-        attn: "Flux2Attention",
         hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: 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))
+        query = query.unflatten(-1, (self.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
+        key = key.unflatten(-1, (self.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        value = value.unflatten(-1, (self.heads, -1))
 
-        query = attn.norm_q(query)
+        query = self.norm_q(query)
-        key = attn.norm_k(key)
+        key = self.norm_k(key)
 
-        if attn.added_kv_proj_dim is not None:
+        if self.added_kv_proj_dim is not None:
-            encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
 
-            encoder_query = attn.norm_added_q(encoder_query)
+            encoder_query = self.norm_added_q(encoder_query)
-            encoder_key = attn.norm_added_k(encoder_key)
+            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(
+        hidden_states = self.attn(query, key, value)
-            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)
 
@@ -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 = self.to_out[0](hidden_states)
-        hidden_states = attn.to_out[1](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(
+        self.norm_q = RMSNorm(dim_head, eps=eps)
-            dim_head, eps=eps, elementwise_affine=elementwise_affine
+        self.norm_k = RMSNorm(dim_head, eps=eps)
-        )
-        self.norm_k = torch.nn.RMSNorm(
-            dim_head, eps=eps, elementwise_affine=elementwise_affine
-        )
 
         # 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:
+        self.attn = USPAttention(
-            processor = self._default_processor_cls()
+            num_heads=num_heads,
-        self.set_processor(processor)
+            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(
+        # Parallel in (QKV + MLP in) projection
-            inspect.signature(self.processor.__call__).parameters.keys()
+        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:
+        # Handle the attention logic
-            logger.warning(
+        query, key, value = qkv.chunk(3, dim=-1)
-                f"joint_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+
+        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}
+            key = _apply_rotary_emb(
-        return self.processor(self, hidden_states, attention_mask, freqs_cis, **kwargs)
+                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

python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py

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

python/sglang/multimodal_gen/test/server/perf_baselines.json

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

python/sglang/multimodal_gen/test/server/testcase_configs.py

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