P5613. 分组查询注意力

场次：2026年10月10日 · 来源：CodeFun2000 P5613

题目描述

分组查询注意力（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$。

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

---

解题思路

第一步：理解问题本质

先看单头缩放点积注意力在做什么：

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

• $Q$ 的第 $i$ 行是“位置 $i$ 想找什么”，$K$ 的第 $j$ 行是“位置 $j$ 能提供什么线索”，两者的点积就是匹配分数；
• 除以 $\sqrt{d_k}$ 是缩放。$d_k$ 越大，点积的方差越大，softmax 会迅速变成“几乎只看分数最高的那一个”，梯度也随之消失；
• $mask[i][j]=1$ 表示位置 $i$ 不允许看位置 $j$，做法是给分数加 $-\infty$，让它的权重变成 $0$；
• softmax 是逐行做的：每个位置对所有位置分配一组和为 $1$ 的权重；
• 最后用这组权重对 $V$ 的各行加权求和，得到这个位置的新表示。

多头注意力把这个过程并行做 $num\_q\_heads$ 次，每个头用自己的 $Q_h$、$K_h$、$V_h$，最后把各头的输出拼起来。

分组查询注意力（GQA）只改一处：KV 头比 Q 头少。设

$$g=\frac{num\_q\_heads}{num\_kv\_heads}$$

则第 $0$ 组是 $Q_0,\dots,Q_{g-1}$ 共享 $K_0,V_0$，第 $1$ 组是 $Q_g,\dots,Q_{2g-1}$ 共享 $K_1,V_1$，依此类推。

这样做的动机很实际：推理时每个已生成的 token 都要缓存它的 $K$、$V$（KV cache），缓存大小正比于 KV 头数。让 $g$ 个 Query 头共享一组 KV，缓存和 KV 投影的计算量都降到 $1/g$，而 Query 头数不变，表达能力基本保留。

两个极端情况可以帮助理解：

• $num\_kv\_heads=num\_q\_heads$（即 $g=1$）：退化成标准多头注意力 MHA；
• $num\_kv\_heads=1$：所有 Query 头共享同一组 KV，即 MQA。

本题不涉及训练，也不涉及参数的随机初始化，输入输出和每一步公式都给定，所以本质是“按维度把公式翻译成循环”。难点不在算法，而在维度对齐。

第二步：暴力解法：每个头现算

最直接的做法是不预先拆头，每次要用到某个头时再从整行里切片：

for h in range(num_q_heads):
    kv = h // g                     # 这个 Query 头属于哪一组 KV
    for i in range(n):
        for j in range(n):
            # 每次都要从整行里切出第 h 个头的 d_k 个分量
            scores[j] = dot(Q[i][h*d_k:(h+1)*d_k], K[j][kv*d_k:(kv+1)*d_k]) / sqrt(d_k)
        ...

这个写法结果正确，但有两个浪费：

1. 切片发生在 $n^2$ 次内层循环里，同一个头的列被反复截取；
2. 同一组 KV 被 $g$ 个 Query 头各用一遍，$K$、$V$ 的列切片要重复取 $g$ 次。

第三步：优化解法：先拆头，再复用 KV

把三件事提到循环外面：

1. 预拆分。 输入矩阵的每一行是一个位置的全部头拼接起来的特征，所以拆头是按列切：

q_heads = split_heads(q_mat, num_q_heads, d_k)   # num_q_heads 个 (n, d_k)
k_heads = split_heads(k_mat, num_kv_heads, d_k)  # num_kv_heads 个 (n, d_k)
v_heads = split_heads(v_mat, num_kv_heads, d_v)  # num_kv_heads 个 (n, d_v)

第 $h$ 个头占据列区间 $[h\cdot width,\ (h+1)\cdot width)$，切完就是一张 $(n,width)$ 的独立矩阵，后面所有循环都只在这个小矩阵里跑。

