| 论文 | 博客 | 路线图 |

---

新闻 🔥

• [2024/2] 前瞻解码（Lookahead Decoding）论文现已发布于 arXiv。现已支持 采样（Sampling） 和 FlashAttention。用于提升 token 预测效果的高级功能已更新。

---

简介

我们提出了前瞻解码（lookahead decoding）：

• 一种用于加速大语言模型（Large Language Model, LLM）推理的并行解码算法。
• 无需草稿模型（draft model）或数据存储（data store）。
• 在每个解码步骤所用计算量（FLOPs）对数增长的情况下，线性减少解码步数。

下图展示了前瞻解码加速 LLaMa-2-Chat 7B 生成的演示：

背景：使用 Jacobi 迭代进行并行 LLM 解码

前瞻解码受 Jacobi 解码 启发，该方法将自回归解码视为求解非线性方程组，并使用不动点迭代法同时解码所有未来 token。下图展示了一个 Jacobi 解码示例。

然而，Jacobi 解码在实际 LLM 应用中几乎无法带来实际运行时间（wall-clock time）上的加速。

前瞻解码：让 Jacobi 解码变得可行

前瞻解码利用 Jacobi 解码的能力，通过收集并缓存从 Jacobi 迭代轨迹中生成的 n-gram（连续 n 个 token 的序列）来实现加速。

下图展示了通过 Jacobi 解码收集 2-gram 并验证它们以加速解码的过程。

为了提升该过程的效率，每个前瞻解码步骤被划分为两个并行分支：前瞻分支（lookahead branch）和验证分支（verification branch）。前瞻分支维护一个固定大小的二维窗口，从 Jacobi 迭代轨迹中生成 n-gram；同时，验证分支选择并验证有希望的 n-gram 候选。

前瞻分支与验证分支

前瞻分支的目标是生成新的 N-gram。该分支通过两个参数定义的二维窗口进行操作：

• 窗口大小 W（Window size W）：我们在未来 token 位置上向前看多远以进行并行解码。
• N-gram 大小 N（N-gram size N）：我们回溯 Jacobi 迭代轨迹多少步以检索 n-gram。

在验证分支中，我们识别出首 token 与当前输入最后一个 token 匹配的 n-gram。这通过简单的字符串匹配确定。一旦识别成功，这些 n-gram 就会被附加到当前输入，并通过一次 LLM 前向传播（forward pass）进行验证。

我们将这两个分支实现在一个注意力掩码（attention mask）中，以进一步利用 GPU 的并行计算能力。

实验结果

我们的研究表明，前瞻解码在单 GPU 上显著降低了延迟，在不同数据集上实现了 1.5 倍至 2.3 倍的加速。见下图。

目录

• 简介
• 目录
• 安装
  • 使用 Pip 安装
  • 从源码安装
  • 使用前瞻解码进行推理
  • 在您自己的代码中使用
• 引用
• 指导

安装

使用 pip 安装

pip install lade

从源码安装

git clone https://github.com/hao-ai-lab/LookaheadDecoding.git
cd LookaheadDecoding
pip install -r requirements.txt
pip install -e .

使用前瞻解码进行推理

您可以运行最小示例，观察前瞻解码带来的加速效果。

python minimal.py # 不使用前瞻解码
USE_LADE=1 LOAD_LADE=1 python minimal.py # 使用前瞻解码，获得 1.6 倍加速

您也可以使用前瞻解码与自己的聊天机器人进行对话。

USE_LADE=1 python applications/chatbot.py  --model_path meta-llama/Llama-2-7b-chat-hf --debug --chat # 使用前瞻解码进行聊天
USE_LADE=0 python applications/chatbot.py  --model_path meta-llama/Llama-2-7b-chat-hf --debug --chat # 不使用前瞻解码进行聊天


USE_LADE=1 python applications/chatbot.py  --model_path meta-llama/Llama-2-7b-chat-hf --debug # 不聊天，但使用前瞻解码
USE_LADE=0 python applications/chatbot.py  --model_path meta-llama/Llama-2-7b-chat-hf --debug # 不聊天，也不使用前瞻解码

在您自己的代码中使用 Lookahead decoding（前瞻解码）

您只需三行代码即可在自己的项目中导入并使用 Lookahead decoding。同时，您需要在命令行中设置 USE_LADE=1，或在 Python 脚本中设置 os.environ["USE_LADE"]="1"。请注意，目前 Lookahead decoding 仅支持 LLaMA 模型。

import lade
lade.augment_all()
lade.config_lade(LEVEL=5, WINDOW_SIZE=7, GUESS_SET_SIZE=7, DEBUG=0) 
# LEVEL、WINDOW_SIZE 和 GUESS_SET_SIZE 是 Lookahead decoding 中三个重要的配置参数（N, W, G），请参考我们的博客！
# 您可以通过在自己的设备上调整 LEVEL/WINDOW_SIZE/GUESS_SET_SIZE 来获得更好的性能。

随后即可加速解码过程。以下是一个使用贪心搜索（greedy search）的示例：

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=torch_device)
model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)
greedy_output = model.generate(**model_inputs, max_new_tokens=1024) # 获得加速效果

以下是一个使用采样（sampling）的示例：

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=torch_device)
model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)
sample_output = model.generate(**model_inputs, max_new_tokens=1024, temperature=0.7) # 获得加速效果

