[diffusion] improve: apply tp optim to cross-attn for wan2.2 (#16788)

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

@@ -139,18 +139,20 @@ class WanSelfAttention(nn.Module):
         self.qk_norm = qk_norm
         self.eps = eps
         self.parallel_attention = parallel_attention
+        self.tp_size = get_tensor_model_parallel_world_size()
 
         # layers
-        self.to_q = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.to_q = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_k = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.to_k = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_v = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.to_v = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.to_out = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
         self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.tp_rmsnorm = self.tp_size > 1 and qk_norm
 
         # Scaled dot product attention
         self.attn = USPAttention(
-            num_heads=num_heads,
+            num_heads=num_heads // self.tp_size,
             head_size=self.head_dim,
             dropout_rate=0,
             softmax_scale=None,
@@ -171,7 +173,7 @@ class WanSelfAttention(nn.Module):
 
 class WanT2VCrossAttention(WanSelfAttention):
 
-    def forward(self, x, context, context_lens, crossattn_cache=None):
+    def forward(self, x, context, context_lens):
         r"""
         Args:
             x(Tensor): Shape [B, L1, C]
@@ -179,23 +181,24 @@ class WanT2VCrossAttention(WanSelfAttention):
             context_lens(Tensor): Shape [B]
         """
         b, n, d = x.size(0), self.num_heads, self.head_dim
+        num_heads_per_rank = n // self.tp_size
 
-        # compute query, key, value
+        q, _ = self.to_q(x)
-        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
+        if self.tp_rmsnorm:
-
+            q = tensor_parallel_rms_norm(q, self.norm_q)
-        if crossattn_cache is not None:
-            if not crossattn_cache["is_init"]:
-                crossattn_cache["is_init"] = True
-                k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
-                v = self.to_v(context)[0].view(b, -1, n, d)
-                crossattn_cache["k"] = k
-                crossattn_cache["v"] = v
-            else:
-                k = crossattn_cache["k"]
-                v = crossattn_cache["v"]
         else:
-            k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
+            q = self.norm_q(q)
-            v = self.to_v(context)[0].view(b, -1, n, d)
+        q = q.view(b, -1, num_heads_per_rank, d)
+
+        k, _ = self.to_k(context)
+        if self.tp_rmsnorm:
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            k = self.norm_k(k)
+        k = k.view(b, -1, num_heads_per_rank, d)
+
+        v, _ = self.to_v(context)
+        v = v.view(b, -1, num_heads_per_rank, d)
 
         # compute attention
         x = self.attn(q, k, v)
@@ -227,10 +230,9 @@ class WanI2VCrossAttention(WanSelfAttention):
             supported_attention_backends=supported_attention_backends,
         )
 
-        self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=False)
-        self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=True)
+        self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=False)
         self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
     def forward(self, x, context, context_lens):
         r"""
@@ -242,15 +244,36 @@ class WanI2VCrossAttention(WanSelfAttention):
         context_img = context[:, :257]
         context = context[:, 257:]
         b, n, d = x.size(0), self.num_heads, self.head_dim
+        num_heads_per_rank = n // self.tp_size
+
+        q, _ = self.to_q(x)
+        if self.tp_rmsnorm:
+            q = tensor_parallel_rms_norm(q, self.norm_q)
+        else:
+            q = self.norm_q(q)
+        q = q.view(b, -1, num_heads_per_rank, d)
+
+        k, _ = self.to_k(context)
+        if self.tp_rmsnorm:
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            k = self.norm_k(k)
+        k = k.view(b, -1, num_heads_per_rank, d)
+
+        v, _ = self.to_v(context)
+        v = v.view(b, -1, num_heads_per_rank, d)
+
+        k_img, _ = self.add_k_proj(context_img)
+        if self.tp_rmsnorm:
+            k_img = tensor_parallel_rms_norm(k_img, self.norm_added_k)
+        else:
+            k_img = self.norm_added_k(k_img)
+        k_img = k_img.view(b, -1, num_heads_per_rank, d)
+
+        v_img, _ = self.add_v_proj(context_img)
+        v_img = v_img.view(b, -1, num_heads_per_rank, d)
 
-        # compute query, key, value
-        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
-        v = self.to_v(context)[0].view(b, -1, n, d)
-        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d)
-        v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
         img_x = self.attn(q, k_img, v_img)
-        # compute attention
         x = self.attn(q, k, v)
 
         # output