题目描述

分组查询注意力（GQA）是标准多头注意力的变体，核心思想是让多个 Query 头共享一组 KV 头，在保证模型效果的同时降低计算开销。请基于缩放点积注意力实现分组查询注意力。

核心数学逻辑

GQA 由基础缩放点积注意力扩展而来，关键差别是头的分组共享：

1. 将 Query 矩阵拆分为 $num\_q\_heads$ 个独立的 Query 头子矩阵，记为 $Q_0, Q_1, \ldots, Q_{num\_q\_heads-1}$。
2. 将 Key、Value 矩阵拆分为 $num\_kv\_heads$ 个独立的 KV 头子矩阵，记为 $K_0, K_1, \ldots, K_{num\_kv\_heads-1}$ 与 $V_0, V_1, \ldots, V_{num\_kv\_heads-1}$。
3. $num\_q\_heads$ 必须是 $num\_kv\_heads$ 的整数倍。每组包含 $g = num\_q\_heads / num\_kv\_heads$ 个 Query 头，共享 1 组 KV 头：
   • 第 $0$ 组：$Q_0, Q_1, \ldots, Q_{g-1}$ 共享 $K_0, V_0$。
   • 第 $1$ 组：$Q_g, Q_{g+1}, \ldots, Q_{2g-1}$ 共享 $K_1, V_1$。
   • 第 $num\_kv\_heads-1$ 组：$Q_{(num\_kv\_heads-1)g}, \ldots, Q_{num\_q\_heads-1}$ 共享 $K_{num\_kv\_heads-1}, V_{num\_kv\_heads-1}$。
4. 对每组内的每个 Query 头，与共享的 KV 头做基础缩放点积注意力，得到该 Query 头的输出。
5. 将全部 Query 头的输出按顺序拼接，得到最终的 GQA 输出矩阵。

基础缩放点积注意力要求

单组 $Q, K, V$ 按下列规则计算。

输入输出

输入：$Q(n, d_k)$、$K(n, d_k)$、$V(n, d_v)$、$mask(n, n)$。

输出：$Output(n, d_v)$。

核心公式

$Attention(Q, K, V, mask) = Softmax\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$

当 $mask[i][j] = 1$ 时 $M[i][j] = -\infty$，否则 $M[i][j] = 0$。实现中 $-\infty$ 可取 $-10^{30}$。

Softmax（逐行）

$A[i][j] = \frac{\exp(S[i][j] - \max_t S[i][t])}{\sum_{t=0}^{n-1} \exp(S[i][t] - \max_t S[i][t])}$

计算步骤

1. 转置：对 $K$ 转置得到 $K^T$。
2. 手动矩阵乘法：$S[i][j] = \sum_{t=0}^{d_k-1} Q[i][t] \cdot K[j][t]$。
3. 缩放：$S[i][j] = S[i][j] / \sqrt{d_k}$。
4. 掩码：若 $mask[i][j] = 1$，则 $S[i][j] = -\infty$。
5. Softmax：按上述公式逐行计算注意力权重 $A$。
6. 加权求和：$Output[i][k] = \sum_{t=0}^{n-1} A[i][t] \cdot V[t][k]$。

算法实现流程

步骤1：输入矩阵拆分

• 原始 $Q$ 的形状为 $(n, d_{total\_q})$，其中 $d_{total\_q} = num\_q\_heads \times d_k$。
• 原始 $K$ 的形状为 $(n, d_{total\_k})$，其中 $d_{total\_k} = num\_kv\_heads \times d_k$。
• 原始 $V$ 的形状为 $(n, d_{total\_v})$，其中 $d_{total\_v} = num\_kv\_heads \times d_v$。

拆分规则：

• $Q$ 拆成 $num\_q\_heads$ 个子矩阵，每个形状为 $(n, d_k)$。
• $K$ 拆成 $num\_kv\_heads$ 个子矩阵，每个形状为 $(n, d_k)$。
• $V$ 拆成 $num\_kv\_heads$ 个子矩阵，每个形状为 $(n, d_v)$。

步骤2：分组计算注意力

• 每组 Query 头数为 $g = \lfloor num\_q\_heads / num\_kv\_heads \rfloor$，即 $num\_q\_heads$ 除以 $num\_kv\_heads$ 后向下取整。
• 对第 $m$ 组（$0 \le m < num\_kv\_heads$）：
  1. 取该组 KV 头 $K_m, V_m$。
  2. 取该组 $g$ 个 Query 头 $Q_{m \cdot g}, Q_{m \cdot g+1}, \ldots, Q_{(m+1) \cdot g-1}$。
  3. 对每个 Query 头 $Q_i$ 计算 $Output_i = Attention(Q_i, K_m, V_m, mask)$。

步骤3：拼接输出

将 $Output_0, Output_1, \ldots, Output_{num\_q\_heads-1}$ 按列拼接，得到最终输出矩阵，形状为 $(n, num\_q\_heads \times d_v)$，即序列长度 ×（Query 头数 × 单个 Value 头维度）。

输入描述

第一行：$n$ $num\_q\_heads$ $num\_kv\_heads$ $d_k$ $d_v$。

字段含义依次为：序列长度、Query 头总数、Key/Value 头总数、单个 Query/Key 头维度、单个 Value 头维度。

• $num\_q\_heads$：Query 头总数，必须是 $num\_kv\_heads$ 的整数倍。
• $num\_kv\_heads$：Key 与 Value 的头总数，二者始终相等。
• $d_k$：单个 Query 头与单个 Key 头的特征维度，二者相同。
• $d_v$：单个 Value 头的特征维度。

