Attention with Linear Biases

原始题目：LeetGPU - Attention with Linear Biases

题目描述

实现论文 "Train Short, Test Long" 中提出的 ALiBi（Attention with Linear Biases）注意力机制。给定查询矩阵 $Q$（$M \times d$）、键矩阵 $K$（$N \times d$）和值矩阵 $V$（$N \times d$），计算：

$$\text{Attention}_{\text{ALiBi}}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^T}{\sqrt{d}} + \alpha \cdot \Delta\right) V$$

其中 $\alpha$ 是控制线性偏置的斜率，$\Delta = i - j$ 表示查询 $i$ 与键 $j$ 之间的相对位置。softmax 按行应用，所有数据均为 float32。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在输出矩阵 output 中。

示例

示例 1

Input:  Q (2×4): [[1,0,0,0],[0,1,0,0]]
        K (3×4): [[1,0,0,0],[0,1,0,0],[0,0,1,0]]
        V (3×4): [[1,2,3,4],[5,6,7,8],[9,10,11,12]]
        α = 0.5
Output: (2×4): [[3.05,4.05,6.05,7.05],[3.93,4.93,5.93,6.93]]

示例 2

Input:  Q (1×2): [[1,2]], K (2×2): [[1,0],[0,1]], V (2×2): [[3,4],[5,6]], α=0.8
Output: (1×2): [[3.95,4.95]]

约束条件

• $Q$ 为 $M \times d$，$K$ 和 $V$ 为 $N \times d$。
• $1 \le M, N \le 2{,}048$，$1 \le d \le 1{,}024$。
• $-1.0 \le \alpha \le 1.0$。
• 性能测试在 $M = N = 2{,}048$ 的规模下进行。

解题思路

ALiBi 用简单的线性偏置替代了位置编码，让模型能够外推到训练时未见过的序列长度。实现上在标准 softmax attention 的 $QK^T/\sqrt{d}$ 基础上加上一个下三角矩阵 $\alpha \cdot (i - j)$（对因果掩码的场景）。偏置矩阵可以预计算并存储在常量内存中。欢迎在 GitHub Discussions 分享你的解法。