[Diffusion] model: fix zimage tp (#16719)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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

@@ -6,6 +6,7 @@ import torch.nn as nn
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm
 from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
 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, apply_qk_norm
@@ -118,9 +119,23 @@ class ZImageAttention(nn.Module):
         self.num_kv_heads = num_kv_heads
         self.qk_norm = qk_norm
 
-        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        tp_size = get_tp_world_size()
-        self.to_k = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False)
+        assert (
-        self.to_v = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False)
+            num_heads % tp_size == 0
+        ), f"num_heads {num_heads} must be divisible by tp world size {tp_size}"
+        assert (
+            num_kv_heads % tp_size == 0
+        ), f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}"
+        self.local_num_heads = num_heads // tp_size
+        self.local_num_kv_heads = num_kv_heads // tp_size
+
+        self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False)
+        self.to_k = ColumnParallelLinear(
+            dim, self.head_dim * num_kv_heads, bias=False, gather_output=False
+        )
+        self.to_v = ColumnParallelLinear(
+            dim, self.head_dim * num_kv_heads, bias=False, gather_output=False
+        )
 
         if self.qk_norm:
             self.norm_q = RMSNorm(self.head_dim, eps=eps)
@@ -134,9 +149,9 @@ class ZImageAttention(nn.Module):
         )
 
         self.attn = USPAttention(
-            num_heads=num_heads,
+            num_heads=self.local_num_heads,
             head_size=self.head_dim,
-            num_kv_heads=num_kv_heads,
+            num_kv_heads=self.local_num_kv_heads,
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
@@ -150,10 +165,9 @@ class ZImageAttention(nn.Module):
         q, _ = self.to_q(hidden_states)
         k, _ = self.to_k(hidden_states)
         v, _ = self.to_v(hidden_states)
-
+        q = q.view(*q.shape[:-1], self.local_num_heads, self.head_dim)
-        q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
+        k = k.view(*k.shape[:-1], self.local_num_kv_heads, self.head_dim)
-        k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim)
+        v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim)
-        v = v.view(*v.shape[:-1], self.num_kv_heads, self.head_dim)
 
         if self.qk_norm:
             if (