以1.2k行纯Python代码,理解PagedAttention、Continuous Batching和Prefix Caching的工程实现
vLLM代码超十万行,新手门槛高,需要轻量级入门路径
请求处理四步流程与框架三大核心组件逐一拆解
Scheduler调度逻辑、Model Runner执行流程、KV cache显存计算
BlockManager分配与回收,理解推理框架资源管理的核心闭环
文字转token ids,多个请求拼接成batch
Prefill阶段,分配逻辑块,model前向运算更新KV cache
Decode阶段,逐轮生成新token,blocks动态增加
De-Tokenization,token ids还原成词汇返回给用户
slot是显存管理的最小单位,block是分配的基本单位
物理块中每个slot的数据位置全局唯一,通过blockID乘以blockSize加上loc计算得出
维护running和waiting队列,组织并下发每一步需要执行的请求
基于PagedAttention原理管理KV cache显存,实现逻辑块到物理块的映射
加载并运行模型,完成每个请求的前向计算,支持TP张量并行
这正是操作系统虚拟内存分页思想在GPU显存管理中的复现
优势在于:充分利用显存、减少显存碎片、减少物理显存的反复申请与释放
持续向GPU送入请求进行推理,一个请求结束后立即下发新请求,有效提升GPU利用效率
通过跨请求复用KV cache减少计算量,请求1可直接复用请求0已计算的前两个token块
支持单节点内TP方式运行,TP大于1时主进程启动多个Model Runner实例共同完成前向计算
它创建Scheduler和Model Runner实例,将请求编码、资源分配、前向运算、结果解码四大步骤串成闭环
内部维护running队列和waiting队列,多个请求在两个队列之间轮转
当收到新prefill请求时可打断当前decode,被中断的请求移入waiting队列
前向计算完成后进行后处理,释放已结束请求占用的KV cache资源
创建context数据,为flash attention算子提供入参
rank 0将请求写入SharedMemory,其他rank通过loop函数读取后启动执行
logits经温度、softmax和gumbel max计算后完成采样,返回token ids
对比vLLM和SGLang的十万行级别代码量,Nano-vLLM将核心机制压缩到可逐行阅读的规模
| 时刻 | 操作 | free_block_ids | used_block_ids | 说明 |
|---|---|---|---|---|
| T0 | 为seq0分配2个block | 2,3,4 | 0,1 | seq0占用block 0和1 |
| T1 | 为seq1分配2个block | 4 | 0,1,2,3 | seq1占用block 2和3 |
| T2 | 释放seq0的block | 0,1,4 | 2,3 | FIFO原则,0和1回到队列尾部 |
| T3 | 为seq2分配2个block | 4 | 2,3,0,1 | 优先分配队列头部的0和1 |
推理框架的核心任务是高效的请求调度与资源分配,传统方式无法满足LLM推理需求
PagedAttention通过逻辑块、物理块和block table三层结构管理KV cache,类似虚拟内存分页
Continuous Batching持续提升GPU利用率,Prefix Caching通过跨请求复用KV cache减少prefill计算量
BlockManager采用FIFO空闲队列管理blocks,有利于历史数据保存和缓存复用
作者将关键函数解耦为可独立运行的Notebook示例,作为本文配套演示