AI 推理优化系列—vLLM PagedAttention 解析：显存利用率从 40% 提升到 90% 的秘密

一、问题大模型推理的显存碎片困局在 vLLM 出现之前大模型推理的显存利用率通常只有20%-40%。这不是因为算法不够好而是因为一个朴素的问题KV Cache 的内存分配方式太粗暴了。1.1 传统 KV Cache 分配方式传统推理框架如 HuggingFace Transformers 的generate()为每个请求预分配一段连续的GPU 显存空间存放 KV Cache# 传统方式为每个请求预分配 max_seq_len 大小的连续显存 # 伪代码示意 past_key_values torch.empty( num_layers, 2, batch_size, max_seq_len, num_heads, head_dim, devicecuda )问题在于请求的实际生成长度是不确定的。你为每个请求预留了 2048 token 的空间但它可能只生成了 50 token 就结束了。预分配的剩余空间全部浪费。1.2 三种显存碎片碎片类型产生原因占比典型场景内部碎片预分配 max_seq_len 但实际只用一部分40%-60%外部碎片请求频繁创建/释放导致显存空洞10%-20%预留碎片为 beam search 等多路径预留的副本空间5%-10%三者叠加实际可用于推理的有效显存可能只占总显存的 30%-40%。1.3 为什么不能简单地按需分配因为 Transformer 的注意力计算需要随机访问整个 KV Cache。如果 KV Cache 分布在不连续的显存块中标准 CUDA kernel 无法高效处理要么性能暴跌要么需要额外的内存拷贝。这就是 PagedAttention 要解决的核心矛盾既要非连续分配消除碎片又要高效随机访问保证性能。二、PagedAttention 核心思想借鉴 OS 虚拟内存2.1 操作系统的解决方案操作系统的虚拟内存机制早就解决了类似问题物理内存被划分为固定大小的页通常 4KB进程使用的是虚拟地址连续的页表Page Table将虚拟页映射到物理页物理页可以不连续但对进程透明PagedAttention 做了一模一样的事情只是把物理内存换成了GPU 显存把进程地址空间换成了序列的逻辑 KV Cache。2.2 PagedAttention 的映射关系OS 虚拟内存 PagedAttention ───────────── ────────────── 物理内存页 (4KB) → KV Cache Block (16 tokens) 虚拟地址空间 → 序列的逻辑 KV Cache 页表 (Page Table) → Block Table 物理内存 → GPU 显存 (KV Cache Pool) 缺页中断 → Block 分配/换入换出2.3 核心数据结构python# vLLM 中的 Block 概念简化版 # 每个 Block 存储 block_size 个 token 的 KV Cache block_size 16 # 每个 block 存 16 个 token # KV Cache Pool预分配的统一显存池 # 形状: [num_blocks, num_layers, 2, block_size, num_kv_heads, head_dim] # 2 Key 和 Value kv_cache_pool torch.empty( num_blocks, num_layers, 2, block_size, num_kv_heads, head_dim, devicecuda, dtypetorch.float16 ) # Block Table每个序列维护一张表逻辑 block → 物理 block # 例: [5, 2, 8, 3] 表示该序列的 4 个逻辑块分别映射到物理块 5, 2, 8, 3 block_table [5, 2, 8, 3]关键点物理块可以不连续但通过 Block Table 的映射序列看到的 KV Cache 是逻辑连续的。三、源码走读关键模块解析以下基于 vLLM 核心架构聚焦 PagedAttention 的实现路径。版本参考 vLLM v0.6.x。3.1 整体架构┌─────────────────────────────────────────────────┐ │ vLLM Engine │ │ ┌──────────────┐ ┌──────────────────────────┐ │ │ │ │ │ BlockSpaceManager │ │ │ │ Scheduler │──│ ┌────────────────────┐ │ │ │ │ │ │ │ GPU Block Allocator│ │ │ │ │ - Waiting │ │ │ - free_blocks │ │ │ │ │ - Running │ │ │ - used_blocks │ │ │ │ │ - Swapped │ │ └────────────────────┘ │ │ │ └──────┬───────┘ └──────────────────────────┘ │ │ │ │ │ ┌──────▼───────────────────────────────────────┐│ │ │ Model Runner ││ │ │ ┌─────────────┐ ┌────────────────────────┐ ││ │ │ │ Attention │ │ KV Cache Pool (GPU) │ ││ │ │ │ Kernel │──│ [block_0][block_1]... │ ││ │ │ │ (Paged) │ │ 物理 Block 数组 │ ││ │ │ └─────────────┘ └────────────────────────┘ ││ │ └───────────────────────────────────────────────┘│ └─────────────────────────────────────────────────┘3.2 BlockSpaceManager显存的操作系统BlockSpaceManager是 PagedAttention 的大管家负责所有 Block 的分配、回收和映射。pythonclass BlockSpaceManager: 管理 GPU 和 CPU 上的 Block 分配 def __init__( self, block_size: int, # 每个 block 的 token 数通常 16 num_gpu_blocks: int, # GPU 上总 block 数 num_cpu_blocks: int, # CPU 上总 block 数用于 swap watermark: float 0.01, # 水位线低于此比例触发抢占 ): self.block_size block_size self.watermark watermark # GPU Block 分配器 self.gpu_allocator GPUBlockAllocator(num_gpu_blocks) # CPU Block 分配器用于显存不足时的换出 self.cpu_allocator CPUBlockAllocator(num_cpu_blocks) # 每个序列的 Block Table self.block_tables: Dict[int, List[int]] {} def allocate(self, seq_id: int, token_ids: List[int]) - None: 为序列分配 KV Cache Block # 计算需要多少个 block num_blocks len(token_ids) // self.block_size if len(token_ids) % self.block_size 0: num_blocks 1 # 从 free pool 中取出 block block_table [] for _ in range(num_blocks): block_number self.gpu_allocator.allocate() block_table.append(block_number) self.block_tables[seq_id] block_table # 将 token 写入对应 block self._write_tokens_to_blocks(seq_id, token_ids) def append_token(self, seq_id: int, token_id: int) - None: 序列生成新 token 时追加到 KV Cache block_table self.block_tables[seq_id] logical_pos self.get_seq_len(seq_id) # 当前序列长度 # 计算该 token 属于哪个 block 的哪个位置 block_idx logical_pos // self.block_size block_offset logical_pos % self.block_size # 如果当前 block 已满分配新 block if block_idx len(block_table): new_block self.gpu_allocator.allocate() block_table.append(new_block) # 将 token 的 KV 写入物理 block self._write_kv_to_block( block_table[block_idx], block_offset, token_id )设计精妙之处allocate()按需分配 Block不预分配整个max_seq_lenappend_token()只在需要时才分配新 Block类似 lazy allocationBlock 来自统一的 free pool没有外部碎片每个 Block 大小固定block_size16内部碎片最多浪费 15 个 token 的空间3.3 Scheduler请求调度与显存抢占Scheduler 决定哪些请求可以在当前 step 执行。当显存不足时它会抢占低优先级请求。pythonclass Scheduler: def __init__(self, config, cache_config, scheduler_config): self.block_manager BlockSpaceManager( block_sizecache_config.block_size, num_gpu_blockscache_config.num_gpu_blocks, num_cpu_blockscache_config.num_cpu_blocks, ) # 三个队列 self.waiting: Deque[SequenceGroup] deque() # 等待 Prefill self.running: Deque[SequenceGroup] deque() # 正在 Decode self.swapped: Deque[SequenceGroup] deque() # 被换出到 CPU def _schedule(self) - SchedulerOutputs: 核心调度逻辑 # 1. 优先调度 running 队列继续 decode running list(self.running) swapped list(self.swapped) if not self.swapped else [] # 2. 检查显存是否足够 while self._check_memory_pressure(): # 显存不足抢占最后加入的 running 序列 if self.running: preempted self.running.pop() # 两种抢占策略 if self.scheduler_config.preemption_mode swap: # 策略 A: 将 KV Cache 换出到 CPU self.block_manager.swap_out(preempted) self.swapped.append(preempted) else: # 策略 B: 直接丢弃重新计算Recomputation self.block_manager.free(preempted) self.waiting.appendleft(preempted) else: break # 3. 从 waiting 队列补充新请求 while self.waiting and self._can_allocate(self.waiting[0]): seq_group self.waiting.popleft() self.block_manager.allocate(seq_group) self.running.append(seq_group) return SchedulerOutputs(...)两种抢占策略对比策略原理优点缺点适用场景SwapKV Cache 拷贝到 CPU 内存恢复快拷贝回 GPU 即可需要 CPU 内存空间显存压力大但不极端Recomputation丢弃 KV Cache重新 Prefill不占 CPU 内存实现简单恢复慢要重新计算 Prefill显存极端不足3.4 PagedAttention Kernel分页注意力计算这是最核心的部分。标准的 Attention Kernel 假设 KV Cache 是连续的PagedAttention 修改了 Kernel使其能通过 Block Table 访问非连续的 KV Cache。python# PagedAttention 的核心修改后的 Attention 计算流程 # 简化版伪代码展示核心逻辑 def paged_attention( query, # [num_heads, head_dim] - 当前 token 的 Q block_table, # [num_blocks] - 该序列的 Block Table kv_cache, # [num_blocks, 2, block_size, num_kv_heads, head_dim] seq_len, # 序列当前长度 ): 核心修改Q 和 K 的点积分多步进行每次只处理一个 Block num_blocks len(block_table) output torch.zeros_like(query) max_score float(-inf) for block_idx in range(num_blocks): # 通过 Block Table 获取物理 Block 编号 physical_block block_table[block_idx] # 取出该 Block 中的 K 和 V keys kv_cache[physical_block, 0] # [block_size, num_kv_heads, head_dim] values kv_cache[physical_block, 1] # [block_size, num_kv_heads, head_dim] # 计算当前 Block 内的 attention scores scores torch.matmul(query, keys.transpose(-1, -2)) scores scores / math.sqrt(head_dim) # 处理 Block 内的 padding最后一个 Block 可能未满 block_start block_idx * block_size block_end min(block_start block_size, seq_len) valid_tokens block_end - block_start scores[valid_tokens:] float(-inf) # mask 掉 padding # 在线 Softmax分块计算的关键 block_max scores.max(dim-1).values max_score torch.maximum(max_score, block_max) # 使用 online softmax 公式更新输出 # exp(score - max) * value 的累加 exp_scores torch.exp(scores - max_score) output output * torch.exp(prev_max - max_score) output torch.matmul(exp_scores, values) # 最终归一化 output output / output.sum(dim-1, keepdimTrue) return output关键设计Online Softmax为什么不能简单地逐 Block 计算后拼接因为 Softmax 需要全局最大值。PagedAttention 使用了Online Softmax算法也叫 FlashAttention 的核心技巧允许在不预先知道全局最大值的情况下分块计算 Softmaxpython# Online Softmax 的数学原理 # 传统: softmax(x_i) exp(x_i - max(x)) / sum(exp(x_j - max(x))) # 需要先遍历一次求 max再遍历一次求 exp 和 sum # Online: 维护 running max 和 running sum # 当处理新 block 时: # new_max max(old_max, block_max) # correction exp(old_max - new_max) # 修正因子 # new_sum old_sum * correction sum(exp(block - new_max)) # new_output old_output * correction matmul(exp(block - new_max), V)这样每个 Block 的计算可以流式处理不需要等所有 Block 都算完。3.5 连续批处理PagedAttention 的红利有了 PagedAttentionvLLM 还实现了Continuous Batching连续批处理这是传统推理框架做不到的传统 Static Batching: ┌──────────────────────────────────────┐ │ Request A: [] │ ← 生成长一直占着 │ Request B: [] │ ← 早就结束了但 slot 空着 │ Request C: [] │ ← 也结束了 │ Request D: [] │ └──────────────────────────────────────┘ 时间 → B/C 结束后必须等 A 完成才能处理新请求 vLLM Continuous Batching: ┌──────────────────────────────────────┐ │ Request A: [] │ │ Request B: [] │ │ Request C: [] │ │ Request E: [] │ ← B 一结束E 立刻插入 │ Request D: [] │ │ Request F: [] │ ← C 一结束F 立刻插入 └──────────────────────────────────────┘ 时间 → 每个位置始终在处理有效请求PagedAttention 使得这种动态插入/移除成为可能因为新请求的 KV Cache 可以从 free pool 中按需分配 Block不需要预留连续空间。四、性能数据实际效果对比4.1 显存利用率对比指标传统框架 (HF Transformers)vLLM (PagedAttention)提升KV Cache 显存利用率20%-40%80%-96%2-3 倍最大并发请求数 (A100 80G, Llama-2-7B)~8~324 倍吞吐量 (tokens/s)~1,200~5,8004.8 倍4.2 不同模型规模下的显存对比以下数据基于 A100 80GB GPUbatch_size32max_seq_len2048模型KV Cache 总大小传统预分配PagedAttention 实际使用节省显存Llama-7B~1.1 GB/序列35.2 GB (32 × 1.1)8.5 GB (实际平均 512 token)26.7 GBLlama-13B~1.7 GB/序列54.4 GB12.3 GB42.1 GBLlama-70B~5.2 GB/序列166 GB (放不下)31.2 GB134.8 GB注传统方式下 Llama-70B 在 A100 80G 上 batch_size32 直接 OOM而 PagedAttention 可以正常运行。4.3 吞吐量基准测试使用 Llama-2-7BA100 80GB输入 512 token输出 128 token┌──────────────────┬────────────┬────────────┬────────────┐ │ 并发数 │ HF generate│ vLLM │ 加速比 │ ├──────────────────┼────────────┼────────────┼────────────┤ │ 1 │ 42 t/s │ 48 t/s │ 1.14x │ │ 8 │ 180 t/s │ 520 t/s │ 2.89x │ │ 16 │ 280 t/s │ 1050 t/s │ 3.75x │ │ 32 │ OOM │ 2100 t/s │ N/A │ │ 64 │ OOM │ 3800 t/s │ N/A │ │ 128 │ OOM │ 5800 t/s │ N/A │ └──────────────────┴────────────┴────────────┴────────────┘关键结论并发越高PagedAttention 的优势越明显。低并发时两者接近因为瓶颈在计算而非显存高并发时传统方案直接 OOM。五、Block Size 的选择不可忽视的调参点block_size是 PagedAttention 唯一的核心超参数直接影响性能Block Size显存利用率Kernel 效率适用场景8最高碎片最小低block 数多循环多显存极度紧张16高最优vLLM 默认通用推荐32中高block 数少循环少长序列生成64低最高极长序列4Kpython# vLLM 启动时指定 block_size from vllm import LLM llm LLM( modelmeta-llama/Llama-2-7b-chat-hf, # block_size 默认 16一般不需要改 # 只有在显存极度紧张或序列特别长时才调整 )为什么 16 是最优因为 GPU 的 shared memory 和 thread block 的访存模式在这个粒度下效率最高。太小如 4会导致 kernel launch 开销占比过高太大如 64会增加最后一个 block 的内部碎片。六、与 FlashAttention 的关系经常有人混淆 PagedAttention 和 FlashAttention它们解决的是不同层面的问题维度FlashAttentionPagedAttention解决的问题Attention 计算的 IO 瓶颈KV Cache 的显存碎片优化层面Kernel 内部的内存访问显存分配与管理策略核心技术Tiling Online Softmax分页 Block Table关系PagedAttention 的 Kernel 借鉴了 FlashAttention 的 Online Softmax 技巧可否单独使用可以只加速计算可以只优化显存组合使用vLLM v0.6 默认组合使用效果最佳一句话总结FlashAttention 让计算更快PagedAttention 让显存更省两者正交互补。七、实际部署建议7.1 什么时候用 vLLM场景推荐度原因高并发 API 服务★★★★★PagedAttention Continuous Batching 的最大优势单条长文本处理★★★☆☆低并发时优势不明显但不会更差批量离线推理★★★★☆吞吐量优势明显流式输出★★★★☆vLLM 0.4 支持流式体验好极低延迟场景★★★☆☆首 token 延迟不如 TensorRT-LLM7.2 快速启动代码pythonfrom vllm import LLM, SamplingParams # 初始化引擎 llm LLM( modelmeta-llama/Llama-2-7b-chat-hf, tensor_parallel_size1, # GPU 数量 gpu_memory_utilization0.90, # 最多使用 90% 显存 max_model_len4096, # 最大序列长度 ) # 批量推理 prompts [ 请解释什么是 PagedAttention, 写一个 Python 快速排序, RAG 系统的核心组件有哪些, ] * 20 # 60 个请求 sampling_params SamplingParams( temperature0.7, top_p0.9, max_tokens512, ) # vLLM 自动进行 Continuous Batching outputs llm.generate(prompts, sampling_params) for output in outputs: print(output.outputs[0].text)7.3 监控显存使用python# 部署时建议监控 KV Cache 使用情况 import torch def print_kv_cache_stats(llm): engine llm.llm_engine block_manager engine.block_manager gpu_allocator block_manager.gpu_allocator total gpu_allocator.num_blocks free len(gpu_allocator.free_blocks) used total - free print(fKV Cache Block 使用情况:) print(f 总 Block 数: {total}) print(f 已使用: {used} ({used/total*100:.1f}%)) print(f 空闲: {free} ({free/total*100:.1f}%)) print(f Block Size: {block_manager.block_size} tokens) print(f 总 KV Cache 显存: {total * block_manager.block_size * 2 * 4096 * 32 * 128 * 2 / 1024**3:.2f} GB)八、总结PagedAttention 的核心贡献不是发明了什么新算法而是把操作系统的成熟经验搬到了 GPU 显存管理上分页机制将 KV Cache 分为固定大小的 Block通过 Block Table 映射消除外部碎片按需分配不再预分配 max_seq_len而是生成过程中逐步分配消除内部碎片在线 Softmax修改 Attention Kernel使其能分块计算兼顾非连续访问和计算效率连续批处理基于分页的灵活分配实现请求的动态加入/移除最大化 GPU 利用率这些组合在一起让显存利用率从 40% 跳到 90%让同样一张卡能服务 4 倍的并发请求。下一篇预告我们将深入 vLLM 的 Continuous Batching 调度器源码解析它如何在不中断生成的情况下动态插入新请求以及不同抢占策略的性能 trade-off。如果觉得有帮助点个赞和收藏这是对我最大的鼓励。关注我的专栏「AI大模型大数据硬件编程」专栏每周更新大模型推理部署的深度实战内容。