diff --git a/python/sglang/multimodal_gen/configs/models/encoders/mistral.py b/python/sglang/multimodal_gen/configs/models/encoders/mistral.py deleted file mode 100644 index 2ab6bbeac..000000000 --- a/python/sglang/multimodal_gen/configs/models/encoders/mistral.py +++ /dev/null @@ -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" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 287b4b484..94e5d7196 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/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 ( - 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( diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 166ab24d5..ec8f680a7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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]) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index ab3145051..fbb752429 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index adde9cd67..3ff593f81 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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 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 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py index 5554f005d..fef6ece6c 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 6661f495f..88d3872bf 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 2bdc766ea..df705db93 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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(