RoPE核心代码手撕(逐行解释)

RoPE核心代码手撕# RoPE核心代码手撕# 导入Pytorch核心张量运算库importtorch# 导入pytorch神经网络基类模块用于自定义层继承importtorch.nnasnn# 自定义旋转位置编码层继承nn.moduleclassRoPE(nn.Module):def__init__(self,head_dim,max_len512,base10000,*args,**kwargs):# 调用父类nn.Module的初始化方法必须执行super().__init__(*args,**kwargs)# 保存单个注意力头的特征维度旋转是按头独立进行的self.head_dimhead_dim# 旋转频率的基数控制频率衰减的速度默认10000是行业通用值self.basebase# 预计算逆频率向量长度为head_dim/2,每个元素对应一组维度的旋转频率# 公式inv_freq[i] 1 / base^(2i / head_dim)# 维度索引越大频率越低负责刻画长距离的全局位置关系对应全局位置信息# 维度索引越小频率越高对应捕捉精细的短距离位置差异对应精细位置inv_freq1.0/(base**(torch.arange(0,head_dim,2).float()/self.head_dim))# 注册为buffer:随模型自动迁移CPU/GPU设备不属于可训练参数# persistent False表示不写入模型权重文件推理时可以随时重建节省存储空间self.register_buffer(inv_freq,inv_freq,persistentFalse)# 初始化时预计算最大长度内的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保持一致确保跨设备兼容ttorch.arange(seq_len,deviceself.inv_freq.device).type_as(self.inv_freq)# 爱因斯坦求和实现外积每个位置索引×每个逆频率每个位置每个维度的旋转角度# 输入t形状[seq_len], inv_freq形状[head_dim // 2]# 输出freqs形状[seq_len, head_dim // 2]freqstorch.einsum(i,j - ij,t,self.inv_freq)# 前后半维度复用相同的频率拼接成完整的head_dim长度# 对应半向量拆分旋转的实现方式输出形状[seq_len, head_dim]embtorch.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,:,:],persistentFalse)self.register_buffer(sin_cached,emb.sin()[None,None,:,:],persistentFalse)def_rotate_half(self,x): 内部工具函数构造旋转的正交分量 将特征拆分为前后两半输出[-x2, x1],对应二维旋转公式的正交部分 # 取特征拆为前后两半x1x[...,:x.shape[-1]//2]x2x[...,x.shape[-1]//2:]# 拼接得到正交分量对应二维旋转中的垂直方向returntorch.cat((-x2,x1),dim-1)defforward(self,q,k,seq_lenNone):前向传播对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_lenq.shape[2]# 动态扩容如果当前序列长度超过了缓存的最大长度重新计算更长的缓存ifseq_lenself.cos_cached.shape[2]:self._set_cos_sin_cache(seq_len)# 截取当前序列长度对应的cos、sin缓存cosself.cos_cached[:,:,:seq_len,:]sinself.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_dim8ropeRoPE(head_dimhead_dim,max_len16)# 构造一个固定的特征向量v所有测试都用它保证初始向量唯一v1torch.randn(1,1,1,head_dim)# [batch1, heads1, seq_len1, dim8]v2torch.randn(1,1,1,head_dim)# [batch1, heads1, seq_len1, dim8]# 构造Q序列位置0放v1位置1放v1其余位置不影响qtorch.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其余位置不影响ktorch.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_rotrope(q,k)# 第0位Q · 第2位K → 相对距离 2score_0_2torch.sum(q_rot[0,0,0,:]*k_rot[0,0,2,:])# 第1位Q · 第3位K → 相对距离 2score_1_3torch.sum(q_rot[0,0,1,:]*k_rot[0,0,3,:])print(f相对距离2的两组注意力分数)print(fscore_0_2 {score_0_2.item():.6f})print(fscore_1_3 {score_1_3.item():.6f})print(f差值浮点误差内应为0{abs(score_0_2-score_1_3).item():.8f})print(✅ 若差值极小1e-6级别说明RoPE相对位置不变性成立代码正确)