GQA - Group Query Attention

MHA의 한계

• 각 Head가 자신만의 K와 V를 가지기 때문에, 모델이 새로운 token을 생성할 때마다 이전의 모든 token에 해당하는 Key-Value cache를 메모리에서 불러와야 함.
• 이로 인해 병목 현상이 생기고, 높은 메모리 대역을 필요로 했다.

MQA (Multi Query Attention)

• 하나의 context에 대해 Haed 개수 만큼의 Q를 만든다.
• 그 다음 공통의 K와 V에 대해 Q의 attention score를 계산한다.
• 이 방식을 통해 key-value cache의 크기를 줄였다.
• 메모리 로딩량이 감소하여 추론 속도가 크게 증가했다.
• 하지만 모델의 표현량과 학습 안정성이 떨어져 전체적인 성능도 떨어졌다.

GQA (Qrouped-Query Attention)

MAH와 MQA의 중간지점

• 하나의 context에 대해 Head 개수 만큼 서로 다른 관점의 Query를 생성한다.
• 이후 $M$개의 Key-Value space를 만든다.
  • $M = H/G$ - 여기서 $H$는 Head의 개수이고, $G$는 임의의 Group의 개수이다.
  • $M$개의 Group 내에서 attention score를 구한다.
• 만약 8개의 Q가 있고, 4개의 group을 사용한다면, 1~2번 Q가 첫 번째 K-V를 공유하고, 3~4번 Q가 두 번째 K-V를 공유하는 식이다.
• G가 1이면, MHA이고, G가 H와 같으면 MQA가 된다.

class GroupQueryAttention(nn.Module):
    def __init__(self, d_in, d_model, num_heads, num_kv_groups, dtype):
        super().__init__()
        assert num_heads % num_kv_groups == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        
        # reshape 시 (B, d_in, num_kv_groups, head_dim) 으로 변환됨
        self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=None)
        self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim)
        
        self.W_query = nn.Linear(d_in, d_model, bias=False, dtype=None)
        self.out_proj = nn.Linear(d_model, d_model, bias=False, dtype=None)
        self.num_kv_groups = num_kv_groups
        
        self.group_size = num_heads // num_kv_groups
        
    def forward(self, x, mask=None, cos=None, sin=None):
        b, num_tokens, d_in = x.shape
        
        Q = self.W_query(x)
        K = self.W_key(x)
        V = self.W_value(x)
        
        Q = Q.view(b, self.num_heads, num_tokens, self.head_dim)
        K = K.view(b, self.num_kv_groups, num_tokens, self.head_dim)
        V = V.view(b, self.num_kv_groups, num_tokens, self.head_dim)
        
        K = K.repeat_interleave(self.group_size, dim=1)
        V = V.repeat_interleave(self.group_size, dim=1)
        
        attn_scores = Q @ K.transpose(2,3)
        
        if mask in None:
            mask = torch.triu(torch.ones(num_tokens, num_tokens, deviec=x.device, dtype=torch.bool), diagonal=1)
        attn_scores.masked_fill_(mask, -torch.inf)
        attn_weights = torch.softmax(attn_scores/K.shape[-1]**0.5, dim=-1)
        
        context_vec = (attn_weights @ V).trannspose(1,2)
        
        context_vec = context_vec.reshape(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec)
        
        return context_vec