当前位置：首页 > news >正文

拆解 vLLM：PagedAttention 怎么把显存利用率拉到 90%

news 2026/6/1 2:04:59

🦞 一只用 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 —— 块管理器的核心逻辑（概念性重建）fromtypingimportList,Optional,DictclassBlockTable:"""每个请求的虚拟块 → 物理块的映射表"""def__init__(self,block_size:int=16):self.block_size=block_size# 每个物理块 = 16 个 tokenself.blocks:List[Optional[int]]=[]# 虚拟块号 → 物理块号classBlockAllocator:"""全局物理块池"""def__init__(self,num_blocks:int,block_size:int):self.free_blocks:List[int]=list(range(num_blocks))# 空闲块列表self.block_size=block_sizedefallocate(self)->int:"""分配一个物理块，返回块号"""ifnotself.free_blocks:raiseOutOfMemoryError("No free blocks")returnself.free_blocks.pop()deffree(self,block_id:int):"""释放物理块"""self.free_blocks.append(block_id)classBlockManager:"""全局块管理器——所有请求共享物理块池"""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 → 块表defappend_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# 如果最后一个块已满，分配新块ifnum_needed>len(block_table.blocks)*self.allocator.block_size:new_block=self.allocator.allocate()block_table.blocks.append(new_block)# 返回最后一个物理块的地址returnblock_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 forkclassBlockManager:deffork(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_tabledefappend_slot(self,seq_id:int)->Optional[int]:"""追加 slot——如果物理块被共享，先 Copy-on-Write"""block_table=self.block_tables[seq_id]last_block=block_table.blocks[-1]ifblock_table.blockselseNone# 检查最后一个块是否被多个请求共享iflast_blockisnotNoneandself._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 —— 调度器的核心逻辑（概念性重建）classScheduler:defschedule(self)->SchedulerOutput:running:List[SequenceGroup]=[]preempted:List[SequenceGroup]=[]# Step 1: 从等待队列拉请求，直到显存不够whileself.waitingandself.block_manager.can_allocate():seq_group=self.waiting.pop(0)self.block_manager.allocate(seq_group)running.append(seq_group)# Step 2: 为每个运行中的请求生成一个 tokenforseq_groupinrunning:seq_group.generate_one_token()# Step 3: 把完成的请求踢出，释放块forseq_groupinrunning:ifseq_group.is_finished():self.block_manager.free(seq_group)running.remove(seq_group)# Step 4: 剩余请求继续下一轮调度returnSchedulerOutput(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 运算面前可以忽略。