| import math |
| from typing import Optional, Sequence, Tuple, Union |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.nn.attention.flex_attention import ( |
| _DEFAULT_SPARSE_BLOCK_SIZE, |
| create_block_mask, |
| create_mask, |
| flex_attention, |
| ) |
| from torch.nn.attention.flex_attention import flex_attention, _vmap_for_bhqkv |
|
|
| try: |
| from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex |
| except ImportError: |
| from torch._higher_order_ops.flex_attention import TransformGetItemToIndex |
| |
| from torch._dynamo import disable |
|
|
|
|
| def generate_alibi_bias(H=12): |
| alibi_bias = [] |
| for h in range(H): |
| alibi_bias.append(-((h + 1) / H)) |
| alibi_bias = torch.tensor(alibi_bias) |
| alibi_bias = torch.exp2(alibi_bias) |
| return alibi_bias |
|
|
|
|
| def get_rel_bias_func(scale, coords=None, qk_scale=1.0): |
| def patch_coords_rel_bias(score, b, h, q_idx, kv_idx): |
| if coords is None: |
| return score |
| with torch.no_grad(): |
| dx = coords[b, q_idx][0] - coords[b, kv_idx][0] |
| dy = coords[b, q_idx][1] - coords[b, kv_idx][1] |
| dist = torch.sqrt(dx * dx + dy * dy) |
| dist = dist.clamp(max=1000) |
| dist = torch.log1p(dist) |
| bias = dist * scale[h] * qk_scale |
| return score - bias |
| return patch_coords_rel_bias |
|
|
|
|
| def key_padding_mask(mask): |
| def padding_mask(b, h, q_idx, kv_idx): |
| return ~mask[b, kv_idx] |
| return padding_mask |
|
|
|
|
| class FlexCore(nn.Module): |
| """ |
| For using "forward hook" |
| """ |
| def forward(self, q, k, v, score_mod=None, block_mask=None, return_lse=False): |
| """ |
| "return_lse=True" should be used with ATTN_MAP_VIS wrapper. |
| Though return_lse is "True", _flex_attention(...) only have an output (attention output, not attention scores). |
| """ |
| return flex_attention(q, k, v, score_mod=score_mod, block_mask=block_mask, return_lse=return_lse) |
|
|
|
|
| class Flex_Attention(nn.Module): |
| def __init__( |
| self, |
| dim: int, |
| num_heads: int = 12, |
| qkv_bias: bool = True, |
| proj_drop: float = 0., |
| use_rel_bias: bool = True, |
| ): |
| super().__init__() |
| assert dim % num_heads == 0, 'dim should be divisible by num_heads' |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
|
|
| self.scale = self.head_dim ** -0.5 |
| self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale) |
| if qkv_bias: |
| self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3)) |
| else: |
| self.in_proj_bias = None |
| |
| self.f_attn = FlexCore() |
|
|
| self.out_proj = nn.Linear(dim, dim) |
| self.out_drop = nn.Dropout(proj_drop) |
| |
| self.max_distance=16 |
| |
| def build_rel_bias(self, coords): |
| return torch.log1p(torch.cdist(coords, coords, p=2)) |
|
|
| def forward(self, x, coords=None, attn_mask: Optional[torch.Tensor] = None, return_attn_score=False): |
| N, L, C = x.shape |
| |
| x_proj = F.linear(x, self.in_proj_weight, self.in_proj_bias).contiguous() |
| q, k, v = [t.contiguous() for t in x_proj.chunk(3, -1)] |
|
|
| q = q.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous() |
| k = k.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous() |
| v = v.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous() |
| |
| if attn_mask is not None: |
| maks_func = create_block_mask( |
| key_padding_mask(attn_mask), N, self.num_heads, L, L |
| ) |
| |
| qk_scale = q.size(-1) ** -0.5 |
| x = self.f_attn( |
| q, k, v, |
| score_mod = get_rel_bias_func(generate_alibi_bias(self.num_heads).to(coords.device), coords, qk_scale) if coords is not None else None, |
| block_mask = maks_func if attn_mask is not None else None, |
| return_lse=return_attn_score, |
| ) |
| |
| x = x.permute(0, 2, 1, 3).contiguous() |
| x = x.reshape(N, L, C).contiguous() |
|
|
| x = self.out_proj(x) |
| x = self.out_drop(x) |
| return x |
|
|
