diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 436e60667..0082d51a3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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) 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 c64e47fdd..d45136d60 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 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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index f8acd57bd..ea7e0be65 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index c81a6f054..3aa2437bb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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)