LLM生成代码后，如何一键合并到源代码中（FastApply技术研究）

okkkk · 前天 19:00

背景

在大语言模型越来越火的今天，越来越多的应用场景开始使用大语言模型来解决实际问题。而辅助编程可以算是大语言模型应用得最成功的场景之一了。早先的时候，更多使用的还是代码补全的能力，但是现在，各家产品都开始支持Chat和Agent的能力了。
之前一直有个疑问，生成的代码明明只是片段，也没有一个很好的规则能直接定位到源文件的位置，甚至有些生成的代码和现有代码没有任何重叠的部分，那这些代码是怎么精准地合并到源代码中的呢？今天就带着大家一起看一下，在Chat、Agent的场景下，如何将生成的代码精准且快速地合并到现有的代码文件中。
从全量重写到Planning+Applying的转变

一个最暴力的方法就是，每次在Chat/Agent里，都让模型生成完整的代码，然后直接全量替换，这样就不用考虑代码合并的问题了。
但是，这种方法的缺点也很明显：

成本高：每次模型都需要输出大量的代码，如果源文件有1000行，那每次模型生成代码就需要输出1000行，这样一次就用掉了大量的token，成本是不可接受的。
速度慢：大模型的输出模式是一个一个token地输出，如果源文件有1万token，那每次模型生成代码就需要走1万次的Decoding的过程，这是极其耗时的。

显然，现在的产品都没有使用这种模式，我们在Chat/Agent界面中看到的都是代码片段。所以，我们先将这个问题进行拆解一下，分成两个步骤：

Planning：大模型先生成代码片段
Applying：将代码片段合并到原始代码中

可能的代码片段形式

接下来我们会从这个原始代码出发，讲一下几种可能的代码片段形式。
import jsondef main(args): # show a greeting print("Hello!") returnif __name__ == '__main__': main("")Code Diff Patch

--- a/greeting.py+++ b/greeting.py@@ -10,4 +10,4 @@ def main(args):    # show a greeting- print("Hello!")+ print("Goodbye!")    return让模型生成完整的Code Diff Patch，然后直接通过Patch程序合并到原始代码中。
这样做的好处是，Applying的过程非常简单，只需要调用Patch程序合并即可。
缺点也很明显，Patch程序对于输入的格式要求非常严格，如果模型生成的Code Diff格式不正确，那合并就会失败。大语言模型对于数字的敏感程度不够，很容易产生幻觉，插入的位置有时候即使是偏了一行，合并出来的代码也就不可用了。
Unified Diff

@@ ... @@ def main(args):    # show a greeting- print("Hello!")+ print("Goodbye!")    return这个Diff的格式来自于Aider。和完整的Diff相比，Unified Diff删除了位置信息，只保留了代码的修改。合并的方式也很简单，只要将以空格和减号开始的行，替换成以空格和加号开始的行即可。这样，就可以有效避免模型关于数字的幻觉了。
这两种代码片段有一个共同的缺点，就是他们的数据都属于不常见的类型。code diff更多还是我们平时用git diff命令看一眼用的，并不会将其存成文件留存下来，模型自然也就很少见到这种数据，所以也就很难准确生成类似的数据了。模型更喜欢的还是生成一段完整的代码片段。
Lazy Format

# ... existing code ... def main(args):    # show a greeting    print("Goodbye!")    return# ... existing code ...相信如果大家让模型生成过代码，就会发现模型生成代码的时候会更偏向于这种Lazy Format的形式。就是用# ... existing code ...来表示代码的上下文，然后只生成中间相关的代码。
这种形式的好处是，模型可以更专注于生成代码，而不需要受到代码变更的干扰，也就是能提升生成的代码的准确性。
缺点也很明显，Applying就变成了一个比较复杂的过程。
基于Lazy Format的Applying

基于AST的代码块替换

对此，Continue的解决方案是借助AST来分解代码，进行代码块的替换：

如上图所示，原始代码被拆分成了3部分，分别是：

import代码块
main代码块
Python入口代码块

而生成的代码也被分为了3部分，分别是：

existing代码块，用...表示
main'代码块(注意这个代码块和原始文件的代码块不是完全相同的)
existing代码块，用...表示

替换步骤则是：

先用...来匹配任意代码块，也就是import代码块
再用main'代码块替换main代码块
最后用...来匹配任意代码块，也就是Python入口代码块

这里只是举了一个简单的例子，详细的算法可以参考Continue的文章。
这个方法的好处是，Applying的过程非常快，也很准确。
缺点则是，有些生成的代码，通过这种方法是没法替换的。比如：
# ... existing code ... # show a greeting print("Goodbye!") return# ... existing code ...压根就没有提供main函数的签名，那替换要从何找起呢？
基于全量代码生成的替换

