vllm源码剖析

一、vllm的预热学习

1）vLLM的分块显存管理（传统kv cache问题点、解决办法、解决问题思路）

• 问题点（传统kv cache管理的问题）
  KV Cache都是按照最长推理长度seq_len去申请，导致占用高时间长且内存碎片严重，实际可并行处理的batch size被大幅压缩，直接影响推理吞吐量
  

• vLLM的处理方法（画图解释）
  
  使用图片介绍，分段后依据token来存放记录（其实存放的是token IDs,如下图的四个ID）
  （block_size（也叫seq_len_per_block）指的是每个缓存块（block）所能容纳的 token 数量上限。）
  

• 请求容量计算：
  2：K和V
  dim：K和V的维度
  dtype：参数类型
  

• 解决方法的思路（也就是pageAttention的思路）
  

2）kv cache的初始化流程（显存分配策略、计算可用KVcache的显存来根据block_size计算GPU块数量、vllm中llama架构举例、显存池）

• 计算显存分配
  假设用户有32G显存，推理大概使用80%，也就是25.6G，那么：
  1）权重已知
  2）预留激活值=推理后的显存占用 - 推理前的显存占用
  3）kvcache = 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中的第几个token）
  GPU 块分配数：指的是整个显存池里一共能划分出多少个这样的块，对应图片里的 num_blocks

• vllm中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，下次用的时候才会重新设置hash
  2）计算前缀匹配：（遍历一个请求中的所有块）
  

• hash计算方式（计算hash的时候，还包含了前面block的信息）
  

• hash淘汰流程
  

• Block池
  为了高效管理 KV Cache 并避免直接操作物理显存块，vLLM 引入了基于 PagedAttention 的内存管理机制，其核心是 BlockPool 类。在系统初始化时，所有可用于 KV Cache 的 GPU 显存会被划分为固定大小的显存块。
  KVCacheManager 通过以下方式管理这些块：

1. 使用 req_to_blocks 字典维护每个请求（Request）与其占用的若干显存块之间的映射关系。
2. 利用 BlockPool 中的 cached_block_hash_to_block 字典支持前缀缓存（prefix caching）。该字典将每个逻辑块的哈希值（基于 token 内容及其前驱块哈希计算）映射到对应的逻辑块，从而实现跨请求的 KV Cache 共享。当新请求到达时，系统会按块粒度计算其提示词（prompt）的块哈希，并在 cached_block_hash_to_block 中查找匹配项，以确定可复用的前缀长度。

• Vllm BlockPool 内部维护两个核心数据结构：

1. 当前系统中可用的空闲物理块列表 self.free_block_queue；
2. 一个从哈希值到对应显存逻辑块的映射表 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中的第几个token）

slotmapping：某次请求中的某个token对应一整个block中的偏移

• 怎么求token的offset
  
  途中请求1的token_ids有1到9，block_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

• 小结：关键变量参数解释
  

4）vLLM引擎模块与流式执行（整个流程介绍、ZeroMQ介绍、推理引擎接受数据）

• 整个流程
  
• ZeroMQ的介绍
  ①REQ/REP通信模式
  特点：一问一答，严格区分客户端和服务端
  
  客户端代码举例
  

②DEALER/ROUTER全双工（比较常用）

客户端

有身份标识，其实就是每个用户登录有唯一的token和uuid可以绑定到唯一的用户


服务器

③PULL/PUSH请求和订阅（有方向的单工通信）
client只能push到服务器，另一端只能pull拉回消息

• 推理引擎接受数据
  
  发送消息数据给客户端
  

debug调试

lanuch json换成attach模式

整体流程

客户端接受数据的流程

5）vLLM Worker和Executor组件的协作(vllm整体架构图)

• vllm整体架构图
  

• worker上报注册服务
  
  

6）PageAttention的基石（显存分布、每层显存块数量计算）

• 显存分布
  ①模型参数
  ②输入、输出、推理时候的激活值
  ③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*2=page_size(乘以2是因为kv是两个)
  ④每个层可分配的显存块数量
  num_blocks = avaliable_memory / (page_size*num_layers)
  

• 备注
  像llama这样的模型中，所有注意力层的KV Cache配置保持相同，即每个层分配的存储空间是一致的

二、vllm的正式学习

1）vLLM 核心组件-调度器如何实现

原理

在 AsyncLLMEngine 中，核心推理任务由EngineCoreProc（运行在独立进程）负责执行。由于 vLLM 支持并发处理多个请求，当前时间步调度哪个请求就成为关键问题。
我们遵循4条调度原则：

