vLLM用虚拟内存分页思想管理KV Cache

可视化演示

展示vLLM的PagedAttention如何借鉴OS虚拟内存分页思想管理KV Cache，从传统连续分配的碎片问题到分页映射、共享物理块的内存优化全流程。

KV Cache显存碎片问题的根源分析

在大语言模型（LLM）推理过程中，Transformer的每一层都需要缓存自注意力机制产生的Key和Value张量，即KV Cache。随着序列长度动态增长，KV Cache的显存占用不断膨胀。在批处理（batching）场景下，不同请求的序列长度各不相同，导致每个请求所需的KV Cache大小差异巨大。传统的连续显存分配策略要求为每个请求预分配一整块连续显存来存放KV Cache，且必须按照最大可能的序列长度来预留空间。

这种预分配策略导致了严重的显存碎片化问题。内部碎片方面，由于预分配空间通常大于实际使用的空间，大量显存被浪费在已分配但未使用的区域。外部碎片方面，不同请求释放的显存块大小不一、位置分散，难以被后续的新请求复用。研究表明，在典型的推理负载下，传统方案中60%~80%的显存被浪费在碎片和预留空间上，极大地限制了系统的并发处理能力和实际吞吐量。

显存碎片化还引发了一系列连锁问题。当显存碎片严重时，即使总可用显存充足，系统也无法为新到达的请求分配足够大的连续空间，导致请求排队等待甚至被拒绝服务。在动态批处理（dynamic batching）场景下，碎片化使得批处理效率大打折扣，GPU的计算单元无法被充分利用。此外，频繁的显存分配和释放操作本身也带来了不可忽视的CPU-GPU同步开销，进一步降低了推理延迟表现。

碎片问题的根本原因在于序列长度的不确定性与传统连续分配模式之间的矛盾。LLM推理的输入长度和输出长度在运行时才能确定，而传统显存管理器要求在分配时就确定大小。这种'静态分配应对动态需求'的范式，在内存管理领域早已被证明是低效的，这也为后来借鉴操作系统虚拟内存分页机制提供了理论基础。

虚拟内存分页机制到KV Cache的映射

操作系统中的虚拟内存分页机制是解决内存碎片问题的经典方案。其核心思想是将物理内存划分为固定大小的页框（page frame），将虚拟地址空间划分为同样大小的页（page），通过页表（page table）维护虚拟页到物理页框的映射关系。这样，一个虚拟页可以映射到任意一个空闲的物理页框，无需连续的物理内存，从根本上消除了外部碎片问题。

PagedAttention巧妙地将这一机制迁移到KV Cache的管理中。具体而言，将KV Cache的存储空间划分为固定大小的块（block），每个块可以容纳固定数量token的KV值。例如，一个block可能存储16个token在所有层、所有注意力头的Key和Value张量。逻辑上，每个请求的KV Cache仍然是一个连续的token序列，但物理上它被分散存储在多个不连续的block中。

块表（block table）是这一映射的核心数据结构，类似于操作系统中的页表。每个请求维护一个独立的块表，记录其逻辑token位置到物理block的映射。当序列增长需要新的KV存储空间时，系统从空闲block池中分配一个物理block，并在块表中添加相应的映射条目。当序列结束或token被驱逐时，对应的block被释放回空闲池。这种机制使得显存利用率接近理论最优，内部碎片仅限于最后一个block中未使用的token槽位。

从类比角度看：虚拟页面对应逻辑KV序列的固定大小分段，物理页框对应GPU显存中的block，页表对应block table，缺页中断对应需要分配新block时的处理逻辑，而TLB（Translation Lookaside Buffer）的优化思路则对应于如何加速block table的查找。这种映射并非简单复制，而是针对GPU显存的特性和注意力计算的特点进行了深度适配。

PagedAttention的注意力计算流程

PagedAttention的核心创新在于重新设计了注意力计算流程，使其能够高效地处理分散在不同物理block中的KV值。在标准注意力计算中，查询向量Q与所有Key向量做点积得到注意力分数，再用该分数对Value向量加权求和。PagedAttention需要在此基础上增加block table的间接寻址逻辑。

具体计算流程分为以下步骤：首先，对于第i个请求的第l层注意力头，从该请求的块表中获取所有物理block的地址列表。然后，遍历这些block，从每个block中读取Key向量，计算当前查询token与这些Key的点积注意力分数。接着，使用softmax归一化得到注意力权重。最后，再次遍历block，用注意力权重对对应的Value向量进行加权求和，得到最终的注意力输出。整个过程通过一次kernel调用完成，避免了多次内存拷贝。

