算法步骤

步骤 1：线性投影

用矩阵乘法计算：

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

此时 $Q,K,V$ 形状都是 $(T, d_{\text{model}})$。

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

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

• 先 reshape 为 $(T, h_q, d_k)$
• 再转置为 $(h_q, T, d_k)$，方便按 head 索引

步骤 3：拆分 Key/Value 为 $h_k$ 个 group，再切子块

每个 group 覆盖 $\text{group\_size}=h_q/h_k$ 个 Query head。

• K reshape 成 $(T, h_k, d_{\text{group}})$ 再转成 $(h_k, T, d_{\text{group}})$
• 再把最后一维拆成 $(\text{group\_size}, d_k)$，得到 $(h_k, T, \text{group\_size}, d_k)$
• 再转置为 $(h_k, \text{group\_size}, T, d_k)$

V 同理得到 $(h_k, \text{group\_size}, T, d_k)$

这样我们就可以用 (g, r) 直接取到与某 head 对齐的 $K_{g,r},V_{g,r}$。

步骤 4：逐 head 计算注意力

对每个 head $i$：

1. 计算 $g=i//\text{group\_size},\ r=i\%\text{group\_size}$

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

3. 打分矩阵：

   $$\text{scores}=(Q_iK_{g,r}^\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,r}$$

步骤 5：拼接所有 head

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

步骤 6：输出层

$$O = HW_O + b_O$$

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
        - 每个 group 的 K/V 宽度是 d_group = d_k * group_size
        - 为了与每个 Q head 对齐计算注意力，将每个 group 的 K/V 再切成 group_size 个子块，每个子块宽度 d_k
        """
        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
        d_k = d_model // h_q
        group_size = h_q // h_k
        d_group = d_k * group_size
        inv_sqrt = 1.0 / np.sqrt(d_k)

        # 1) 线性投影
        Q = Xn @ WQ + bQ
        K = Xn @ WK + bK
        V = Xn @ WV + bV

        # 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 个 group，再在组内切成 group_size 个子块: (h_k, group_size, T, d_k)
        Kg = (
            K.reshape(T, h_k, d_group)
            .transpose(1, 0, 2)
            .reshape(h_k, T, group_size, d_k)
            .transpose(0, 2, 1, 3)
        )
        Vg = (
            V.reshape(T, h_k, d_group)
            .transpose(1, 0, 2)
            .reshape(h_k, T, group_size, d_k)
            .transpose(0, 2, 1, 3)
        )

        # 4) 逐 head 注意力
        Hh = np.empty((h_q, T, d_k), dtype=np.float64)
        for i in range(h_q):
            g = i // group_size
            r = i % group_size
            Qi = Qh[i]      # (T, d_k)
            Ki = Kg[g, r]   # (T, d_k)
            Vi = Vg[g, r]   # (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 分组数：

  $$h_k,\quad h_k < h_q$$

• 每个 Query head 的维度：

  $$d_k = \frac{d_{\text{model}}}{h_q}$$

你需要实现 分组查询注意力（Grouped Query Attention, GQA），流程如下。

1. 线性投影

$$Q = X W_Q + b_Q,\quad K = X W_K + b_K,\quad V = X W_V + b_V$$

2. 拆分 Q 为多个 Query heads

将 $Q$ 均匀拆分成 $h_q$ 个 Query heads：

$$Q_i \in \mathbb{R}^{T \times d_k}$$

3. 将 K 和 V 拆分为 h_k 组（group）

每个组的维度为：

$$d_{\text{group}} = d_k \cdot \frac{h_q}{h_k}$$

拆分得到：

$$K_j,\ V_j \in \mathbb{R}^{T \times d_{\text{group}}},\quad j=1,\dots,h_k$$

4. 每个 Query head 选择其对应的 K/V 组

令组索引：

$$g(i) = \left\lfloor \frac{i}{h_q/h_k} \right\rfloor$$

对每个 Query head 计算注意力：

$$H_i = \text{softmax}\left(\frac{Q_i K_{g(i)}^\top}{\sqrt{d_k}}\right)V_{g(i)}$$

5. 拼接所有 Query head 的输出

$$H = \text{concat}(H_1,\dots,H_{h_q})$$

6. 输出层线性变换

$$O = H W_O + b_O$$

输入参数

• X：形状为 $T \times d_{\text{model}}$ 的输入序列
• W_Q, W_K, W_V：形状为 $d_{\text{model}} \times d_{\text{model}}$ 的投影矩阵
• b_Q, b_K, b_V：长度为 $d_{\text{model}}$ 的偏置项
• h_q：Query head 数
• h_k：Key/Value 分组数，需满足 $h_k < h_q$
• 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, 0],
 [0, 1, 1, 0],
 [1, 0, 0, 1],
 [0, 0, 1, 1]]

b_K = [0, 0, 0, 0]

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

b_V = [0, 0, 0, 0]

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 =
[[3.46, 4.49, 3.99, 3.96], 
 [3.76, 3.0, 2.5, 4.26]]

提示

• 输入范围： $-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$
• head 数要求： $1 \le h_k < h_q \le d_{\text{model}}$ 且 $d_{\text{model}} \bmod h_q = 0$
• softmax 输出满足： $0 \le \text{softmax}(z_i) \le 1$，且 $\sum_i \text{softmax}(z_i) = 1$