步骤 1：线性投影（注意 K/V 输出维度已压缩）

用矩阵乘法计算：

• $Q = XW_Q + b_Q$
• $K = XW_K + b_K$
• $V = XW_V + b_V$

其中：

题目描述

设输入序列为 $X \in \mathbb{R}^{T \times d_{\text{model}}}$。

给定：

• Query head 数：$h_q$
• Key/Value head 数（分组数）：$h_k$，满足 $1 \le h_k < h_q$
• 每个 Query head 的维度：$d_k = \frac{d_{\text{model}}}{h_q}$

并要求：

• $d_{\text{model}} \bmod h_q = 0$
• $h_q \bmod h_k = 0$

你需要实现分组查询注意力（Grouped Query Attention, GQA）。其核心特点是：Q 有 $h_q$ 个 head，而 K/V 只有 $h_k$ 个 head，因此多个 Query head 共享同一组 K/V。

1. 线性投影（K/V 的投影维度被压缩）

其中参数形状为：

• $W_Q \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$，$b_Q \in \mathbb{R}^{d_{\text{model}}}$
• $W_K \in \mathbb{R}^{d_{\text{model}} \times (h_k d_k)}$，$b_K \in \mathbb{R}^{h_k d_k}$
• $W_V \in \mathbb{R}^{d_{\text{model}} \times (h_k d_k)}$，$b_V \in \mathbb{R}^{h_k d_k}$

因此：

• $Q \in \mathbb{R}^{T \times d_{\text{model}}}$
• $K, V \in \mathbb{R}^{T \times (h_k d_k)}$

2. 多头拆分

将 $Q$ 拆分为 $h_q$ 个 Query heads：

将 $K,V$ 拆分为 $h_k$ 个 heads：

3. Query head 到 K/V head 的映射

令每组包含的 Query head 数为：

对第 $i$ 个 Query head，其对应的 K/V head 索引为：

4. 分组注意力计算

对每个 Query head $i$，使用共享的 $K_{g(i)}, V_{g(i)}$：

其中 softmax 沿最后一维归一化。

5. 拼接所有 Query heads 的输出

6. 输出层线性变换

其中 $W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$，$b_O \in \mathbb{R}^{d_{\text{model}}}$，输出 $O \in \mathbb{R}^{T \times d_{\text{model}}}$。

输入参数

• $X$：形状为 $T \times d_{\text{model}}$
• $W_Q$：形状为 $d_{\text{model}} \times d_{\text{model}}$，$b_Q$：长度为 $d_{\text{model}}$
• $W_K,W_V$：形状为 $d_{\text{model}} \times (h_k d_k)$
• $b_K,b_V$：长度为 $h_k d_k$
• $h_q$：Query head 数
• $h_k$：Key/Value head 数（分组数），满足 $1 \le h_k < h_q$ 且 $h_q \bmod h_k = 0$
• $W_O$：形状为 $d_{\text{model}} \times d_{\text{model}}$，$b_O$：长度为 $d_{\text{model}}$

返回值

• $O$：形状为 $T \times d_{\text{model}}$

示例

输入：

X =
[[1, 0, 1, 2],
 [0, 1, 1, 0]]

W_Q =
[[1, 0, 0, 1],
 [0, 1, 1, 0],
 [1, 0, 1, 0],
 [0, 1, 0, 1]]

b_Q = [0, 0, 0, 0]

W_K =
[[1, 1],
 [0, 1],
 [1, 0],
 [0, 1]]

b_K = [0, 0]

W_V =
[[1, 0],
 [1, 1],
 [0, 1],
 [0, 0]]

b_V = [0, 0]

h_q = 2

h_k = 1

W_O =
[[1, 0, 1, 0],
 [0, 1, 0, 1],
 [1, 0, 0, 1],
 [0, 1, 1, 0]]

b_O = [0, 0, 0, 0]

输出：

O =
[[2.00, 2.02, 2.01, 2.01],
 [2.00, 2.30, 2.20, 2.11]]

提示

• 输入范围： $-1000 \le X[i,j] \le 1000$
• 权重范围： $-10 \le W_Q,\ W_K,\ W_V,\ W_O \le 10$
• 偏置范围： $-5 \le b_Q,\ b_K,\ b_V,\ b_O \le 5$