接下来 $n$ 行：原始 Query 矩阵 $Q$，形状为 $(n, num\_q\_heads \times d_k)$，每行 $num\_q\_heads \times d_k$ 个浮点数。

接下来 $n$ 行：原始 Key 矩阵 $K$，形状为 $(n, num\_kv\_heads \times d_k)$，每行 $num\_kv\_heads \times d_k$ 个浮点数。

接下来 $n$ 行：原始 Value 矩阵 $V$，形状为 $(n, num\_kv\_heads \times d_v)$，每行 $num\_kv\_heads \times d_v$ 个浮点数。

接下来 $n$ 行：掩码矩阵，形状为 $(n, n)$。$0$ 表示可见，$1$ 表示屏蔽；若全为 $0$，则表示无掩码。

输入约束：

• $1 \le n \le 100$
• $1 \le num\_q\_heads \le 16$，$1 \le num\_kv\_heads \le 8$，且 $num\_q\_heads \bmod num\_kv\_heads = 0$
• $1 \le d_k, d_v \le 10$
• 矩阵元素取值范围：$[-10.0, 10.0]$

输出描述

输出一个二维浮点数矩阵，表示分组查询注意力的计算结果。

形状：$(n, num\_q\_heads \times d_v)$（序列长度 ×（Query 头数 × 单个 Value 头维度））。

输出格式：共 $n$ 行，每行 $num\_q\_heads \times d_v$ 个浮点数，每个数保留 $2$ 位小数。

样例

样例1

输入

2 4 2 1 1
1.0 0.0 0.0 1.0
0.0 1.0 1.0 0.0
0.0 1.0
1.0 0.0
2.0 4.0
8.0 6.0
0 0
0 0

输出

6.39 5.00 5.00 4.54
5.00 6.39 4.54 5.00

说明

$n=2$，$num\_q\_heads=4$，$num\_kv\_heads=2$，$d_k=1$，$d_v=1$，故 $g=2$。掩码全为 $0$，$\sqrt{d_k}=1$。

第 $0$ 组 $Q_0, Q_1$ 共享 $K_0, V_0$；第 $1$ 组 $Q_2, Q_3$ 共享 $K_1, V_1$。按列拆开后：

• $Q_0$ 两行为 $1, 0$，$Q_1$ 两行为 $0, 1$，$K_0$ 两行为 $0, 1$，$V_0$ 两行为 $2, 8$。
• $Q_2$ 两行为 $0, 1$，$Q_3$ 两行为 $1, 0$，$K_1$ 两行为 $1, 0$，$V_1$ 两行为 $4, 6$。

得分相差 $1$ 时，较大位置的 Softmax 权重为 $e/(1+e) \approx 0.731059$，较小位置为 $1/(1+e) \approx 0.268941$。两行得分相同时，权重各为 $0.5$。

• $Q_0$：第 $0$ 行得分为 $0, 1$，输出 $0.268941 \times 2 + 0.731059 \times 8 = 6.386354$，保留两位小数为 $6.39$；第 $1$ 行得分为 $0, 0$，输出 $0.5 \times 2 + 0.5 \times 8 = 5.00$。
• $Q_1$：第 $0$ 行得分为 $0, 0$，输出 $5.00$；第 $1$ 行得分为 $0, 1$，输出 $6.39$。
• $Q_2$：第 $0$ 行得分为 $0, 0$，输出 $0.5 \times 4 + 0.5 \times 6 = 5.00$；第 $1$ 行得分为 $1, 0$，输出 $0.731059 \times 4 + 0.268941 \times 6 = 4.537882$，保留两位小数为 $4.54$。
• $Q_3$：第 $0$ 行得分为 $1, 0$，输出 $4.54$；第 $1$ 行得分为 $0, 0$，输出 $5.00$。

四个头按列拼接，即得到输出。

样例2

输入

2 2 2 2 1
1.0 0.0 0.0 1.0
0.0 1.0 1.0 0.0
1.0 0.0 0.0 1.0
0.0 1.0 1.0 0.0
2.0 6.0
4.0 8.0
0 1
0 0

输出

2.00 6.00
3.34 7.34

说明

$n=2$，$num\_q\_heads=2$，$num\_kv\_heads=2$，$d_k=2$，$d_v=1$，故 $g=1$。$Q_0$ 只使用 $K_0, V_0$，$Q_1$ 只使用 $K_1, V_1$。$\sqrt{d_k}=\sqrt{2}\approx 1.414214$。掩码第 $0$ 行第 $1$ 列为 $1$，该位置在缩放之后改为 $-\infty$。

$Q_0$ 两行是 $(1, 0)$ 与 $(0, 1)$，$K_0$ 两行是 $(1, 0)$ 与 $(0, 1)$，$V_0$ 两行是 $2, 4$。点积矩阵两行分别为 $1, 0$ 与 $0, 1$，缩放后第 $0$ 行约为 $0.707107, 0$；掩码使第二项为 $-\infty$，权重为 $1, 0$，输出 $2.00$。第 $1$ 行缩放后约为 $0, 0.707107$，权重约为 $0.330238, 0.669762$，输出 $0.330238 \times 2 + 0.669762 \times 4 = 3.339524$，保留两位小数为 $3.34$。

$Q_1$ 两行是 $(0, 1)$ 与 $(1, 0)$，$K_1$ 两行是 $(0, 1)$ 与 $(1, 0)$，$V_1$ 两行是 $6, 8$。缩放后的得分与 $Q_0$ 相同，掩码也相同，因此第 $0$ 行输出 $6.00$，第 $1$ 行输出 $0.330238 \times 6 + 0.669762 \times 8 = 7.339524$，保留两位小数为 $7.34$。

两头按列拼接，即得到输出。