在GPU实现层面，PagedAttention的CUDA kernel被精心设计以最大化内存访问效率。由于每个block内的KV数据是连续存储的，kernel可以利用GPU的共享内存和合并内存访问（coalesced memory access）特性。对于每个block，kernel先将Key/Value数据加载到共享内存，然后在共享内存内完成点积计算。不同block之间的计算可以流水线化执行，充分利用GPU的并行计算能力。

对于prefill阶段和decode阶段的处理策略有所不同。Prefill阶段需要一次性处理所有输入token，计算量大但只需一次前向传播，block table的查找开销被大量计算摊薄。Decode阶段每次只生成一个token，需要增量地从KV Cache中读取历史Key/Value，此时block table的高效查找对延迟影响更为显著。vLLM通过精心设计的kernel和缓存策略，将block table查找的开销控制在可接受的范围内。

写时复制与块共享机制

PagedAttention借鉴了操作系统中写时复制（Copy-on-Write, CoW）的思想，实现了KV Cache block的共享机制。在大模型推理中，多个场景需要共享KV Cache：束搜索（beam search）的多个候选共享公共前缀的KV Cache；并行采样（parallel sampling）的多个输出共享相同的prompt KV Cache；链式提示（chain-of-thought）的不同分支共享推理前缀。通过块共享，这些场景的显存占用可以大幅降低。

写时复制机制的工作流程如下：当多个序列需要共享某段KV Cache时，它们对应的块表条目指向同一个物理block，同时该block的引用计数（reference count）加1。所有共享者都可以读取该block中的KV数据。当某个序列需要修改（追加新token导致block内容变化）一个共享block时，系统首先检查该block的引用计数。如果引用计数为1，直接修改；如果引用计数大于1，则先复制该block到一个新的物理block，然后在新block上进行修改，并更新块表映射。原block的引用计数减1。

块共享机制在束搜索中的应用最为典型。假设束搜索的宽度为k，初始prompt的KV Cache只需要存储一份。随着解码的进行，k个候选各自产生不同的后续token，它们从分叉点开始各自分配独立的block，但公共前缀部分始终共享。这相比为每个候选独立分配完整的KV Cache，显存节省比例随公共前缀的长度增加而增大。在典型的束搜索场景中，显存节省可达40%~70%。

写时复制的安全性由引用计数机制保证。在任何时刻，系统都能准确追踪每个block被多少序列共享。当序列完成推理被清理时，其引用的所有block的引用计数减1，引用计数降为0的block被释放回空闲池。这种机制确保了显存资源的正确回收，避免了内存泄漏。同时，由于复制操作仅在真正需要修改时发生，避免了不必要的数据拷贝，保持了高效的显存使用。

显存管理与调度策略设计

vLLM的显存管理系统采用分层架构设计。最上层是请求调度器（Scheduler），负责决定每个调度步骤中哪些请求参与执行、是执行prefill还是decode。中间层是KV Cache管理器，维护block table和空闲block池。最底层是GPU显存分配器，负责物理block的分配和回收。这种分层设计使得各层职责清晰，便于独立优化。

调度器的核心策略包括：首先，采用先来先服务（FCFS）与最短作业优先的混合策略，在保证公平性的同时优化整体延迟。其次，实现动态批处理（continuous batching），允许不同阶段的请求（prefill和decode）在同一个batch中混合执行，最大化GPU利用率。当新请求到达时，调度器检查是否有足够的空闲block来满足该请求的最低需求（至少能处理当前prompt的block数量），如果可以则加入执行队列，否则放入等待队列。

显存水线（watermark）机制是防止OOM（Out of Memory）的关键设计。系统预留一部分显存（如5%）作为安全缓冲，当空闲block比例低于水线时，调度器暂停接收新请求，优先处理当前batch中的请求以释放block。当某些请求完成推理后，其占用的block被释放，空闲block比例回升到水线以上，系统恢复接收新请求。这种机制确保了系统在高负载下的稳定性。

此外，vLLM还实现了显存驱逐（eviction）策略。当显存压力持续较大时，系统可以选择将某些等待队列中请求的已分配block驱逐（将其KV Cache数据转移到CPU内存），为优先级更高的请求腾出空间。被驱逐的请求在重新被调度时，需要将KV Cache数据从CPU内存传回GPU显存。这种GPU-CPU协同的显存管理策略进一步提升了系统在极端负载下的表现。

PagedAttention的工程实践与局限

在工程实践中，PagedAttention的性能表现受多种因素影响。block大小（block size）是一个关键的超参数：较小的block减少内部碎片但增加block table的大小和查找开销；较大的block减少查找开销但增加内部碎片。vLLM的默认block size通常为16个token，这在大多数场景下取得了较好的平衡。实际部署中，需要根据模型大小、序列长度分布、并发量等因素进行调优。

