[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 (
-    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(