[diffusion] fix: fix pack qkv opt break tensor parallel (#15225)

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

@@ -36,10 +36,7 @@ from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 
 # from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
-from sglang.multimodal_gen.runtime.layers.linear import (
-    QKVParallelLinear,
-    ReplicatedLinear,
-)
+from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     NDRotaryEmbedding,
@@ -102,13 +99,8 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
         self.norm_q = RMSNorm(dim_head, eps=eps)
         self.norm_k = RMSNorm(dim_head, eps=eps)
 
-        # Use QKVParallelLinear for fused QKV projections
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=query_dim,
-            head_size=dim_head,
-            total_num_heads=num_heads,
-            bias=bias,
-        )
+        # Use ReplicatedLinear for fused QKV projections
+        self.to_qkv = ReplicatedLinear(query_dim, self.inner_dim * 3, bias=bias)
 
         if not self.pre_only:
             self.to_out = torch.nn.ModuleList([])
@@ -121,12 +113,9 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
         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)
-            # Use QKVParallelLinear for added (encoder) QKV projections
-            self.to_added_qkv = QKVParallelLinear(
-                hidden_size=added_kv_proj_dim,
-                head_size=dim_head,
-                total_num_heads=num_heads,
-                bias=added_proj_bias,
+            # Use ReplicatedLinear for added (encoder) QKV projections
+            self.to_added_qkv = ReplicatedLinear(
+                added_kv_proj_dim, self.inner_dim * 3, bias=added_proj_bias
             )
             self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
 

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 QKVParallelLinear
+from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     NDRotaryEmbedding,
     _apply_rotary_emb,
@@ -123,13 +123,8 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
         self.added_kv_proj_dim = added_kv_proj_dim
         self.added_proj_bias = added_proj_bias
 
-        # Use QKVParallelLinear for fused QKV projections
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=query_dim,
-            head_size=dim_head,
-            total_num_heads=num_heads,
-            bias=bias,
-        )
+        # Use ReplicatedLinear for fused QKV projections
+        self.to_qkv = ReplicatedLinear(query_dim, self.inner_dim * 3, bias=bias)
 
         # QK Norm
         self.norm_q = RMSNorm(dim_head, eps=eps)
@@ -142,12 +137,9 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
         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)
-            # Use QKVParallelLinear for added (encoder) QKV projections
-            self.to_added_qkv = QKVParallelLinear(
-                hidden_size=added_kv_proj_dim,
-                head_size=dim_head,
-                total_num_heads=num_heads,
-                bias=added_proj_bias,
+            # Use ReplicatedLinear for added (encoder) QKV projections
+            self.to_added_qkv = ReplicatedLinear(
+                added_kv_proj_dim, self.inner_dim * 3, bias=added_proj_bias
             )
             self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
 

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

@@ -16,10 +16,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
-from sglang.multimodal_gen.runtime.layers.linear import (
-    QKVParallelLinear,
-    ReplicatedLinear,
-)
+from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
     apply_rotary_embedding,
     fuse_scale_shift_kernel,
@@ -261,13 +258,9 @@ class QwenImageCrossAttention(nn.Module):
         self.parallel_attention = parallel_attention
         self.added_kv_proj_dim = added_kv_proj_dim
 
-        # Use QKVParallelLinear for fused QKV projections
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=head_dim,
-            total_num_heads=num_heads,
-            bias=True,
-        )
+        # Use ReplicatedLinear for fused QKV projections
+        qkv_dim = num_heads * head_dim * 3
+        self.to_qkv = ReplicatedLinear(dim, qkv_dim, bias=True)
 
         if self.qk_norm:
             self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
@@ -277,13 +270,8 @@ class QwenImageCrossAttention(nn.Module):
         self.inner_kv_dim = self.inner_dim
 
         if added_kv_proj_dim is not None:
-            # Use QKVParallelLinear for added (encoder) QKV projections
-            self.to_added_qkv = QKVParallelLinear(
-                hidden_size=added_kv_proj_dim,
-                head_size=head_dim,
-                total_num_heads=num_heads,
-                bias=True,
-            )
+            # Use ReplicatedLinear for added (encoder) QKV projections
+            self.to_added_qkv = ReplicatedLinear(added_kv_proj_dim, qkv_dim, bias=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)

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

@@ -8,12 +8,7 @@ from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
 from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
 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 (
-    MergedColumnParallelLinear,
-    QKVParallelLinear,
-    ReplicatedLinear,
-    RowParallelLinear,
-)
+from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@@ -79,16 +74,9 @@ class TimestepEmbedder(nn.Module):
 class FeedForward(nn.Module):
     def __init__(self, dim: int, hidden_dim: int):
         super().__init__()
-        self.w13 = MergedColumnParallelLinear(
-            input_size=dim,
-            output_sizes=[hidden_dim] * 2,
-            bias=False,
-        )
-        self.w2 = RowParallelLinear(
-            input_size=hidden_dim,
-            output_size=dim,
-            bias=False,
-        )
+        # Use ReplicatedLinear for gate and up projection (fused)
+        self.w13 = ReplicatedLinear(dim, hidden_dim * 2, bias=False)
+        self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False)
         self.act = SiluAndMul()
 
     def forward(self, x):
@@ -114,13 +102,9 @@ class ZImageAttention(nn.Module):
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
 
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            total_num_kv_heads=num_kv_heads,
-            bias=False,
-        )
+        # Use ReplicatedLinear for QKV projection (fused)
+        qkv_dim = dim + 2 * (num_kv_heads * self.head_dim)
+        self.to_qkv = ReplicatedLinear(dim, qkv_dim, bias=False)
 
         if self.qk_norm:
             self.norm_q = RMSNorm(self.head_dim, eps=eps)