post/hands-on-rope-position-embedding

[giscus[bot]]

post/hands-on-rope-position-embedding

深入讲解 RoPE 旋转位置编码的核心原理与 PyTorch 实现。从 2D 旋转矩阵推导相对位置编码，逐行手写代码实现 LLaMA Qwen 风格 RoPE，附热力图可视化帮助理解。适合想彻底搞懂 RoPE 位置编码的开发者。

https://yuanchaofa.com/post/hands-on-rope-position-embedding

[Sudun]

博主，rope的rotate half好像有问题，比如[x1, x2, x3, x4] -> [-x3, -x4, x1, x2]，q_embed = (q * cos) + (rotate_half(q) * sin)的结果不是[x1cos - x3sin, x2cos - x4sin, ......]吗？

[JiYueBo666]

我也有这个疑惑，laama的实现和原论文实现似乎不一样

[wendfor]

我看源码是通过复数形式实现旋转的

[HailToThee]

是为了利用GPU批处理的性质来进行加速吧，决定因素还是m,n的相对位置。而token内部dim的两两配对则似乎并不那么重要，相比于速度来说，由于cos, sin stack后是[theta1, theta2, theta3, theta1, theta2, theta3], 如果要适配则需要做[d1, d2, d3, d4, d5, d6] -> [d1, d3, d5, d2, d4, d6] [-d2, -d4, -d6, d1, d3, d5] 这样的交换很浪费时间。

[fhzzz]

我也认为是错的，我用一些能力可能没那么强的AI给出的代码和文章里作者给出的代码是一样的，但是一些能力可能比较强的AI给出的代码使用torch.stack，我自己推导验证了这两种写法，我还是认为作者给出的代码是错误的，在后续进行q * freqs_cos + rotate_half(q) * freqs_sin计算时是错误的

[bbruceyuan]

刚收到邮件，看到这么多人有疑惑，我统一回复一下. @Sudun @JiYueBo666 @wendfor @HailToThee @fhzzz

1. 关于 rotate_half 的两种写法

首先，@HailToThee 对效率方面的解释是有道理的。

RoPE 存在两种常见的维度排列方式。

一种是论文中更直观的相邻维度配对：

$$ (x_0,x_1),\quad (x_2,x_3),\quad \ldots $$

对应的实现通常是：

def rotate_half(x):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack([-x2, x1], dim=-1).flatten(-2)

这种情况下，cos/sin 也需要按照相邻维度重复排列：

$$ [c_0,c_0,c_1,c_1,\ldots] $$

另一种是 LLaMA/Hugging Face 使用的排列方式。它把前半部分和后半部分进行配对：

$$ (x_0,x_{d/2}),\quad (x_1,x_{d/2+1}),\quad \ldots $$

对应的实现是：

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

同时，cos/sin 按照下面的方式排列：

$$ [c_0,c_1,\ldots,c_{d/2-1}, c_0,c_1,\ldots,c_{d/2-1}] $$

也就是：

cos = torch.cat([cos, cos], dim=-1)
sin = torch.cat([sin, sin], dim=-1)

从计算效率角度看，LLaMA 这种写法避免了显式的交错维度重排，因此实现更加适合批量计算。这也是 LLaMA 采用这种写法的一个原因。

2. 关于复数实现

RoPE 也可以使用复数形式实现，本质上是把每一对二维向量看作一个复数，然后乘以对应位置的旋转因子：

$$(x_{2i}+ix_{2i+1})e^{im\theta_i}$$

不过由于我对 PyTorch 的复数 API 不够熟悉，所以文章中采用了实数形式实现。实数形式和复数形式在数学上是等价的。

3. 关于官方实现

官方代码不代表 RoPE 只有这一种写法。

我这里采用的是 LLaMA/Hugging Face 中的实现：

Hugging Face LLaMA RoPE 实现

其中核心代码是：

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

以及：

q_embed = (q * cos) + (rotate_half(q) * sin)

Hugging Face 的源码也特别注明了：

# Different from paper, but it uses a different permutation
# in order to obtain the same calculation

也就是说，LLaMA 的实现与论文中常见的维度排列不同，但通过不同的维度排列获得了等价的计算结果。

文章前面为了讲清楚 RoPE 的数学原理，采用了苏神文章中较直观的相邻维度配对方式；后面写代码时，则采用了我比较熟悉的 LLaMA 实, 我最早就是看的这个代码。这里没有提前说明两种排列方式的区别，确实容易让读者产生疑惑，我会更新文章中的相关表述。

预计本周我重新 review 一下本篇文章

[fhzzz]

哇好的，竟然如此之快收到回复，看到关于效率的回复之后我确实还没有细究这个问题，只是从代码忠于方法原理实现的角度感到疑惑。非常感谢作者的回复，这让我对RoPE的原理和工程实现之间的异同印象更加深刻，祝作者越来越好~

…

