| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.nn.attention import sdpa_kernel, SDPBackend |
| from .RotaryPositionalEmbedding import RotaryPositionalEmbedding |
|
|
| _SDPA_BACKENDS = [ |
| SDPBackend.CUDNN_ATTENTION, |
| SDPBackend.FLASH_ATTENTION, |
| SDPBackend.EFFICIENT_ATTENTION, |
| SDPBackend.MATH, |
| ] |
|
|
| class GroupedQueryAttention(nn.Module): |
| def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, dropout: float): |
| super().__init__() |
| assert d_model % num_heads == 0 |
| assert num_heads % num_kv_heads == 0, \ |
| "num_heads phải chia hết cho num_kv_heads" |
|
|
| self.num_heads = num_heads |
| self.num_kv_heads = num_kv_heads |
| self.num_groups = num_heads // num_kv_heads |
| self.d_k = d_model // num_heads |
| self.dropout_rate = dropout |
| self.q_dim = num_heads * self.d_k |
| self.kv_dim = num_kv_heads * self.d_k |
| self.wqkv = nn.Linear(d_model, self.q_dim + 2 * self.kv_dim, bias=False) |
| self.wo = nn.Linear(self.q_dim, d_model, bias=False) |
|
|
| def _project_qkv(self, x: torch.Tensor): |
| B, T, _ = x.shape |
| qkv = self.wqkv(x) |
| q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) |
| q = q.view(B, T, self.num_heads, self.d_k) |
| k = k.view(B, T, self.num_kv_heads, self.d_k) |
| v = v.view(B, T, self.num_kv_heads, self.d_k) |
| return q, k, v |
|
|
| def _merge(self, out: torch.Tensor) -> torch.Tensor: |
| B, _, T, _ = out.shape |
| return self.wo(out.transpose(1, 2).contiguous().view(B, T, self.num_heads * self.d_k)) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attn_mask=None) -> torch.Tensor: |
| q, k, v = self._project_qkv(x) |
|
|
| q = RotaryPositionalEmbedding.apply_rope(q, cos, sin) |
| k = RotaryPositionalEmbedding.apply_rope(k, cos, sin) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| dropout_p = self.dropout_rate if self.training else 0.0 |
|
|
| with sdpa_kernel(_SDPA_BACKENDS): |
| out = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=attn_mask, is_causal=(attn_mask is None), dropout_p=dropout_p, |
| enable_gqa=True, |
| ) |
| return self._merge(out) |
|
|
| def prefill(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): |
| q, k, v = self._project_qkv(x) |
|
|
| q = RotaryPositionalEmbedding.apply_rope(q, cos, sin) |
| k = RotaryPositionalEmbedding.apply_rope(k, cos, sin) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| with sdpa_kernel(_SDPA_BACKENDS): |
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True) |
| return self._merge(out), (k, v) |
|
|
| def forward_with_cache(self, x: torch.Tensor, past_kv, cache_len: int, cos: torch.Tensor, sin: torch.Tensor): |
| B, T, _ = x.shape |
| q, k, v = self._project_qkv(x) |
|
|
| q = RotaryPositionalEmbedding.apply_rope(q, cos, sin) |
| k = RotaryPositionalEmbedding.apply_rope(k, cos, sin) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| past_kv[0][:B, :, cache_len:cache_len + T, :] = k |
| past_kv[1][:B, :, cache_len:cache_len + T, :] = v |
|
|
| k_full = past_kv[0][:B, :, :cache_len + T, :] |
| v_full = past_kv[1][:B, :, :cache_len + T, :] |
|
|
| with sdpa_kernel(_SDPA_BACKENDS): |
| out = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=False, enable_gqa=True) |
| return self._merge(out) |
|
|
