Rotary Positional Embedding

agicy2026/6/6大约 2 分钟

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 中以减少内存往返。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void rope_kernel(const float* Q, const float* cos, const float* sin,
    float* output, int M, int D) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int half = D / 2;
    if (i < M) {
        for (int d = 0; d < half; d++) {
            float x1 = Q[i*D + d];
            float x2 = Q[i*D + d + half];
            float c = cos[i*D + d], s = sin[i*D + d];
            output[i*D + d] = x1*c - x2*s;
            output[i*D + d + half] = x1*s + x2*c;
        }
    }
}
extern "C" void solve(const float* Q, const float* cos, const float* sin,
    float* output, int M, int D) {
    rope_kernel<<<(M+255)/256, 256>>>(Q, cos, sin, output, M, D);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def rope_kernel(Q_ptr, cos_ptr, sin_ptr, out_ptr, M: tl.constexpr, D: tl.constexpr, BLOCK: tl.constexpr):
    i = tl.program_id(0)*BLOCK + tl.arange(0, BLOCK)
    mask = i < M
    half = D // 2
    d_half = tl.arange(0, half)
    x1 = tl.load(Q_ptr + i[:,None]*D + d_half[None,:], mask=mask[:,None])
    x2 = tl.load(Q_ptr + i[:,None]*D + d_half[None,:] + half, mask=mask[:,None])
    c  = tl.load(cos_ptr + i[:,None]*D + d_half[None,:], mask=mask[:,None])
    s  = tl.load(sin_ptr + i[:,None]*D + d_half[None,:], mask=mask[:,None])
    tl.store(out_ptr + i[:,None]*D + d_half[None,:], x1*c - x2*s, mask=mask[:,None])
    tl.store(out_ptr + i[:,None]*D + d_half[None,:] + half, x1*s + x2*c, mask=mask[:,None])