From f7f5c3896d3fdb54fae3146879365000b916f0a1 Mon Sep 17 00:00:00 2001 From: triple-mu Date: Fri, 9 Jan 2026 10:40:31 +0900 Subject: [PATCH] [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> --- .../runtime/layers/layernorm.py | 19 +++++++ .../runtime/models/dits/wanvideo.py | 50 +++++++++++++------ 2 files changed, 53 insertions(+), 16 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 3cf61d54f..a9f5811e2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index ee63a77ce..35e15935e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) - self.to_v = 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=False) + 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_k = RMSNorm(dim_head, eps=eps) + self.norm_q = RMSNorm(self.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: - query = self.norm_q(query) + 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: - key = self.norm_k(key) + 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)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.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, (-1, self.dim_head)) + value = value.squeeze(1).unflatten(2, (-1, self.dim_head)) # Apply rotary embeddings cos, sin = freqs_cis