FlashAttention 支持

安装原始的 FlashAttention：

pip install flash-attn==2.3.3 # 原始 FlashAttention

有两种方式安装专为 Lookahead Decoding 优化的 FlashAttention：

1. 从 https://github.com/Viol2000/flash-attention-lookahead/releases/tag/v2.3.3 下载预编译包并安装（快速，推荐）。
   例如，我的环境是 cuda==11.8、python==3.9、torch==2.1，则应执行以下命令：

wget https://github.com/Viol2000/flash-attention-lookahead/releases/download/v2.3.3/flash_attn_lade-2.3.3+cu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
pip install flash_attn_lade-2.3.3+cu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

2. 从源码安装（较慢，不推荐）：

git clone https://github.com/Viol2000/flash-attention-lookahead.git
cd flash-attention-lookahead && python setup.py install

以下是使用 FlashAttention 运行模型的示例脚本：

python minimal-flash.py # 不使用 Lookahead decoding，启用 FlashAttention
USE_LADE=1 LOAD_LADE=1 python minimal-flash.py # 使用 Lookahead decoding，启用 FlashAttention，相比不使用 FlashAttention 提升约 20% 速度

在您自己的代码中，调用 config_lade 时需设置 USE_FLASH=True，并在调用 AutoModelForCausalLM.from_pretrained 时设置 attn_implementation="flash_attention_2"：

import lade
lade.augment_all()
lade.config_lade(LEVEL=5, WINDOW_SIZE=7, GUESS_SET_SIZE=7, USE_FLASH=True, DEBUG=0) 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=torch_device, attn_implementation="flash_attention_2")
model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)
greedy_output = model.generate(**model_inputs, max_new_tokens=1024) # 获得加速效果

我们将把 FlashAttention 直接集成到本仓库中，以简化安装和使用流程。

引用

@article{fu2024break,
  title={Break the sequential dependency of llm inference using lookahead decoding},
  author={Fu, Yichao and Bailis, Peter and Stoica, Ion and Zhang, Hao},
  journal={arXiv preprint arXiv:2402.02057},
  year={2024}
}

使用指南

核心实现在 decoding.py 中。Lookahead decoding 需要针对每个具体模型进行适配，示例见 models/llama.py。

LookaheadDecoding 快速上手指南

LookaheadDecoding 是一种用于加速大语言模型（LLM）推理的并行解码算法，无需草稿模型或外部数据存储，已在 LLaMA 系列模型上验证有效。

---

环境准备

• 操作系统：Linux（推荐 Ubuntu 20.04+）
• Python 版本：≥ 3.8
• PyTorch 版本：≥ 2.0（建议 2.1+）
• CUDA：≥ 11.7（若使用 GPU 加速）
• 模型支持：目前仅支持 LLaMA / LLaMA-2 架构（包括 meta-llama/Llama-2-7b-chat-hf 等 Hugging Face 模型）

💡 建议使用国内镜像源加速依赖安装，例如：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple

---

安装步骤

方法一：通过 pip 安装（推荐）

pip install lade

方法二：从源码安装

git clone https://github.com/hao-ai-lab/LookaheadDecoding.git
cd LookaheadDecoding
pip install -r requirements.txt
pip install -e .

（可选）启用 FlashAttention 支持（进一步提速）

1. 安装通用 FlashAttention：

   pip install flash-attn==2.3.3
   

2. 推荐：安装专为 LookaheadDecoding 优化的预编译版本（根据你的 CUDA、PyTorch 和 Python 版本选择）
   示例（CUDA 11.8 + PyTorch 2.1 + Python 3.9）：

   wget https://github.com/Viol2000/flash-attention-lookahead/releases/download/v2.3.3/flash_attn_lade-2.3.3+cu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   pip install flash_attn_lade-2.3.3+cu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   

---

基本使用

1. 最简示例（命令行）

# 不使用 LookaheadDecoding
python minimal.py

# 启用 LookaheadDecoding（约 1.6 倍加速）
USE_LADE=1 LOAD_LADE=1 python minimal.py

2. 在自己的代码中启用（仅需 3 行）

import os
os.environ["USE_LADE"] = "1"  # 或在命令行设置 USE_LADE=1

import lade
lade.augment_all()
lade.config_lade(LEVEL=5, WINDOW_SIZE=7, GUESS_SET_SIZE=7, DEBUG=0)

# 正常加载模型和生成文本
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "meta-llama/Llama-2-7b-chat-hf"
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.float16, device_map=torch_device
)

input_text = "Hello, how are you?"
model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)

# 自动生成时自动加速（支持 greedy 和 sampling）
output = model.generate(**model_inputs, max_new_tokens=128, temperature=0.7)
print(tokenizer.decode(output[0], skip_special_tokens=True))

3. 启用 FlashAttention（如已安装）

在 config_lade 中设置 USE_FLASH=True，并在加载模型时指定 attn_implementation="flash_attention_2"：

lade.config_lade(LEVEL=5, WINDOW_SIZE=7, GUESS_SET_SIZE=7, USE_FLASH=True, DEBUG=0)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map=torch_device,
    attn_implementation="flash_attention_2"
)

🔧 提示：LEVEL、WINDOW_SIZE、GUESS_SET_SIZE 对应论文中的 N、W、G 参数，可根据硬件和任务微调以获得最佳性能。