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

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

Browse Source
This commit is contained in:
Adarsh Shirawalmath
2026-01-16 13:46:17 +05:30
committed by GitHub
parent daa4841e86
commit 7c39ea68f3
11 changed files with 676 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
docs/supported_models/diffusion_models.md
View File

@@ -4,7 +4,7 @@ SGLang Diffusion is an inference framework for accelerated image and video gener
## Key Features
- **Broad Model Support**: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, and more
- **Broad Model Support**: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image, and more
- **Fast Inference**: Optimized kernels from sgl-kernel, efficient scheduler loop, and Cache-DiT acceleration
- **Ease of Use**: OpenAI-compatible API, CLI, and Python SDK
- **Multi-Platform**: NVIDIA GPUs (H100, H200, A100, B200, 4090) and AMD GPUs (MI300X, MI325X)
@@ -110,12 +110,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
@@ -1004,13 +1008,14 @@ All Cache-DiT parameters can be set via the following environment variables:
SGLang Diffusion x Cache-DiT supports almost all models originally supported in SGLang Diffusion:
| Model Family | Example Models |
|--------------|-----------------------------|
| Wan | Wan2.1, Wan2.2 |
| Flux | FLUX.1-dev, FLUX.2-dev |
| Z-Image | Z-Image-Turbo |
| Qwen | Qwen-Image, Qwen-Image-Edit |
| Hunyuan | HunyuanVideo |
| Model Family | Example Models |
|--------------|-------------------------------------------|
| Wan | Wan2.1, Wan2.2 |
| Flux | FLUX.1-dev, FLUX.2-dev, FLUX.2-Klein |
| Z-Image | Z-Image-Turbo |
| Qwen | Qwen-Image, Qwen-Image-Edit |
| GLM | GLM-Image |
| Hunyuan | HunyuanVideo |
## Performance Tips

2
python/sglang/multimodal_gen/README.md
View File

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

2
python/sglang/multimodal_gen/configs/models/encoders/__init__.py
View File

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

85
python/sglang/multimodal_gen/configs/models/encoders/qwen3.py Normal file
View File

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

8
python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
View File

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

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

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

16
python/sglang/multimodal_gen/docs/support_matrix.md
View File

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

21
python/sglang/multimodal_gen/registry.py
View File

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

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

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

422
python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py Normal file
View File

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

8
python/sglang/multimodal_gen/runtime/pipelines/flux_2_klein.py Normal file
View File

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