123 lines
5.3 KiB
Python
123 lines
5.3 KiB
Python
import math
|
|
import torch
|
|
from einops import rearrange
|
|
from torch import Tensor
|
|
|
|
try:
|
|
from flash_attn import flash_attn_func
|
|
except ImportError:
|
|
flash_attn_func = None
|
|
print("Flash Attention 2 is not available. Falling back to PyTorch's native attention implementation.")
|
|
|
|
# Efficient implementation equivalent to the following:
|
|
def scaled_dot_product_attention(query, key, value=None, attn_mask=None, dropout_p=0.0,
|
|
is_causal=False, scale=None, enable_gqa=False, return_qk_logits=False) -> torch.Tensor:
|
|
L, S = query.size(-2), key.size(-2)
|
|
scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
|
|
attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
|
|
if is_causal:
|
|
assert attn_mask is None
|
|
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
|
|
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
|
attn_bias.to(query.dtype)
|
|
|
|
if attn_mask is not None:
|
|
if attn_mask.dtype == torch.bool:
|
|
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
|
|
else:
|
|
attn_bias = attn_mask + attn_bias
|
|
|
|
if enable_gqa:
|
|
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
|
|
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
|
|
|
|
attn_weight = query @ key.transpose(-2, -1) * scale_factor
|
|
attn_weight += attn_bias
|
|
attn_weight = torch.softmax(attn_weight, dim=-1)
|
|
attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
|
|
if return_qk_logits:
|
|
return attn_weight
|
|
else:
|
|
return attn_weight @ value
|
|
|
|
|
|
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, task_register_embeddings: Tensor = None, use_flash_attention: bool = False) -> Tensor:
|
|
q, k = apply_rope(q, k, pe)
|
|
|
|
if task_register_embeddings is not None:
|
|
k_mask_type_embeds, v_mask_type_embeds = task_register_embeddings[0], task_register_embeddings[1]
|
|
k = torch.cat((k, k_mask_type_embeds), dim=2)
|
|
v = torch.cat((v, v_mask_type_embeds), dim=2)
|
|
|
|
# Use Flash Attention 2 for acceleration (if available)
|
|
try:
|
|
if use_flash_attention and flash_attn_func is not None:
|
|
# Check if the number of heads match for GQA (Grouped Query Attention)
|
|
q_heads = q.size(1)
|
|
k_heads = k.size(1)
|
|
v_heads = v.size(1)
|
|
|
|
# Handle the case where k,v have fewer heads than q (Grouped Query Attention)
|
|
if q_heads % k_heads == 0 and q_heads % v_heads == 0:
|
|
# For GQA, we need to repeat k and v
|
|
k = k.repeat_interleave(q_heads // k_heads, dim=1)
|
|
v = v.repeat_interleave(q_heads // v_heads, dim=1)
|
|
elif k_heads != q_heads or v_heads != q_heads:
|
|
# If heads don't match and can't use GQA, fall back to standard attention
|
|
raise ValueError(f"Incompatible head dimensions: q={q_heads}, k={k_heads}, v={v_heads}")
|
|
|
|
# Input shape: [B, H, L, D] but flash_attn_func expects [B, L, H, D]
|
|
# Need to transpose dimensions 1 and 2
|
|
q_trans = q.transpose(1, 2) # [B, L, H, D]
|
|
k_trans = k.transpose(1, 2) # [B, L, H, D]
|
|
v_trans = v.transpose(1, 2) # [B, L, H, D]
|
|
|
|
# causal=False means non-autoregressive
|
|
x_trans = flash_attn_func(q_trans, k_trans, v_trans, causal=False) # [B, L, H, D]
|
|
|
|
# Convert back to original dimension order
|
|
x = x_trans.transpose(1, 2) # [B, H, L, D]
|
|
|
|
# Output shape: [B, H, L, D], need to convert to [B, L, H*D]
|
|
x = rearrange(x, "B H L D -> B L (H D)")
|
|
else:
|
|
# Fallback to native implementation
|
|
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
x = rearrange(x, "B H L D -> B L (H D)")
|
|
except Exception as e:
|
|
# Fallback to native implementation if flash attention fails
|
|
print(f"Flash attention failed with error: {e}. Falling back to native implementation.")
|
|
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 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, xk, freqs_cis):
|
|
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
|
# print(xq_.shape, freqs_cis.shape)
|
|
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
|
|
if xk is not None:
|
|
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
|
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)
|
|
else:
|
|
return xq_out.reshape(*xq.shape).type_as(xq), None
|
|
|
|
|
|
def apply_learnable_pos_emb(q: Tensor, k: Tensor, embeddings: Tensor) -> Tensor:
|
|
if embeddings.ndim == 3:
|
|
embeddings = embeddings.unsqueeze(1)
|
|
q = q + embeddings
|
|
k = k + embeddings
|
|
return q, k
|