RoPE核心代码手撕
# RoPE核心代码手撕# 导入Pytorch核心张量运算库importtorch# 导入pytorch神经网络基类模块，用于自定义层继承importtorch.nnasnn# 自定义旋转位置编码层，继承nn.moduleclassRoPE(nn.Module):def__init__(self,head_dim,max_len=512,base=10000,*args,**kwargs):# 调用父类nn.Module的初始化方法，必须执行super().__init__(*args,**kwargs)# 保存单个注意力头的特征维度，旋转是按头独立进行的self.head_dim=head_dim# 旋转频率的基数，控制频率衰减的速度，默认10000是行业通用值self.base=base# 预计算逆频率向量：长度为head_dim/2,每个元素对应一组维度的旋转频率# 公式：inv_freq[i] = 1 / base^(2i / head_dim)# 维度索引越大，频率越低，负责刻画长距离的全局位置关系，对应全局位置信息；# 维度索引越小，频率越高，对应捕捉精细的短距离位置差异，对应精细位置inv_freq=1.0/(base**(torch.arange(0,head_dim,2).float()/self.head_dim))# 注册为buffer:随模型自动迁移CPU/GPU设备，不属于可训练参数# persistent = False表示不写入模型权重文件，推理时可以随时重建，节省存储空间self.register_buffer("inv_freq",inv_freq,persistent=False)# 初始化时预计算最大长度内的cos/sin值，缓存起来避免重复计算self._set_cos_sin_cache(max_len)def_set_cos_sin_cache(self,seq_len):"""内部工具函数：预计算所有位置的cos、sin旋转角度值并缓存"""# 生成位置索引序列[0, 1, 2, ..., seq_len-1]# 设备、数据类型与inv_freq保持一致，确保跨设备兼容t=torch.arange(seq_len,device=self.inv_freq.device).type_as(self.inv_freq)# 爱因斯坦求和实现外积：每个位置索引×每个逆频率=每个位置每个维度的旋转角度# 输入：t形状[seq_len], inv_freq形状[head_dim // 2]# 输出：freqs形状[seq_len, head_dim // 2]freqs=torch.einsum("i,j -> ij",t,self.inv_freq)# 前后半维度复用相同的频率，拼接成完整的head_dim长度# 对应半向量拆分旋转的实现方式，输出形状[seq_len, head_dim]emb=torch.cat((freqs,freqs),dim=-1)# 计算所有位置的cos值，扩展2个维度适配[batch, heads, seq_len, dim]的广播机制# 最终缓存形状：[1, 1, seq_len, head_dim]self.register_buffer("cos_cached",emb.cos()[None,None,:,:],persistent=False)self.register_buffer("sin_cached",emb.sin()[None,None,:,:],persistent=False)def_rotate_half(self,x):""" 内部工具函数：构造旋转的正交分量 将特征拆分为前后两半，输出[-x2, x1],对应二维旋转公式的正交部分 """# 取特征拆为前后两半x1=x[...,:x.shape[-1]//2]x2=x[...,x.shape[-1]//2:]# 拼接得到正交分量，对应二维旋转中的垂直方向returntorch.cat((-x2,x1),dim=-1)defforward(self,q,k,seq_len=None):"""前向传播：对Q和K同时施加旋转位置编码 Args: q:查询向量，形状[batch, num_heads, seq_len, head_dim] k:键向量，形状[batch, num_heads, seq_len, head_dim] seq_len: 可选，指定序列长度，默认自动从q中读取 Returns: q_rot, k_rot:旋转后的Q、K向量，形状与输入完全一致 """# 未传入seq_len时，从q的第3维读取当前序列长度ifseq_lenisNone:seq_len=q.shape[2]# 动态扩容：如果当前序列长度超过了缓存的最大长度，重新计算更长的缓存ifseq_len>self.cos_cached.shape[2]:self._set_cos_sin_cache(seq_len)# 截取当前序列长度对应的cos、sin缓存cos=self.cos_cached[:,:,:seq_len,:]sin=self.sin_cached[:,:,:seq_len,:]# 核心旋转公式：x' = x * cos + rotate_half(x) * sin# 数学上等价于二维旋转变换，向量化并行计算所有位置、所有头q_rot=(q*cos)+(self._rotate_half(q)*sin)k_rot=(k*cos)+(self._rotate_half(k)*sin)returnq_rot,k_rot# 验证用例# 性质：两个向量的注意力点积结果，仅和它们的相对位置差有关，和绝对位置无关if__name__=="__main__":head_dim=8rope=RoPE(head_dim=head_dim,max_len=16)# 构造一个固定的特征向量v，所有测试都用它，保证初始向量唯一v1=torch.randn(1,1,1,head_dim)# [batch=1, heads=1, seq_len=1, dim=8]v2=torch.randn(1,1,1,head_dim)# [batch=1, heads=1, seq_len=1, dim=8]# 构造Q序列：位置0放v1，位置1放v1，其余位置不影响q=torch.zeros(1,1,4,head_dim)q[0,0,0,:]=v1 q[0,0,1,:]=v1# print("v1_shape:", v1.shape)# 构造K序列：位置2放v2，位置3放v2，其余位置不影响k=torch.zeros(1,1,4,head_dim)k[0,0,2,:]=v2 k[0,0,3,:]=v2# print("v2_shape:", v2.shape)# 执行RoPE旋转q_rot,k_rot=rope(q,k)# 第0位Q · 第2位K → 相对距离 = 2score_0_2=torch.sum(q_rot[0,0,0,:]*k_rot[0,0,2,:])# 第1位Q · 第3位K → 相对距离 = 2score_1_3=torch.sum(q_rot[0,0,1,:]*k_rot[0,0,3,:])print(f"相对距离=2的两组注意力分数：")print(f"score_0_2 ={score_0_2.item():.6f}")print(f"score_1_3 ={score_1_3.item():.6f}")print(f"差值（浮点误差内应为0）：{abs(score_0_2-score_1_3).item():.8f}")print("✅ 若差值极小（1e-6级别），说明RoPE相对位置不变性成立，代码正确")
