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的自定义实现，笔者也会利用业余时间学习，补一篇文章进行介绍。
LG Wing暗示下一款发布滑出屏手机
爱立信携手“钢铁巨头”打造工业领域最大5G专网，助力企业达成脱碳目标
TE提供的组件以及整套传感器解决方案 成为医疗器械应用的卓越选择
全场景智能车芯加速智能驾驶落地，芯驰科技获评智驾优秀企业
回顾2020年比亚迪的发展状况
LLMEngine下一层级的模块内如何实现各自功能接口
英特尔推出Kapoho Point开发板 降低神经拟态开发的门槛
怎么样提高变频调速系统的可靠性
合肥工厂探秘：蔚来ES8产能未来将提至平均3分钟一台，9月前完成1万台交付
多用户多回路宿舍用电管理解决方案
苹果公开可折叠iPhone新专利
如何利用智能技术提升重型移动设备的性能与效率
人工智能将渗透各个行业 1/5劳动将仰赖人工智能
国家重点研发计划启动实施“物联网与智慧城市关键技术及示范”重点专项
江森自控以132亿美元出售汽车电池的动力解决方案业务Power Solutions
人工气候室具有哪些应用优势
齐鲁工大：研发石墨烯/碳纳米管气凝胶的柔性应变传感器
人工智能系统BioMind在15分钟内诊断脑肿瘤的准确性达到了87%
谷歌公布并开源Pigweed 嵌入式库的集合
苹果业绩下滑 库克薪酬减少155万美元
回到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的自定义实现，笔者也会利用业余时间学习，补一篇文章进行介绍。
LG Wing暗示下一款发布滑出屏手机
爱立信携手“钢铁巨头”打造工业领域最大5G专网，助力企业达成脱碳目标
TE提供的组件以及整套传感器解决方案 成为医疗器械应用的卓越选择
全场景智能车芯加速智能驾驶落地，芯驰科技获评智驾优秀企业
回顾2020年比亚迪的发展状况
LLMEngine下一层级的模块内如何实现各自功能接口
英特尔推出Kapoho Point开发板 降低神经拟态开发的门槛
怎么样提高变频调速系统的可靠性
合肥工厂探秘：蔚来ES8产能未来将提至平均3分钟一台，9月前完成1万台交付
多用户多回路宿舍用电管理解决方案
苹果公开可折叠iPhone新专利
如何利用智能技术提升重型移动设备的性能与效率
人工智能将渗透各个行业 1/5劳动将仰赖人工智能
国家重点研发计划启动实施“物联网与智慧城市关键技术及示范”重点专项
江森自控以132亿美元出售汽车电池的动力解决方案业务Power Solutions
人工气候室具有哪些应用优势
齐鲁工大：研发石墨烯/碳纳米管气凝胶的柔性应变传感器
人工智能系统BioMind在15分钟内诊断脑肿瘤的准确性达到了87%
谷歌公布并开源Pigweed 嵌入式库的集合
苹果业绩下滑 库克薪酬减少155万美元