为了解决这个问题，只好又求助与大语言模型了，让大语言模型根据原始代码和生成的代码，直接全量生成。这个时候大家可能会问了，之前不是说了全量生成成本高吗，为什么又要提出这种方法呢？
这就要仰仗于我们前面提到的Planning+Applying的思路了。在我们把问题拆解成两部分后，就可以对每一部分进行优化了，Planning的过程更关注于用户问题的理解，精准的代码生成，这通常依赖于大模型的通用能力，要求模型是一个比较强的模型。所以我们只能使用GPT-4o/DeepSeek R1/Claude 3.7 Sonnet这种大模型来生成代码，如果让它们全量生成代码，那成本确实就会非常高了。但是Applying的过程，只是把两个代码合并到一起而已，这对模型的要求就低很多了，所以我们就可以用小模型来处理该任务。
蒸馏大模型

第一个想法就是用大模型蒸馏小模型，FastApply-7B-v1.0就是专门做代码合并用的小模型，通过让Claude Sonnet 3.5 (70%)和GPT-4 (30%) 生成合并代码，并微调Qwen2.5-Coder-7B的模型，得到了该模型，细节可以参考该github仓库。
该模型在DeepSeek的评估下，准确率达到了99.95%，可以说是非常高了。
基于该模型，我用A6000的显卡进行了实际的速度测试。未量化的版本在vllm的加持下，可以得到45tokens/s的速度。我随机找了某个代码库统计了一下数据分布，每个文件的token中值在939左右，平均在1150左右。按照1000token来算，那么该模型就需要22s左右的时间才能合并完所有代码。换算成Qwen2.5官网提到的A100的速度（85tokens/s），那也需要12s左右的时间，依然是不可用的。
Speculative Decoding

好在，我们还有Speculative Decoding技术，这种技术的原理就是用小模型来打草稿（生成可能的Decoding token），然后用大模型来验证草稿。这样只要小模型生成出来的token被接受一部分，就可以显著减少大模型Decoding的时间了。具体的过程可以参考下图：

一般情况下，Speculative Decoding的加速比可以达到2-3倍，也就是说，用A6000来推理的话，可以达到120tokens/s左右的速度，这样得到的结果就是8s左右的时间，虽然快了挺多的，但是还是不大可用。
Prompt Lookup Decoding

又好在，代码合并这个任务有很好的性质:

生成的内容全部来自于原始代码和代码片段。这样我们就可以用“规则”来替换小模型，省掉小模型生成token的时间。
一旦正确定位到了某个位置，后续的一大段token都是可以被接受的。可以一次“生成”大量的token草稿，让大模型验证，并且验证成功的概率也很大。

基于这两点，我们就可以用Prompt Lookup Decoding来加速代码合并的过程。
原理其实很简单，先通过n-gram匹配Decoding的token和输入部分（原始代码和代码片段）的token（下图中的import），然后“摘抄”k个token（下图中Decoding中虚框部分），让模型来验证，模型会一次输出所有token的概率。前N个绿色框的token的概率都是最大的，所以可以被接受，，从第一个被拒绝的token(("Hello"))开始，后面的token都需要拒绝。因为将该token替换成("Goodbye")（概率最大）后，后面的token就应该基于("Goodbye")token来生成了。（注：这里拆分的token和模型的tokenizer拆分的token并不是完全一样的，只是用于让示例更直观）

vllm已经支持了PLD，经过测试，用上PLD之后，将生成的token数设置成100，1000个token的场景，加速比来到了12倍，也就是用A6000来推理可以达到550tokens/s左右的速度，耗时2s左右，看起来已经可以接受了。
PLD+

虽然PLD的加速比已经很高了，但是还有没有更进一步的可能性呢？我找到了一篇叫PLD+的论文，论文中提出了PLD+的技术，基本的思想如下：如果ngram一次匹配上了多个候选的token，那到底该选哪一个呢？按照PLD的话，就是直接选最新的一个，这样就有概率选错。而刚好我们的模型的中间输出可以给我们提供参考。模型推理的时候会生成每个token的hidden states和attention，这俩都可以用于帮助选择token。hidden states可以通过计算最新的一个token和所有候选token之间的余弦相似度，选相似度最高的一个。针对attention则可以选attention score最大的一个token。
参考下图：

基于vllm，我实现了一个简单版本的PLD+，论文中提到要使用9层hidden states来计算，但是vllm在推理的时候，只返回最后一层，所以我就直接用最后一层来计算了。得到的结果是：有提升，但是提升的幅度并不大。从12倍提升到了13倍，这起码验证了PLD+的思路是可行的，后续如果更进一步优化，应该可以得到更大的提升。另外，如果是大规模应用，那这多一倍的提升，也意味着很大的收益了。
最终的探索结果：在1000个token的场景下，用PLD+的加速比达到13倍，也就是用A6000来推理可以达到600tokens/s左右的速度，耗时1.7s左右，非常不错了。
未来的探索方向

基于测试的结果，如果将生成的token数设置成100，当token数很大的时候，加速比并不会提升，意味着耗时会线性增长，对于更大的原始代码，将生成的token数设置得大一些可能会取得更好的加速比。
PLD+提到的实现如果能在vllm上完整复现，可以期待一下更大的提升。
还有一些背景知识并没有被利用，比如：可以通过定位existing code的token来动态设置生成的token数。
在更大的模型上验证一下加速比。

参考