minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
co-authored by Mick
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
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+69 -33
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.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm 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 ( from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding, NDRotaryEmbedding,
apply_flashinfer_rope_qk_inplace, apply_flashinfer_rope_qk_inplace,
@@ -83,14 +83,18 @@ class Flux2FeedForward(nn.Module):
dim_out = dim_out or dim dim_out = dim_out or dim
# Flux2SwiGLU will reduce the dimension by half # 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.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: 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.act_fn(x)
x = self.linear_out(x) x, _ = self.linear_out(x)
return x return x
@@ -123,31 +127,52 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
self.added_kv_proj_dim = added_kv_proj_dim self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias self.added_proj_bias = added_proj_bias
self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) self.to_q = ColumnParallelLinear(
self.to_k = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) query_dim, self.inner_dim, bias=bias, gather_output=True
self.to_v = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) )
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 # QK Norm
self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps)
self.to_out = torch.nn.ModuleList([]) 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)) self.to_out.append(torch.nn.Dropout(dropout))
if added_kv_proj_dim is not None: if added_kv_proj_dim is not None:
self.norm_added_q = RMSNorm(dim_head, eps=eps) self.norm_added_q = RMSNorm(dim_head, eps=eps)
self.norm_added_k = RMSNorm(dim_head, eps=eps) self.norm_added_k = RMSNorm(dim_head, eps=eps)
self.add_q_proj = ReplicatedLinear( self.add_q_proj = ColumnParallelLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
) )
self.add_k_proj = ReplicatedLinear( self.add_k_proj = ColumnParallelLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
) )
self.add_v_proj = ReplicatedLinear( self.add_v_proj = ColumnParallelLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias 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( self.attn = USPAttention(
num_heads=num_heads, num_heads=num_heads,
@@ -212,9 +237,9 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
], ],
dim=1, 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) hidden_states = self.to_out[1](hidden_states)
if encoder_hidden_states is not None: if encoder_hidden_states is not None:
@@ -265,10 +290,11 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
self.mlp_mult_factor = mlp_mult_factor self.mlp_mult_factor = mlp_mult_factor
# Fused QKV projections + MLP input projection # Fused QKV projections + MLP input projection
self.to_qkv_mlp_proj = torch.nn.Linear( self.to_qkv_mlp_proj = ColumnParallelLinear(
self.query_dim, self.query_dim,
self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
bias=bias, bias=bias,
gather_output=True,
) )
self.mlp_act_fn = Flux2SwiGLU() self.mlp_act_fn = Flux2SwiGLU()
@@ -277,8 +303,11 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
self.norm_k = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps)
# Fused attention output projection + MLP output projection # Fused attention output projection + MLP output projection
self.to_out = torch.nn.Linear( self.to_out = ColumnParallelLinear(
self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias self.inner_dim + self.mlp_hidden_dim,
self.out_dim,
bias=out_bias,
gather_output=True,
) )
self.attn = USPAttention( self.attn = USPAttention(
@@ -297,7 +326,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
**kwargs, **kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
# Parallel in (QKV + MLP in) projection # 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( qkv, mlp_hidden_states = torch.split(
hidden_states, hidden_states,
[3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor], [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 # Concatenate and parallel output projection
hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) 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 return hidden_states
@@ -557,14 +586,16 @@ class Flux2Modulation(nn.Module):
super().__init__() super().__init__()
self.mod_param_sets = mod_param_sets 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() self.act_fn = nn.SiLU()
def forward( def forward(
self, temb: torch.Tensor self, temb: torch.Tensor
) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
mod = self.act_fn(temb) mod = self.act_fn(temb)
mod = self.linear(mod) mod, _ = self.linear(mod)
if mod.ndim == 2: if mod.ndim == 2:
mod = mod.unsqueeze(1) mod = mod.unsqueeze(1)
@@ -646,9 +677,11 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
) )
# 4. Input projections # 4. Input projections
self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False) self.x_embedder = ColumnParallelLinear(
self.context_embedder = nn.Linear( in_channels, self.inner_dim, bias=False, gather_output=True
joint_attention_dim, self.inner_dim, bias=False )
self.context_embedder = ColumnParallelLinear(
joint_attention_dim, self.inner_dim, bias=False, gather_output=True
) )
# 5. Double Stream Transformer Blocks # 5. Double Stream Transformer Blocks
@@ -689,8 +722,11 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin):
eps=eps, eps=eps,
bias=False, bias=False,
) )
self.proj_out = nn.Linear( self.proj_out = ColumnParallelLinear(
self.inner_dim, patch_size * patch_size * self.out_channels, bias=False self.inner_dim,
patch_size * patch_size * self.out_channels,
bias=False,
gather_output=True,
) )
self.layer_names = ["transformer_blocks", "single_transformer_blocks"] 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] single_stream_mod = self.single_stream_modulation(temb)[0]
# 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states) # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
hidden_states = self.x_embedder(hidden_states) hidden_states, _ = self.x_embedder(hidden_states)
encoder_hidden_states = self.context_embedder(encoder_hidden_states) encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
# 3. Calculate RoPE embeddings from image and text tokens # 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 # 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 # 6. Output layers
hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states) output, _ = self.proj_out(hidden_states)
return output return output
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
+13
View File
@@ -10,6 +10,7 @@ from functools import lru_cache
from pathlib import Path from pathlib import Path
from typing import Any, Dict, Optional from typing import Any, Dict, Optional
import torch
from dateutil.tz import UTC from dateutil.tz import UTC
import sglang import sglang
@@ -149,6 +150,12 @@ class StageProfiler:
self.logger.info(f"[{self.stage_name}] started...") self.logger.info(f"[{self.stage_name}] started...")
if (self.enabled and self.timings) or self.simple_log: 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() self.start_time = time.perf_counter()
return self return self
@@ -157,6 +164,12 @@ class StageProfiler:
if not ((self.enabled and self.timings) or self.simple_log): if not ((self.enabled and self.timings) or self.simple_log):
return False 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 execution_time_s = time.perf_counter() - self.start_time
if exc_type: if exc_type:
python/sglang/multimodal_gen/test/server/perf_baselines.json
+50 -50
View File
@@ -1508,56 +1508,56 @@
"DecodingStage": 321.94 "DecodingStage": 321.94
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 66.94, "0": 129.07,
"1": 511.02, "1": 437.16,
"2": 537.24, "2": 437.7,
"3": 516.92, "3": 437.67,
"4": 537.18, "4": 437.84,
"5": 520.79, "5": 438.03,
"6": 531.95, "6": 438.09,
"7": 522.69, "7": 437.65,
"8": 532.77, "8": 437.95,
"9": 522.36, "9": 438.31,
"10": 530.91, "10": 437.99,
"11": 526.67, "11": 438.54,
"12": 530.23, "12": 438.47,
"13": 526.16, "13": 438.2,
"14": 528.55, "14": 438.56,
"15": 524.96, "15": 438.69,
"16": 526.36, "16": 438.69,
"17": 527.71, "17": 438.98,
"18": 528.84, "18": 437.96,
"19": 528.73, "19": 438.9,
"20": 527.95, "20": 438.87,
"21": 525.1, "21": 438.04,
"22": 528.21, "22": 437.88,
"23": 528.06, "23": 439.09,
"24": 527.11, "24": 438.61,
"25": 530.93, "25": 437.68,
"26": 526.27, "26": 439.2,
"27": 528.66, "27": 439.63,
"28": 526.33, "28": 438.65,
"29": 530.99, "29": 439.32,
"30": 528.51, "30": 439.01,
"31": 528.53, "31": 438.84,
"32": 528.73, "32": 438.72,
"33": 527.02, "33": 439.09,
"34": 527.49, "34": 438.3,
"35": 528.25, "35": 439.48,
"36": 529.96, "36": 438.2,
"37": 528.43, "37": 439.67,
"38": 528.5, "38": 440.65,
"39": 525.72, "39": 439.96,
"40": 529.65, "40": 439.0,
"41": 525.62, "41": 439.2,
"42": 526.24, "42": 439.37,
"43": 528.77, "43": 439.98,
"44": 526.68, "44": 438.6,
"45": 528.06, "45": 439.58,
"46": 524.92, "46": 440.23,
"47": 529.5, "47": 440.1,
"48": 524.22, "48": 440.21,
"49": 528.21 "49": 439.22
}, },
"expected_e2e_ms": 27624.8, "expected_e2e_ms": 27624.8,
"expected_avg_denoise_ms": 518.23, "expected_avg_denoise_ms": 518.23,