File size: 1,098 Bytes
2f9282b |
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 |
# -*- coding: utf-8 -*-
from typing import Optional, Tuple
import torch
from einops import rearrange
from fla.ops.linear_attn.utils import normalize_output
def naive_chunk_linear_attn(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
scale: Optional[float] = None,
normalize: bool = False
) -> Tuple[torch.Tensor, torch.Tensor]:
if scale is None:
scale = q.shape[-1] ** -0.5
chunk_size = 64
q = rearrange(q, 'b h (n c) d -> b h n c d', c=chunk_size) * scale
k = rearrange(k, 'b h (n c) d -> b h n c d', c=chunk_size)
v = rearrange(v, 'b h (n c) d -> b h n c d', c=chunk_size)
kv = k.transpose(-1, -2) @ v
kv = kv.cumsum(2)
kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2)
inter = q @ kv
intra = ((
q @ k.transpose(-1, -2)).masked_fill_(
torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1),
0
)) @ v
o = inter + intra
if normalize:
o = normalize_output(q * scale, k, o)
return rearrange(o, 'b h n c d -> b h (n c) d')
|