[diffusion] feat: support tp for qwen-image-edit-2511 (#18619)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -218,7 +218,7 @@ class GPUWorker:
 
             # Save output to file and return file path only if requested. Avoid the serialization
             # and deserialization overhead between scheduler_client and gpu_worker.
-            if req.save_output and req.return_file_paths_only:
+            if req.save_output and req.return_file_paths_only and self.rank == 0:
                 output_paths = save_outputs(
                     output_batch.output,
                     req.data_type,

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

@@ -9,7 +9,6 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
@@ -25,7 +24,10 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     apply_layernorm_only,
     apply_qk_norm,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
@@ -60,6 +62,81 @@ def _get_qkv_projections(
     return img_query, img_key, img_value, txt_query, txt_key, txt_value
 
 
+class GELU(nn.Module):
+    r"""
+    GELU activation function with tanh approximation support with `approximate="tanh"`.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(
+        self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True
+    ):
+        super().__init__()
+        self.proj = ColumnParallelLinear(
+            dim_in, dim_out, bias=bias, gather_output=False
+        )
+        self.approximate = approximate
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        return F.gelu(hidden_states[0], approximate=self.approximate)
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        activation_fn: str = "geglu",
+        inner_dim=None,
+        bias: bool = True,
+    ):
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim, bias=bias)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+        else:
+            raise NotImplementedError(
+                f"activation_fn '{activation_fn}' is not supported."
+            )
+
+        self.net = nn.ModuleList([])
+        self.net.append(act_fn)
+        self.net.append(nn.Identity())
+        self.net.append(
+            RowParallelLinear(inner_dim, dim_out, bias=True, input_is_parallel=True)
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
 class QwenTimestepProjEmbeddings(nn.Module):
     def __init__(self, embedding_dim, use_additional_t_cond=False):
         super().__init__()
@@ -484,34 +561,44 @@ class QwenImageCrossAttention(nn.Module):
         # Use separate Q/K/V projections
         self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
         self.inner_kv_dim = self.inner_dim
-        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_q = ColumnParallelLinear(
-        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+            dim, self.inner_dim, bias=True, gather_output=True
-        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        )
+        self.to_k = ColumnParallelLinear(
+            dim, self.inner_dim, bias=True, gather_output=True
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, self.inner_dim, bias=True, gather_output=True
+        )
 
         if self.qk_norm:
             self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
             self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         if added_kv_proj_dim is not None:
-            self.add_q_proj = ReplicatedLinear(
+            self.add_q_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True
+                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
             )
-            self.add_k_proj = ReplicatedLinear(
+            self.add_k_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True
+                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
             )
-            self.add_v_proj = ReplicatedLinear(
+            self.add_v_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True
+                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
             )
 
         if context_pre_only is not None and not context_pre_only:
-            self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
+            self.to_add_out = ColumnParallelLinear(
+                self.inner_dim, self.dim, bias=out_bias, gather_output=True
+            )
         else:
             self.to_add_out = None
 
         if not pre_only:
             self.to_out = nn.ModuleList([])
             self.to_out.append(
-                ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
+                ColumnParallelLinear(
+                    self.inner_dim, self.dim, bias=out_bias, gather_output=True
+                )
             )
         else:
             self.to_out = None
@@ -644,8 +731,8 @@ class QwenImageTransformerBlock(nn.Module):
         # Image processing modules
         self.img_mod = nn.Sequential(
             nn.SiLU(),
-            nn.Linear(
+            ColumnParallelLinear(
-                dim, 6 * dim, bias=True
+                dim, 6 * dim, bias=True, gather_output=True
             ),  # For scale, shift, gate for norm1 and norm2
         )
         self.img_norm1 = LayerNormScaleShift(
@@ -669,8 +756,8 @@ class QwenImageTransformerBlock(nn.Module):
         # Text processing modules
         self.txt_mod = nn.Sequential(
             nn.SiLU(),
-            nn.Linear(
+            ColumnParallelLinear(
-                dim, 6 * dim, bias=True
+                dim, 6 * dim, bias=True, gather_output=True
             ),  # For scale, shift, gate for norm1 and norm2
         )
         self.txt_norm1 = LayerNormScaleShift(
@@ -790,8 +877,8 @@ class QwenImageTransformerBlock(nn.Module):
         modulate_index: Optional[List[int]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Get modulation parameters for both streams
-        img_mod_params = self.img_mod[1](temb_img_silu)  # [B, 6*dim]
+        img_mod_params = self.img_mod[1](temb_img_silu)[0]  # [B, 6*dim]
-        txt_mod_params = self.txt_mod[1](temb_txt_silu)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod[1](temb_txt_silu)[0]  # [B, 6*dim]
 
         # Split modulation parameters for norm1 and norm2
         img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
@@ -837,7 +924,7 @@ class QwenImageTransformerBlock(nn.Module):
             residual_x=hidden_states,
         )
         img_mlp_output = self.img_mlp(img_modulated2)
-        hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)
+        hidden_states = self.fuse_mul_add(img_mlp_output[0], img_gate2, hidden_states)
 
         # Process text stream - norm2 + MLP
         txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
@@ -851,7 +938,7 @@ class QwenImageTransformerBlock(nn.Module):
         txt_gate2 = txt_gate2_raw.unsqueeze(1)
         txt_mlp_output = self.txt_mlp(txt_modulated2)
         encoder_hidden_states = self.fuse_mul_add(
-            txt_mlp_output, txt_gate2, encoder_hidden_states
+            txt_mlp_output[0], txt_gate2, encoder_hidden_states
         )
 
         # Clip to prevent overflow for fp16

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -76,8 +76,6 @@ class ImageEncodingStage(PipelineStage):
     def move_to_device(self, device):
         fields = [
             "image_processor",
-            "image_encoder",
-            "text_encoder",
         ]
         for field in fields:
             processor = getattr(self, field, None)