题目内容

给定三个矩阵 $Q$、$K$、$V$ 以及一个可选的屏蔽矩阵 $mask$，需要按以下步骤完成计算：

1. 相似度矩阵：计算 $S = \frac{Q K^T}{\sqrt{d}}$，其中 $d$ 为 $Q$ 和 $K$ 的列数。
2. 屏蔽处理：若 $mask$ 不为空（非 null），则令 $S \leftarrow S + mask$；若 $mask$ 为 null，则视为与 $S$ 形状相同的全 $0$ 矩阵。$mask$ 中的元素取值为 $0$ 或 $-10^9$，用于屏蔽特定位置。
3. 数值稳定的行级指数归一化：对 $S$ 的每一行，先减去该行的最大值，即 $S_{ij} \leftarrow S_{ij} - \max_k S_{ik}$；然后计算 $P_{ij} = \exp(S_{ij}) / \sum_k \exp(S_{ik})$，得到权重矩阵 $P$。
4. 加权求和：计算输出矩阵 $O = P V$，其形状为 $(m, d)$。
5. 展平：将 $O$ 按行优先顺序展开成一个一维数组。

约束条件：