File size: 1,516 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 | import numpy as np
import torch
def segsum(x):
T = x.size(-1)
x_cumsum = torch.cumsum(x, dim=-1)
x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :]
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool))
x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
return x_segsum
def construct_level_mask(level, L):
T = L.size(-1)
if level == 0:
return torch.diag_embed(L[..., level, :])
indices = torch.cartesian_prod(torch.arange(T), torch.arange(T)).to(L.device)
mask = torch.where(
torch.logical_and(
torch.logical_and(
indices[:, 0] % (1 << level) >= (1 << (level - 1)),
indices[:, 1] + (1 << (level - 1))
>= indices[:, 0] - (indices[:, 0] % (1 << (level - 1))),
),
indices[:, 1] < indices[:, 0] - (indices[:, 0] % (1 << (level - 1))),
).view(T, T),
L[..., level, :].unsqueeze(-1).expand(*([-1] * (len(L.shape) - 2)), T, T),
0,
)
return mask
def construct_H_matrix(a, L):
T = a.size(-1)
A = torch.exp(segsum(a))
H = torch.zeros_like(A)
for level in range(int(np.ceil(np.log2(T))) + 1):
mask = construct_level_mask(level, L)
H += A * mask
return H
def naive_log_linear_attn(q, k, v, g, level_scales):
H = construct_H_matrix(g.permute(0, 2, 1), level_scales.permute(0, 2, 3, 1))
M = torch.einsum("bhlc,blhn,bchn->bhlc", H, q, k)
return torch.einsum("bhlc,bchp->blhp", M, v)
|