|
|
|
|
|
import torch |
|
from einops import rearrange |
|
|
|
|
|
def naive_chunk_rwkv6( |
|
q: torch.Tensor, |
|
k: torch.Tensor, |
|
v: torch.Tensor, |
|
w: torch.Tensor, |
|
u: torch.Tensor, |
|
chunk_size: int = 32 |
|
): |
|
assert q.shape[-2] % chunk_size == 0 |
|
orig_dtype = q.dtype |
|
num_chunk = q.shape[-2] // chunk_size |
|
u = u.unsqueeze(0) |
|
|
|
q, k, v, w = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size).float(), (q, k, v, w)) |
|
|
|
w_cumsum = w.cumsum(-2) |
|
|
|
kw = k * (w_cumsum[..., -1, None, :] - w_cumsum).exp() |
|
wkv = kw.transpose(-1, -2) @ v |
|
|
|
wkv_new = torch.zeros_like(wkv) |
|
|
|
for i in range(num_chunk - 1): |
|
wkv_new[:, :, i+1] = (wkv_new[:, :, i] * w_cumsum[:, :, i, -1, :, None].exp()) + wkv[:, :, i] |
|
|
|
o_inter = torch.einsum('b h n d p, b h n c d -> b h n c p', wkv_new, (q * (w_cumsum - w).exp())) |
|
|
|
o_intra = torch.zeros_like(o_inter) |
|
for i in range(chunk_size): |
|
attn = (q[:, :, :, i, None] * k * (w_cumsum[:, :, :, i, None] - w[:, :, :, i, None] - w_cumsum).exp()).sum(-1) |
|
mask = (torch.arange(0, chunk_size) < i).to(attn.device) |
|
attn.masked_fill_(~mask, 0) |
|
intra_inter_o = (attn.unsqueeze(-1) * v).sum(-2) |
|
intra_intra_o = (q[:, :, :, i] * u.unsqueeze(2) * k[:, :, :, i]).sum(-1).unsqueeze(-1) * v[:, :, :, i] |
|
o_intra[:, :, :, i] = intra_inter_o + intra_intra_o |
|
o = o_inter + o_intra |
|
return rearrange(o, 'b h n c d -> b h (n c) d').to(orig_dtype) |
|
|
