chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
triple-mu
2026-01-09 10:40:31 +09:00
committed by GitHub
parent 9e3a032ad6
commit f7f5c3896d
2 changed files with 53 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

@@ -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)

50
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
View File

@@ -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