vLLM 核心技术 PagedAttention 原理详解

本文是 vLLM 系列文章的第二篇,介绍 vLLM 核心技术 PagedAttention 的设计理念与实现机制。

vLLM PagedAttention 论文精读视频可以在这里观看:https://www.bilibili.com/video/BV1GWjjzfE1b

往期文章:

1 引言:为什么大模型推理的内存管理如此关键?

随着大语言模型(LLM)在聊天机器人、代码补全、智能问答等场景中的广泛应用,越来越多的公司开始将其作为核心服务进行部署。但运行这类模型的成本极高。相比传统的关键词搜索请求,处理一次 LLM 推理的代价可能高出十倍以上,而背后的主要成本之一,正是 GPU 内存的使用效率。

在大多数主流 LLM 中,推理过程需要缓存每一步生成过程中的 Key 和 Value 向量(即 KV Cache),以便后续生成阶段引用。这部分缓存并不会随模型权重一起常驻 GPU,而是随着用户请求的长度动态增长和释放。在高并发场景下,不合理的 KV Cache 管理方式会导致大量内存碎片和资源浪费,最终限制可并发处理的请求数量,拉低整体吞吐量。

为了解决这一瓶颈,vLLM 引入了一个全新的注意力机制 —— PagedAttention。它借鉴了操作系统中的虚拟内存分页技术,将 KV Cache 分块存储在非连续的内存地址中,配合 block-level 的共享与 copy-on-write 机制,极大提升了内存利用率,从而显著提高了模型的吞吐能力。

vLLM 团队将 vLLM 的推理吞吐量与 HuggingFace Transformers(HF)HuggingFace Text Generation Inference(TGI) 进行了对比。评估在两种硬件设置下进行:在 NVIDIA A10G GPU 上运行 LLaMA-7B 模型,以及在 NVIDIA A100(40GB)GPU 上运行 LLaMA-13B 模型。实验结果表明,与 HF 相比,vLLM 的吞吐量最高可达 24 倍,与 TGI 相比,吞吐量最高可达 3.5 倍。

本文将结合 PagedAttention 的论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》,深入解析 PagedAttention 的设计理念与实现细节,并说明它是如何有效缓解内存瓶颈,显著提升大模型推理性能的。

2 背景知识:LLM 推理中 KV Cache 的角色

在大语言模型(LLM)如 GPT、OPT、LLaMA 的推理过程中,一个关键机制是自回归生成(autoregressive generation)。这意味着模型会基于用户提供的 prompt(提示词),逐步生成下一个 token,每一步都依赖之前生成的 token。这种生成方式的效率,极大依赖于 Transformer 架构中的自注意力机制(self-attention)

在 self-attention 中,模型为每个 token 计算三个向量:Query(Q)Key(K)Value(V)。每生成一个新 token,模型会将当前的 Query 向量与之前所有 token 的 Key 向量进行点积计算注意力分数,再据此对 Value 向量做加权求和。这种计算需要频繁访问此前的 token 信息。

为了避免重复计算这些历史 Key 和 Value 向量,推理系统通常会将它们缓存下来,称为 KV Cache(Key-Value 缓存)。这不仅节省了大量重复计算,也显著提升了推理效率。

推理过程可以划分为两个阶段:

  • prefill 阶段:模型接收完整的用户输入 prompt,并一次性并行计算所有 token 的 Query、Key 和 Value 向量。这一阶段是高度并行的,能够充分利用 GPU 的算力资源,因此属于 compute-bound(计算受限),瓶颈在于算力而非内存。

  • decode 阶段:此阶段模型开始逐个 token 地生成输出。每一步仅处理一个新的 token,模型首先计算其对应的 Query、Key 和 Value 向量。当前 token 的 Query 会与历史及当前的 Key 向量进行点积计算,通过

