Multi-Head Attention

原始题目：LeetGPU - Multi-Head Attention

题目描述

实现多头自注意力（Multi-Head Self-Attention）。给定查询矩阵 $Q$、键矩阵 $K$ 和值矩阵 $V$（大小均为 $N \times d_{model}$），计算：

$$\text{MultiHead}(Q, K, V) = \text{Concat}(head_1, \dots, head_h)$$

其中每个头计算：

$$head_i = \text{softmax}\!\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i$$

$d_k = d_{model} / h$，$Q_i, K_i, V_i$ 是输入矩阵的第 $i$ 个头分区。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

示例

N=2, d_model=4, h=2
Q=[[1,0,2,3],[4,5,6,7]], K=[[1,2,3,4],[5,6,7,8]], V=[[0.5,1,1.5,2],[2.5,3,3.5,4]]
Output (2×4): [[2.39,2.89,3.50,4.00],[2.50,3.00,3.50,4.00]]

约束条件

• $1 \le N \le 10{,}000$，$2 \le d_{model} \le 1{,}024$，$1 \le h \le d_{model}$，$d_{model} \bmod h = 0$。
• 性能测试在 $N=1{,}024, d_{model}=1{,}024$ 下进行。

解题思路

MHA 是 Transformer 的核心算子。每个头独立计算 softmax attention，最后拼接输出。GPU 上高效实现的关键是将多个小矩阵乘法 batch 化：将 $Q, K, V$ 的重塑视为 batched GEMM。每个头的数据量较小（$d_k$ 一般 64–128），可以多个头并行处理。FlashAttention 通过分块计算 $QK^T$ 避免完整 attention 矩阵的显存分配，是生产级实现的标准方案。欢迎在 GitHub Discussions 分享你的解法。