聊聊vLLM并行推理框架的设计

背景

实现并行/批量推理需要解决如下三个问题：

对Early-finished Requests的处理。不同请求所生成的文本长度不一致，可能差别很大，并且不易预测。如果没有一个将已生成结束的请求从Batch中移除并提前返回结果的机制，那么只能等一个Batch内所有请求都完成生成后才返回生成结果，导致生成短文本的用户则需要多“陪跑”数秒到数十秒才能得到结果，这对于服务响应时间是不利的；
对Late-joining Requests的处理。完整生成一段文本需要长达数秒或数十秒的时间，是漫长的。所以如果没有一个将新请求插入到推理Batch的机制，那么只能像CV业务那样，等前面的请求都完成推理了才进行后续请求的推理。这会导致请求需要在系统中长时间等待排队，表现为服务响应时间过长甚至不可接受；
Batching an arbitrary set of requests。每个请求对应的QKV Tensor的Length维度各不相同，在批量计算Attention时，需要处理此问题。诚然Padding+Masking的方法可以解决，但严重浪费算力和显存，对于算力和显存均有限的推理GPU是不利的。

根据vLLM的官方文档，它向用户提供了两种调用它的方法，分别是：

OpenAI-Compatible API Server（官方推荐）：兼容了OpenAI请求格式的server，包括OpenAI Completions API和OpenAI Chat API。
Simple Demo API Server（测试开发用，官方不推荐，相关脚本也不再维护）

vLLM的两种调用方式与内核引擎LLMEngine的关系如下（图片来自vLLM团队2023 first meetup PPT）:

PagedAttention的设计灵感来自操作系统的虚拟内存分页管理技术。主要来说是如下两点：

PagedAttention允许在非连续的内存空间中存储连续的键和值。具体来说，PagedAttention 将每个序列的 KV 缓存划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。
PagedAttention支持高效的内存重用共享；对于相同数据(输入的prompt)对应的KV cache，能复用则尽量复用；无法复用时，再开辟新的物理空间。

调度是一种常见的设计，vLLM中通过队列与调度协调工作实现了请求的批量处理。调度的原则就是最优的协调GPU资源、内存资源与批量请求的平衡。

vLLM对请求的调度处理流程大概如下：

当请求批量打到服务器上，将其存放在队列中，按照队列的特性，按顺序处理请求。
多种”状态”队列；running队列、waiting队列、swapped队列协调工作处理请求数据的调度进出。
在每1个推理阶段，决定要把哪些数据送给模型做推理，同时负责给这些模型分配KV Cache物理块。
当GPU资源不足时，暂停请求的调度，将其相关的KV cache物理块全部都先swap（置换、卸载）在cpu上，等后续gpu显存充足时，再把它们加载回gpu上继续做相关请求的推理。

总的来说，调度的处理对象是：GPU资源与请求prompt。