【llm对话系统】 LLM 大模型推理python实现：vLLM 框架

在 LLM 的应用中，推理 (Inference) 阶段至关重要。它指的是利用训练好的 LLM 模型，根据输入 (Prompt) 生成文本的过程。然而，LLM 的推理速度往往较慢，尤其是在处理长序列或高并发请求时，效率瓶颈尤为突出。

为了解决这个问题，vLLM 应运而生！vLLM 是一个专为 LLM 设计的高吞吐、低延迟的推理和服务引擎，它能够显著提升 LLM 的推理速度，让你的应用如虎添翼！

今天，我们就来一起探索 vLLM 的奥秘，看看它是如何实现高效推理的！

一、vLLM 是什么？

vLLM 是一个开源的 LLM 推理和服务引擎，由 UC Berkeley 的研究团队开发。它采用了 PagedAttention 算法，能够有效地管理注意力键和值的内存，从而显著提高推理速度和吞吐量。

vLLM 的主要特点：

• 速度快： 比 HuggingFace Transformers 快 14x-24x，比 Text Generation Inference (TGI) 快 2.2x-2.5x！
• 高效的内存管理： 采用 PagedAttention 算法，有效减少内存碎片，提高显存利用率。
• 支持多种模型： 支持各种 HuggingFace Transformers 模型，包括 LLaMA, Mistral, Falcon, MPT 等流行模型。
• 易于使用： 提供简单易用的 API，可以轻松集成到你的应用中。
• 灵活部署： 支持多种部署方式，包括单机、多卡、多节点等。

二、vLLM 的核心技术：PagedAttention

vLLM 之所以如此高效，关键在于它采用了 PagedAttention 算法。在介绍 PagedAttention 之前，我们先来回顾一下传统的注意力机制 (Attention)。

1. 传统 Attention 的内存瓶颈

在传统的 Attention 机制中，每个 token 的键 (Key) 和值 (Value) 向量都需要连续存储在显存中。随着序列长度的增加，KV 缓存 (Key-Value Cache) 的大小也会线性增长，最终可能导致显存不足 (OOM)。

问题所在：

• 内存碎片： 由于请求的长度不同，KV 缓存中会产生大量的内存碎片，降低显存利用率。
• 内存浪费： 即使是同一个请求，不同 token 的 KV 缓存大小也可能不同，例如 masked 的 token 只需要很少的内存，而传统 Attention 仍然会为它们分配与其他 token 相同大小的内存。

2. PagedAttention：分页机制，解决内存难题

PagedAttention 借鉴了操作系统中虚拟内存和分页的思想，将每个请求的 KV 缓存分成多个块 (Block)，每个块可以存储固定数量的 token 的键和值。这些块不需要在连续的内存空间中，可以像分页一样按需分配。

PagedAttention 的优势：

• 减少内存碎片： 通过分页机制，可以有效减少内存碎片，提高显存利用率。
• 动态分配内存： 根据实际需要分配内存，避免内存浪费。
• 支持更长的序列： 由于减少了内存占用，PagedAttention 可以支持更长的序列推理。
• 高效的内存共享： 类似于操作系统的 Copy-on-Write 机制，PagedAttention 可以在不同请求之间共享内存块，例如对于并行的采样请求，其共享相同的 Prompt 部分的 KV 缓存，那么就可以共享最前面 Prompt 部分的 Block。

通俗理解：

想象一下，你有一本书 (LLM)，这本书有很多页 (Block)。传统的 Attention 机制要求你把这本书的所有页都平铺在一张桌子上 (连续的显存空间)，这样会占用很大的空间，而且不方便查找。

PagedAttention 则允许你将这本书分成不同的章节 (Block)，并将这些章节分别放到不同的抽屉里 (非连续的显存空间)。当你需要阅读某个章节时，只需要打开对应的抽屉即可。

三、vLLM 的安装和使用

1. 安装

pip install vllm

2. 使用 vLLM 进行离线推理 (Offline Inference)

离线推理指的是一次性处理一批输入，并生成所有输出。

from vllm import LLM, SamplingParams

# 准备输入 prompts
prompts = [
    "Hello, my name is",
    "The capital of France is",
    "The future of AI is",
]

# 定义采样参数
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=64)

# 初始化 LLM
llm = LLM(model="facebook/opt-125m")  # 可以替换成你想要使用的模型，支持大部分 HuggingFace 模型，比如 mistralai/Mistral-7B-v0.1，meta-llama/Llama-2-7b-hf 等

# 生成文本
outputs = llm.generate(prompts, sampling_params)

# 打印输出
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

代码解释：

1. 我们使用 vllm 库来进行推理。
2. prompts 变量存储了多个输入 prompt。
3. sampling_params 变量定义了采样参数，例如 temperature 和 top_p 用于控制生成文本的多样性，max_tokens 控制生成文本的最大长度。
4. 我们使用 LLM 类初始化一个 LLM 对象，并指定要使用的模型。
5. 我们调用 llm.generate() 方法生成文本。
6. 最后，我们打印每个 prompt 和生成的文本。

3. 使用 vLLM 进行在线推理 (Online Inference / Serving)

在线推理指的是搭建一个服务，实时处理请求并生成输出。

启动 OpenAI 兼容的 API 服务：

python -m vllm.entrypoints.openai.api_server --model facebook/opt-125m # 指定模型

使用 curl 命令发送请求：

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "facebook/opt-125m",
        "prompt": "The capital of France is",
        "max_tokens": 64,
        "temperature": 0.8
    }'

使用 Python 的 openai 库发送请求：

import openai

openai.api_base = "http://localhost:8000/v1" # API 地址
openai.api_key = "EMPTY"  # 不需要 API 密钥

response = openai.Completion.create(
  model="facebook/opt-125m",
  prompt="The capital of France is",
  max_tokens=64,
  temperature=0.8,
)

print(response.choices[0].text)