题目内容

给定三个矩阵 $Q$、$K$、$V$ 和一个可选的忽略矩阵 $\textit{mask}$，现需按以下步骤计算加权结果并展平输出：

1. 计算相似度矩阵 $S = \frac{Q \, K^{\mathsf{T}}}{\sqrt{d}}$，其中 $d$ 为 $Q$ 的列数，也是 $K$ 的列数。
2. 若 $\textit{mask}$ 不为 $\text{null}$，则令 $S \leftarrow S + \textit{mask}$；否则 $S$ 保持不变（等价于加上全 $0$ 矩阵）。矩阵 $\textit{mask}$ 的元素为 $0$ 或 $-10^9$，使得下一步中被标记的位置权重趋于 $0$。
3. 对 $S$ 的每一行进行数值稳定的指数归一化：$$S \leftarrow S - \max_{\text{row}} S,$$ $$P_{ij} = \frac{e^{S_{ij}}}{\sum_{k} e^{S_{ik}}}.$$
4. 计算加权聚合 $O = P \, V$，结果 $O$ 的形状为 $(m, d)$。
5. 将 $O$ 按行优先顺序展平为一维数组。