LLMEngine下一层级的模块内如何实现各自功能接口

最近业余时间在看新番vLLM，在读源码过程中，对其显存管理原理有了清晰的认识。vLLM系统主要是基于python+cuda实现的，很多其他python项目实现都很混乱（各种重复代码、语意不明/模糊的抽象设计），但vLLM的系统设计却特别工整，为怕遗忘，特别开启本篇，top down的记录一下vLLM框架结构。

回到vLLM这个项目，vLLM针对GPT类模型推理过程中KVCache这个显存占用大头专门设计了block_table，将KVCache分段成多个block存储在GPU中。一方面，这种设计可以共用beam search多batch之间share prompt sequence（的KVCache），减少显存占用。另一方面，在gpu显存和cpu内存间调度这些block，可以在有限的gpu显存空间下同时推理更大batch的sequence，换句话说，就是尽可能拉满gpu显存使用率，提高吞吐。

本篇文章将会按top down的方式介绍整个系统。先总览整个框架包含的基本类型，包括类型之间的关系、各类职责。然后，针对系统入口LLMEngine，介绍各个类之间如何通过接口互相组织完成推理过程，加深各个类功能的抽象理解。更进一步的，分析LLMEngine下一层级的模块内如何实现各自功能接口。（后续也会抽时间专门开一篇介绍vLLM中用到的cuda ops源码，特别是PageAttention部分，敬请期待）

框架概览

vLLM类关系图

整个框架核心的模块关系如上：

LLMEngine：是整个系统的入口，add_request负责输入prompt请求，step迭代推理，最终返回LLM生成的结果。其内部组合了一个Scheduler和一组Worker。

Scheduler：在每个推理步调度可处理的Sequence输入信息，其组合包含了一个BlockSpaceManager

BlockSpaceManager：维护gpu显存和cpu内存的使用情况，以及Sequence对应Cache的BlockTable信息。

Worker：在每个推理步执行LlamaForCausalLM推理，并返回采样后结果。除一个LLM模型外，其另一个核心组件是CacheEngine。

CacheEngine：负责执行相关gpu、cpu空间的换入、换出、拷贝等操作。

LLMEngine

LLMEngine实现细节

从图中可以看到，从上到下按先后顺序LLMEngine分别进行了__init__、add_request、step。

在构造LLMEngine时，LLMEngine就会调用Worker中的CacheEngine，初始化gpu、cpu空间，计算能容纳多少个block。每个block包含block_size个token对应的各层KVCache大小。在后续的组织中都会将Sequence对应的KVCache分成block_size大小的cache block，以方便管理对应block的空间。

add_request接口执行多次，接收多个待处理的prompt，将prompt处理成对应token的Sequence。每个输入prompt构造一个SequenceGroup， 其中包含了多个重复的Sequence为后续beam search做准备。SequenceGroup会最终保存到Scheduler中，以进行后续的调度。

step执行一个推理步。首先Scheduler会调度一组SequenceGroup和相关信息作为当前推理步的执行输入，除了输入外，还会包含当前SequenceGroup所需KVCache的换入换出信息。然后，Worker会将执行一次LLM推理（当然会先执行CacheEngine先准备KVCache）。Worker采样的输出结果会再次更新到Scheduler中的SequenceGroup内，以更新其内部的状态。最后，多次step循环，直到所有prompt对应的SequenceGroup都生成结束。

Scheduler & BlockSpaceManager

Scheduler

Scheduler中包含了三个队列：waitting、running、swapped。每当新增一个SequenceGroup时，添加至waitting队列中。

这三个队列之间的关系如下：

waitting：等待计算KVCache的SequenceGroup（也就是prompt序列）

running：执行推理的SequenceGroup，会在当前step中作为输入，一共包含两类：

prompt：来自waitting，未计算KVCache的SequenceGroup

generate token：计算过KVCache的SequenceGroup，准备生成下一个token

swapped：KVCache暂时换出到cpu内存的SequenceGroup

在每次schedule执行时，会调度几个队列之间的SequenceGroup，维护队列间的状态，使得当前执行推理尽可能占满显存空间。详细逻辑如上图中的数字标号顺序所示，值得注意的是，通过调度能实现两种解决显存不足的策略，一个是换出到cpu中，另一个就是重计算了（只有在SequenceGroup内只有一个Sequence的情况才能使用）。

当SequenceGroup推理新增了token时，update接口会更新一遍SequenceGroup内的状态。如下图所示，SequenceGroup内包含一组beam search的seq，每次执行推理的时候，每个seq采样s次，那么久会生成n*s个生成的token，根据这里面token保留置信度topn个生成结果。下图所示的结果就是n=4的情况，可以看到topn保留的output里seq1和seq3都来自原始输入seq1（parent_seq=1），此时需要BlockSpaceManager将原始的seq3释放（因为保留的output里没有seq3的输出），然后从seq1拷贝/fork到seq3，再讲新token添加到各个seq中。