PagedAttention与多种推理优化技术具有良好的兼容性。它可以与张量并行（tensor parallelism）结合，在多GPU上分布式存储KV Cache；可以与量化（quantization）结合，对block中的KV值进行低精度存储以减少显存占用；可以与投机解码（speculative decoding）结合，加速自回归生成过程。这种兼容性使得PagedAttention成为现代LLM推理引擎的基础组件。

然而，PagedAttention也存在一些局限性。首先，block table的间接寻址引入了额外的内存访问开销，在短序列场景下（序列长度小于一个block大小），这种开销可能抵消分页带来的收益。其次，对于需要频繁修改历史KV Cache的场景（如滑动窗口注意力的变体），分页机制增加了数据管理的复杂性。此外，block table本身也需要占用显存，在超大batch size和超长序列场景下，block table的显存开销不可忽略。

未来的发展方向包括：自适应block大小策略，根据序列长度动态调整block大小；层次化block table设计，减少超大序列场景下的block table显存开销；与硬件特性的深度结合，如利用NVIDIA H100的异步拷贝特性优化block数据的搬运；以及探索更精细的KV Cache压缩技术与分页机制的协同，进一步突破显存瓶颈。

AI 应用场景

Python 演示

class BlockTable:    def __init__(self, block_size=16, num_gpu_blocks=1024):        self.block_size = block_size        self.num_gpu_blocks = num_gpu_blocks        self.free_blocks = list(range(num_gpu_blocks))        self.tables = {}        self.ref_counts = {}     def allocate_block(self, seq_id):        if not self.free_blocks:            raise MemoryError("No free blocks available")        phys = self.free_blocks.pop(0)        if seq_id not in self.tables:            self.tables[seq_id] = []        self.tables[seq_id].append(phys)        self.ref_counts[phys] = 1        return phys     def free_block(self, seq_id, logical_idx):        if seq_id not in self.tables:            return        table = self.tables[seq_id]        if logical_idx >= len(table):            return        phys = table[logical_idx]        self.ref_counts[phys] -= 1        if self.ref_counts[phys] == 0:            self.free_blocks.append(phys)            del self.ref_counts[phys]        table[logical_idx] = -1     def copy_on_write(self, src_seq, dst_seq, logical_idx):        src_table = self.tables[src_seq]        phys = src_table[logical_idx]        if self.ref_counts[phys] <= 1:            if dst_seq not in self.tables:                self.tables[dst_seq] = list(src_table)            else:                self.tables[dst_seq].append(phys)            self.ref_counts[phys] += 1            return phys        new_phys = self.allocate_block(dst_seq)        return new_phys     def get_physical_blocks(self, seq_id):        return self.tables.get(seq_id, [])     def simulated_attention(self, seq_id, query):        blocks = self.get_physical_blocks(seq_id)        scores = []        for block_phys in blocks:            if block_phys == -1:                continue            block_score = sum(query[i] * (block_phys + 1) * 0.01 for i in range(min(len(query), 4)))            scores.append(block_score)        if not scores:            return 0.0        max_s = max(scores)        exp_scores = [__import__("math").exp(s - max_s) for s in scores]        total = sum(exp_scores)        weights = [e / total for e in exp_scores]        output = sum(w * (b + 1) * 0.1 for w, b in zip(weights, blocks) if b != -1)        return output

1. class BlockTable:定义BlockTable类，封装分页KV Cache的块表管理逻辑，包括分配、释放、写时复制和注意力模拟
2. def __init__(self, block_size=16, num_gpu_blocks=1024):初始化方法：设置block大小（每block容纳的token数）和GPU总block数量，维护空闲block列表、每个序列的块表和block引用计数
3. self.tables = {}allocate_block方法：为指定序列分配一个物理block，从空闲列表中取出一个block并追加到该序列的块表中
4. if seq_id not in self.tables:free_block方法：释放指定序列的逻辑block，将引用计数减1，当引用计数降为0时回收至空闲池
5. returncopy_on_write方法：实现写时复制语义，当引用计数为1时直接共享物理block，否则分配新block并复制数据
6. else:get_physical_blocks方法：获取指定序列的所有物理block地址列表，即该序列的块表内容
7. return physsimulated_attention方法：模拟PagedAttention的注意力计算，通过块表间接寻址访问分散的物理block中的KV数据
8. def get_physical_blocks(self, seq_id):遍历每个物理block计算注意力分数，模拟从非连续物理block中读取Key并与query做点积的过程
9. scores = []使用softmax归一化注意力分数，先减去最大值保证数值稳定性，再计算指数和归一化权重
10. continue使用归一化后的注意力权重对Value进行加权求和，Value值通过物理block地址模拟，最终返回注意力输出

常见问答

常见误区

参考来源

• webpagevLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention
• webpageEfficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023)
• webpagevLLM GitHub Repository
• webpageOperating System Concepts - Virtual Memory Paging
• webpageFlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness