File size: 3,348 Bytes
b66f552 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import torch
from einops import rearrange
def naive_recurrent_kda(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
):
dtype = v.dtype
B, T, H, K, V = *q.shape, v.shape[-1]
if scale is None:
scale = K ** -0.5
q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
q = q * scale
S = k.new_zeros(B, H, K, V).to(q)
if initial_state is not None:
S += initial_state
o = torch.zeros_like(v)
for i in range(0, T):
q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
S = S * g_i[..., None].exp()
S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
if not output_final_state:
S = None
return o.to(dtype), S
def naive_chunk_kda(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
chunk_size: int = 64,
):
dtype = v.dtype
B, T, H, K, V = *q.shape, v.shape[-1]
BT = chunk_size
NT = T // BT
if scale is None:
scale = K ** -0.5
assert T % BT == 0
q, k, v, g, beta = map(lambda x: rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float), [q, k, v, g, beta])
q = q * scale
g = g.cumsum(-2)
# note that diagonal is masked.
mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)
A = torch.zeros(*q.shape[:-1], BT, dtype=torch.float, device=q.device)
for i in range(BT):
k_i = k[..., i, :]
g_i = g[..., i:i+1, :]
A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i)
A = A * beta[..., None]
A = -A.masked_fill(mask, 0)
for i in range(1, BT):
A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta[..., None, :]
w = A @ (g.exp() * k)
u = A @ v
S = k.new_zeros(B, H, K, V).to(q)
if initial_state is not None:
S += initial_state
o = torch.zeros_like(v)
mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
for i in range(0, NT):
# [B, H, BT, ...]
q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i]
A = torch.zeros(B, H, BT, BT, dtype=torch.float, device=q.device)
for j in range(BT):
k_j = k[:, :, i, j]
g_j = g[:, :, i, j:j+1, :]
A[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j)
A = A.masked_fill(mask, 0)
v_i = u_i - w_i @ S
o[:, :, i] = (q_i * g_i.exp()) @ S + A @ v_i
S = S * rearrange(g_i[:, :, -1].exp(), 'b h k -> b h k 1')
S += rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, 'b h c k -> b h k c') @ v_i
if not output_final_state:
S = None
return rearrange(o, 'b h n c d -> b (n c) h d').to(dtype), S
|