摘要

本文探讨了在代码生成领域中，开发者实时编辑代码并请求大型语言模型（LLM）即时预测下一个令牌或下一行代码的典型场景。传统的LLM需要重新编码整个KV缓存以提供精确预测，这在处理长序列时计算成本高昂。本文提出了一种名为Positional Integrity Encoding（PIE）的新方法，该方法基于旋转位置编码（RoPE），通过移除引起时间混淆的旋转矩阵并重新应用正确的旋转矩阵，确保了令牌间的位置关系正确，仅需一轮矩阵乘法操作。实验证明，PIE在RepoBench-C-8k数据集上显著减少了计算开销，同时保持了模型性能。

原理

PIE的核心在于优化LLM在实时编辑场景中的效率和准确性。传统的LLM在处理代码编辑时，需要重新计算整个KV缓存，这在长序列处理中非常耗时。PIE通过以下步骤解决这一问题：

1. 移除旋转矩阵：首先移除导致时间混淆的旋转矩阵，这些矩阵在KV缓存中引入了不正确的位置信息。
2. 重新应用旋转矩阵：然后重新应用正确的旋转矩阵，确保编辑后的序列中令牌的位置关系是连续且唯一的。
3. 单轮矩阵乘法：整个过程仅需一轮矩阵乘法操作，极大地减少了计算开销。 通过这种方式，PIE确保了LLM在处理实时编辑时的效率和准确性，无需重新计算整个KV缓存。

流程

PIE的工作流程如下：

1. 编辑检测：检测到代码编辑操作后，识别出编辑的起始和结束位置。
2. 缓存更新：更新与编辑区域相关的KV缓存，移除旧的旋转矩阵并应用新的旋转矩阵。
3. 预测生成：使用更新后的KV缓存进行下一个令牌的预测。 例如，在代码插入场景中，PIE首先识别插入位置，更新该位置及其后的KV缓存，然后使用这些更新后的缓存进行预测。

应用

PIE的应用前景广泛，特别是在需要实时交互的代码编辑助手和开发工具中。通过显著减少计算开销，PIE使得LLM能够更快速、更准确地响应开发者的编辑操作，提高了开发效率和用户体验。未来，PIE还可以与其他优化技术结合，进一步扩展其应用范围。