[diffusion] model: wan tp+usp optimize (#16720)

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

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -11,6 +11,11 @@ import torch.nn.functional as F
 from sgl_kernel import fused_add_rmsnorm, rmsnorm
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
 from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
     fuse_scale_shift_kernel,
@@ -469,3 +474,17 @@ def apply_qk_norm(
     q_out = q_norm(q.view(-1, head_dim)).view(q_shape)
     k_out = k_norm(k.view(-1, head_dim)).view(k_shape)
     return q_out, k_out
+
+
+def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
+    tp_rank = get_tensor_model_parallel_rank()
+    tp_size = get_tensor_model_parallel_world_size()
+    src_dtype = x.dtype
+    weight = norm.weight.tensor_split(tp_size)[tp_rank].float()
+    x_fp32 = x.float()
+    variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+    variance = get_tp_group().all_reduce(
+        variance, op=torch._C._distributed_c10d.ReduceOp.AVG
+    )
+    output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight
+    return output.to(dtype=src_dtype)

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

@@ -11,7 +11,11 @@ import torch.nn as nn
 
 from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
 from sglang.multimodal_gen.configs.sample.wan import WanTeaCacheParams
-from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
+from sglang.multimodal_gen.runtime.distributed import divide
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_world_size,
+    get_tensor_model_parallel_world_size,
+)
 from sglang.multimodal_gen.runtime.layers.attention import (
     MinimalA2AAttnOp,
     SparseLinearAttention,
@@ -24,8 +28,12 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     RMSNorm,
     ScaleResidual,
     ScaleResidualLayerNormScaleShift,
+    tensor_parallel_rms_norm,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     NDRotaryEmbedding,
@@ -273,11 +281,11 @@ class WanTransformerBlock(nn.Module):
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
-        self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
 
-        self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
+        self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True)
         if attention_type == "sla":
             self.attn1 = MinimalA2AAttnOp(
                 SparseLinearAttention(
@@ -286,7 +294,7 @@ class WanTransformerBlock(nn.Module):
             )
         else:
             self.attn1 = USPAttention(
-                num_heads=num_heads,
+                num_heads=divide(num_heads, get_tensor_model_parallel_world_size()),
                 head_size=dim // num_heads,
                 causal=False,
                 supported_attention_backends=supported_attention_backends,
@@ -295,10 +303,10 @@ class WanTransformerBlock(nn.Module):
 
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
-        dim_head = dim // num_heads
+        self.dim_head = dim // num_heads
         if qk_norm == "rms_norm":
-            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_q = RMSNorm(self.dim_head, eps=eps)
-            self.norm_k = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(self.dim_head, eps=eps)
         elif qk_norm == "rms_norm_across_heads":
             # LTX applies qk norm across all heads
             self.norm_q = RMSNorm(dim, eps=eps)
@@ -307,6 +315,7 @@ class WanTransformerBlock(nn.Module):
             logger.error("QK Norm type not supported")
             raise Exception
         assert cross_attn_norm is True
+        self.qk_norm = qk_norm
         self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -388,15 +397,24 @@ class WanTransformerBlock(nn.Module):
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
-
+        tp_rmsnorm = (
+            self.qk_norm == "rms_norm_across_heads"
+            and get_tensor_model_parallel_world_size() > 1
+        )
         if self.norm_q is not None:
+            if tp_rmsnorm:
+                query = tensor_parallel_rms_norm(query, self.norm_q)
+            else:
                 query = self.norm_q(query)
         if self.norm_k is not None:
+            if tp_rmsnorm:
+                key = tensor_parallel_rms_norm(key, self.norm_k)
+            else:
                 key = self.norm_k(key)
 
-        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        query = query.squeeze(1).unflatten(2, (-1, self.dim_head))
-        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (-1, self.dim_head))
-        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (-1, self.dim_head))
 
         # Apply rotary embeddings
         cos, sin = freqs_cis