[diffusion] model: support flux Klein (#17173)

python/sglang/multimodal_gen/README.md

@@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus
 ## Key Features
 
 SGLang Diffusion has the following features:
-  - Broad model support: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux
+  - Broad model support: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image
   - Fast inference speed: enpowered by highly optimized kernel from sgl-kernel and efficient scheduler loop
   - Ease of use: OpenAI-compatible api, CLI, and python sdk support
   - Multi-platform support: NVIDIA GPUs (H100, H200, A100, B200, 4090) and AMD GPUs (MI300X, MI325X)

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

@@ -11,6 +11,7 @@ from sglang.multimodal_gen.configs.models.encoders.clip import (
     CLIPVisionConfig,
 )
 from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
+from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
 from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
 
 __all__ = [
@@ -21,5 +22,6 @@ __all__ = [
     "CLIPTextConfig",
     "CLIPVisionConfig",
     "LlamaConfig",
+    "Qwen3TextConfig",
     "T5Config",
 ]

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

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3 text encoder configuration for SGLang diffusion models."""
+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 Qwen3TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for Qwen3 text encoder.
+
+    Qwen3 is similar to LLaMA but with QK-Norm (RMSNorm on Q and K before attention).
+    """
+
+    vocab_size: int = 151936
+    hidden_size: int = 2560
+    intermediate_size: int = 9728
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 40960
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 151643
+    bos_token_id: int = 151643
+    eos_token_id: int = 151645
+    tie_word_embeddings: bool = True
+    rope_theta: float = 1000000.0
+    rope_scaling: dict | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int = 128
+    text_len: int = 512
+    output_hidden_states: bool = True  # Klein needs hidden states from layers 9, 18, 27
+
+    # Stacked params for weight loading with tensor parallelism
+    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),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+    )
+
+    # FSDP sharding conditions for CPU offload
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
+    )
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Qwen3TextConfig(TextEncoderConfig):
+    """Top-level config for Qwen3 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3TextArchConfig)
+    prefix: str = "qwen3"

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

@@ -7,7 +7,11 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
 from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import (
     DiffusersGenericPipelineConfig,
 )
-from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.flux import (
+    Flux2KleinPipelineConfig,
+    Flux2PipelineConfig,
+    FluxPipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
     Flux2FinetunedPipelineConfig,
 )
@@ -29,6 +33,8 @@ __all__ = [
     "HunyuanConfig",
     "FastHunyuanConfig",
     "FluxPipelineConfig",
+    "Flux2PipelineConfig",
+    "Flux2KleinPipelineConfig",
     "Flux2FinetunedPipelineConfig",
     "PipelineConfig",
     "SlidingTileAttnConfig",

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

@@ -15,6 +15,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.qwen3 import Qwen3TextConfig
 from sglang.multimodal_gen.configs.models.encoders.qwen_image import (
     _is_transformer_layer,
 )
@@ -340,6 +341,20 @@ def flux2_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Te
     return prompt_embeds
 
 
+def flux2_klein_postprocess_text(
+    outputs: BaseEncoderOutput, _text_inputs
+) -> torch.Tensor:
+    hidden_states_layers: list[int] = [9, 18, 27]
+
+    out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1)
+    batch_size, num_channels, seq_len, hidden_dim = out.shape
+    prompt_embeds = out.permute(0, 2, 1, 3).reshape(
+        batch_size, seq_len, num_channels * hidden_dim
+    )
+
+    return prompt_embeds
+
+
 @dataclass
 class Flux2MistralTextArchConfig(TextEncoderArchConfig):
     stacked_params_mapping: list[tuple[str, str, str]] = field(
@@ -607,3 +622,61 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         # remove noise over input image
         noise = noise[:, : latents.size(1) :]
         return noise
+
+
+@dataclass
+class Flux2KleinPipelineConfig(Flux2PipelineConfig):
+    # Klein is distilled, so no guidance embeddings
+    should_use_guidance: bool = False
+    task_type: ModelTaskType = ModelTaskType.T2I
+
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Qwen3TextConfig(),)
+    )
+
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text,),
+    )
+
+    postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (flux2_klein_postprocess_text,)
+    )
+
+    def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
+        if prompts and isinstance(prompts[0], list):
+            prompts = [p for prompt in prompts for p in prompt]
+
+        def _apply_chat_template(prompt: str) -> str:
+            messages = [{"role": "user", "content": prompt}]
+            try:
+                return tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    enable_thinking=False,
+                )
+            except TypeError:
+                return tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                )
+
+        texts = [_apply_chat_template(prompt) for prompt in prompts]
+
+        tok_kwargs = dict(tok_kwargs or {})
+        max_length = tok_kwargs.pop("max_length", 512)
+        padding = tok_kwargs.pop("padding", "max_length")
+        truncation = tok_kwargs.pop("truncation", True)
+        return_tensors = tok_kwargs.pop("return_tensors", "pt")
+
+        return tokenizer(
+            texts,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            **tok_kwargs,
+        )

python/sglang/multimodal_gen/docs/support_matrix.md

@@ -38,12 +38,16 @@ default parameters when initializing and generating videos.
 
 ### Image Generation Models
 
-| Model Name      | HuggingFace Model ID           | Resolutions    |
-|:----------------|:-------------------------------|:---------------|
-| FLUX.1-dev      | `black-forest-labs/FLUX.1-dev` | Any resolution |
-| FLUX.2-dev      | `black-forest-labs/FLUX.2-dev` | Any resolution |
-| Qwen Image      | `Qwen/Qwen-Image`              | Any resolution |
-| Qwen Image Edit | `Qwen/Qwen-Image-Edit`         | Any resolution |
+| Model Name       | HuggingFace Model ID                    | Resolutions    |
+|:-----------------|:----------------------------------------|:---------------|
+| FLUX.1-dev       | `black-forest-labs/FLUX.1-dev`          | Any resolution |
+| FLUX.2-dev       | `black-forest-labs/FLUX.2-dev`          | Any resolution |
+| FLUX.2-Klein     | `black-forest-labs/FLUX.2-klein-4B`     | Any resolution |
+| Z-Image-Turbo    | `Tongyi-MAI/Z-Image-Turbo`              | Any resolution |
+| GLM-Image        | `zai-org/GLM-Image`                     | Any resolution |
+| Qwen Image       | `Qwen/Qwen-Image`                       | Any resolution |
+| Qwen Image 2512  | `Qwen/Qwen-Image-2512`                  | Any resolution |
+| Qwen Image Edit  | `Qwen/Qwen-Image-Edit`                  | Any resolution |
 
 ## Verified LoRA Examples
 

python/sglang/multimodal_gen/registry.py

@@ -38,7 +38,10 @@ from sglang.multimodal_gen.configs.pipeline_configs import (
     ZImagePipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
-from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.flux import (
+    Flux2KleinPipelineConfig,
+    Flux2PipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.glm_image import (
     GlmImagePipelineConfig,
 )
@@ -534,13 +537,27 @@ def _register_configs():
         ],
         model_detectors=[lambda hf_id: "flux.1" in hf_id.lower()],
     )
+    register_configs(
+        sampling_param_cls=FluxSamplingParams,
+        pipeline_config_cls=Flux2KleinPipelineConfig,
+        hf_model_paths=[
+            "black-forest-labs/FLUX.2-klein-4B",
+            "black-forest-labs/FLUX.2-klein-9B",
+        ],
+        model_detectors=[
+            lambda hf_id: "flux.2-klein" in hf_id.lower()
+            or "flux2-klein" in hf_id.lower()
+        ],
+    )
     register_configs(
         sampling_param_cls=FluxSamplingParams,
         pipeline_config_cls=Flux2PipelineConfig,
         hf_model_paths=[
             "black-forest-labs/FLUX.2-dev",
         ],
-        model_detectors=[lambda hf_id: "flux.2" in hf_id.lower()],
+        model_detectors=[
+            lambda hf_id: "flux.2" in hf_id.lower() and "klein" not in hf_id.lower()
+        ],
     )
     register_configs(
         sampling_param_cls=ZImageSamplingParams,

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

@@ -550,7 +550,11 @@ class Flux2TransformerBlock(nn.Module):
 
 class Flux2TimestepGuidanceEmbeddings(nn.Module):
     def __init__(
-        self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False
+        self,
+        in_channels: int = 256,
+        embedding_dim: int = 6144,
+        bias: bool = False,
+        guidance_embeds: bool = True,
     ):
         super().__init__()
 
@@ -561,24 +565,32 @@ class Flux2TimestepGuidanceEmbeddings(nn.Module):
             in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias
         )
 
-        self.guidance_embedder = TimestepEmbedding(
-            in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias
-        )
+        if guidance_embeds:
+            self.guidance_embedder = TimestepEmbedding(
+                in_channels=in_channels,
+                time_embed_dim=embedding_dim,
+                sample_proj_bias=bias,
+            )
+        else:
+            self.guidance_embedder = None
 
-    def forward(self, timestep: torch.Tensor, guidance: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
         timesteps_proj = self.time_proj(timestep)
         timesteps_emb = self.timestep_embedder(
             timesteps_proj.to(timestep.dtype)
         )  # (N, D)
 
-        guidance_proj = self.time_proj(guidance)
-        guidance_emb = self.guidance_embedder(
-            guidance_proj.to(guidance.dtype)
-        )  # (N, D)
-
-        time_guidance_emb = timesteps_emb + guidance_emb
-
-        return time_guidance_emb
+        if guidance is not None and self.guidance_embedder is not None:
+            guidance_proj = self.time_proj(guidance)
+            guidance_emb = self.guidance_embedder(
+                guidance_proj.to(guidance.dtype)
+            )  # (N, D)
+            time_guidance_emb = timesteps_emb + guidance_emb
+            return time_guidance_emb
+        else:
+            return timesteps_emb
 
 
 class Flux2Modulation(nn.Module):
@@ -650,8 +662,10 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
         axes_dims_rope: Tuple[int, ...] = config.axes_dims_rope
         rope_theta: int = config.rope_theta
         eps: float = config.eps
+        guidance_embeds: bool = getattr(config, "guidance_embeds", True)
         self.out_channels = out_channels or in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
+        self.guidance_embeds = guidance_embeds
 
         # 1. Sinusoidal positional embedding for RoPE on image and text tokens
         self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
@@ -661,6 +675,7 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
             in_channels=timestep_guidance_channels,
             embedding_dim=self.inner_dim,
             bias=False,
+            guidance_embeds=guidance_embeds,
         )
 
         # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.)
@@ -767,7 +782,8 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
 
         # 1. Calculate timestep embedding and modulation parameters
         timestep = timestep.to(hidden_states.dtype)
-        guidance = guidance.to(hidden_states.dtype)
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype)
 
         temb = self.time_guidance_embed(timestep, guidance)
 

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

@@ -0,0 +1,422 @@
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+
+
+class Qwen3MLP(nn.Module):
+    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen3Attention(nn.Module):
+    """Qwen3 attention with QK-Norm and tensor parallelism.
+
+    Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
+    """
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 1000000.0,
+        rope_scaling: dict[str, Any] | None = None,
+        max_position_embeddings: int = 40960,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.head_dim = getattr(
+            config, "head_dim", self.hidden_size // self.total_num_heads
+        )
+        self.rotary_dim = self.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        # QKV projection with tensor parallelism
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        # Output projection
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # QK-Norm: Key difference from LLaMA
+        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+        # Rotary embeddings
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=int(rope_theta),
+            rope_scaling=rope_scaling,
+            is_neox_style=True,
+        )
+
+        # Attention with FlashAttention/SageAttn support
+        self.attn = LocalAttention(
+            self.num_heads,
+            self.head_dim,
+            self.num_kv_heads,
+            softmax_scale=self.scaling,
+            causal=True,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # QKV projection
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        # Reshape for QK-norm
+        batch_size, seq_len = q.shape[0], q.shape[1]
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        # Apply QK-Norm (key difference from LLaMA)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Reshape back for rotary embeddings
+        q = q.reshape(batch_size, seq_len, -1)
+        k = k.reshape(batch_size, seq_len, -1)
+
+        # Apply rotary embeddings
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Reshape for attention
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        # Attention
+        attn_output = self.attn(q, k, v)
+        attn_output = attn_output.reshape(batch_size, seq_len, -1)
+
+        # Output projection
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3DecoderLayer(nn.Module):
+    """Qwen3 transformer decoder layer."""
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 1000000.0)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        max_position_embeddings = getattr(config, "max_position_embeddings", 40960)
+        attention_bias = getattr(config, "attention_bias", False)
+
+        self.self_attn = Qwen3Attention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(
+                config, "num_key_value_heads", config.num_attention_heads
+            ),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+
+        # MLP
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class Qwen3ForCausalLM(TextEncoder):
+    """Qwen3 causal language model for text encoding in diffusion models.
+
+    Features:
+    - Tensor parallelism support
+    - FlashAttention/SageAttn/SDPA support via LocalAttention
+    - QK-Norm for better training stability
+    - FSDP sharding for CPU offload
+    """
+
+    def __init__(self, config: Qwen3TextConfig) -> None:
+        super().__init__(config)
+
+        self.config = config
+        self.quant_config = config.quant_config
+
+        # Embedding layer with tensor parallelism
+        if config.lora_config is not None:
+            max_loras = getattr(config.lora_config, "max_loras", 1)
+            lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1)
+            lora_vocab = lora_vocab_size * max_loras
+        else:
+            lora_vocab = 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=config.quant_config,
+        )
+
+        # Transformer layers
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(
+                    config=config,
+                    quant_config=config.quant_config,
+                    prefix=f"{config.prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        # Final layer norm
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+
+        residual = None
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1], device=hidden_states.device
+            ).unsqueeze(0)
+
+        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (
+                    (hidden_states,)
+                    if residual is None
+                    else (hidden_states + residual,)
+                )
+            hidden_states, residual = layer(position_ids, hidden_states, residual)
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        # Add hidden states from the last decoder layer
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights with support for tensor parallelism and weight remapping."""
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # Strip 'model.' prefix from HuggingFace Qwen3 weights
+            if name.startswith("model."):
+                name = name[6:]  # len("model.") == 6
+
+            # Skip rotary embedding weights
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                continue
+
+            # Handle KV scale remapping
+            if "scale" in name:
+                kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
+                if kv_scale_name is None:
+                    continue
+                else:
+                    name = kv_scale_name
+
+            # Handle stacked params mapping (qkv_proj, gate_up_proj)
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                # Skip loading extra bias for GPTQ models
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+EntryClass = Qwen3ForCausalLM

python/sglang/multimodal_gen/runtime/pipelines/flux_2_klein.py

@@ -0,0 +1,8 @@
+from sglang.multimodal_gen.runtime.pipelines.flux_2 import Flux2Pipeline
+
+
+class Flux2KleinPipeline(Flux2Pipeline):
+    pipeline_name = "Flux2KleinPipeline"
+
+
+EntryClass = Flux2KleinPipeline