Rotary Positional Embedding

原始题目：LeetGPU - Rotary Positional Embedding

题目描述

编写一个 GPU 程序，计算一批查询向量的旋转位置编码（RoPE）。RoPE 是 Transformer 模型中编码位置信息的一种方法，通过使用预计算的余弦和正弦分量对查询和键向量进行旋转。

数学上，给定查询向量 $x$ 及对应的余弦和正弦向量：

$$\text{RoPE}(x) = x \odot \cos + \text{rotate\_half}(x) \odot \sin$$

其中 $\odot$ 表示逐元素乘法。$\text{rotate\_half}(x)$ 操作交换向量的前后两半并取反前半部分。对于维度为 $d$ 的向量：

$$\text{rotate\_half}([x_1, \dots, x_{d/2}, x_{d/2+1}, \dots, x_d]) = [-x_{d/2+1}, \dots, -x_d, x_1, \dots, x_{d/2}]$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 输入张量 $Q$、$cos$、$sin$ 形状均为 $(M, D)$（$M$ 为 token 数，$D$ 为头维度，$D$ 保证为偶数）。
• 最终结果必须以 $(M, D)$ 的形状存储在 output 变量中。

示例

Input:  Q   = [[1,2,3,4],[1,1,1,1]]
        Cos = [[1,1,1,1],[0,0,0,0]]
        Sin = [[0,0,0,0],[1,1,1,1]]
Output: (Row 0 恒等，Row 1 旋转): [[1,2,3,4],[-1,-1,1,1]]

约束条件

• $Q$、$cos$、$sin$ 维度相同，$D \bmod 2 = 0$。
• $1 \le M, D \le 10{,}000$。
• 性能测试在 $D = 128,\ M = 1{,}048{,}576$ 的规模下进行。

解题思路

RoPE 的计算是纯逐元素操作：$\cos$ 和 $\sin$ 是预计算的，$\text{rotate\_half}$ 只需重新排列向量元素和取反。每个 token 的每对相邻维度 $(2k, 2k+1)$ 独立旋转，完全数据并行。这是典型的带宽受限内核。LLaMA、Mistral 等模型广泛使用 RoPE，高效的 CUDA 实现通常将其融合到 attention kernel 中以减少内存往返。欢迎在 GitHub Discussions 分享你的解法。