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

用矩阵乘法计算：

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

其中：

• $Q$ 形状为 $(T, d_{\text{model}})$
• $K,V$ 形状为 $(T, h_k d_k)$，其中 $d_k = d_{\text{model}}/h_q$

步骤 2：拆分 Query 为 $h_q$ 个 head

把 $Q$ 的最后一维按 $h_q$ 均分成 $h_q$ 段，每段长度 $d_k$：

• reshape 为 $(T, h_q, d_k)$
• 转置为 $(h_q, T, d_k)$

步骤 3：拆分 Key/Value 为 $h_k$ 个 head（不再组内切子块）

把 $K,V$ reshape 为：

• $K \rightarrow (T, h_k, d_k)$，再转置为 $(h_k, T, d_k)$
• $V \rightarrow (T, h_k, d_k)$，再转置为 $(h_k, T, d_k)$

步骤 4：逐 head 计算注意力（每组共享同一个 K/V head）

令 $group_{size}=h_q/h_k$，对每个 Query head $i$：

1. 计算共享索引：

   $$g(i) = \left\lfloor \frac{i}{group_size} \right\rfloor$$

2. 取 $Q_i \in (T,d_k)$、$K_{g(i)} \in (T,d_k)$、$V_{g(i)} \in (T,d_k)$

3. 打分矩阵：

   $$\text{scores}=\frac{Q_iK_{g(i)}^\top}{\sqrt{d_k}}$$

4. 稳定 softmax（对每行）：

   • scores -= scores.max(axis=1, keepdims=True)
   • attn = exp(scores) / exp(scores).sum(axis=1, keepdims=True)

5. 输出：

   $$H_i=\text{attn}\cdot V_{g(i)}$$

步骤 5：拼接所有 head

把 $H_i$ 按 head 维拼回 $(T,d_{\text{model}})$。

步骤 6：输出层

Python 参考实现

import numpy as np
from typing import List


class Solution:
    def grouped_query_attention(
        self,
        X: List[List[float]],
        W_Q: List[List[float]],
        b_Q: List[float],
        W_K: List[List[float]],
        b_K: List[float],
        W_V: List[List[float]],
        b_V: List[float],
        h_q: int,
        h_k: int,
        W_O: List[List[float]],
        b_O: List[float],
    ) -> List[List[float]]:
        """
        标准分组查询注意力（GQA）。

        关键要点：
        - d_k = d_model / h_q
        - group_size = h_q / h_k
        - K/V head 数为 h_k（压缩），每个 head 维度仍为 d_k
        - 每 group_size 个 Query heads 共享同一个 K/V head：
          g(i) = floor(i / group_size)
        """
        Xn = np.asarray(X, dtype=np.float64)
        WQ = np.asarray(W_Q, dtype=np.float64)
        WK = np.asarray(W_K, dtype=np.float64)
        WV = np.asarray(W_V, dtype=np.float64)
        WO = np.asarray(W_O, dtype=np.float64)
        bQ = np.asarray(b_Q, dtype=np.float64)
        bK = np.asarray(b_K, dtype=np.float64)
        bV = np.asarray(b_V, dtype=np.float64)
        bO = np.asarray(b_O, dtype=np.float64)

        T, d_model = Xn.shape
        if d_model % h_q != 0:
            raise ValueError("d_model must be divisible by h_q")
        if not (1 <= h_k < h_q):
            raise ValueError("require 1 <= h_k < h_q")
        if h_q % h_k != 0:
            raise ValueError("require h_q divisible by h_k")

        d_k = d_model // h_q
        group_size = h_q // h_k
        inv_sqrt = 1.0 / np.sqrt(d_k)

        # WK/WV 输出维度为 h_k*d_k
        hkdk = h_k * d_k
        if WK.shape != (d_model, hkdk) or WV.shape != (d_model, hkdk):
            raise ValueError("W_K/W_V must have shape (d_model, h_k*d_k)")
        if bK.shape != (hkdk,) or bV.shape != (hkdk,):
            raise ValueError("b_K/b_V must have length h_k*d_k")

        # 1) 线性投影
        Q = Xn @ WQ + bQ              # (T, d_model)
        K = Xn @ WK + bK              # (T, h_k*d_k)
        V = Xn @ WV + bV              # (T, h_k*d_k)

        # 2) Q 拆成 h_q 个 head: (h_q, T, d_k)
        Qh = Q.reshape(T, h_q, d_k).transpose(1, 0, 2)

        # 3) K/V 拆成 h_k 个 head: (h_k, T, d_k)
        Kh = K.reshape(T, h_k, d_k).transpose(1, 0, 2)
        Vh = V.reshape(T, h_k, d_k).transpose(1, 0, 2)

        # 4) 逐 head 注意力
        Hh = np.empty((h_q, T, d_k), dtype=np.float64)
        for i in range(h_q):
            g = i // group_size
            Qi = Qh[i]      # (T, d_k)
            Ki = Kh[g]      # (T, d_k)
            Vi = Vh[g]      # (T, d_k)

            scores = (Qi @ Ki.T) * inv_sqrt
            scores = scores - scores.max(axis=1, keepdims=True)
            attn = np.exp(scores)
            attn = attn / attn.sum(axis=1, keepdims=True)

            Hh[i] = attn @ Vi

        # 5) 拼接 head: (T, d_model)
        H = Hh.transpose(1, 0, 2).reshape(T, d_model)

        # 6) 输出层
        O = H @ WO + bO
        return O.tolist()

题目描述

设输入序列为 $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$