通过多层推测实现闪电般快速的代码编辑

🚀 我们提出了 Blazedit，这是一种极其简单但通用的推测解码方法，可将整个文件代码编辑的速度提高多达 7.7 倍，涵盖全面的编辑场景。我们的代码库可在 ise-uiuc/blazedit 获取，我们期待进一步的优化和集成。

背景

代码编辑

大型语言模型（LLM）已广泛应用于当今的软件工程。代码补全等应用已广泛集成到现代 IDE 中。最近，代码编辑已成为 LLM 辅助开发人员的一种更基础的方法。例如，只需一个简短的用户提示，LLM 就可以自动编辑以修复错误或提高大型代码文件的代码质量。

然而，使用 LLM 编辑大型文件并非易事。主要有两种代码编辑格式：

1. 整个文件编辑：从头开始生成一个新文件——虽然这种方法有效，但由于需要生成大量 token，因此既昂贵又缓慢，不适合交互式编程。
2. 基于差异的编辑：生成一个紧凑的差异——这种方法效率更高，但对于在纯代码文本上训练的 LLM 来说具有挑战性。

在本文中，我们重点优化整个文件代码编辑。

推测解码

推测解码是一种众所周知的技术，通过权衡计算量来加速 LLM 生成，从而降低延迟。在常规的自回归解码中，每个 token 都需要 LLM 进行一次完整的正向传播，其效率受内存带宽限制。推测解码通过以下两个步骤加速常规解码：

• 草稿生成：使用更廉价的方法（例如较小的 LLM 或查找之前的 token）推测草稿 token 序列
• 验证：通过 token 预填充，在一次模型正向传播中验证整个草稿 token 序列

由于可以在一次正向传播中生成多个 token，因此解码速度可以显著提高。

更具体地说，Huggingface Transformers🤗 中有两种不需要修改原始模型的草稿生成方法：

• 辅助解码 通过调用较小的 LLM 来提出草稿 token。当前的默认实现（2025 年 2 月 13 日）还会动态调整草稿长度以优化接受率。
• 提示查找解码 (PLD) 通过复制粘贴前缀与提示后缀匹配的子提示来提出草稿 token，这几乎是免费的。该策略在代码编辑中非常高效，因为编辑后的代码通常与编辑前的代码大量重叠。

Blazedit：多层推测解码

我们提出了 Blazedit，一个简单的框架，结合了辅助解码和 PLD，以进一步加速整个文件代码编辑。

我们首先介绍现有代码编辑推测解码方法的局限性。

• 草稿模型中的高开销：草稿模型可以在实际编辑期间生成有意义的草稿 token，而不是简单地复制，从而导致更高的接受率。尽管如此，草稿生成仍然是自回归的，因此开销不可忽略，尤其是在草稿长度较长时。
• PLD 中的低接受率：PLD 效率很高，因为草稿生成成本可以忽略不计。然而，“复制”机制可能导致在目标模型进行实际编辑时，验证步骤中的接受率非常低。

Blazedit 使用优雅的多层推测解码策略解决了这些限制。在高级别上，与辅助解码类似，Blazedit 使用草稿模型来提出草稿 token，并由目标模型验证，以获得良好的接受率。同时，Blazedit 使用 PLD 来加速草稿模型，减少草稿模型生成的开销。具体来说，PLD 步骤会执行多次以累积草稿 token，然后才调用目标模型的前向传播。这允许草稿模型提出自适应数量的草稿 token，从而优化目标模型的接受率。

• 当 PLD 层获得高接受率时，它会检测到复制密集型场景，从而草稿模型可以提出更多草稿 token。
• 当 PLD 层获得低接受率时，它会检测到编辑密集型场景，从而草稿模型可以提出更少的草稿 token。

下面显示了更详细的算法描述：

评估

我们在一组代码编辑场景中对 Blazedit 和基线进行了基准测试。具体来说，我们要求它们编辑 90 个长度约为 5K 字符的源文件，其预期编辑比率均匀分布在 10% 到 90% 之间。为了展示最佳性能，我们对 A100 GPU 上的不同方法进行了网格搜索，以找到最佳配置。

我们首先演示运行 Qwen2.5-Coder-32B 和 DeepSeekCoder-33B 指令调优模型时的端到端每秒字符数和加速。我们看到与常规解码相比有显著加速，DeepSeekCoder-33B 的加速高达 7.7 倍。与最佳基线 (PLD) 相比，我们还看到了 1.15-1.17 倍的加速。Blazedit 还将最差情况（90% 分位数）比最佳基线 (PLD) 提高了 1.29-1.43 倍。

下图绘制了不同方法的每秒字符数和编辑比率分布。总体而言，Blazedit 和 PLD 都对编辑比率敏感，编辑比率越小，生成速度越快。Blazedit 数据点的分布总体优于 PLD。相比之下，辅助解码总体上对编辑比率不敏感，尽管它使用自适应草稿窗口来实现高接受率，这意味着草稿模型生成的开销可能是一个主要影响。

下面的案例研究了一个随机编辑样本。X 轴显示目标模型的前向传播步骤（越少越好），Y 轴显示建议（灰色）和接受（彩色）的 token 数量。与 PLD 相比，Blazedit 可以自适应地调整草稿 token 调度以优化接受率。与 动态辅助解码 相比，Blazedit 中的草稿模型生成速度快得多，因为它使用了 PLD 层。

展望

目前，用户可以通过两种方式使用 Blazedit：

1. 猴子补丁。 只需按照此 脚本 劫持草稿模型生成。
2. 自定义 Transformers 分支。 安装分叉的 分支，只需在 generate 方法中传递一个 blazedit_config 参数。

model_output = target_model.generate(
    input_ids,
    generation_config=generation_config,
    assistant_model=draft_model,
    # to trigger the custom candidate generation
    blazedit_config=dict(
        micro_draft_tokens=80, max_num_run=4, max_matching_ngram_size=10
    )
)

下一步，我们期待在其他推理框架中实现和集成 Blazedit。我们也将很快发布一篇包含更多细节和实验的论文。

虽然我们展示了使用 PLD 和辅助解码的多层推测的有效性，但探索不同层中更高级的推测方法组合也很有趣。