4、预计算与缓存策略：如何避免推理时的重复计算。

各位好，我是老张。今天咱们聊聊 RoPE 在端侧部署时，一个非常实在的问题——重复计算。

你想想看，端侧设备就那么点算力，电池就那么点电。如果每次推理都去重新算一遍旋转位置编码，那简直是浪费。我刚开始做端侧优化时，就踩过这个坑。当时模型跑在手机上，发热严重，一查才发现 RoPE 计算占了将近 15% 的耗时。嗯，这必须得优化。

4.1 为什么会有重复计算？

RoPE 的核心公式其实很简单：

q_rotated = q * cos(θ) + rotate_half(q) * sin(θ)
k_rotated = k * cos(θ) + rotate_half(k) * sin(θ)

这里面的 cos(θ) 和 sin(θ)，只跟位置 pos 和维度索引 i 有关。跟输入的内容 q、k 完全没关系。

说白了，只要序列长度固定，这些三角函数值就是固定的。每次推理都重新算一遍，就是做无用功。

4.2 预计算策略：把算过的存起来

我个人习惯的做法是：在模型初始化时，一次性把所有可能用到的 cos/sin 值算好，存起来。

具体来说，假设模型支持的最大序列长度是 max_seq_len，隐藏层维度是 dim。那么预计算的缓存就是一个 [max_seq_len, dim] 的矩阵。

代码实现大概是这样的：

class RotaryEmbedding:
    def __init__(self, dim, max_seq_len=2048):
        # 预计算所有位置的 cos 和 sin
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
        positions = torch.arange(max_seq_len)
        # 外积得到 [max_seq_len, dim/2]
        angles = positions[:, None] * inv_freq[None, :]
        # 拼接成 [max_seq_len, dim]
        self.cos_cached = torch.cos(angles).repeat(1, 2)
        self.sin_cached = torch.sin(angles).repeat(1, 2)

    def forward(self, q, k, pos):
        # 直接从缓存取，不用再算
        cos = self.cos_cached[pos]
        sin = self.sin_cached[pos]
        return apply_rope(q, k, cos, sin)

4.3 缓存策略：什么时候该清，什么时候该留？

预计算解决了「算一次」的问题。但端侧推理还有一个场景——流式推理。

比如你做一个语音助手，用户说一句话，模型要一个字一个字地生成。这时候，cos/sin 缓存是随着位置递增的。

我建议的策略是：

• 静态缓存：模型加载时一次性算好，整个生命周期不变。适合固定长度推理。
• 动态扩展：流式推理时，按需扩展缓存。比如当前最大位置到了 100，缓存只有 64，那就追加计算 64~100 的部分。

我曾经在项目里犯过一个错：每次生成新 token 都重新算整个序列的 RoPE。结果推理速度慢得离谱。后来改成动态扩展，速度提升了 3 倍。

4.4 量化后的缓存处理

端侧部署经常要做量化。RoPE 的 cos/sin 值范围是 [-1, 1]，量化后精度损失很小。

我一般这样处理：

1. 预计算 float32 的 cos/sin 缓存
2. 量化为 int8 或 fp16
3. 推理时直接使用量化后的值

实测下来，int8 量化后精度损失不到 0.1%，但内存占用减少 75%。对于端侧设备来说，这笔买卖很划算。

4.5 知识体系总览

下面这张图，是我自己总结的 RoPE 缓存策略全流程。你可以对照着看：

4.6 实际部署中的避坑指南

最后，分享几个我踩过的坑：

• 缓存对齐问题：有些 NPU 要求内存地址 16 字节对齐。预计算时要注意 padding，否则推理时会 crash。
• 多 batch 场景：如果同时处理多个序列，每个序列的位置可能不同。缓存要支持按位置索引，而不是按 batch 索引。
• 精度问题：fp16 的 cos/sin 在极端位置（比如 pos=0）没问题，但 pos 很大时误差会累积。我建议用 fp32 预计算，推理时再转成 fp16。

---

公众号:蓝海资料掘金营，微信deep3321