diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 3b88d99bd..411f97ada 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 6b7da272d..8435571bb 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index c7948bac7..6f1adeb6e 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -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,