# vLLM推理引擎性能翻倍的秘密:PagedAttention技术详解与ModelScope适配实战 如果你最近在折腾大语言模型的本地部署,大概率已经听过vLLM这个名字。这个由加州大学伯克利分校团队推出的推理引擎,凭借“吞吐量提升24倍”的宣传口号,迅速在开发者社区中走红。但当你真正尝试在ModelScope生态里部署一个通义千问模型时,可能会发现,那些简单的Docker命令背后,隐藏着一套相当精妙的内存管理哲学。 今天我们不谈那些一键部署的脚本,而是深入vLLM最核心的**PagedAttention**技术,看看它到底如何解决大模型推理中最头疼的显存碎片问题。更重要的是,我会结合在ModelScope平台上适配不同模型的实战经验,分享一些官方文档里不会写的调优细节和避坑指南。无论你是想优化线上服务的并发能力,还是单纯想在本地机器上更高效地跑起大模型,理解这些底层原理都能让你少走很多弯路。 ## 1. 传统注意力机制的显存困境与PagedAttention的破局思路 要理解PagedAttention的价值,得先看看在没有它的时候,大模型推理面临什么样的挑战。 当你用传统方式(比如直接使用Hugging Face的Transformers库)运行一个大语言模型时,每次推理过程都需要为当前序列的所有历史token存储键(Key)和值(Value)缓存,这就是常说的**KV Cache**。假设你正在处理一个长度为8192的序列,模型隐藏层维度为4096,使用FP16精度,那么单个token的KV缓存大小大约是 `2 * 4096 * 2 bytes = 16KB`。对于8192个token,这个缓存就会膨胀到128MB。 这听起来似乎还能接受,但问题在于实际服务场景中的复杂性: 1. **请求长度不均**:不同用户的输入提示(prompt)长度差异巨大,从几十个token到上万个token都有可能。 2. **并发处理**:服务端需要同时处理多个请求,每个请求都有自己的KV缓存。 3. **生成过程动态**:模型是自回归生成的,每个新token的产生都会扩展对应序列的KV缓存。 在传统的连续批处理(continuous batching)中,系统需要为每个请求分配一块连续的显存来存储其KV缓存。由于请求长度不确定且动态增长,这会导致严重的**显存碎片化**——显存中散布着许多大小不一的空闲块,虽然总空闲显存可能还很多,但当你需要为一个新请求或一个增长的序列分配一块较大的连续空间时,却可能找不到合适的位置。 > 提示:你可以把这种情况想象成硬盘的文件存储。如果你频繁创建和删除不同大小的文件,硬盘上就会留下很多无法利用的小空隙,即使总剩余空间很大,也可能无法存入一个稍大的新文件。 更糟糕的是,当多个请求共享相同的前缀(比如系统提示词)时,传统方案会在每个请求的KV缓存中重复存储这部分内容,造成显存的浪费。vLLM团队在他们的论文中指出,这种碎片化和重复存储可能浪费**超过50%** 的显存空间。 那么PagedAttention是如何解决这个问题的呢?它的核心思想借鉴了操作系统中的**虚拟内存分页**机制。 ### 1.1 PagedAttention的基本工作原理 PagedAttention将每个请求的KV缓存从逻辑上划分为固定大小的“块”(block),每个块包含固定数量的token(比如128个)。这些物理块在显存中不需要连续存放,系统维护一个类似“页表”的映射关系,将请求的逻辑块序列映射到物理块上。 当注意力计算需要访问某个位置的KV缓存时,系统通过这个映射表找到对应的物理块,然后读取数据。这种方式带来了几个关键优势: * **消除外部碎片**:因为所有块大小相同,显存分配变得像“乐高积木”一样规整,完全避免了因为请求大小不一而产生的碎片。 * **共享缓存**:如果多个请求有相同的提示前缀,它们可以映射到相同的物理块上,显存中只存储一份,多个请求共享。 * **高效的内存分配**:系统可以预先分配一个大的“块池”,所有请求都从这个池中按需取用和归还块,分配和释放的开销极小。 为了更直观地展示这种差异,我们用一个简单的表格对比两种方案在处理两个不同长度请求时的显存布局: | 特性 | 传统连续分配方案 | PagedAttention分块方案 | | :--- | :--- | :--- | | **请求A (512 tokens)** | 分配一块连续的512-token空间 | 分配4个128-token的块(可能不连续) | | **请求B (256 tokens)** | 分配一块连续的256-token空间 | 分配2个128-token的块 | | **碎片情况** | 请求A和B之间可能产生无法利用的空隙 | 无外部碎片,所有块大小相同,可互换 | | **共享前缀(128 tokens)** | 在两个请求的缓存中各存一份,共256 tokens | 两个请求映射到同一个物理块,只存128 tokens | | **扩展性** | 请求增长可能需要重新分配更大连续空间,成本高 | 只需分配新的块并添加到映射表,成本低 | 这种设计使得vLLM在服务变长请求时,能够将GPU显存的利用率推向极限,这也是其吞吐量实现数量级提升的根本原因。 ## 2. 在ModelScope生态中集成vLLM:从模型下载到服务启动 理解了原理,我们来看看如何在实际的ModelScope环境中利用vLLM。虽然网络上有不少基于Docker的一键部署教程,但为了更深入地理解各个环节,我们不妨从更底层的Python API开始,这能让你对配置参数有更强的掌控力。 ### 2.1 环境准备与模型获取 首先,你需要一个支持CUDA的Python环境。我个人的习惯是使用conda创建独立的环境,避免包冲突。 ```bash # 创建并激活conda环境 conda create -n vllm-demo python=3.10 -y conda activate vllm-demo # 安装vLLM,这里指定从官方PyPI安装 pip install vllm # 安装ModelScope,用于下载模型 pip install modelscope ``` 接下来,我们需要从ModelScope社区下载模型。以通义千问2.5系列的一个小规模模型`Qwen2.5-0.5B-Instruct`为例,这个模型参数量小,适合在消费级显卡上快速实验。 ```python # 这是一个可选的模型下载脚本,你可以提前下载模型 # 但vLLM也支持运行时自动下载 import os from modelscope import snapshot_download model_name = "Qwen/Qwen2.5-0.5B-Instruct" # 指定模型缓存目录,方便管理 model_dir = snapshot_download(model_name, cache_dir="./models") print(f"模型已下载至: {model_dir}") ``` 如果你身处网络环境复杂的地区,可能会遇到下载缓慢的问题。ModelScope支持通过环境变量设置镜像源,可以尝试以下方式加速: ```bash # 在终端中设置临时环境变量 export MODELSCOPE_CACHE=./models export MODELSCOPE_ENDPOINT=https://mirror.modelscope.cn ``` ### 2.2 使用vLLM Python API启动推理服务 模型准备好后,就可以用vLLM加载并启动服务了。vLLM提供了两种主要的接口:离线的批量推理`LLM.generate()`和在线的API服务`vllm serve`。我们先看离线的批量推理,这能帮你快速验证模型加载是否成功。 创建一个名为`test_vllm.py`的脚本: ```python from vllm import LLM, SamplingParams import os # 关键一步:告诉vLLM使用ModelScope作为模型源 os.environ['VLLM_USE_MODELSCOPE'] = 'True' # 定义采样参数,控制生成文本的随机性和多样性 sampling_params = SamplingParams( temperature=0.8, # 温度值,越高越随机,越低越确定 top_p=0.95, # 核采样参数,累积概率超过此值的token会被考虑 max_tokens=256, # 生成的最大token数 ) # 初始化LLM引擎 # 注意:对于Qwen等需要trust_remote_code的模型,必须设置此参数 llm = LLM( model="Qwen/Qwen2.5-0.5B-Instruct", # 直接使用ModelScope模型ID trust_remote_code=True, # 允许执行远程代码(模型自定义层需要) gpu_memory_utilization=0.85, # GPU显存利用率,建议保留一些余量 max_model_len=8192, # 模型支持的最大上下文长度 ) # 准备一批提示词 prompts = [ "请用一句话解释人工智能是什么。", "法国的首都是哪里?", "写一个关于夏天的五言绝句。", ] # 执行批量生成 outputs = llm.generate(prompts, sampling_params) # 打印结果 for i, output in enumerate(outputs): print(f"提示 {i+1}: {output.prompt}") print(f"生成文本: {output.outputs[0].text}") print("-" * 50) ``` 运行这个脚本,如果一切顺利,你应该能看到模型生成的回答。这个过程中,vLLM在后台已经启用了PagedAttention来管理这些并发请求的KV缓存。 > 注意:第一次运行时会下载模型(如果未提前下载)并编译一些CUDA内核,可能会花费几分钟时间。后续运行会快很多。 ## 3. 关键参数调优:让vLLM在ModelScope模型上发挥全力 直接跑起来只是第一步,要让vLLM真正高效地工作,还需要根据你的硬件资源和业务场景调整一系列参数。这些参数就像引擎的调校旋钮,拧对了地方,性能提升立竿见影。 ### 3.1 显存与并行配置 这是影响性能最直接的一组参数,主要取决于你的GPU硬件。 * `gpu_memory_utilization`:这个参数控制vLLM可以使用的GPU显存比例。**不建议设置为1.0**,因为需要为系统和其他进程留出空间。对于独占的GPU,0.85-0.9是比较安全的选择。设置后,vLLM会预先分配好显存块池。 * `tensor_parallel_size`:**张量并行度**。如果你的模型很大,单张GPU放不下,或者你想利用多张GPU来加速单个请求的推理,就需要设置这个参数。例如,对于一个70B的模型,你可能需要设置为4或8,并在多张GPU上运行。 * `max_model_len`:模型支持的最大上下文长度。**这个值不能超过模型本身训练时的最大长度**,否则会出现不可预知的行为。对于Qwen2.5-0.5B,通常是32768,但你可以设置得更小以节省显存。 ```python # 多GPU配置示例(假设你有4张GPU) llm = LLM( model="Qwen/Qwen-7B-Chat", tensor_parallel_size=4, # 使用4张GPU进行张量并行 gpu_memory_utilization=0.88, max_model_len=8192, trust_remote_code=True, ) ``` ### 3.2 批处理与吞吐量优化 vLLM的吞吐量优势很大程度上来自于其先进的批处理策略。你需要根据服务类型调整相关参数。 * `max_num_batched_tokens`:**单个批处理步骤中处理的最大token数**。这是控制吞吐量和延迟平衡的关键参数。较大的值可以提高吞吐量(GPU利用率更高),但会增加每个请求的等待时间(因为要等批处理满)。对于高并发在线服务,可以设置小一些(如256-1024);对于离线批量任务,可以设置大一些(如2048-8192)。 * `enforce_eager`:是否强制使用Eager执行模式。默认为False,vLLM会使用CUDA图(CUDA Graphs)来捕获和重放计算图,这能显著减少内核启动开销。**只有在遇到兼容性问题(如某些自定义算子不支持)时,才需要设置为True**。 一个常见的性能调优策略是,在服务启动后,使用不同长度的请求进行压测,观察吞吐量(tokens/second)和延迟(time to first token, TTFT)的变化,找到`max_num_batched_tokens`的最佳值。 ### 3.3 与ModelScope模型相关的特殊参数 由于ModelScope上的一些模型结构比较特殊,可能需要额外注意: * `trust_remote_code`:**对于几乎所有ModelScope上的国产大模型(Qwen, Baichuan, ChatGLM等),这个参数都必须设置为True**。因为这些模型通常使用了自定义的注意力实现或模型层,需要动态加载代码。 * `revision`:指定模型的版本号。ModelScope上的模型可能会更新,如果你需要锁定某个特定版本,可以指定`revision="v1.0.0"`之类的标签。 * `dtype`:模型加载的数据类型。默认是`auto`,vLLM会自动选择。对于支持BF16的GPU(如Ampere架构及以上),使用`dtype="bfloat16"`可以在保持数值范围的同时减少显存占用。如果遇到精度问题,可以回退到`dtype="float16"`。 下面是一个综合了上述考量的配置示例,适用于在A100上部署Qwen-7B-Chat模型提供在线服务: ```python llm = LLM( model="Qwen/Qwen-7B-Chat", trust_remote_code=True, dtype="bfloat16", gpu_memory_utilization=0.9, max_model_len=16384, max_num_batched_tokens=512, # 针对低延迟在线服务优化 tensor_parallel_size=1, # 单卡运行 # 启用CUDA图优化(默认) enforce_eager=False, ) ``` ## 4. 实战:构建一个高性能的OpenAI兼容API服务 对于生产环境,我们通常不会直接运行Python脚本,而是启动一个常驻的API服务。vLLM提供了开箱即用的OpenAI兼容API,这意味你可以直接使用`openai`这个几乎成为标准的Python库来调用你的私有模型。 ### 4.1 启动API服务器 最直接的方式是使用vLLM的命令行工具。在终端中执行: ```bash # 设置使用ModelScope export VLLM_USE_MODELSCOPE=True # 启动API服务器 vllm serve Qwen/Qwen2.5-0.5B-Instruct \ --trust-remote-code \ --gpu-memory-utilization 0.85 \ --max-model-len 8192 \ --port 8000 ``` 这个命令会在本地的8000端口启动一个服务。你可以看到输出中包含了API的访问地址和Swagger UI地址(通常是`http://localhost:8000/docs`)。 ### 4.2 使用Python客户端进行调用 服务启动后,在另一个终端或Python脚本中,你可以像调用OpenAI API一样调用它: ```python from openai import OpenAI # 配置客户端,指向本地的vLLM服务 client = OpenAI( api_key="token-abc123", # 如果启动时设置了--api-key,这里需要匹配 base_url="http://localhost:8000/v1", # vLLM的OpenAI兼容端点 ) # 调用聊天补全接口 response = client.chat.completions.create( model="Qwen/Qwen2.5-0.5B-Instruct", # 模型名称,需与启动时一致 messages=[ {"role": "system", "content": "你是一个乐于助人的助手。"}, {"role": "user", "content": "请给我推荐三本值得阅读的科幻小说。"} ], temperature=0.7, max_tokens=500, stream=False, # 设置为True可以启用流式输出 ) print(response.choices[0].message.content) ``` ### 4.3 高级部署:结合FastChat提供Web UI 如果你想要一个类似ChatGPT的交互界面,可以结合FastChat。FastChat是一个开源项目,它提供了一个分布式的模型服务框架和漂亮的Web UI。虽然vLLM本身性能更强,但FastChat的UI体验更完善。 部署方式稍微复杂一些,需要启动三个组件:控制器(controller)、模型工作器(model worker,使用vLLM后端)和Web服务器(web server)。 ```bash # 1. 启动控制器(在一个终端) python -m fastchat.serve.controller --host 0.0.0.0 --port 21001 # 2. 启动vLLM模型工作器(在另一个终端) # 注意FASTCHAT_USE_MODELSCOPE环境变量 FASTCHAT_USE_MODELSCOPE=True python -m fastchat.serve.vllm_worker \ --model-path Qwen/Qwen2.5-0.5B-Instruct \ --trust-remote-code \ --controller http://localhost:21001 \ --worker http://localhost:21002 \ --port 21003 # 3. 启动Web UI(在第三个终端) python -m fastchat.serve.gradio_web_server --controller http://localhost:21001 --host 0.0.0.0 --port 7860 ``` 完成以上步骤后,打开浏览器访问`http://localhost:7860`,就能看到一个功能完整的聊天界面了。这种组合方案既享受了vLLM的高性能,又获得了友好的用户交互界面,非常适合内部工具或演示场景。 ## 5. 性能监控、问题排查与进阶技巧 服务跑起来之后,工作还没结束。你需要知道它运行得怎么样,以及在出现问题时如何快速定位。 ### 5.1 监控关键指标 vLLM提供了一些内置的监控方式。启动API服务时,添加`--metrics-interval`参数可以定期在日志中输出性能指标: ```bash vllm serve Qwen/Qwen2.5-0.5B-Instruct \ --trust-remote-code \ --metrics-interval 10 # 每10秒打印一次指标 ``` 你会看到类似下面的输出,其中包含了吞吐量、缓存利用率、队列深度等关键信息: ``` INFO 07-15 14:30:10 metrics.py:150] Throughput: 125.3 tokens/s, Running: 4 reqs, Waiting: 2 reqs INFO 07-15 14:30:10 metrics.py:151] KV Cache Usage: 45.3% (56 blocks used, 68 blocks free) INFO 07-15 14:30:10 metrics.py:152] GPU Memory: 78.5% used, 21.5% free ``` 对于生产环境,建议将这些指标接入到Prometheus + Grafana等监控系统中。vLLM支持通过`--prometheus-port`参数暴露Prometheus格式的指标。 ### 5.2 常见问题与解决方案 在实际使用中,你可能会遇到一些典型问题: * **问题:模型加载失败,提示“无法找到对应架构”** * **原因**:vLLM没有原生支持该模型架构。 * **解决**:确保`trust_remote_code=True`。如果还不行,可以尝试设置`model_impl="transformers"`,强制vLLM使用Hugging Face Transformers的后端来加载模型,虽然性能可能略有损失,但兼容性最好。 ```python llm = LLM(model="一些冷门模型", model_impl="transformers", trust_remote_code=True) ``` * **问题:生成结果乱码或重复** * **原因**:采样参数设置不当,或者模型本身在长文本生成上不稳定。 * **解决**:调整`temperature`(降低)、`repetition_penalty`(增加到1.1-1.2)、或启用`top_k`采样。对于Qwen系列模型,可以尝试以下组合: ```python sampling_params = SamplingParams(temperature=0.7, top_p=0.9, repetition_penalty=1.05) ``` * **问题:服务运行一段时间后显存缓慢增长(内存泄漏)** * **原因**:可能是由于请求失败后资源没有正确释放,或者是CUDA图缓存积累。 * **解决**:首先检查vLLM版本是否为最新。可以尝试在启动参数中添加`--disable-custom-all-reduce`(如果使用了张量并行)。对于CUDA图问题,可以设置`enforce_eager=True`作为临时排查手段,但这会牺牲性能。 ### 5.3 进阶技巧:使用vLLM进行模型量化推理 如果你的GPU显存紧张,或者想进一步提升吞吐量,可以考虑使用量化模型。vLLM支持AWQ(Activation-aware Weight Quantization)和GPTQ等量化格式。ModelScope社区也提供了许多模型的量化版本。 加载量化模型非常简单,只需要在初始化时指定`quantization`参数: ```python llm = LLM( model="Qwen/Qwen-7B-Chat-Int4", # ModelScope上的Int4量化模型 quantization="awq", # 量化类型,也可以是'gptq' trust_remote_code=True, gpu_memory_utilization=0.7, # 量化后显存需求降低,可以调低 ) ``` 根据我的实测,使用4-bit量化的Qwen-7B模型,显存占用可以从原始的约14GB降低到约5GB,使得在RTX 3090(24GB)这类消费级显卡上运行70B以下的模型成为可能,而吞吐量相比FP16版本通常只有10%-20%的损失,性价比极高。 最后,别忘了vLLM和ModelScope都处于快速迭代中。我遇到过一个坑是早期版本对Qwen2.5的Flash Attention 2支持有问题,导致生成速度很慢,更新到最新版后就解决了。多关注项目的GitHub仓库和ModelScope社区的公告,能帮你及时避开这些已知问题。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值