GPT-2 Transformer Block

原始题目：LeetGPU - GPT-2 Transformer Block

题目描述

实现单个 GPT-2 Transformer 解码器块。给定输入张量 $x$（形状 $(seq\_len, 768)$）和包含所有块参数的打包权重缓冲区，使用 pre-norm 架构计算输出：

$$\begin{aligned} x' &= x + \text{MHA}(\text{LN}_1(x)) \\ output &= x' + \text{FFN}(\text{LN}_2(x')) \end{aligned}$$

架构参数：$d_{model}=768$，12 头，$d_k=64$，FFN 维度 3072，GELU 使用 tanh 近似：

$$\text{GELU}(x) = 0.5x\left(1 + \tanh\!\left(\sqrt{\frac{2}{\pi}}(x + 0.044715 x^3)\right)\right)$$

LayerNorm 使用 $\epsilon = 10^{-5}$。注意力无因果掩码（双向）。

实现要求

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

示例

seq_len=4, x 均匀采样 [-1,1], 权重随机初始化
weights (打包): 7,087,872 个 float32
output: shape (4, 768)

约束条件

• $1 \le seq\_len \le 4{,}096$。
• 性能测试在 $seq\_len = 1{,}024$ 下进行。

解题思路

GPT-2 块包含两个大矩阵乘法（QKV 投影 $768 \to 2304$ 和 FFN 第一层 $768 \to 3072$）以及两个相对较小的投影。QKV 可以合并为一次矩阵乘法后 split。注意力层和 FFN 层之间通过残差连接，需要额外的内存分配或原地操作优化。GELU 激活可融合到前一个 matmul kernel 中。欢迎在 GitHub Discussions 分享你的解法。