vllm源码剖析
文章目录一、vllm的预热学习1vLLM的分块显存管理传统kv cache问题点、解决办法、解决问题思路2kv cache的初始化流程显存分配策略、计算可用KVcache的显存来根据block_size计算GPU块数量、vllm中llama架构举例、显存池3请求和显存块的映射(请求映射到一维数据、token和请求块的映射)4vLLM引擎模块与流式执行整个流程介绍、ZeroMQ介绍、推理引擎接受数据5vLLM Worker和Executor组件的协作(vllm整体架构图)6PageAttention的基石显存分布、每层显存块数量计算二、vllm的正式学习1vLLM 核心组件-调度器如何实现原理执行过程总结2vLLM模型权重加载一、vllm的预热学习1vLLM的分块显存管理传统kv cache问题点、解决办法、解决问题思路问题点传统kv cache管理的问题KV Cache都是按照最长推理长度seq_len去申请导致占用高时间长且内存碎片严重实际可并行处理的batch size被大幅压缩直接影响推理吞吐量vLLM的处理方法画图解释使用图片介绍分段后依据token来存放记录其实存放的是token IDs,如下图的四个IDblock_size也叫seq_len_per_block指的是每个缓存块block所能容纳的 token 数量上限。请求容量计算2K和VdimK和V的维度dtype参数类型解决方法的思路也就是pageAttention的思路2kv cache的初始化流程显存分配策略、计算可用KVcache的显存来根据block_size计算GPU块数量、vllm中llama架构举例、显存池计算显存分配假设用户有32G显存推理大概使用80%也就是25.6G那么1权重已知2预留激活值推理后的显存占用 - 推理前的显存占用3kvcache 25.6G - 推理后的显存占用计算GPU块分配数量①根据block_szie以及每个block占用的显存大小来计算出GPU块数量②根据GPU块数量构建显存池BlockPool主要功能包括①维护系统中空闲的显存块②维护哈希到显存快的映射表便于前缀匹配时快速命中已有块,避免重复分配。命中时只需要将对应块的引用计数加1③块的淘汰机制当引用计数为0重新挂入空闲队列并在重新分配的时候再清楚哈希映射关系page_size指的是单个 GPU 块block占用的显存大小单位是字节一个block最多可存放block_size个token的KV Cache,而每个token占用连续的一行空间因此要定位某个token在其block的位置还必须知道它在该block的偏移量offset即该token是当前block中的第几个tokenGPU 块分配数指的是整个显存池里一共能划分出多少个这样的块对应图片里的 num_blocksvllm中llama架构举例假设llama有6层attenton每层大小为per_layer_size那么总共占用的显存是6xper_layer_size为什么要reshape原始张量刚分配时是按 [num_blocks, block_size, num_kv_heads, head_size] 这样的逻辑维度存储的这更贴合 “块” 的管理视角。reshape 后会把前两个维度拼接最终得到三维张量例如 [num_blocks × block_size, num_kv_heads, head_size]。这样做是为了让每个 block 里的 token 数据在显存里是连续存放的当后续需要按 token 粒度访问或进行注意力计算时可以直接定位到连续的显存地址大幅提升访问效率。具体代码函数显存池管理1池子里面用过的block有hash下次用的时候才会重新设置hash2计算前缀匹配遍历一个请求中的所有块hash计算方式计算hash的时候还包含了前面block的信息hash淘汰流程Block池为了高效管理 KV Cache 并避免直接操作物理显存块vLLM 引入了基于 PagedAttention 的内存管理机制其核心是 BlockPool 类。在系统初始化时所有可用于 KV Cache 的 GPU 显存会被划分为固定大小的显存块。KVCacheManager 通过以下方式管理这些块使用 req_to_blocks 字典维护每个请求Request与其占用的若干显存块之间的映射关系。利用 BlockPool 中的 cached_block_hash_to_block 字典支持前缀缓存prefix caching。该字典将每个逻辑块的哈希值基于 token 内容及其前驱块哈希计算映射到对应的逻辑块从而实现跨请求的 KV Cache 共享。当新请求到达时系统会按块粒度计算其提示词prompt的块哈希并在 cached_block_hash_to_block 中查找匹配项以确定可复用的前缀长度。Vllm BlockPool 内部维护两个核心数据结构当前系统中可用的空闲物理块列表 self.free_block_queue一个从哈希值到对应显存逻辑块的映射表 self.cached_block_hash_to_block用于支持 KV Cache 的快速查找与复用。另外还有两个重要的字段就是block_id和ref_cnt以及它的hash值一个是一个唯一化的id另一个是这个block被使用了多少次当一个block的ref_cnt被请求使用的次数减到0之后就要放还free_block_queue中以下图为例block1被两个请求所使用所以它的ref_cnt就是2如果这里的任意一个请求结束那么block1中对应中的ref_cnt就要减去1。3请求和显存块的映射(请求映射到一维数据、token和请求块的映射)请求映射到一维数据里面一维向量也统计了不同请求之间在一维向量数据的起始位置和长度token和请求块的映射一个block最多可存放block_size个token的KV Cache,而每个token占用连续的一行空间因此要定位某个token在其block的位置还必须知道它在该block的偏移量offset即该token是当前block中的第几个tokenslotmapping某次请求中的某个token对应一整个block中的偏移怎么求token的offset途中请求1的token_ids有1到9block_size是4那么1到4属于block_id为0,5到8属于block_id为1位移offset可由模得出78模出2和3怎么计算一次请求占用的显存量一次请求所占用显存token数量*每个token的显存占用 token数量*num_heads X head_size X 2(KV) X dtype小结关键变量参数解释4vLLM引擎模块与流式执行整个流程介绍、ZeroMQ介绍、推理引擎接受数据整个流程ZeroMQ的介绍①REQ/REP通信模式特点一问一答严格区分客户端和服务端客户端代码举例②DEALER/ROUTER全双工比较常用客户端有身份标识其实就是每个用户登录有唯一的token和uuid可以绑定到唯一的用户服务器③PULL/PUSH请求和订阅有方向的单工通信client只能push到服务器另一端只能pull拉回消息推理引擎接受数据发送消息数据给客户端debug调试lanuch json换成attach模式整体流程客户端接受数据的流程5vLLM Worker和Executor组件的协作(vllm整体架构图)vllm整体架构图worker上报注册服务6PageAttention的基石显存分布、每层显存块数量计算显存分布①模型参数②输入、输出、推理时候的激活值③kv cache传统静态推理假设创建4096个max len的动态扩容(按需增长)vllm的做法①按需分配②节省kv cache块使用③前缀匹配免去多余计算显存块token需要的显存①一个token需要显存 num_heads x head_size x 1 x data_type (这里的1表示一个token就是1个block_size)②vllm会模拟跑一遍用推理后的显存减去客户端允许的显存大小就能得到可以分配的kv cache显存分配大小③1 block 1 token需要的显存 x block_size*2page_size(乘以2是因为kv是两个)④每个层可分配的显存块数量num_blocks avaliable_memory / (page_size*num_layers)备注像llama这样的模型中所有注意力层的KV Cache配置保持相同即每个层分配的存储空间是一致的二、vllm的正式学习1vLLM 核心组件-调度器如何实现原理在 AsyncLLMEngine 中核心推理任务由EngineCoreProc运行在独立进程负责执行。由于 vLLM 支持并发处理多个请求当前时间步调度哪个请求就成为关键问题。我们遵循4条调度原则等待时间优先优先处理等待最久的请求避免用户因长时间无响应而流失。已经开始输出的请求优先正在生成输出的请求应保持流畅避免卡顿影响体验。新请求次之尚未开始处理的请求可适当延后但需控制最大等待时间以防饿死。抢占规则如果需要抢占请求那么优先选择最晚达到且优先级低的。请求从客户端通过ZMQ ROUTER-DEALER 模式异步发送至EngineCoreProc由其接收并放入self.input_queue。随后_process_input_queue从队列中消费请求调用 EngineCore.add_request()最终将请求注册到调度器 self.scheduler —— 至此请求正式进入调度生命周期。执行过程调度器调度running队列在上述步骤中系统已经接收到来自客户端的一个新请求。接下来在调度函数中我们需要对当前所有待处理的请求进行统一调度依据一定的优先级策略选择出在当前时间步最适合执行的一个或多个请求。需要注意的是一个时间步并非只能执行单个请求——多个请求可以并行推进只要资源允许。这些请求可能处于不同的执行阶段有些已经生成了部分输出即用户已收到部分响应而另一些则仍在等待首个 token 的生成。这里的资源限制主要是以下的两个部分显存是否有足够的物理 Block 来存放这些请求生成的 Key/Value 缓存。系统限制单步处理的 Token 总数不能超过 GPU 的饱和负载或设定的阈值例如 max_num_batched_tokens。如图所示当前系统中的请求可分为两种状态正在执行中的请求(decode)这类请求已经生成了至少一个输出 token用户端已有部分内容返回。如请求1其 prompt 已在之前的步骤中被完整处理当前只需为其生成下一个 token即自回归解码阶段。为了方便起见我们将它叫做running请求。总的来说它是一种访存密集型IO Bound任务算力需求低但对显存带宽读取 KV Cache要求极高。等待首 token 的请求prefill这类请求尚未开始生成输出通常是因为其输入 prompt 还未被完全处理或尚未被调度执行。如请求2其prompt例如Hi!对应的 token 序列均尚未被处理调度时需对其进行首次前向计算 —— 可选择一次性处理完整 prompt或分批处理取决于系统资源和调度策略同理将它叫做waiting请求。总的来说GPU 需要一次性并行计算整个 Prompt 的所有 Token它能极大地压榨 GPU 的算力但也会瞬间消耗大量的 KV Cache 显存块。优先级之前提到的系统会优先保障 Running 请求因为如果它们被中断用户会明显感觉到大模型对话输出“卡顿”。waiting 队列还没开始算Prefill 阶段 ← 低优先级 running 队列已经在算了Decode 阶段 ← 高优先级先我们需要调度的是已经在执行的请求也就是在running队列中的请求。在上文我们提到调度器可在单个时间步内并行调度多个请求的序列。这些序列可能处于不同阶段有的在执行 prompt 预填充prefill有的在生成下一个 tokendecode。为控制系统负载与显存压力vLLM 支持通过参数如token_budget限制单步可处理的最大 token 总数。调度器会动态组合多个请求使总 token 数不超过该阈值同时尽可能提升 GPU 利用率。例如假设当前系统设置最大批处理 token 数为 3请求1decode 阶段只需生成 1 个 token。请求2prefill 阶段需处理 prompt “Hi!”共 2 个 token调度器仍可将两者同时调度。详细说明链接限制 1不能超过模型的最大记忆长度 模型最多能记住 4096 个 tokenmax_model_len 已经算了 1000 个 tokennum_computed_tokens 这一步最多能算4096 - 1 - 1000 3095 个 就像一张纸只能写 4096 个字已经写了 1000 个剩下最多还能写 3095 个。 限制 2不能超过系统当前的Token 预算token_budget 系统同时处理 10 个请求总预算是 512 个 token/步 已经给其他请求分配了 400 个 这个请求最多只能分到512 - 400 112 个 就像食堂今天只有 512 份饭10 个班级一起来打你们班来晚了只剩 112 份。 限制 3长请求要切片long_prefill_token_threshold 如果一个请求的输入特别长比如 10000 个 token 系统规定每步最多处理 2048 个 那这一步只处理前 2048 个下一步再处理后面的 就像搬家公司每次最多搬 2 吨10 吨的东西要分 5 趟搬。 三个条件取最小值就是最终的 num_new_tokens num_new_tokens min( max_model_len - 1 - num_computed_tokens, # 限制1 token_budget, # 限制2 long_prefill_token_threshold # 限制3 )总结请求接入与初始状态客户端请求通过 ZMQ 异步发送至 EngineCoreProc经由 input_queue 被消费后调用 add_request加入调度器的 waiting 队列状态设为 WAITING。调度优先级原则遵循“已激活请求优先、等待时间最长优先、新请求次之”的策略保障生成连续性并防止请求饿死。Running 队列调度优先处理正在运行的请求decode 阶段按 FCFS 顺序为其分配 KV Cache block若资源不足则抢占 running 队列末尾请求释放显存并将其移回 waiting 头部。Waiting 队列调度仅当无抢占且 token_budget 0 时从 waiting 中取出请求完成 block 分配后转入 running 队列状态更新为 RUNNING。资源与负载控制通过 token_budget 限制单步总 token 数支持超长 prompt 切分处理确保 GPU 利用率与显存使用的平衡。调度输出与状态更新生成 SchedulerOutput包含待执行请求、token 数量、block 分配等信息随后更新各请求的 num_computed_tokens完成调度闭环。2vLLM模型权重加载