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

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

Browse Source
This commit is contained in:
wxy
2026-01-10 21:25:59 +08:00
committed by GitHub
parent d1ec93e3ac
commit 76d4881794
1 changed files with 54 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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_k = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_v = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_out = ColumnParallelLinear(dim, dim, gather_output=True)
self.to_q = ColumnParallelLinear(dim, dim, gather_output=False)
self.to_k = ColumnParallelLinear(dim, dim, gather_output=False)
self.to_v = ColumnParallelLinear(dim, dim, gather_output=False)
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.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
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"]
q, _ = self.to_q(x)
if self.tp_rmsnorm:
q = tensor_parallel_rms_norm(q, self.norm_q)
else:
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)
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)
# 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_v_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=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