2. 组内复用。 第 $h$ 个 Query 头对应的 KV 组号是 $h\ //\ g$，同一个组里的 $g$ 个头用的是同一对 $K_m,V_m$：

kv_index = h // group_size

3. 跳过零权重。 被掩码屏蔽的位置权重是 $0$，加权求和时可以直接 continue，省掉一次长度为 $d_v$ 的内层循环。

计算顺序也固定下来，一次成型：

点积 -> 除以 sqrt(d_k) -> 掩码置 -1e30 -> 逐行 softmax -> 乘 V

第四步：完整实现细节

1. 各个矩阵的形状

名称	形状	说明
$Q$	$n\times(num\_q\_heads\cdot d_k)$	每行是一个位置的全部 Query 头
$Q_h$	$n\times d_k$	第 $h$ 个 Query 头
$K$	$n\times(num\_kv\_heads\cdot d_k)$	每行是一个位置的全部 Key 头
$K_m$	$n\times d_k$	第 $m$ 个 KV 组的 Key
$V$	$n\times(num\_kv\_heads\cdot d_v)$	每行是一个位置的全部 Value 头
$V_m$	$n\times d_v$	第 $m$ 个 KV 组的 Value
$S$、$A$	$n\times n$	分数矩阵与注意力权重矩阵
$Output_h$	$n\times d_v$	第 $h$ 个 Query 头的输出
最终输出	$n\times(num\_q\_heads\cdot d_v)$	各头按列拼接

注意输出的宽度是 $num\_q\_heads\cdot d_v$，与输入的 $Q$ 宽度 $num\_q\_heads\cdot d_k$ 不一定相同，因为 $d_v$ 和 $d_k$ 是两个独立的参数。

2. 掩码的位置

掩码是在缩放之后、softmax 之前生效的。顺序不能换：

• 先缩放再加 $-\infty$，与“先加 $-\infty$ 再缩放”在结果上等价（$-\infty$ 缩放仍是 $-\infty$），但和 softmax 的先后不能错；
• 如果错写成“softmax 之后再清零权重”，分母里已经算进了被屏蔽项的指数，其余权重之和不再是 $1$，结果会偏小。

实现里用题目建议的 $-10^{30}$ 代替 $-\infty$：

if mask_i[j] == 1:
    scores[j] = -1e30

这样在减最大值、取指数之后，被屏蔽项的指数是 $e^{-10^{30}}$，直接下溢成 $0.0$，不会出现 nan。

3. softmax 的数值稳定性

逐行减去该行最大值再取指数：

$$A[i][j]=\frac{\exp(S[i][j]-\max_t S[i][t])}{\sum_t \exp(S[i][t]-\max_t S[i][t])}$$

分子分母同乘 $e^{-\max}$，结果不变，但指数的输入不会超过 $0$，不可能上溢。

4. 全屏蔽行的行为

如果某一行的所有位置都被屏蔽，这一行的分数全是 $-10^{30}$，减去最大值后全为 $0$，权重退化成 $\frac{1}{n},\dots,\frac{1}{n}$。这是用 $-10^{30}$ 哨兵值的必然结果（用真正的 $-\infty$ 会算出 $0/0$）。本题的正常数据每行都有可见位置，这一条只作为边界认知。

5. 拼接顺序

第 $h$ 个头的 $d_v$ 个分量占据输出的第 $h\cdot d_v$ 到 $(h+1)\cdot d_v-1$ 列，顺序与输入中头的编号一致。拼接顺序和拆头顺序必须互为逆操作，否则投影矩阵的列会对不上。

6. 输出格式

每个数保留 $2$ 位小数。浮点结果可能是 $-0.001$ 这种绝对值很小但符号为负的数，格式化前把 $|x|<0.005$ 归零，避免输出 -0.00。

---

完整代码实现

import math
import sys

MASKED = -1e30   # 用题面建议的 -1e30 代替 -inf，避免 nan 传播


