Understanding RoPE from an engineer's perspective — from the math to the CUDA kernel.
旋转编码是苏剑林老师
因为只靠 Attention 模块无法捕捉输入的位置信息。换句话说，将输入随意做 permute，再对 Attention 的输出做 unpermute，获得的结果都是一样的。这是一件比较糟糕的事情，为此，我们需要引入位置编码。RoPE 的核心思想是”通过绝对位置编码的方式实现相对位置编码”。这句话可能还是有一点点抽象，让笔者详细解释一下：相对位置编码的目标是让第 i 个词和第 j 个词之间距离只和 (i - j) 相关，传统相对位置编码的实现方式主要是想办法对 Attention 计算中间过程里 QK 这个矩阵做一点操作，从而保证 QK 矩阵里的每一行的第 i 个元素和第 j 个元素之间有一个和 (i - j) 相关的差；绝对位置编码的意思就是不去修改 Attention 里的矩阵，只修改 Attention 里的输入。而 RoPE 就只修改了输入，同时还能保留相对位置信息。
既然把”只修改输入前提下保留相对位置”这句话翻译成数学的语言就是下面这行公式。
求解公式 $\eqref{eq:rope-core}$，得到下面这个函数是上面这个方程的一个解，OK，我们就得到了 RoPE 了。（点击上面这个公式右侧的编号可以跳转回来。）
上面这个函数又等价于下面这个函数，下面这个函数里的矩阵都是 Dense 的，因此算起来更高效，所以在后文计算的时候都以下面这个矩阵为准。其中 $ \theta_i=10000^{-2 i / d} $，这点和最经典的 Sinusoidal 保持一致。
\[\left(\begin{array}{c} q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \cos m \theta_0 \\ \cos m \theta_0 \\ \cos m \theta_1 \\ \cos m \theta_1 \\ \vdots \\ \cos m \theta_{d / 2-1} \\ \cos m \theta_{d / 2-1} \end{array}\right)+\left(\begin{array}{c} -q_1 \\ q_0 \\ -q_3 \\ q_2 \\ \vdots \\ -q_{d-1} \\ q_{d-2} \end{array}\right) \otimes\left(\begin{array}{c} \sin m \theta_0 \\ \sin m \theta_0 \\ \sin m \theta_1 \\ \sin m \theta_1 \\ \vdots \\ \sin m \theta_{d / 2-1} \\ \sin m \theta_{d / 2-1} \end{array}\right)\]另外，在实现上为了更高效，Llama 其实是这么实现的，下面也以这种实现为基准。
\[\left(\begin{array}{c} q_0 \\ q_1 \\ \vdots \\ q_{d/2-1} \\ q_{d/2} \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \cos m \theta_0 \\ \cos m \theta_1 \\ \vdots \\ \cos m \theta_{d / 2-1} \\ \cos m \theta_{0} \\ \vdots \\ \cos m \theta_{d / 2-2} \\ \cos m \theta_{d / 2-1} \end{array}\right)+\left(\begin{array}{c} -q_0 \\ -q_1 \\ \vdots \\ -q_{d/2-1} \\ q_{d/2} \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \sin m \theta_0 \\ \sin m \theta_1 \\ \vdots \\ \sin m \theta_{d / 2-1} \\ \sin m \theta_{0} \\ \vdots \\ \sin m \theta_{d / 2-2} \\ \sin m \theta_{d / 2-1} \end{array}\right)\]That’s It! 下面就让我们来实现一版吧！
让我们以 2025 年 1 月初 Megatron-LM 里的一个实现作为 baseline，链接。下面简单写一个伪代码。
torch.arange(seq_len) 生成。torch.arange 的结果完全可以复用，因此可以 Cache。inv_freq，其结果也可以复用。_rotate_half 生成最后那个公式里右边的那部分。def _rotate_half(x):
x1, x2 = torch.chunk(x, 2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2) / dim))
freqs = torch.outer(torch.arange(seq_len), inv_freq)
cos_ = torch.cos(freqs)
sin_ = torch.sin(freqs)
t = (t * cos_) + (_rotate_half(t) * sin_)
Apex、TransformerEngine、FlashAttention 都提供了 API，具体依然可以参考 Megatron 里的代码。本文的目的既不是想探究 API 怎么用（这部分看 Megatron 代码就好了），也懒得自己写一个 CUDA 代码了，因此直接跑一个实验，看看这几个 API 实现的效果如何，从而大概了解一下这种类型的函数应该怎么实现。
首先， 是不同 API 在 A100 上的实验结果。可以看到 FlashAttention 的实现其实还差一点，Apex 和 TransformerEngine 的实现是最优的。恰好笔者对 TransformerEngine 这个库比较熟，就直接看看这里是怎么实现的吧。
| 实现 | 耗时 (ms) |
|---|---|
| RotaryEmbeddingMegatron | 73.51 |
| RotaryEmbeddingApex | 14.40 |
| RotaryEmbeddingTE | 14.08 |
| RotaryEmbeddingFlash | 21.12 |
具体代码是 FusedRoPEFunc 这个函数，一上来就是巨大的槽点。也就是说，不管输入是什么类型，TransformerEngine 都会先把它转成 Float32。那么很自然的大家就会想到，如果输入是 BFloat16，正确性还能保证吗？笔者做了实验，发现把 Input 的 Dtype 从 Float 改成 BFloat16 之后，FlashAttention 和 TransformerEngine 都和 baseline 对不齐。emmm 只能说世界就是巨大的草台班子，大家调用这种 API 之前最好还是本地实测一下。至于 Megatron 很机智的用了 Apex 的实现做了默认的实现，不知道是不是出于这个正确性的考虑。
if freqs.dtype != torch.float32:
freqs = freqs.float()
让我们再看看具体 CUDA 实现，毕竟这个实现性能还是很好的。让我们无视掉每个变量具体的含义，快速地看一些 High level Idea，思想很简单，就是每个线程去做 Load/Store。这个函数是典型的 Memory Bound 的函数，因此 sincosf 这个计算绝对是 Free 的。这也解释了，为什么 FlashAttention 的实现是提前算好 sin/cos，而 TransformerEngine 是每次计算 sin/cos，反而是后者的性能更好，答案是因为计算根本就是 free 的。我看了 FlashAttention 的代码，他是用 Triton 实现的，那确实性能比 CUDA 差是符合预期的。
#pragma unroll
for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {
float v_cos, v_sin;
sincosf(freqs[s_id * d2 + d_id], &v_sin, &v_cos);
#pragma unroll
for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
int offset_src = offset_block + h_id * stride_h + d_id * stride_d;
int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
float v_src = src[offset_src];
float v_src_rotate = (d_id + d2 / 2 < d2)
? -static_cast<float>(src[offset_src + (d2 / 2) * stride_d])
: static_cast<float>(src[offset_src + (d2 / 2 - d2) * stride_d]);
dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;
}
}
本文就到此结束了，代码都在笔者的 这个 GitHub 仓库 里。如果想深究、追求一个极致，还可以想办法通过 float4 这种方式来加速 Load/Store。或许日后某天会补上吧。
RoPE was proposed by Jianlin Su
Attention alone cannot capture positional information. In other words, if you arbitrarily permute the inputs and then un-permute the attention outputs, you get the exact same result. That is undesirable, so we need to introduce position encodings. The core idea of RoPE is “to realize relative position encoding through absolute position encoding.” That may still be a little abstract, so let me explain: the goal of relative position encoding is to make the distance between token $i$ and token $j$ depend only on $(i - j)$. Traditional approaches mostly operate on the $QK$ matrix in the middle of attention, ensuring that within each row of $QK$ the difference between elements $i$ and $j$ depends on $(i - j)$. Absolute position encoding, by contrast, leaves the attention matrices untouched and only modifies the attention inputs. RoPE modifies only the inputs, while still preserving relative-position information.
Translating “preserve relative position while only modifying the inputs” into math gives the equation below.
Solving this equation, the function below is one solution — and that gives us RoPE.
The function above is equivalent to the one below, whose matrices are all dense and therefore more efficient to compute — so all later computation uses this form. Here $ \theta_i=10000^{-2 i / d} $, consistent with the classic Sinusoidal encoding.
\[\left(\begin{array}{c} q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \cos m \theta_0 \\ \cos m \theta_0 \\ \cos m \theta_1 \\ \cos m \theta_1 \\ \vdots \\ \cos m \theta_{d / 2-1} \\ \cos m \theta_{d / 2-1} \end{array}\right)+\left(\begin{array}{c} -q_1 \\ q_0 \\ -q_3 \\ q_2 \\ \vdots \\ -q_{d-1} \\ q_{d-2} \end{array}\right) \otimes\left(\begin{array}{c} \sin m \theta_0 \\ \sin m \theta_0 \\ \sin m \theta_1 \\ \sin m \theta_1 \\ \vdots \\ \sin m \theta_{d / 2-1} \\ \sin m \theta_{d / 2-1} \end{array}\right)\]For further efficiency, Llama actually implements it as follows, and this is the form used as our baseline below.
\[\left(\begin{array}{c} q_0 \\ q_1 \\ \vdots \\ q_{d/2-1} \\ q_{d/2} \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \cos m \theta_0 \\ \cos m \theta_1 \\ \vdots \\ \cos m \theta_{d / 2-1} \\ \cos m \theta_{0} \\ \vdots \\ \cos m \theta_{d / 2-2} \\ \cos m \theta_{d / 2-1} \end{array}\right)+\left(\begin{array}{c} -q_0 \\ -q_1 \\ \vdots \\ -q_{d/2-1} \\ q_{d/2} \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{array}\right) \otimes\left(\begin{array}{c} \sin m \theta_0 \\ \sin m \theta_1 \\ \vdots \\ \sin m \theta_{d / 2-1} \\ \sin m \theta_{0} \\ \vdots \\ \sin m \theta_{d / 2-2} \\ \sin m \theta_{d / 2-1} \end{array}\right)\]That’s it! Let’s implement one.
Let’s take an implementation from Megatron-LM around early January 2025 as our baseline (link). Here is some simplified pseudocode.
torch.arange(seq_len). Its result can be fully reused, so it can be cached.inv_freq, whose result can also be reused.outer takes two 1-D vectors and outputs a 2-D matrix — the Cartesian product of the inputs._rotate_half produces the right-hand part of the final formula.def _rotate_half(x):
x1, x2 = torch.chunk(x, 2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2) / dim))
freqs = torch.outer(torch.arange(seq_len), inv_freq)
cos_ = torch.cos(freqs)
sin_ = torch.sin(freqs)
t = (t * cos_) + (_rotate_half(t) * sin_)
Apex, TransformerEngine, and FlashAttention all provide APIs — again, see the Megatron code for details. The point of this post is neither to explore how to call the APIs (just read Megatron for that) nor to write my own CUDA kernel; instead I’ll just run an experiment to see how these implementations perform, and get a rough sense of how this kind of function should be implemented.
First, here are results for the different APIs on an A100. You can see FlashAttention’s implementation lags a bit, while Apex and TransformerEngine are the best. Since I happen to be familiar with TransformerEngine, let’s look at how it’s implemented there.
RotaryEmbeddingMegatron take 73.508257 ms while size = 20971520
RotaryEmbeddingApex take 14.397789 ms while size = 20971520
RotaryEmbeddingTE take 14.080192 ms while size = 20971520
RotaryEmbeddingFlash take 21.117988 ms while size = 20971520
The relevant code is FusedRoPEFunc (this function), and the very first thing it does is a big red flag: no matter the input type, TransformerEngine first casts it to Float32. Naturally you’d wonder — if the input is BFloat16, is correctness still guaranteed? I ran an experiment and found that after changing the input dtype from Float to BFloat16, both FlashAttention and TransformerEngine failed to match the baseline. Well — the world is one big improvised stage; you’re better off testing these APIs locally before trusting them. As for Megatron, it cleverly uses the Apex implementation as the default — perhaps out of exactly this correctness concern.
if freqs.dtype != torch.float32:
freqs = freqs.float()
Let’s look at the actual CUDA implementation, since its performance is quite good. Ignoring the exact meaning of each variable, let’s quickly grasp the high-level idea: it’s simple — each thread just does a Load/Store. This is a textbook memory-bound function, so the sincosf computation is essentially free. That also explains why FlashAttention precomputes sin/cos while TransformerEngine recomputes them every time — yet the latter is faster, because the computation is basically free. FlashAttention’s code is written in Triton, so it being slower than CUDA is expected.
#pragma unroll
for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) {
float v_cos, v_sin;
sincosf(freqs[s_id * d2 + d_id], &v_sin, &v_cos);
#pragma unroll
for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) {
int offset_src = offset_block + h_id * stride_h + d_id * stride_d;
int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d;
float v_src = src[offset_src];
float v_src_rotate = (d_id + d2 / 2 < d2)
? -static_cast<float>(src[offset_src + (d2 / 2) * stride_d])
: static_cast<float>(src[offset_src + (d2 / 2 - d2) * stride_d]);
dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin;
}
}
That’s the end of this post. All the code is in this GitHub repo. If you want to go deeper and squeeze out every last bit, you could try speeding up the Load/Store with something like float4. Maybe I’ll add that someday.