深入解析vLLM：加速大模型推理的高效框架

秋声studio

已于 2024-12-22 19:35:05 修改

阅读量3.9k

点赞数 29

分类专栏：口语化解析文章标签：人工智能模型调优 nlp vLLM

于 2024-12-11 08:17:51 首次发布

作品来源于秋声工作室

本文链接：https://blog.csdn.net/TrueYSH/article/details/144372246

版权

口语化解析专栏收录该内容

23 篇文章

订阅专栏

深入解析vLLM：加速大模型推理的高效框架

引言

在现代深度学习应用中，大型语言模型（如GPT、GLM、LLaMA、Qwen等）的推理速度和资源利用率成为关键挑战。vLLM (Virtual Large Language Model) 是由伯克利大学LMSYS组织开源的大语言模型高速推理框架，旨在极大地提升实时场景下的语言模型服务的吞吐量与内存使用效率。本文将深入探讨vLLM的工作原理及其优化技术。

一、vLLM概述

vLLM简介

vLLM是一个快速且易于使用的库，用于LLM推理和服务，可以与HuggingFace无缝集成。它利用了全新的注意力算法「PagedAttention」，有效地管理注意力键和值。vLLM的核心理念是通过高效的内存管理和资源调度，优化大语言模型的部署和执行。

降低GPU资源消耗：通过更高效的资源管理减少VRAM占用。
提升模型对话的并发量：允许更多的用户同时与模型交互，而不影响性能。

二、核心原理：KV Cache与Page Attention

vLLM的核心优化技术是基于KV Cache和Page Attention思想，这些技术显著提高了GPU资源的利用效率。

KV Cache (Key-Value Cache)

所有生成式大模型，无论是GPT、GLM、LLaMA还是Qwen，底层都是Decoder结构。解码器类似于Encoder，但底层使用Multi-head Self-Attention机制，涉及K (Key)、Q (Query)、V (Value) 矩阵运算。不同的是，解码器在计算Q乘以K转置后，会加上一个Mask Matrix，确保每个词仅受其前面词的影响，从而实现Causal Attention。

在K-Q-V计算过程中，后续的词会用到前面词的K和V矩阵。因此，在推理过程中，如果能存储前面词的K和V值，就不需要重新计算，从而提升了推理速度。这是一种典型的“用空间换时间”的做法。

然而，KV Cache本身也占用了大量VRAM。例如，如果大模型占59%的VRAM，KV Cache可能占用另外31%，这使得合理利用这31%的VRAM显得尤为重要。

Page Attention

为了解决KV Cache占用大量VRAM的问题，vLLM引入了Page Attention技术。具体来说，vLLM将KV Cache划分为多个小块（pages）。根据用户输入token的数量，动态分配这些小块的空间。未被占用的空间可以供其他任务使用，从而避免显存浪费。

例如，如果用户输入的句子较短，vLLM只会分配必要的KV Cache空间，而不是预分配整个缓存空间。这使得其他任务能够共享剩余的VRAM资源。

实现细节

Page Attention的工作方式，可以参考以下代码示例：

import torch
from vllm import KVCacheManager

# 初始化KV Cache Manager
kv_cache_manager = KVCacheManager(max_pages=1024, page_size=256)

# 动态分配KV Cache
def allocate_kv_cache(tokens):
    pages_needed = len(tokens) // kv_cache_manager.page_size + 1
    kv_cache_manager.allocate(pages_needed)
    return kv_cache_manager.get_allocated_pages()

# 示例：分配KV Cache
tokens = [torch.randn(768)] * 512  # 假设有512个tokens
allocated_pages = allocate_kv_cache(tokens)
print(f"Allocated Pages: {allocated_pages}")

三、vLLM的组织结构

vLLM的运作以Step函数为中心，Step函数的功能是进行一次单步推理，encode或decode。以下是vLLM的主要组件：

LLMEngine

LLMEngine是vLLM的离线推理引擎，负责处理离线推理任务。它的设计使得vLLM能够在后台高效地管理推理过程。

SequenceGroup

SequenceGroup是将用户请求抽象化的类，一个用户请求对应一个SequenceGroup。这样可以更好地管理和调度用户的推理请求。

Scheduler

Scheduler主要负责SequenceGroup级别的调度。它使用三个队列来管理请求：

Waiting Queue：存放可以用于下一轮推理的请求。
Running Queue：存放当前Step正在处理的请求。
Swapped Queue：暂时被交换到CPU中等待的请求。

vLLM中所有的存储设备都被划分成等大的块，涉及到存储设备的调度会调用Scheduler的Block_manager来实现。

Worker

Worker是模型及其计算的抽象类，一个worker对应一个GPU。Worker负责模型具体的计算，例如Attention。一个模型是否支持主要取决于Worker中是否实现了该模型的计算框架。具体来说就是将模型中的计算更换为vLLM的算子。

四、性能评估与调优

在实际应用中，选择合适的参数和优化策略对于提高vLLM的性能至关重要。以下是一些常见的调优建议：

调整聚类数量（nlist）：增加聚类数量可以提高检索精度，但也可能降低检索速度。
调整nprobe值：控制每次检索查询的聚类数量，权衡检索速度和精度。
选择合适的相似性度量：根据应用场景选择最合适的相似性度量方法。

五、使用层面的优势

Loading Local Models (加载本地模型)

vLLM提供了简便的接口来加载本地部署的大模型。只需要从vLLM中导入LLM类，并传入模型路径即可。

from vllm import LLM

# 加载本地模型
model = LLM(model_path='/path/to/your/model')

API Packaging and Testing (API封装与测试)

vLLM还支持快速封装大模型API，只需一行代码即可实现：

from vllm.entrypoints.openai.api_server import run_api_server

# 启动API服务器
run_api_server(model)

在API测试方面，可以采用以下两种方法：

使用API测试工具（如Postman）：
- 使用Postman等工具发送HTTP POST请求测试API。
使用Python发送请求：
- 通过Python脚本发送HTTP请求测试API。

import requests

url = 'http://localhost:8000/v1/completions'
data = {
    "prompt": "你好，世界",
    "max_tokens": 50
}

response = requests.post(url, json=data)
print(response.json())