| 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) |
|
|
