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

[Diffusion] Flux2 tp support (#16219)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-01-03 17:02:27 +08:00
committed by GitHub
parent 65b0b5b248
commit d0fb24ee7b
3 changed files with 132 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,7 +23,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous
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.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
apply_flashinfer_rope_qk_inplace,
@@ -83,14 +83,18 @@ class Flux2FeedForward(nn.Module):
dim_out = dim_out or dim
# Flux2SwiGLU will reduce the dimension by half
self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias)
self.linear_in = ColumnParallelLinear(
dim, inner_dim * 2, bias=bias, gather_output=True
)
self.act_fn = Flux2SwiGLU()
self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias)
self.linear_out = ColumnParallelLinear(
inner_dim, dim_out, bias=bias, gather_output=True
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.linear_in(x)
x, _ = self.linear_in(x)
x = self.act_fn(x)
x = self.linear_out(x)
x, _ = self.linear_out(x)
return x
@@ -123,31 +127,52 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias
self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_k = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_v = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_q = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_k = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
self.to_v = ColumnParallelLinear(
query_dim, self.inner_dim, bias=bias, gather_output=True
)
# 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(
ColumnParallelLinear(
self.inner_dim, self.out_dim, bias=out_bias, gather_output=True
)
)
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 = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
self.add_q_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_k_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
self.add_k_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.add_v_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
self.add_v_proj = ColumnParallelLinear(
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
)
self.to_add_out = ColumnParallelLinear(
self.inner_dim, query_dim, bias=out_bias, gather_output=True
)
self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
self.attn = USPAttention(
num_heads=num_heads,
@@ -212,9 +237,9 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
],
dim=1,
)
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)
hidden_states = self.to_out[0](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:
@@ -265,10 +290,11 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
self.mlp_mult_factor = mlp_mult_factor
# Fused QKV projections + MLP input projection
self.to_qkv_mlp_proj = torch.nn.Linear(
self.to_qkv_mlp_proj = ColumnParallelLinear(
self.query_dim,
self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
bias=bias,
gather_output=True,
)
self.mlp_act_fn = Flux2SwiGLU()
@@ -277,8 +303,11 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
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
self.to_out = ColumnParallelLinear(
self.inner_dim + self.mlp_hidden_dim,
self.out_dim,
bias=out_bias,
gather_output=True,
)
self.attn = USPAttention(
@@ -297,7 +326,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
**kwargs,
) -> torch.Tensor:
# Parallel in (QKV + MLP in) projection
hidden_states = self.to_qkv_mlp_proj(hidden_states)
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],
@@ -335,7 +364,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
# Concatenate and parallel output projection
hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
hidden_states = self.to_out(hidden_states)
hidden_states, _ = self.to_out(hidden_states)
return hidden_states
@@ -557,14 +586,16 @@ class Flux2Modulation(nn.Module):
super().__init__()
self.mod_param_sets = mod_param_sets
self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias)
self.linear = ColumnParallelLinear(
dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True
)
self.act_fn = nn.SiLU()
def forward(
self, temb: torch.Tensor
) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
mod = self.act_fn(temb)
mod = self.linear(mod)
mod, _ = self.linear(mod)
if mod.ndim == 2:
mod = mod.unsqueeze(1)
@@ -646,9 +677,11 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
)
# 4. Input projections
self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
self.context_embedder = nn.Linear(
joint_attention_dim, self.inner_dim, bias=False
self.x_embedder = ColumnParallelLinear(
in_channels, self.inner_dim, bias=False, gather_output=True
)
self.context_embedder = ColumnParallelLinear(
joint_attention_dim, self.inner_dim, bias=False, gather_output=True
)
# 5. Double Stream Transformer Blocks
@@ -689,8 +722,11 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
eps=eps,
bias=False,
)
self.proj_out = nn.Linear(
self.inner_dim, patch_size * patch_size * self.out_channels, bias=False
self.proj_out = ColumnParallelLinear(
self.inner_dim,
patch_size * patch_size * self.out_channels,
bias=False,
gather_output=True,
)
self.layer_names = ["transformer_blocks", "single_transformer_blocks"]
@@ -740,8 +776,8 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
single_stream_mod = self.single_stream_modulation(temb)[0]
# 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
hidden_states = self.x_embedder(hidden_states)
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
hidden_states, _ = self.x_embedder(hidden_states)
encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
# 3. Calculate RoPE embeddings from image and text tokens
# NOTE: the below logic means that we can't support batched inference with images of different resolutions or
@@ -773,7 +809,7 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
# 6. Output layers
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
output, _ = self.proj_out(hidden_states)
return output

13
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
View File

@@ -10,6 +10,7 @@ from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Optional
import torch
from dateutil.tz import UTC
import sglang
@@ -149,6 +150,12 @@ class StageProfiler:
self.logger.info(f"[{self.stage_name}] started...")
if (self.enabled and self.timings) or self.simple_log:
if (
os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
and self.stage_name.startswith("denoising_step_")
and torch.cuda.is_available()
):
torch.cuda.synchronize()
self.start_time = time.perf_counter()
return self
@@ -157,6 +164,12 @@ class StageProfiler:
if not ((self.enabled and self.timings) or self.simple_log):
return False
if (
os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
and self.stage_name.startswith("denoising_step_")
and torch.cuda.is_available()
):
torch.cuda.synchronize()
execution_time_s = time.perf_counter() - self.start_time
if exc_type:

100
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -1508,56 +1508,56 @@
"DecodingStage": 321.94
},
"denoise_step_ms": {
"0": 66.94,
"1": 511.02,
"2": 537.24,
"3": 516.92,
"4": 537.18,
"5": 520.79,
"6": 531.95,
"7": 522.69,
"8": 532.77,
"9": 522.36,
"10": 530.91,
"11": 526.67,
"12": 530.23,
"13": 526.16,
"14": 528.55,
"15": 524.96,
"16": 526.36,
"17": 527.71,
"18": 528.84,
"19": 528.73,
"20": 527.95,
"21": 525.1,
"22": 528.21,
"23": 528.06,
"24": 527.11,
"25": 530.93,
"26": 526.27,
"27": 528.66,
"28": 526.33,
"29": 530.99,
"30": 528.51,
"31": 528.53,
"32": 528.73,
"33": 527.02,
"34": 527.49,
"35": 528.25,
"36": 529.96,
"37": 528.43,
"38": 528.5,
"39": 525.72,
"40": 529.65,
"41": 525.62,
"42": 526.24,
"43": 528.77,
"44": 526.68,
"45": 528.06,
"46": 524.92,
"47": 529.5,
"48": 524.22,
"49": 528.21
"0": 129.07,
"1": 437.16,
"2": 437.7,
"3": 437.67,
"4": 437.84,
"5": 438.03,
"6": 438.09,
"7": 437.65,
"8": 437.95,
"9": 438.31,
"10": 437.99,
"11": 438.54,
"12": 438.47,
"13": 438.2,
"14": 438.56,
"15": 438.69,
"16": 438.69,
"17": 438.98,
"18": 437.96,
"19": 438.9,
"20": 438.87,
"21": 438.04,
"22": 437.88,
"23": 439.09,
"24": 438.61,
"25": 437.68,
"26": 439.2,
"27": 439.63,
"28": 438.65,
"29": 439.32,
"30": 439.01,
"31": 438.84,
"32": 438.72,
"33": 439.09,
"34": 438.3,
"35": 439.48,
"36": 438.2,
"37": 439.67,
"38": 440.65,
"39": 439.96,
"40": 439.0,
"41": 439.2,
"42": 439.37,
"43": 439.98,
"44": 438.6,
"45": 439.58,
"46": 440.23,
"47": 440.1,
"48": 440.21,
"49": 439.22
},
"expected_e2e_ms": 27624.8,
"expected_avg_denoise_ms": 518.23,