109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
import torch
|
|
import torch.cuda.amp as amp
|
|
|
|
from .fuser import (get_sequence_parallel_rank,
|
|
get_sequence_parallel_world_size, get_sp_group,
|
|
xFuserLongContextAttention)
|
|
|
|
|
|
def pad_freqs(original_tensor, target_len):
|
|
seq_len, s1, s2 = original_tensor.shape
|
|
pad_size = target_len - seq_len
|
|
padding_tensor = torch.ones(
|
|
pad_size,
|
|
s1,
|
|
s2,
|
|
dtype=original_tensor.dtype,
|
|
device=original_tensor.device)
|
|
padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
|
|
return padded_tensor
|
|
|
|
|
|
@amp.autocast(enabled=False)
|
|
@torch.compiler.disable()
|
|
def rope_apply(x, grid_sizes, freqs):
|
|
"""
|
|
x: [B, L, N, C].
|
|
grid_sizes: [B, 3].
|
|
freqs: [M, C // 2].
|
|
"""
|
|
s, n, c = x.size(1), x.size(2), x.size(3) // 2
|
|
# split freqs
|
|
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
|
|
|
# loop over samples
|
|
output = []
|
|
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
|
seq_len = f * h * w
|
|
|
|
# precompute multipliers
|
|
x_i = torch.view_as_complex(x[i, :s].to(torch.float32).reshape(
|
|
s, n, -1, 2))
|
|
freqs_i = torch.cat([
|
|
freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
|
|
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
|
|
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
|
|
],
|
|
dim=-1).reshape(seq_len, 1, -1)
|
|
|
|
# apply rotary embedding
|
|
sp_size = get_sequence_parallel_world_size()
|
|
sp_rank = get_sequence_parallel_rank()
|
|
freqs_i = pad_freqs(freqs_i, s * sp_size)
|
|
s_per_rank = s
|
|
freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
|
|
s_per_rank), :, :]
|
|
x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
|
|
x_i = torch.cat([x_i, x[i, s:]])
|
|
|
|
# append to collection
|
|
output.append(x_i)
|
|
return torch.stack(output)
|
|
|
|
|
|
def rope_apply_qk(q, k, grid_sizes, freqs):
|
|
q = rope_apply(q, grid_sizes, freqs)
|
|
k = rope_apply(k, grid_sizes, freqs)
|
|
return q, k
|
|
|
|
|
|
def usp_attn_flashhead_forward(self,
|
|
x,
|
|
seq_lens,
|
|
grid_sizes,
|
|
freqs,
|
|
dtype=torch.bfloat16,
|
|
t=0):
|
|
"""
|
|
FlashHead specific self-attention forward with xFuser support.
|
|
Similar to WanSelfAttention but with sequence parallel.
|
|
"""
|
|
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
|
half_dtypes = (torch.float16, torch.bfloat16)
|
|
|
|
def half(x):
|
|
return x if x.dtype in half_dtypes else x.to(dtype)
|
|
|
|
# query, key, value function
|
|
def qkv_fn(x):
|
|
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
|
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
|
v = self.v(x).view(b, s, n, d)
|
|
return q, k, v
|
|
|
|
q, k, v = qkv_fn(x)
|
|
q, k = rope_apply_qk(q, k, grid_sizes, freqs)
|
|
|
|
# xFuser attention
|
|
x = xFuserLongContextAttention()(
|
|
None,
|
|
query=half(q),
|
|
key=half(k),
|
|
value=half(v),
|
|
window_size=self.window_size)
|
|
|
|
# output
|
|
x = x.flatten(2)
|
|
x = self.o(x)
|
|
return x
|