vLLM Prefix 前缀缓存详解：复用 System Prompt 的计算

讲解 vLLM 的 Prefix 特性：把每次推理都重复计算的 System Prompt 前缀缓存下来以避免自注意力阶段的冗余计算，配合 FIFO 或贪心的动态调度策略提升吞吐并降低延迟，并给出使用示例。

1. 介绍

在大语言模型（LLM）的实际应用中，通常需要在给模型具体任务信息之前，先介绍任务背景或上下文，以提升模型在特定任务上的表现。这部分内容通常被称为 System Prompt，它在每次推理过程中都会被重复计算，导致计算资源的浪费。为了解决这一问题，vLLM 提出了 Prefix 特性，通过缓存重复的前缀计算结果，减少自注意力阶段的冗余计算，从而提高吞吐量并降低延迟。

2. 核心思想

vLLM-Prefix 的核心在于利用 Prefix 的共享特性来优化推理过程。具体来说：

Prefix 缓存：将重复的前缀计算结果缓存起来，避免在每次推理时重复计算。
动态调度：通过调度策略（如 FIFO 或贪心算法）管理 Prefix 缓存，确保显存的有效利用。
多级 Prefix：利用 Trie 数据结构维护 Prefix，支持多级共享，进一步减少显存占用。
3. 技术实现

vLLM-Prefix 的实现基于 vLLM 的核心特性，包括 PagedAttention 和 Triton 优化：

PagedAttention：借鉴操作系统中的分页思想，将显存分页管理，支持灵活的 Token 存储方式，减少显存浪费。
Triton 优化：使用 Triton 语言实现高效的 GPU 并行计算，结合 FlashAttention 思想，优化 Self-Attention 的计算效率。
在 Prefix 的具体实现中：
PrefixPool 和 PrefixTrie：通过 PrefixPool 管理 Prefix 缓存，使用 PrefixTrie 数据结构实现 Prefix 的高效存储和检索。
调度策略：支持多种调度策略，如 FIFO 和贪心调度。贪心调度策略通过优先处理显存占用较大的 Prefix，优化推理效率。
4. 使用方法

vLLM-Prefix 的使用需要通过指定 prefix_pos 参数来标记前缀的长度。以下是一个简单的使用示例：

from vllm import LLM, SamplingParams

# 初始化模型
model_path = "/models/Llama-2-70b-chat-hf"
llm = LLM(model=model_path, tensor_parallel_size=1)

# 设置采样参数
sampling_params = SamplingParams(temperature=0, max_tokens=100)

# 定义 Prompt 和 Prefix 长度
prompts = ["这是一个带有前缀的 Prompt 示例。", "这是另一个带有相同前缀的 Prompt。"]
prefix_len = [16, 16]  # 前缀长度

# 生成文本
outputs = llm.generate(prompts, sampling_params=sampling_params, prefix_pos=prefix_len)
for output in outputs:
    print(output)

5. 实验效果

根据实验数据，vLLM-Prefix 在推理加速方面表现显著：

在 Baichuan2-13B-chat 模型上，约 700 个 Token 的 Prompt 中，前 596 个 Token 作为 Prefix，推理速度提升了 2-3 倍。

与 SGLang 等其他优化框架相比，vLLM-Prefix 在某些场景下略逊一筹，但在稳定性上更具优势。

# 获取数据集
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json

# 执行benchmark_throughput评测脚本
python benchmark_throughput.py --model meta-llama/Llama-2-70b-chat-hf --tokenizer meta-llama/Llama-2-70b-chat-hf --temsor-parallel-size 8 --dataset ShareGPT_V3_unfiltered_cleaned_split.json

benchmarks执行日志

6. 未来展望

vLLM-Prefix 作为一种高效的推理加速技术，未来仍有优化空间：

Prefix 删除机制：目前的 Prefix 缓存删除策略仍有待完善，以更好地适应动态推理场景。
调度策略优化：进一步优化 Prefix 调度策略，以适应更多复杂的推理任务。
跨框架兼容性：探索与其他推理框架（如 SGLang）的兼容性，提升整体推理效率。
7. 总结

vLLM-Prefix 通过缓存重复的前缀计算结果，显著提高了大模型的推理效率。它结合了 PagedAttention 和 Triton 优化，实现了高效的显存管理和计算加速。通过合理的调度策略和数据结构设计，vLLM-Prefix 在实际应用中展现了强大的性能优化潜力，为大模型的高效部署提供了新的思路。

1. 介绍

2. 核心思想

3. 技术实现

4. 使用方法

5. 实验效果

6. 未来展望

7. 总结