旋转矩阵计算
rotary_emb 对应 L l a m a R o t a r y E m b e d d i n g LlamaRotaryEmbedding LlamaRotaryEmbedding层，其中内置 i n i t init init 初始化方法和 f o r w a r d forward forward 前向调用，负责生成旋转矩阵中的 c o s cos cos 和 s i n sin sin。
代码
class LlamaRotaryEmbedding(torch.nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim)) self.register_buffer("inv_freq", inv_freq) # Build here to make `torch.jit.trace` work. self.max_seq_len_cached = max_position_embeddings t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case. if seq_len > self.max_seq_len_cached: self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1).to(x.device) self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) return ( self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), ) __init__函数关键代码
根据公式计算 θ \theta θ
源码：inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim)) 公式：θ i = 1000 0 − 2 i d \theta_i = 10000^{\frac {-2i}d} θi=10000d−2i example
假设base=10000, dim=8, device=“cpu”
dim, base, device=8, 10000, 'cpu' inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim)) inv_freq inv_freq为size=torch.Size([dim//2])的tensor
tensor([1.0000, 0.1000, 0.0100, 0.0010]) 生成所有位置对应的ID
源码：t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype) example:假设max_position_embeddings=10，则
max_position_embeddings =10 t = torch.arange(max_position_embeddings, dtype=inv_freq.dtype) t 输出：
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) 计算 m θ m\theta mθ
源码：freqs = torch.einsum("i,j->ij", t, self.inv_freq) example: freqs = torch.einsum("i,j->ij", t, inv_freq) freqs 输出：
tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00], [1.0000e+00, 1.0000e-01, 1.0000e-02, 1.0000e-03], [2.0000e+00, 2.0000e-01, 2.0000e-02, 2.0000e-03], [3.0000e+00, 3.0000e-01, 3.0000e-02, 3.0000e-03], [4.0000e+00, 4.0000e-01, 4.0000e-02, 4.0000e-03], [5.0000e+00, 5.0000e-01, 5.0000e-02, 5.0000e-03], [6.0000e+00, 6.0000e-01, 6.0000e-02, 6.0000e-03], [7.0000e+00, 7.0000e-01, 7.0000e-02, 7.0000e-03], [8.0000e+00, 8.0000e-01, 8.0000e-02, 8.0000e-03], [9.0000e+00, 9.0000e-01, 9.0000e-02, 9.0000e-03]]) 将 m θ m\theta mθ拼接两次
源码：emb = torch.cat((freqs, freqs), dim=-1).to(x.device) example: emb = torch.cat((freqs, freqs), dim=-1) emb tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00], [1.0000e+00, 1.0000e-01, 1.0000e-02, 1.0000e-03, 1.0000e+00, 1.0000e-01, 1.0000e-02, 1.0000e-03], [2.0000e+00, 2.0000e-01, 2.0000e-02, 2.0000e-03, 2.0000e+00, 2.0000e-01, 2.0000e-02, 2.0000e-03], [3.0000e+00, 3.0000e-01, 3.0000e-02, 3.0000e-03, 3.0000e+00, 3.0000e-01, 3.0000e-02, 3.0000e-03], [4.0000e+00, 4.0000e-01, 4.0000e-02, 4.0000e-03, 4.0000e+00, 4.0000e-01, 4.0000e-02, 4.0000e-03], [5.0000e+00, 5.0000e-01, 5.0000e-02, 5.0000e-03, 5.0000e+00, 5.0000e-01, 5.0000e-02, 5.0000e-03], [6.0000e+00, 6.0000e-01, 6.0000e-02, 6.0000e-03, 6.0000e+00, 6.0000e-01, 6.0000e-02, 6.0000e-03], [7.0000e+00, 7.0000e-01, 7.0000e-02, 7.0000e-03, 7.0000e+00, 7.0000e-01, 7.0000e-02, 7.0000e-03], [8.0000e+00, 8.0000e-01, 8.0000e-02, 8.0000e-03, 8.0000e+00, 8.0000e-01, 8.0000e-02, 8.0000e-03], [9.0000e+00, 9.0000e-01, 9.0000e-02, 9.0000e-03, 9.0000e+00, 9.0000e-01, 9.0000e-02, 9.0000e-03]]) 计算sin、cos
源码：self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) rotate_half
def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) 将原始向量从中间分为x1、x2 两部分，然后拼接为 [-x2, x1] ：
[q1,q2,q3,q4,q5,q6,q7,q8,q9,q10] -> [-q6,-q7,-q8,-q9,-q10,q1,q2,q3,q4,q5] apply_rotary_pos_emb
def apply_rotary_pos_emb(q, k, cos, sin, position_ids): gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1] gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3]) cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed 生成index tensor
gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1] gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3]) lookup index获取对应位置的 cos m θ \cos{mθ} cosmθ 和 sin m θ \sin{mθ} sinmθ 值
cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) 分别计算前面的 c o s cos cos 部分，再计算后面的 sin θ ∗ \sinθ * sinθ∗ rotate_half 部分
q_embed = (q * cos) + (rotate_half(q) * sin) 注意
q、cos、sin对应下标顺序与
q = [-q6,-q7,-q8,-q9,-q10,q1,q2,q3,q4,q5] cos = [cosθ1,cosθ2,cosθ3,cosθ4,cosθ5,cosθ1,cosθ2,cosθ3, cosθ4, cosθ5] sin = [sinθ1,sinθ2,sinθ3,sinθ4,sinθ5,sinθ1,sinθ2,sinθ3,sinθ4, sinθ5] 参考
LLM - 旋转位置编码 RoPE 代码详解
RoPE旋转位置编码深度解析：理论推导、代码实现、长度外推
图解RoPE旋转位置编码及其特性
Rotary Positional Embeddings (RoPE)
