以下代码实现了分组查询注意力（GQA），请阅读代码并填写空缺部分。 [1] 和 [2] 处应分别填入：

import torch

import torch.nn as nn

import torch.nn.functional as F

import math

class GroupedQueryAttention(nn.Module):

   """分组查询注意力（GQA）：多个Query头共享同一组Key和Value头"""

   def __init__(self, hidden_dim, num_q_heads, num_kv_heads):

       super().__init__()

       self.num_q_heads = num_q_heads  # Query头的数量，例如32

       self.num_kv_heads = num_kv_heads  # KV头的数量，例如8

       self.head_dim = hidden_dim // num_q_heads

       # [1] 计算每个KV头对应多少个Query头（分组大小）

       __________________________________________________________

       self.q_proj = nn.Linear(hidden_dim, num_q_heads * self.head_dim)

       self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)

       self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)

       self.o_proj = nn.Linear(hidden_dim, hidden_dim)

   def forward(self, x):

       B, L, _ = x.shape

       q = self.q_proj(x).view(B, L, self.num_q_heads, self.head_dim).transpose(1, 2)

       k = self.k_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)

       v = self.v_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)

       # 将 KV 头扩展以匹配 Q 头的数量（分组复制）

       k = k.repeat_interleave(self.num_groups, dim=1)

       v = v.repeat_interleave(self.num_groups, dim=1)

       # 缩放点积注意力

       scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

       attn_weights = F.softmax(scores, dim=-1)

       attn_output = torch.matmul(attn_weights, v)

       # 合并多头输出

       attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, -1)

       output = self.o_proj(attn_output)

       return output