1. 等待时间优先：优先处理等待最久的请求，避免用户因长时间无响应而流失。
2. 已经开始输出的请求优先：正在生成输出的请求应保持流畅，避免卡顿影响体验。
3. 新请求次之：尚未开始处理的请求可适当延后，但需控制最大等待时间以防饿死。
4. 抢占规则：如果需要抢占请求，那么优先选择最晚达到且优先级低的。
   请求从客户端通过ZMQ ROUTER-DEALER 模式异步发送至EngineCoreProc，由其接收并放入self.input_queue。随后，_process_input_queue从队列中消费请求，调用 EngineCore.add_request()，最终将请求注册到调度器 self.scheduler —— 至此，请求正式进入调度生命周期。

执行过程

调度器调度running队列
在上述步骤中，系统已经接收到来自客户端的一个新请求。接下来，在调度函数中，我们需要对当前所有待处理的请求进行统一调度，依据一定的优先级策略，选择出在当前时间步最适合执行的一个或多个请求。需要注意的是，一个时间步并非只能执行单个请求——多个请求可以并行推进，只要资源允许。这些请求可能处于不同的执行阶段：有些已经生成了部分输出（即用户已收到部分响应），而另一些则仍在等待首个 token 的生成。这里的资源限制主要是以下的两个部分：

1. 显存：是否有足够的物理 Block 来存放这些请求生成的 Key/Value 缓存。
2. 系统限制：单步处理的 Token 总数不能超过 GPU 的饱和负载或设定的阈值（例如 max_num_batched_tokens）。
   如图所示，当前系统中的请求可分为两种状态：
3. 正在执行中的请求(decode)：这类请求已经生成了至少一个输出 token，用户端已有部分内容返回。如请求1，其 prompt 已在之前的步骤中被完整处理，当前只需为其生成下一个 token（即自回归解码阶段）。为了方便起见，我们将它叫做running请求。总的来说，它是一种访存密集型（IO Bound）任务，算力需求低，但对显存带宽（读取 KV Cache）要求极高。
4. 等待首 token 的请求（prefill）：这类请求尚未开始生成输出，通常是因为其输入 prompt 还未被完全处理或尚未被调度执行。如请求2，其prompt（例如Hi!对应的 token 序列）均尚未被处理，调度时需对其进行首次前向计算 —— 可选择一次性处理完整 prompt，或分批处理（取决于系统资源和调度策略），同理将它叫做waiting请求。总的来说，GPU 需要一次性并行计算整个 Prompt 的所有 Token，它能极大地压榨 GPU 的算力，但也会瞬间消耗大量的 KV Cache 显存块。
   优先级：之前提到的，系统会优先保障 Running 请求，因为如果它们被中断，用户会明显感觉到大模型对话输出“卡顿”。

waiting 队列：还没开始算（Prefill 阶段） ← 低优先级 running 队列：已经在算了（Decode 阶段） ← 高优先级

先我们需要调度的是已经在执行的请求，也就是在running队列中的请求。
在上文我们提到，调度器可在单个时间步内并行调度多个请求的序列。这些序列可能处于不同阶段：有的在执行 prompt 预填充（prefill），有的在生成下一个 token（decode）。为控制系统负载与显存压力，vLLM 支持通过参数（如token_budget）限制单步可处理的最大 token 总数。调度器会动态组合多个请求，使总 token 数不超过该阈值，同时尽可能提升 GPU 利用率。例如，假设当前系统设置最大批处理 token 数为 3：

1. 请求1（decode 阶段）：只需生成 1 个 token。
2. 请求2（prefill 阶段）：需处理 prompt “Hi!”，共 2 个 token，调度器仍可将两者同时调度。

详细说明链接

限制 1：不能超过模型的"最大记忆长度" 模型最多能记住 4096 个 token（max_model_len） 已经算了 1000 个 token（num_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 )

总结

1. 请求接入与初始状态：客户端请求通过 ZMQ 异步发送至 EngineCoreProc，经由 input_queue 被消费后调用 add_request，加入调度器的 waiting 队列，状态设为 WAITING。
2. 调度优先级原则：遵循“已激活请求优先、等待时间最长优先、新请求次之”的策略，保障生成连续性并防止请求饿死。
3. Running 队列调度：优先处理正在运行的请求（decode 阶段），按 FCFS 顺序为其分配 KV Cache block，若资源不足则抢占 running 队列末尾请求，释放显存并将其移回 waiting 头部。
4. Waiting 队列调度：仅当无抢占且 token_budget > 0 时，从 waiting 中取出请求，完成 block 分配后转入 running 队列，状态更新为 RUNNING。
5. 资源与负载控制：通过 token_budget 限制单步总 token 数，支持超长 prompt 切分处理，确保 GPU 利用率与显存使用的平衡。
6. 调度输出与状态更新：生成 SchedulerOutput，包含待执行请求、token 数量、block 分配等信息；随后更新各请求的 num_computed_tokens，完成调度闭环。

2）vLLM模型权重加载