拆解 vLLM:PagedAttention 怎么把显存利用率拉到 90%
🦞 一只用 AI Agent 搭副业产线的程序员
你跑过一个 70B 模型吗?Q4 量化后大概 40GB,一张 A100(80GB)放得下。但生产环境的问题不是"放不放得下",而是一个请求只用了 2K 上下文,为什么显存就不够处理第二个请求了?
答案是 KV Cache。传统方法预分配了太多永远用不到的显存空间。
vLLM 的 PagedAttention 就是来解决这个问题的。这篇文章我们看它怎么把操作系统的虚拟内存管理思想搬到了 GPU 显存管理上。
项目简介
vLLM(GitHub 40k+ Stars)是 UC Berkeley 开源的 LLM 推理引擎,核心贡献是 PagedAttention——一种把 KV Cache 按"页"管理的算法。它把显存利用率从传统框架的 30-40% 提升到 90% 以上,吞吐量提升 2-4 倍。现在被 LMSYS(Chatbot Arena)和多家公司用于生产环境。
架构全景
┌──────────────────────────────────────────────────────────────┐
│ API 服务层 │
│ OpenAI-compatible: /v1/completions, /v1/chat, │
│ /v1/embeddings, /v1/models │
├──────────────────────────────────────────────────────────────┤
│ 调度器(Scheduler) │
│ Continuous Batching — 不再等整批完成,来一个处理一个 │
│ ┌─────────┐ ┌──────────┐ ┌───────────────┐ │
│ │ 请求队列 │→│ Prefill │→│ Decode 循环 │ │
│ └─────────┘ └──────────┘ └───────────────┘ │
├──────────────────────────────────────────────────────────────┤
│ 块管理器(Block Manager)—— PagedAttention 核心 │
│ ┌──────────┐ ┌───────────┐ ┌─────────────┐ │
│ │ 物理块池 │ │ 块映射表 │ │ Copy-on-Write │ │
│ │ (Physical) │ │ (BlockTable)│ │ (beam search) │ │
│ └──────────┘ └───────────┘ └─────────────┘ │
├──────────────────────────────────────────────────────────────┤
│ CUDA Kernel 层 │
│ PagedAttention · FlashAttention · FP8/INT8 量化 · TP/PP │
└──────────────────────────────────────────────────────────────┘
先理解问题:传统 KV Cache 为什么浪费显存
LLM 推理时,每生成一个 token,都要拿当前的 query 去和之前所有 token 的 key/value 做 attention。为了避免重复计算,程序把每一层的 K 和 V 张量存下来——这就是 KV Cache。
传统框架(FasterTransformer、TGI)的处理方式:
请求 1,上下文 2000 tokens → 预分配 (max_context=4096) × K × V ≈ 2GB
请求 2,上下文 500 tokens → 预分配 (max_context=4096) × K × V ≈ 2GB
请求 3,上下文 8000 tokens → 预分配 (max_context=8192) × K × V ≈ 4GB
─────────────────────────────────────────────────────────────
总占用:8GB(但实际只用到了 (2000+500+8000)/(4096+4096+8192) ≈ 64%)
两个问题:
- 预分配:必须按"最大可能长度"分配,绝大多数请求用不完。
- 碎片化:请求 1 结束后释放 2GB,但紧接着来一个需要 3GB 的请求——那 2GB 的碎片用不上,要 defrag 或者 OOM。
核心问题:KV Cache 是连续分配的。连续分配 = 外部碎片 = 浪费。
关键设计一:分页管理——显存版的虚拟内存
vLLM 的答案是:不按"请求"分配,按"块"分配。每个块固定大小(比如 16 个 token),物理块池化,请求通过"块表"引用物理块。
# vllm/core/block_manager.py —— 块管理器的核心逻辑(概念性重建)
from typing import List, Optional, Dict
class BlockTable:
"""每个请求的虚拟块 → 物理块的映射表"""
def __init__(self, block_size: int = 16):
self.block_size = block_size # 每个物理块 = 16 个 token
self.blocks: List[Optional[int]] = [] # 虚拟块号 → 物理块号
class BlockAllocator:
"""全局物理块池"""
def __init__(self, num_blocks: int, block_size: int):
self.free_blocks: List[int] = list(range(num_blocks)) # 空闲块列表
self.block_size = block_size
def allocate(self) -> int:
"""分配一个物理块,返回块号"""
if not self.free_blocks:
raise OutOfMemoryError("No free blocks")
return self.free_blocks.pop()
def free(self, block_id: int):
"""释放物理块"""
self.free_blocks.append(block_id)
class BlockManager:
"""全局块管理器——所有请求共享物理块池"""
def __init__(self, num_gpu_blocks: int, block_size: int = 16):
self.allocator = BlockAllocator(num_gpu_blocks, block_size)
self.block_tables: Dict[int, BlockTable] = {} # 请求 ID → 块表
def append_slot(self, seq_id: int) -> Optional[int]:
"""为一个请求追加一个 slot(需要时分配新块)"""
block_table = self.block_tables[seq_id]
# 计算需要几个块
num_needed = (len(block_table.blocks) * self.allocator.block_size) + 1
# 如果最后一个块已满,分配新块
if num_needed > len(block_table.blocks) * self.allocator.block_size:
new_block = self.allocator.allocate()
block_table.blocks.append(new_block)
# 返回最后一个物理块的地址
return block_table.blocks[-1]
这个设计的效果:
改前(连续分配):
请求A: ┌──────────4096 tokens──────────┬ 碎片 ┐
请求B: ┌────2048 tokens────┬ 碎片 ┐
改后(分页分配,块大小=16):
物理块池: [A1][B1][A2][空][B2][A3][空][空][A4][B3]...
└── A 的块表: [0, 2, 5, 8] ──┘
└── B 的块表: [1, 4, 9] ──┘
没有外部碎片了——因为所有分配都是固定大小的块。内部碎片最多 15 个 token(最后一个块没装满),在上下文的尺度下可以忽略不计。
设计洞察:这就是操作系统的分页思想,直接搬到 GPU 显存管理。页表 + 物理页框池 + 按需分配。香不香?香。新不新?不新。但能把 60 年前的 OS 思想用到 LLM 推理里并做到生产可用——这就是工程的魅力。
关键设计二:Copy-on-Write——并行生成的零拷贝优化
一个常见场景:用户要求"生成 3 个候选回复"。怎么做?
朴素方案:KV Cache 复制 3 份。一个 4K 上下文的请求 = 2GB KV Cache。3 份 = 6GB。
PagedAttention 方案:共享前缀部分的物理块,只在分叉点复制块表指针。
# vllm/core/block_manager.py —— Copy-on-Write fork
class BlockManager:
def fork(self, parent_seq_id: int, child_seq_id: int):
"""从父请求 fork 一个子请求(beam search / parallel sampling)"""
parent_table = self.block_tables[parent_seq_id]
# 子请求共享父请求的块表(shallow copy)
child_table = BlockTable(block_size=parent_table.block_size)
child_table.blocks = list(parent_table.blocks) # 引用相同的物理块
self.block_tables[child_seq_id] = child_table
def append_slot(self, seq_id: int) -> Optional[int]:
"""追加 slot——如果物理块被共享,先 Copy-on-Write"""
block_table = self.block_tables[seq_id]
last_block = block_table.blocks[-1] if block_table.blocks else None
# 检查最后一个块是否被多个请求共享
if last_block is not None and self._ref_count(last_block) > 1:
# COW: 分配新物理块,复制内容
new_block = self.allocator.allocate()
self._copy_block(last_block, new_block)
self.allocator.free(last_block) # 减少旧块的引用计数
block_table.blocks[-1] = new_block
# 剩下的逻辑跟普通 append 一样
...
这个优化让 parallel sampling(生成 n 个候选回复)的显存开销从 O(n) 降到 O(1),只额外花在分叉后产生差异的 token 上。
设计洞察:Copy-on-Write 的通用性极高——fork 进程用它、Redis 的 BGSAVE 用它、vLLM 的 parallel sampling 也用它。理解一个模式,能用一辈子。
关键设计三:Continuous Batching——请求级别的流水线
传统批处理(Static Batching):等到一批请求全部完成,再处理下一批。
请求A(200 tokens)→ ████████████████████
请求B(50 tokens) → █████ → 等 A 完成 → 空闲
请求C(10 tokens) → ██ → 等 A 和 B 完成 → 空闲
Continuous Batching:一个请求完成立即踢出,把空出来的计算资源给等待队列的下一个。
# vllm/core/scheduler.py —— 调度器的核心逻辑(概念性重建)
class Scheduler:
def schedule(self) -> SchedulerOutput:
running: List[SequenceGroup] = []
preempted: List[SequenceGroup] = []
# Step 1: 从等待队列拉请求,直到显存不够
while self.waiting and self.block_manager.can_allocate():
seq_group = self.waiting.pop(0)
self.block_manager.allocate(seq_group)
running.append(seq_group)
# Step 2: 为每个运行中的请求生成一个 token
for seq_group in running:
seq_group.generate_one_token()
# Step 3: 把完成的请求踢出,释放块
for seq_group in running:
if seq_group.is_finished():
self.block_manager.free(seq_group)
running.remove(seq_group)
# Step 4: 剩余请求继续下一轮调度
return SchedulerOutput(
scheduled=running,
preempted=preempted,
num_waiting=len(self.waiting),
)
调度的核心在 Step 1:“能分配就分配”。不等到"最佳批次大小",而是只要显存有空就拉新请求。
这个设计的关键收益:短请求不用等长请求。50 个 token 的请求,生成完立刻释放 KV Cache 块,给下一个请求腾空间。在混合长短请求的场景下,吞吐量提升最明显。
核心代码拆解:PagedAttention 的 CUDA Kernel 是怎么读取 KV Cache 的
把 KV Cache 分页之后,attention 计算就不能用连续的矩阵乘法了——K 和 V 分散在不同物理块里。vLLM 为此写了一个定制的 CUDA kernel:
// vllm/csrc/attention/paged_attention.cu —— 简化逻辑
__global__ void paged_attention_kernel(
float* output, // [num_tokens, num_heads, head_size]
const float* query, // [num_tokens, num_heads, head_size]
const float* key_cache, // [num_blocks, num_heads, block_size, head_size]
const float* value_cache,// [num_blocks, num_heads, block_size, head_size]
const int* block_table, // [num_requests, max_num_blocks]
const int* context_lens, // [num_requests]
int num_heads,
float scale,
int block_size
) {
int tid = threadIdx.x;
int seq_idx = blockIdx.x; // 每个请求一个 block
int head_idx = blockIdx.y; // 每个 head 一个…嗯,另一个 block
int num_blocks = (context_lens[seq_idx] + block_size - 1) / block_size;
for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
// 关键:通过块表把"虚拟块号"转成"物理块号"
int physical_block = block_table[seq_idx * max_num_blocks + block_idx];
// 用物理块号去读 K 和 V
// key_cache[physical_block * block_stride + head_idx * block_size * head_size + ...]
int block_offset = physical_block * block_stride
+ head_idx * block_size * head_size;
// 计算这个 token 和当前块里所有 token 的 attention score
for (int t = 0; t < block_size; t++) {
float score = 0;
for (int d = 0; d < head_size; d++) {
score += query[q_offset + d] * key_cache[block_offset + t * head_size + d];
}
scores[block_idx * block_size + t] = score * scale;
}
}
// softmax + weighted sum(跟标准 attention 一样)
...
}
kernel 的核心只有一行:physical_block = block_table[seq_idx * max_blocks + block_idx]。这一行就是 PagedAttention 的全部魔法——其余部分都是在做普通的 attention 运算。代价只是一次额外的全局内存读取(读块表),在已经在做 O(n^2) 的 attention 运算面前可以忽略。
你可以抄的作业
1. 分页思路不只属于操作系统
任何"大块连续分配会碎片化"的场景都能用分页。你要做个内存池管理游戏对象?分页。管理网络数据包的缓冲?分页。只要分配单元不固定且碎片化严重,固定大小的块 + 映射表就是标准答案。
2. Copy-on-Write 的通用模式
fork + 共享读取 + 写时复制。vLLM 的 block fork 和 Linux 的 fork() 系统调用本质上是一回事。理解 COW,你就多了一个优化"同源并行操作"的武器。
3. Continuous Batching 就是"能进就进"
Pipeline 优化的本质不是"凑满批次",而是"不要等"。任何批处理系统——ETL 管道、消息队列消费者、API Gateway——都可以用这个思路:有空闲资源就拉,做完立刻放,不等。
4. 好的优化 = 好的数据结构 + 老的计算机思想
PagedAttention 没有发明新数学,它只是选了正确的数据结构(页表)和正确的资源管理策略(按需分配 + COW)。“创新"很多时候就是"把别处已经验证过的思想搬到你的领域”。
最后
vLLM 的故事跟很多"学术项目变成工业标准"的故事一样:不是算法有多新,而是工程优化做得够深。PagedAttention 把显存利用率从 30% 拉到 90%,靠的不是一篇论文、一个公式,而是一个完整的系统设计——块管理、调度器、CUDA kernel 三者协同工作。
理解 vLLM 不只是学一个推理框架,更是学"如何把一个孤立的算法优化做成系统级的解决方案"。
下一讲拆 Dify。一个低代码 AI 应用平台,是怎么设计插件系统让 100+ 个模型和工具能无缝接入的?
本文拆解的 vLLM 版本:v0.6.x。源码地址:github.com/vllm-project/vllm
🦞 一只用 AI Agent 搭副业产线的程序员
全平台同名:虾哥不加班 | 源码:GitHub - lobster-bujiaban
需要定制 AI 工具?来聊聊 → lob_ai
更多推荐


所有评论(0)