________________________________ 发件人: Chaofa Yuan｜袁朝发 ***@***.***> 发送时间: 2026年8月26日 16:34 收件人: bbruceyuan/bbruceyuan.github.io ***@***.***> 抄送: 一粒尘埃 ***@***.***>; Mention ***@***.***> 主题: Re: [bbruceyuan/bbruceyuan.github.io] post/hands-on-rope-position-embedding (Discussion #56) 刚收到邮件，看到这么多人有疑惑，我统一回复一下. @Sudun<https://github.com/Sudun> @JiYueBo666<https://github.com/JiYueBo666> @wendfor<https://github.com/wendfor> @HailToThee<https://github.com/HailToThee> @fhzzz<https://github.com/fhzzz> 1. 关于 rotate_half 的两种写法 ***@***.***<https://github.com/HailToThee> 对效率方面的解释是有道理的。 RoPE 存在两种常见的维度排列方式。 一种是论文中更直观的相邻维度配对： $$ (x_0,x_1),\quad (x_2,x_3),\quad \ldots $$ 对应的实现通常是： def rotate_half(x): x1 = x[..., ::2] x2 = x[..., 1::2] return torch.stack([-x2, x1], dim=-1).flatten(-2) 这种情况下，cos/sin 也需要按照相邻维度重复排列： $$ [c_0,c_0,c_1,c_1,\ldots] $$ 另一种是 LLaMA/Hugging Face 使用的排列方式。它把前半部分和后半部分进行配对： $$ (x_0,x_{d/2}),\quad (x_1,x_{d/2+1}),\quad \ldots $$ 对应的实现是： def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) 同时，cos/sin 按照下面的方式排列： $$ [c_0,c_1,\ldots,c_{d/2-1}, c_0,c_1,\ldots,c_{d/2-1}] $$ 也就是： cos = torch.cat([cos, cos], dim=-1) sin = torch.cat([sin, sin], dim=-1) 从计算效率角度看，LLaMA 这种写法避免了显式的交错维度重排，因此实现更加适合批量计算。这也是 LLaMA 采用这种写法的一个原因。 2. 关于复数实现 RoPE 也可以使用复数形式实现，本质上是把每一对二维向量看作一个复数，然后乘以对应位置的旋转因子： $$(x_{2i}+ix_{2i+1})e^{im\theta_i}$$ 不过由于我对 PyTorch 的复数 API 不够熟悉，所以文章中采用了实数形式实现。实数形式和复数形式在数学上是等价的。 3. 关于官方实现 官方代码不代表 RoPE 只有这一种写法。 我这里采用的是 LLaMA/Hugging Face 中的实现： Hugging Face LLaMA RoPE 实现<https://github.com/huggingface/transformers/blob/e42587f596181396e1c4b63660abf0c736b10dae/src/transformers/models/llama/modeling_llama.py#L173-L188> 其中核心代码是： def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) 以及： q_embed = (q * cos) + (rotate_half(q) * sin) Hugging Face 的源码也特别注明了： # Different from paper, but it uses a different permutation # in order to obtain the same calculation 也就是说，LLaMA 的实现与论文中常见的维度排列不同，但通过不同的维度排列获得了等价的计算结果。 文章前面为了讲清楚 RoPE 的数学原理，采用了苏神文章中较直观的相邻维度配对方式；后面写代码时，则采用了我比较熟悉的 LLaMA 实, 我最早就是看的这个代码。这里没有提前说明两种排列方式的区别，确实容易让读者产生疑惑，我会更新文章中的相关表述。 预计本周我重新 review 一下本篇文章 ― Reply to this email directly, view it on GitHub<#56?email_source=notifications&email_token=A7IX33ZQQLAJDWU4UAZTSMD5L2OJHA5CNFSNUABIM5UWIORPF5TWS5BNNB2WEL2ENFZWG5LTONUW63SDN5WW2ZLOOQXTCOBRGU4TKNRYUZZGKYLTN5XKO3LFNZ2GS33OUVSXMZLOOSWGM33PORSXEX3DNRUWG2Y#discussioncomment-18159568>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/A7IX3363DJLQSIAOX55XN7T5L2OJHAVCNFSNUABHKJSXA33TNF2G64TZHMZDINJXGY3TINZVHNCGS43DOVZXG2LPNY5TSNBVGIZTEOFBOYBA>. Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS<https://github.com/notifications/mobile/ios/A7IX336NU3AAPG47ULN67P35L2OJHA5CNFSNUABIM5UWIORPF5TWS5BNNB2WEL2ENFZWG5LTONUW63SDN5WW2ZLOOQXTCOBRGU4TKNRYUZZGKYLTN5XKO3LFNZ2GS33OUVSXMZLOOSVGM33PORSXEX3JN5ZQ> and Android<https://github.com/notifications/mobile/android/A7IX336D53LMEMNPX7VZFU35L2OJHA5CNFSNUABIM5UWIORPF5TWS5BNNB2WEL2ENFZWG5LTONUW63SDN5WW2ZLOOQXTCOBRGU4TKNRYUZZGKYLTN5XKO3LFNZ2GS33OUVSXMZLOOSXGM33PORSXEX3BNZSHE33JMQ>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

[bbruceyuan]

昨天 (0830) 已经更新, 感兴趣可以看一下~

备注: 关于效率，我之前的表述过强。half-split 在某些 eager PyTorch 实现中可能少一个高层张量操作(这个也是几年前我看各种blog和博主这么说的, 所以本次没有考究就这么写了), 但并不一定会效率更高。实际上，Meta 官方 LLaMA 使用相邻维度的复数实现；Hugging Face 的 Llama 实现则采用 half-split，并在权重转换阶段同步置换 Q/K 权重。两种布局经固定置换后数学等价，实际效率要进行测试, 可能没有什么特别大的差距。