File size: 2,210 Bytes
c8c00f0 | 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 | import math
import torch
from einops import rearrange
from torch import Tensor
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
q, k = apply_rope(q, k, pe)
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
x = rearrange(x, "B H L D -> B L (H D)")
return x
def attention_masked(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, patch_ids: list[int], mask: Tensor, return_weight:bool=False) -> Tensor:
q, k = apply_rope(q, k, pe)
# x = torch.nn.functional.scaled_dot_product_attention(q,k,v,attn_mask=attn_mask)
if return_weight:
x, m = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight)
x = rearrange(x, "B H L D -> B L (H D)")
return x, m
else:
x = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight)
x = rearrange(x, "B H L D -> B L (H D)")
return x
# return x, m
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
assert dim % 2 == 0
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
omega = 1.0 / (theta**scale)
out = torch.einsum("...n,d->...nd", pos, omega)
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
return out.float()
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
def scaled_dot_product_attention_masked(query, key, value, patch_ids, attn_mask, return_weight=False):
scale_factor = 1 / math.sqrt(query.size(-1))
attn_weight = query @ key.transpose(-2, -1) * scale_factor
attn_weight += attn_mask
attn_weight = torch.softmax(attn_weight, dim=-1)
if return_weight:
return attn_weight @ value, attn_weight
return attn_weight @ value
|