def softmax_row(scores):
    """对一行分数做 softmax。

    先减去该行最大值：exp 的输入不会超过 0，避免上溢；
    这一步只是把分子分母同时除以同一个常数，不改变结果。
    """
    row_max = max(scores)
    exps = [math.exp(s - row_max) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def single_head_attention(q, k, v, mask, n, d_k, d_v, scale):
    """单头缩放点积注意力：Attention(Q, K, V, mask)，返回 (n, d_v)。"""
    output = []
    for i in range(n):
        qi = q[i]
        mask_i = mask[i]

        # 第 1、2 步：手动矩阵乘法 S[i][j] = Q[i] · K[j]，再除以 sqrt(d_k)
        scores = []
        for j in range(n):
            kj = k[j]
            dot = 0.0
            for t in range(d_k):
                dot += qi[t] * kj[t]
            scores.append(dot / scale)

        # 第 4 步：掩码在缩放之后生效，被屏蔽的位置直接压成 -inf（这里用 -1e30）
        for j in range(n):
            if mask_i[j] == 1:
                scores[j] = MASKED

        # 第 5 步：逐行 softmax
        weights = softmax_row(scores)

        # 第 6 步：用注意力权重对 V 加权求和，得到 d_v 维输出
        row = [0.0] * d_v
        for j in range(n):
            w = weights[j]
            if w == 0.0:
                continue          # 被屏蔽位置的权重是 0，跳过可以省掉一次内层循环
            vj = v[j]
            for t in range(d_v):
                row[t] += w * vj[t]
        output.append(row)

    return output


def split_heads(matrix, num_heads, width):
    """把 (n, num_heads * width) 的矩阵按列切成 num_heads 个 (n, width) 子矩阵。"""
    return [
        [row[h * width:(h + 1) * width] for row in matrix]
        for h in range(num_heads)
    ]


def format_float(value):
    """保留 2 位小数，并把 -0.00 归一成 0.00。"""
    if abs(value) < 0.005:
        value = 0.0
    return f"{value:.2f}"


def solve():
    """读取输入、计算分组查询注意力并按要求输出。"""
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    pos = 0
    n, num_q_heads, num_kv_heads, d_k, d_v = (int(x) for x in tokens[pos:pos + 5])
    pos += 5

    def read_matrix(rows, cols, cast=float):
        """从当前位置读取 rows x cols 的矩阵。"""
        nonlocal pos
        mat = []
        for _ in range(rows):
            mat.append([cast(x) for x in tokens[pos:pos + cols]])
            pos += cols
        return mat

    q_total = num_q_heads * d_k      # Q 每行的宽度
    k_total = num_kv_heads * d_k     # K 每行的宽度
    v_total = num_kv_heads * d_v     # V 每行的宽度

    q_mat = read_matrix(n, q_total)
    k_mat = read_matrix(n, k_total)
    v_mat = read_matrix(n, v_total)
    mask = read_matrix(n, n, cast=int)

    # 按列拆头：Q 有 num_q_heads 个、K/V 各有 num_kv_heads 个
    q_heads = split_heads(q_mat, num_q_heads, d_k)
    k_heads = split_heads(k_mat, num_kv_heads, d_k)
    v_heads = split_heads(v_mat, num_kv_heads, d_v)

    scale = math.sqrt(d_k)              # 缩放因子，d_k 越大点积的量级越大
    group_size = num_q_heads // num_kv_heads   # 每组共享同一对 KV 的 Query 头数

    # 每个 Query 头独立算一次注意力，KV 头按组复用
    head_outputs = []
    for h in range(num_q_heads):
        kv_index = h // group_size       # 第 h 个 Query 头属于第 h // g 组
        head_outputs.append(
            single_head_attention(
                q_heads[h], k_heads[kv_index], v_heads[kv_index],
                mask, n, d_k, d_v, scale,
            )
        )

    # 按头编号拼接：位置 i 的输出 = [头 0 的 d_v 维 | 头 1 的 d_v 维 | ...]
    for i in range(n):
        row = []
        for h in range(num_q_heads):
            row.extend(head_outputs[h][i])
        print(" ".join(format_float(x) for x in row))


if __name__ == "__main__":
    solve()

NumPy 版本

以下版本使用 NumPy。只有在评测环境允许第三方库时才能提交；CodeFun2000 的普通 Python 环境不一定安装 NumPy。

import sys

import numpy as np

MASKED = -1e30   # 用题面建议的 -1e30 代替 -inf，避免 nan 传播


def gqa(q_mat, k_mat, v_mat, mask, num_q_heads, num_kv_heads, d_k, d_v):
    """分组查询注意力：一次算一整组 Query 头，返回 (n, num_q_heads * d_v)。

    所有头都按「列」拆：reshape 成 (n, heads, width) 后把 heads 维提到前面，
    得到 (heads, n, width) 的三维数组，后面直接做批量矩阵乘法。
    """
    n = q_mat.shape[0]

    q_heads = q_mat.reshape(n, num_q_heads, d_k).transpose(1, 0, 2)    # (hq, n, dk)
    k_heads = k_mat.reshape(n, num_kv_heads, d_k).transpose(1, 0, 2)   # (hkv, n, dk)
    v_heads = v_mat.reshape(n, num_kv_heads, d_v).transpose(1, 0, 2)   # (hkv, n, dv)

    scale = np.sqrt(d_k)                       # 缩放因子，d_k 越大点积的量级越大
    group_size = num_q_heads // num_kv_heads   # 每组共享同一对 KV 的 Query 头数
    blocked = (mask == 1)                      # 布尔掩码，(n, n)

    head_outputs = []
    for m in range(num_kv_heads):
        # 第 m 组的 g 个 Query 头编号连续，共享 K_m、V_m
        q_group = q_heads[m * group_size:(m + 1) * group_size]   # (g, n, dk)
        k_group = k_heads[m]                                     # (n, dk)
        v_group = v_heads[m]                                     # (n, dv)

        # 批量打分 + 缩放：(g, n, dk) @ (dk, n) -> (g, n, n)，g 个头一起算
        scores = q_group @ k_group.T / scale

        # 掩码在缩放之后、softmax 之前生效，被屏蔽的位置压成 -1e30
        scores = np.where(blocked, MASKED, scores)

        # 逐行 softmax：先减该行最大值，指数输入不会超过 0，不会上溢
        scores = scores - scores.max(axis=2, keepdims=True)
        weights = np.exp(scores)
        weights /= weights.sum(axis=2, keepdims=True)

        # 用注意力权重对 V 加权求和：(g, n, n) @ (n, dv) -> (g, n, dv)
        head_outputs.append(weights @ v_group)

    # 各组头的编号本身就是连续的，直接按头维拼起来再还原成按位置的行
    out = np.concatenate(head_outputs, axis=0)          # (hq, n, dv)
    return out.transpose(1, 0, 2).reshape(n, num_q_heads * d_v)


def format_float(value):
    """保留 2 位小数，并把 -0.00 归一成 0.00。"""
    if abs(value) < 0.005:
        value = 0.0
    return f"{value:.2f}"


def solve():
    """读取输入、计算分组查询注意力并按要求输出。"""
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    pos = 0
    n, num_q_heads, num_kv_heads, d_k, d_v = (int(x) for x in tokens[pos:pos + 5])
    pos += 5

    def read_matrix(rows, cols, dtype):
        """从当前位置读取 rows x cols 的矩阵。"""
        nonlocal pos
        mat = np.array(tokens[pos:pos + rows * cols], dtype=dtype).reshape(rows, cols)
        pos += rows * cols
        return mat

    q_total = num_q_heads * d_k      # Q 每行的宽度
    k_total = num_kv_heads * d_k     # K 每行的宽度
    v_total = num_kv_heads * d_v     # V 每行的宽度

    q_mat = read_matrix(n, q_total, np.float64)
    k_mat = read_matrix(n, k_total, np.float64)
    v_mat = read_matrix(n, v_total, np.float64)
    mask = read_matrix(n, n, np.int64)

    output = gqa(q_mat, k_mat, v_mat, mask, num_q_heads, num_kv_heads, d_k, d_v)

    for row in output:
        print(" ".join(format_float(x) for x in row))


if __name__ == "__main__":
    solve()

这一版的三个关键点：

• 拆头用 reshape + transpose：一行里前 $d_k$ 个数是头 0、接着 $d_k$ 个是头 1……所以 reshape(n, heads, width) 之后第 1 维就是头编号，transpose(1, 0, 2) 把它换到最前面，得到 $(heads, n, width)$ 的三维数组。这一步是纯视图变换，不拷贝数据；
• 批量矩阵乘法：第 $m$ 组一次取 $g$ 个 Query 头，q_group @ k_group.T 直接算出 $(g, n, n)$ 的分数矩阵，weights @ v_group 又把它压回 $(g, n, d_v)$，$g$ 个头的打分与加权求和都只走一次 BLAS。掩码用 np.where(blocked, MASKED, scores) 广播到整组头上，softmax 用 max(axis=2, keepdims=True) 逐行减最大值，语义与纯 Python 版完全相同；
• 拼接：各组头的编号连续，np.concatenate(head_outputs, axis=0) 得到的就是 $(h_q, n, d_v)$，再 transpose(1, 0, 2).reshape(n, h_q\cdot d_v) 还原成每个位置一行。顺序与纯 Python 版按列拼接一致。

另外要说明一点精度：这里的点积交给 BLAS，它可能用 FMA（乘加合并成一次舍入）并分块累加，与纯 Python 里逐个 dot += qi[t] * kj[t] 的舍入顺序不完全一样，分数可能相差 $10^{-16}$ 量级。这点差别经过 softmax 后仍在 $10^{-16}$ 量级，本题输出只保留 2 位小数，实测与纯 Python 版逐字一致。

---

示例推演

样例 1：4 个 Query 头共享 2 组 KV

本题 $n=2$、$num\_q\_heads=4$、$num\_kv\_heads=2$、$d_k=1$、$d_v=1$，于是 $g=4/2=2$，$\sqrt{d_k}=1$，掩码全为 $0$，输出宽度 $4\times1=4$。

按列拆头。 $Q$ 的两行是 $(1,0,0,1)$ 与 $(0,1,1,0)$：

头	$Q_h$ 的两行
$Q_0$	$1,\ 0$
$Q_1$	$0,\ 1$
$Q_2$	$0,\ 1$
$Q_3$	$1,\ 0$

$K$ 的两行是 $(0,1)$ 与 $(1,0)$，所以 $K_0$ 是列 0 即 $0,1$，$K_1$ 是列 1 即 $1,0$。
$V$ 的两行是 $(2,4)$ 与 $(8,6)$，所以 $V_0$ 是 $2,8$，$V_1$ 是 $4,6$。

第 0 组（$Q_0,Q_1$ 共享 $K_0,V_0$）。 分数只差 $1$ 时，较大位置的权重是

$$\frac{e}{1+e}\approx0.731059$$

较小位置是 $1-0.731059=0.268941$。

• $Q_0$ 第 0 行：分数 $1\cdot0=0$ 与 $1\cdot1=1$，权重 $0.268941,\ 0.731059$，输出 $0.268941\times2+0.731059\times8=6.386354$，保留两位小数为 6.39；
• $Q_0$ 第 1 行：$Q$ 分量是 $0$，两个分数都是 $0$，权重各 $0.5$，输出 $0.5\times2+0.5\times8=5.00$；
• $Q_1$ 第 0 行：分量是 $0$，输出 $5.00$；
• $Q_1$ 第 1 行：分量是 $1$，分数 $0$ 与 $1$，输出 $6.39$。

第 1 组（$Q_2,Q_3$ 共享 $K_1,V_1$）。 $K_1$ 是 $1,0$，$V_1$ 是 $4,6$：

• $Q_2$ 第 0 行：分量 $0$，两个分数都是 $0$，输出 $0.5\times4+0.5\times6=5.00$；
• $Q_2$ 第 1 行：分量 $1$，分数 $1\cdot1=1$ 与 $1\cdot0=0$，权重 $0.731059,\ 0.268941$，输出 $0.731059\times4+0.268941\times6=4.537882$，保留两位小数为 4.54；
• $Q_3$ 第 0 行：分量 $1$，输出 $4.54$；
• $Q_3$ 第 1 行：分量 $0$，输出 $5.00$。

按头拼接。 第 0 行依次是头 $0,1,2,3$ 的第 0 行：

6.39 5.00 5.00 4.54

第 1 行同理：

5.00 6.39 4.54 5.00

样例 2：掩码如何生效

$n=2,\ num\_q\_heads=num\_kv\_heads=2,\ d_k=2,\ d_v=1$，所以 $g=1$，每个 Query 头独占一组 KV。掩码是 $\begin{bmatrix}0&1\\0&0\end{bmatrix}$，即位置 0 不能看位置 1。

$Q_0$ 的两行是 $(1,0)$、$(0,1)$；$K_0$ 的两行也是 $(1,0)$、$(0,1)$；$V_0$ 是 $2,4$。

• $Q_0$ 第 0 行：点积是 $1\cdot1+0\cdot0=1$ 与 $1\cdot0+0\cdot1=0$，除以 $\sqrt{2}\approx1.414214$ 得 $0.707107$ 与 $0$。第 1 列被屏蔽，分数变 $-\infty$，归一化后权重是 $1,0$，输出 $2.00$。
• $Q_0$ 第 1 行：点积 $0$ 与 $1$，缩放后 $0$ 与 $0.707107$，这一行没有屏蔽。权重为

$$\frac{1}{1+e^{0.707107}}\approx0.330238,\qquad \frac{e^{0.707107}}{1+e^{0.707107}}\approx0.669762$$

输出 $0.330238\times2+0.669762\times4=3.339524$，保留两位小数为 3.34。

$Q_1$ 的列与 $Q_0$ 相反（两行是 $(0,1)$、$(1,0)$），$K_1$ 的两行也是 $(0,1)$、$(1,0)$，$V_1$ 是 $6,8$。缩放后的分数与掩码情况完全一样，因此第 0 行输出 $6.00$，第 1 行输出 $0.330238\times6+0.669762\times8=7.339524$，即 7.34。

拼接后就是样例输出：

2.00 6.00
3.34 7.34

---

复杂度分析

步骤	时间复杂度	空间复杂度	说明
按列拆头	$O\left(n\left(h_qd_k+h_{kv}(d_k+d_v)\right)\right)$	同左	每个元素只搬一次
单头打分与缩放	$O(n^2d_k)$	$O(n^2)$	$S$ 为 $n\times n$
单头 softmax	$O(n^2)$	$O(n)$	每行一次
单头加权求和	$O(n^2d_v)$	$O(nd_v)$	屏蔽项可跳过
全部 $h_q$ 个头	$O\left(h_qn^2(d_k+d_v)\right)$	$O\left(n^2+n\,h_qd_v\right)$	KV 被同组 $g$ 个头共享，但打分仍按 Query 头各自进行
拼接输出	$O\left(n\,h_qd_v\right)$	$O\left(n\,h_qd_v\right)$	逐行拼列

$n\le100$、$num\_q\_heads\le16$、$d_k,d_v\le10$，最坏计算量约为 $16\times10^4\times20=3.2\times10^6$ 次乘加，非常小。空间上最大的开销是每个头的 $n\times n$ 分数矩阵，逐头计算、算完即弃，所以峰值只有一份。

---

易错点总结

1. 拆头是按列切，不是按行切

输入的一行是同一个位置的所有头拼在一起。所以第 $h$ 个头取的是每行的第 $h\cdot d_k$ 到 $(h+1)\cdot d_k-1$ 列。按行切会把不同位置的同一个头当成一个头，结果完全错乱。

2. 组号是 $h\ //\ g$

每组恰好 $g=num\_q\_heads/num\_kv\_heads$ 个 Query 头，且编号连续。所以第 $h$ 个头用第 $h//g$ 组 KV。写成 $h\ \%\ num\_kv\_heads$ 在 $g=1$ 时恰好也对，$g>1$ 时就会错位。

3. 缩放因子是 $\sqrt{d_k}$

用单个 Query/Key 头的维度 $d_k$ 开方，不是 $\sqrt{n}$，也不是 $\sqrt{num\_q\_heads}$，更不是直接除以 $d_k$。示例 2 里 $d_k=2$，所以除 $1.414214$。

4. 掩码必须在 softmax 之前

被屏蔽的位置要先变成 $-\infty$，softmax 之后它的权重才会是 $0$。如果在 softmax 之后再把权重清零，分母仍然包含这些位置的指数，剩下的权重加起来小于 $1$，输出会整体偏小。

5. softmax 是逐行归一化

对第 $i$ 个位置（也就是 $S$ 的第 $i$ 行）在整个序列方向上归一化，方向是列下标 $j$。如果误按列归一化，就变成“每个 key 对所有 query 归一化”，语义完全变了。

6. 用 $-10^{30}$ 而不是真的 $-\infty$

真正的 $-\infty$ 在“减最大值”时会得到 $-\infty-(-\infty)$，结果是 nan，整行输出都会坏掉。用 $-10^{30}$ 时减完是 $0$，指数是 $0$ 或下溢成 $0$，行为稳定。

7. 拼接顺序与拆头顺序必须互逆

输出第 $h\cdot d_v$ 到 $(h+1)\cdot d_v-1$ 列来自第 $h$ 个头。拼反了（比如按 KV 组拼接）在 $g=1$ 时看不出来，$g>1$ 时结果全错。

8. 三块矩阵的宽度各不相同

$Q$ 每行 $num\_q\_heads\cdot d_k$ 个数，$K$ 每行 $num\_kv\_heads\cdot d_k$ 个数，$V$ 每行 $num\_kv\_heads\cdot d_v$ 个数。三者从同一个 token 流里顺序读取，只要有一个宽度写错，后面的掩码就会整体错位，现象往往是“输出全是 0 或者每行都一样”。源码用 read_matrix(n, 宽度) 显式指定列数，就是为了避免这种串行。

9. 输出格式

保留 $2$ 位小数；负零要归零，避免 -0.00。另外 $d_v$ 与 $d_k$ 可能不同，输出宽度是 $num\_q\_heads\cdot d_v$，不要按输入宽度输出。

---

扩展思考

• 三个特例：$num\_kv\_heads=num\_q\_heads$ 是 MHA，$num\_kv\_heads=1$ 是 MQA，中间状态就是 GQA。工程上常取 $g=4$ 或 $8$，在效果和显存之间折中。
• KV cache 的收益：推理时每个 token 的 KV 缓存大小是 $2\cdot num\_kv\_heads\cdot d_k$（乘上层数）。GQA 把 $num\_kv\_heads$ 缩小 $g$ 倍，长上下文场景下省下的显存非常可观。
• 因果掩码：自回归模型用 $mask[i][j]=1\ (j>i)$ 的上三角掩码，保证位置 $i$ 只能看到自己和左边。本题的掩码是任意给定的，逻辑上更一般。
• 不稳定的 softmax：这里逐行存下 $n\times n$ 的分数矩阵。$n$ 很大时可以分块计算、边算边累加指数和（FlashAttention 的思路），把空间从 $O(n^2)$ 降到 $O(n)$。
• 精度：本题数值范围小，双精度足够。真实模型里 softmax 的输入可能到几十，减最大值这一步不能省。

---

相关题目

• 模域双向循环注意力网络前向计算
• 滑动窗口注意力机制
• K-Median 聚类