题目内容

实现 Grouped Query Attention（GQA）算法。

给定输入序列 $X$（$L \times d$）、Query 头数 $h_q$、KV 头数 $h_{kv}$（$h_q$ 为 $h_{kv}$ 的整数倍），以及投影矩阵，按以下步骤计算输出：

1. 线性投影：$Q = X W_q$ 拆分为 $h_q$ 个头，$K = X W_k$、$V = X W_v$ 拆分为 $h_{kv}$ 个头。
2. 分组计算：每 $g = h_q / h_{kv}$ 个 Query 头共享同一个 KV 头。
3. 注意力计算：每个 Query 头与对应 KV 头计算 $\mathrm{softmax}(QK^{\mathsf T}/\sqrt{d_k})V$。
4. 拼接输出：将所有头输出拼接后通过 $W_o$ 投影。

输入描述

第一行包含四个整数 $L, d, h_q, h_{kv}$（$1 \le L \le 8$，$2 \le d \le 8$，$h_{kv} \le h_q \le d$，$h_q \bmod h_{kv} = 0$，$d \bmod h_q = 0$）。

接下来 $L$ 行，每行 $d$ 个浮点数，表示输入序列 $X$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_q$。

接下来 $d$ 行，每行 $d_{kv}$ 个浮点数，表示 $W_k$（$d_{kv} = d \cdot h_{kv} / h_q$）。

接下来 $d$ 行，每行 $d_{kv}$ 个浮点数，表示 $W_v$。

接下来 $d$ 行，每行 $d$ 个浮点数，表示 $W_o$。

输出描述

输出 $L$ 行，每行 $d$ 个浮点数，表示 GQA 输出，保留 $4$ 位小数。

样例 $1$

输入

2 4 4 2
1.0 0.0 1.0 0.0
0.0 1.0 0.0 1.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0
1.0 0.0
0.0 1.0
0.0 0.0
0.0 0.0
0.0 1.0
0.0 0.0
0.0 0.0
0.0 0.0
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0

输出

0.7311 0.5000 0.7311 0.5000
0.5000 0.7311 0.5000 0.7311

说明

$h_q = 4$，$h_{kv} = 2$，分组比例 $g = 2$。

Query 拆分为 $4$ 头，每头维度 $1$。KV 拆分为 $2$ 头，每头维度 $1$。Query 头 $0,1$ 共享 KV 头 $0$；Query 头 $2,3$ 共享 KV 头 $1$。