Tech Architecture

Nano-vLLM 轻量级推理框架详解

以1.2k行纯Python代码,理解PagedAttention、Continuous Batching和Prefix Caching的工程实现

kaiyuankaiyuanInfraTech
微信公众号
Nano-vLLM
02 / 15

这篇文章是怎么往前推进的

起点

vLLM代码超十万行,新手门槛高,需要轻量级入门路径

展开

请求处理四步流程与框架三大核心组件逐一拆解

深入

Scheduler调度逻辑、Model Runner执行流程、KV cache显存计算

落点

BlockManager分配与回收,理解推理框架资源管理的核心闭环

微信公众号
02 / 15
Why Engine

传统方式已无法满足LLM推理需求

第2节
03 / 15
流程 · 请求处理
04 / 15

请求处理的四个步骤

Step 1

文字转token ids,多个请求拼接成batch

Step 2

Prefill阶段,分配逻辑块,model前向运算更新KV cache

Step 3

Decode阶段,逐轮生成新token,blocks动态增加

Step 4

De-Tokenization,token ids还原成词汇返回给用户

第1节
04 / 15
Excerpt · 显存管理
05 / 15
Excerpt
每个token占用相同显存。slot是显存最小单位,多个slot构成block。第1节 · 显存分配
Side Note

slot是显存管理的最小单位,block是分配的基本单位

物理块中每个slot的数据位置全局唯一,通过blockID乘以blockSize加上loc计算得出

第1节
05 / 15
架构 · 三大组件
06 / 15

Nano-vLLM 用三大模块承载推理框架的核心职责

Scheduler

维护running和waiting队列,组织并下发每一步需要执行的请求

Block Manager

基于PagedAttention原理管理KV cache显存,实现逻辑块到物理块的映射

Model Runner

加载并运行模型,完成每个请求的前向计算,支持TP张量并行

第2.1节
06 / 15
Excerpt · PagedAttention
07 / 15
Excerpt
逻辑块与物理块通过块表完成映射。逻辑块连续,物理块不一定连续。第2.1节 · PagedAttention
Side Note

这正是操作系统虚拟内存分页思想在GPU显存管理中的复现

优势在于:充分利用显存、减少显存碎片、减少物理显存的反复申请与释放

第2.1节
07 / 15
特性 · 两大优化
08 / 15

Nano-vLLM 实现的两大核心优化机制

Continuous Batching

持续向GPU送入请求进行推理,一个请求结束后立即下发新请求,有效提升GPU利用效率

Prefix Caching

通过跨请求复用KV cache减少计算量,请求1可直接复用请求0已计算的前两个token块

Tensor Parallelism

支持单节点内TP方式运行,TP大于1时主进程启动多个Model Runner实例共同完成前向计算

第2.1节
08 / 15
判断 · Engine
09 / 15
Engine Core

LLM Engine 是推理框架的调度中枢

它创建Scheduler和Model Runner实例,将请求编码、资源分配、前向运算、结果解码四大步骤串成闭环

第2.2节
09 / 15
调度 · Scheduler
10 / 15

Scheduler 采用prefill优先的抢占式调度

双队列

内部维护running队列和waiting队列,多个请求在两个队列之间轮转

抢占逻辑

当收到新prefill请求时可打断当前decode,被中断的请求移入waiting队列

资源回收

前向计算完成后进行后处理,释放已结束请求占用的KV cache资源

第2.3节
10 / 15
执行 · Model Runner
11 / 15

Model Runner 通过SharedMemory实现多 rank 协作

数据准备

创建context数据,为flash attention算子提供入参

分布式协同

rank 0将请求写入SharedMemory,其他rank通过loop函数读取后启动执行

采样输出

logits经温度、softmax和gumbel max计算后完成采样,返回token ids

第2.4节
11 / 15
数据 · 规模
12 / 15
1.2k

行纯 Python 代码实现完整推理框架

对比vLLM和SGLang的十万行级别代码量,Nano-vLLM将核心机制压缩到可逐行阅读的规模

第2.1节 / 第3节
12 / 15
表格 · BlockManager
13 / 15

BlockManager 分配与回收时序

时刻操作free_block_idsused_block_ids说明
T0为seq0分配2个block2,3,40,1seq0占用block 0和1
T1为seq1分配2个block40,1,2,3seq1占用block 2和3
T2释放seq0的block0,1,42,3FIFO原则,0和1回到队列尾部
T3为seq2分配2个block42,3,0,1优先分配队列头部的0和1
第3.2节
13 / 15
核心总结
14 / 15
01

推理框架的核心任务是高效的请求调度与资源分配,传统方式无法满足LLM推理需求

02

PagedAttention通过逻辑块、物理块和block table三层结构管理KV cache,类似虚拟内存分页

03

Continuous Batching持续提升GPU利用率,Prefix Caching通过跨请求复用KV cache减少prefill计算量

04

BlockManager采用FIFO空闲队列管理blocks,有利于历史数据保存和缓存复用

全文总结
14 / 15
Notebook

从1.2k行代码开始,理解十万行框架

作者将关键函数解耦为可独立运行的Notebook示例,作为本文配套演示

github.com/CalvinXKY/InfraTech
15 / 15