DeepSeekV4中RoPE设计解析

本文作者kaiyuanDeepSeekV4中RoPE设计解析想深耕AI Infra领域欢迎访问InfraTech库内容涵盖大模型基础、PyTorch/vLLM/SGLang框架入门、性能加速等核心方向配套50知识干货及适合初学者的notebook练习:https://github.com/CalvinXKY/InfraTechDeepSeek V4采用RoPE进行位置编码但由于注意力结构升级会带来两个核心问题在CSA/HCA中存在压缩操作多个token会被压缩为一个token位置信息应在压缩前注入还是压缩后注入Attention采用MQA模式K与V共享表示。若直接对KV旋转会将位置信息引入V该如何处理下面围绕这两个问题梳理DSV4的位置编码设计。1 MLA中的RoPE处理回顾在分析V4之前先回顾V2/V3的MLAMulti-headLatentAttention方案因为MLA同样涉及MQA与KV cache压缩问题。在标准RoPE中计算queryq m q_mqm​与keyk n k_nkn​内积时可写为( R m q m ) ⊤ ( R n k n ) q m ⊤ R m ⊤ R n k n q m ⊤ R n − m k n (R_m q_m)^\top(R_n k_n)q_m^\top R_m^\top R_n k_nq_m^\top R_{n-m}k_n(Rm​qm​)⊤(Rn​kn​)qm⊤​Rm⊤​Rn​kn​qm⊤​Rn−m​kn​其中R m R_mRm​、R n R_nRn​是位置m 、 n m、nm、n对应的旋转矩阵。由于旋转矩阵正交满足R ( θ ) ⊤ R ( θ ) − 1 R ( − θ ) R(\theta)^\topR(\theta)^{-1}R(-\theta)R(θ)⊤R(θ)−1R(−θ)因此内积只依赖相对位置n − m n−mn−m。具体介绍参考《彻底搞懂RoPE计算原理》[1]在MLA中KV下采样后K与V共享同一份cache值这样可以节省显存但也带来问题如果给K注入RoPEV会被一并旋转导致V值“掺杂”了位置信息。一种直观做法是将K、V拆开仅对K旋转。但这样需要分别存储K cache和V cache开销会回到接近GQA。MLA采用了一个折中的做法在Q、K隐藏维度中设置一部分专门用于RoPE计算。图片参考https://github.com/CalvinXKY/InfraTech/tree/main/models/deepseek_v3这样既能让K携带位置信息又能避免污染V同时只需额外存储较小的RoPE相关K cache如图中的k_pe远小于完整拆分K/V cache。公式推导参考part2、3[2]2 CSA/HCA中的RoPE处理在DSV4的CSA/HCA中同样存在KV cache压缩与MQA下KV共享的问题。CSA与HCA在RoPE处理上的思路一致下面以HCA为例说明。HCA中涉及RoPE的主要位置包括窗口通道SWA的KV值C128A压缩器输出的压缩KV值上采样后的Q值Attention输出的O值。图片参考https://github.com/CalvinXKY/InfraTech/blob/main/models/deepseek_v4看到这个设计思考问题如下2.1 为什么要对输出O做一次旋转前面在MLA中提到KV共享时直接旋转KV会让V也带上位置信息。HCA中窗口通道与压缩通道都在KV的最后rope_head_dim维度上施加RoPE。对应维度上的计算可写为o n ϕ ( ( R m q m ) ⊤ ( R n k n ) ) R n v n p ⋅ R n v n o_n\phi\big((R_m q_m)^\top(R_n k_n)\big)\,R_n v_np\cdot R_n v_non​ϕ((Rm​qm​)⊤(Rn​kn​))Rn​vn​p⋅Rn​vn​这里多出的R n R_nRn​等价于给输出引入绝对位置信息。绝对位置信息并非一定不可训练但在可扩展性尤其是长上下文外推上通常不如相对位置形式稳定。因此HCA会对输出再做一次逆旋转O ′ R − i ⋅ O R − i ϕ ( ( R m q m ) ⊤ ( R n k n ) ) R n v n p ⋅ R n − i v n OR_{-i}\cdot OR_{-i}\,\phi\big((R_m q_m)^\top(R_n k_n)\big)\,R_n v_np\cdot R_{n-i}v_nO′R−i​⋅OR−i​ϕ((Rm​qm​)⊤(Rn​kn​))Rn​vn​p⋅Rn−i​vn​这样位置项从绝对位置转为相对位置。这里有个小问题是否能采用正向旋转 i ii答案是否定的因为从公式看到结果将仍偏向绝对位置表达。2.2 能否直接给P旋转不行。设V VV末两维为[seq,head_dim]而P PP维度为[seq,seq]。RoPE旋转作用在head_dim维P PP与旋转维度不匹配。 从计算上看P V PVPV可理解为“标量权重乘向量”P PP本身是标量权重集合不具备可旋转的向量维度。2.3 旋转应在压缩前还是压缩后RoPE角度与绝对位置相关θ ( m , i ) m ⋅ θ i m ⋅ 10000 − 2 i / d \theta(m,i)m\cdot\theta_im\cdot 10000^{-2i/d}θ(m,i)m⋅θi​m⋅10000−2i/d其中m mm是token位置索引d为注意力头维度即hidden_size/num_headsi的取值范围是0 , 1 , … , d / 2 − 1 0, 1, \dots, d/2 - 10,1,…,d/2−1。而C128A会将128个KV状态压缩为1个KV状态QK计算使用的是压缩后的K。核心问题是K旋转角度的系数位置m怎么选若在压缩前旋转每个token先带位置再压缩。这样看似直观但位置信息会在序列维累加混合容易破坏RoPE所需的相对位置结构。若在压缩后旋转给每个压缩K值指定一个标定位置即可。该位置可选起始、结束或中点只要映射规则全程一致。HCA采用的是“每128段取起始位置”旋转角度公式θ ( m ′ , i ) ( 128 ⋅ t ) ⋅ 10000 − 2 i / d \theta(m,i)(128\cdot t)\cdot 10000^{-2i/d}θ(m′,i)(128⋅t)⋅10000−2i/d其中t tt是当前压缩K值在压缩序列中的索引。C128A压缩计算中RoPE的位置建议的前置阅读彻底搞懂RoPE计算原理从1D到3D图解DeepSeek V4详细计算流程解析参考[1]https://zhuanlan.zhihu.com/p/2023493768003724514[2]https://spaces.ac.cn/archives/10091InfraTech申明未经允许不得转载