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 倍。