BlockSpaceManager

BlockSpaceManager的功能是管理各个SequenceGroup对应KVCache存储信息。回顾LLMEngine提到过的，每个Sequence的KVCache序列会分成多个block_size长度的cache block，每个cache block的位置信息记录在BlocKspaceManager。如下图所示，BlockSpaceManager包含一个block_tables，其记录cache block到gpu显存或cpu内存物理地址的映射。

SequenceGroup刚加入Scheduler的时候并没有分配cache block空间，第一次进入running的时候需要向BlockSpaceManager申请可用的block空间。如下图所示，BlockSpaceManager分配block空间是以一个SequenceGroup作为一组输入，而且默认分配空间的时候，所有SequenceGroup内的token都是一样的（即是相同的prompt），因此会为所有Sequence都指向同一片cache block区域，该区域被引用数为Sequence数量。

当需要为一个Sequence新增token时，如下图所示，有两种情况：

当前cache block空间不足添加新token，则新增cache block。

当前空间足以添加新token，但last block与其他Sequence共用时（ref_count>1），如果新token还是添加到last block，那么会与共用last block的其他Sequence冲突，则需要释放掉last block（free不是真的释放，只是ref_count-1），分配一个新的last block。最后，返回信息标记原本last block内容需要拷贝到这个新的last block，即所谓的“copy-on-write”。

最后就是BlockSpaceManager其他接口的实现图解了，详细可参加下图：

实际上，BlockSpaceManager只负责维护cache block到gpu/cpu空间的索引，真正进行换入、换出、拷贝操作都需要通过Worker中CacheEngine进行。因此在Scheduler调度的时候，也会收集BlockSpaceManager返回结果，得到当前step所需KVCache的block_to_swap_in、block_to_swap_out、block_to_copy，以供后续CacheEngine操作内存空间。

Worker

Worker负责缓存更新执行和LLM推理执行。关于Worker的这个图比较长，因此这里截断成两张图来看。

如上图所示，Worker在执行时首先进行两个操作。

缓存更新：SchedulerOutputs包含了前面提到的当前所需swap in/swap out/copy的cache block信息，然后通过CacheEngine自定义的ops去执行缓存搬运操作，得到cuda stream的event，后续会在推理LLM各层的时候用到。

准备输入token序列__prepare_input：上图右侧的框内是预处理的过程，将SequenceGroupMetadata包含Scehduler调度得到running的所有SequenceGroup组合一个flatten的token序列，作为LLM的初始输入。Scheduler那节中提到过，running队列中当前执行的SequenceGroup有两类：一类未计算prompt（前缀）的KVCache，这部分需要完整的prompt token输入去计算各层的KVCache（全量推理）。另一类已经计算并缓存前缀的KVCache，因此只需要last token作为输入计算下一个generation token的分布（增量推理）。如上图所示，输入token序列的前半部分是多个prompt的token全量推理序列，后半部分是各个增量推理序列的last token。此外，全量推理的SequenceGroup中多个Sequence共享prompt，因此只需要任意一个Sequence的prompt作用输入就行。

上图是Worker执行LLM模型的过程。由__prepare_input组装的flatten token在各层映射成flatten hidden state。除了线性层、激活层等token独立计算的层以外，attention层的计算涉及不同token的hidden state依赖。上图主要展开了Attention层的四个主要步骤：

prompt全量推理：prompt序列通过xformers的attention算子计算得到下个layer的hidden state。由于这里attention计算的输入是完整的tensor，不是KVCache中分散的cache block，所以可以用第三方的attention算子完成计算。

等待缓存事件：CacheEngine中发送了异步缓存操作，因此只有当前层序列的cache block完成缓存更新，才能进一步获取KVCache或者记录KVCache，这种异步的实现能通过overlap计算和缓存搬运，节省一部分缓存搬运时间。

记录当前KVCache：当前层输入的hidden state作为KVCache通过自定义算子记录到对应cache block内，这里记录所有有效token的hidden state，包括prompt和last token（last token是前几次step中新增的，所以也没有缓存hidden state到KVCache）。

generation token增量推理：vLLM的核心PageAttention即在此实现，这里作者通过一个自定义算子（好像是参考Faster Transformer实现），实现了基于BlockTable分散KVCache的增量attention计算。

最后LLM内的采样器进行采样，将beam_search结果（新token）返回给Worker输出。

碎碎念

至此，笔者基本完成想要表达的的vLLM top down系统架构，相关的框架drawio已上库（图画的都有点挫，文章里可能不方便看。。），希望这篇文章能帮助有意愿在vLLM上做贡献的小伙伴。针对vLLM作者设计的cache_ops、attention_ops的自定义实现，笔者也会利用业余时间学习，补一篇文章进行介绍。

审核编辑：彭菁