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