Attention with Linear Biases (ALiBi)

题面

实现 ALiBi：$\text{Attention\_ALiBi}(Q,K,V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d}} + \alpha \cdot \Delta \right) V$，其中 $\Delta = i - j$ 为相对位置偏置，$\alpha$ 为斜率。按行应用 softmax，$Q, K, V, \text{output}, \alpha$ 为 float32。

Implementation Requirements

• Use only native features (external libraries are not permitted)
• The solve function signature must remain unchanged
• 输出写入 output

Examples

见页面两组示例（含 α 参数）。

Constraints

• $Q \in \mathbb{R}^{M \times d}$; $K, V \in \mathbb{R}^{N \times d}$
• $1 \le M, N \le 2048$; $1 \le d \le 1024$; $-1.0 \le \alpha \le 1.0$
• Performance: $M = 2048, N = 2048$