Chú Ý Nhóm Truy Vấn (Grouped-Query Attention - GQA) LLaMA-3

Mục tiêu kiến thức

• Trong Multi-Head Attention (MHA) truyền thống, số lượng đầu K, V bằng đúng số lượng đầu Q (H_KV = H_Q). Điều này khiến KV Cache phình to chiếm trọn bộ nhớ GPU khi context dài.
• Multi-Query Attention (MQA) chỉ dùng 1 đầu K, V duy nhất (H_KV = 1), tiết kiệm tối đa bộ nhớ nhưng làm suy giảm khả năng biểu diễn.
• Grouped-Query Attention (GQA - Ainslie et al., 2023):
• Chia H_Q đầu Query thành H_KV nhóm (H_Q chia hết cho H_KV, tỷ lệ nhóm G = H_Q // H_KV).
• Mỗi nhóm các đầu Query chia sẻ chung 1 đầu Key và Value.
• Trước khi tính tích vô hướng attention, tensor Key và Value được lặp lại G lần dọc theo chiều head (repeat_kv).
• LLaMA-3 8B dùng H_Q = 32, H_KV = 8 (G = 4), giảm dung lượng KV Cache đi 4 lần!

Yêu cầu

Xây dựng hàm repeat_kv và lớp GroupedQueryAttention(torch.nn.Module):

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # x shape: (B, num_kv_heads, seq_len, head_dim)
    # returns shape: (B, num_kv_heads * n_rep, seq_len, head_dim)
    pass

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model: int, num_query_heads: int, num_kv_heads: int):
        pass

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        pass

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.