vLLM核心技术：PagedAttention机制详解
各位朋友，大家好！今天我们来深入探讨vLLM的核心技术——PagedAttention，一种旨在解决KV Cache显存碎片化与吞吐量瓶颈的创新机制。在大模型推理场景下，KV Cache（Key-Value Cache）用于存储Transformer层中Key和Value的中间表示，是影响推理速度和显存利用率的关键因素。PagedAttention通过巧妙地管理KV Cache，显著提升了大模型的推理性能。
一、KV Cache与性能瓶颈
在传统的Transformer推理中，每当处理一个新的token，都需要将Key和Value向量存储在显存中。随着序列长度的增加，KV Cache的体积也随之线性增长。对于长序列推理，KV Cache很容易占据大量的显存空间，导致OOM（Out Of Memory）错误。
此外，传统的KV Cache管理方式容易造成显存碎片化。例如，当处理不同长度的序列时，会频繁地分配和释放KV Cache空间，导致显存中出现许多不连续的小块空闲空间。这些碎片化的空间无法有效地被利用，进一步降低了显存利用率。
更重要的是，KV Cache的读取成为了性能瓶颈。在自注意力机制中，需要频繁地读取KV Cache中的向量，计算注意力权重。如果KV Cache存储在不连续的显存空间中，会导致大量的随机访存操作，从而降低推理速度。
二、PagedAttention：分页管理的KV Cache
PagedAttention的核心思想是将KV Cache划分为多个固定大小的页面（Page）。每个页面包含固定数量的Key和Value向量，例如16个或32个。这些页面在逻辑上是连续的，但在物理上可以是不连续的。
PagedAttention维护了一个页表（Page Table），用于记录每个token对应的KV Cache页面。页表将token ID映射到其对应的物理页面地址。通过页表，PagedAttention可以快速地定位到任何token的KV Cache，而无需遍历整个KV Cache空间。
2.1 页面分配与回收
PagedAttention采用动态页面分配策略。当处理一个新的token时，如果当前序列的KV Cache空间不足，PagedAttention会分配一个新的页面。如果显存中没有足够的连续空间来分配整个页面，PagedAttention会从空闲页面池中选择一个可用的页面。
当序列结束或被中断时，PagedAttention会将该序列占用的页面放回空闲页面池中，以便后续使用。通过这种方式，PagedAttention可以有效地管理KV Cache空间，避免显存碎片化。
2.2 页面共享
PagedAttention支持页面共享机制，允许多个序列共享同一个KV Cache页面。例如，在并行解码或多任务学习场景下，多个序列可能共享一部分相同的上下文信息。在这种情况下，PagedAttention可以将共享的KV Cache存储在同一个页面中，从而节省显存空间。
2.3 页面交换
当显存空间不足时，PagedAttention可以将一部分KV Cache页面交换到CPU内存中。当需要访问这些页面时，再将它们从CPU内存加载回显存。这种页面交换机制可以有效地扩展KV Cache的容量，允许处理更长的序列。
三、PagedAttention的优势
PagedAttention相比于传统的KV Cache管理方式，具有以下优势：
- 更高的显存利用率： 通过分页管理和页面共享，PagedAttention可以有效地避免显存碎片化，从而提高显存利用率。
- 更快的推理速度： 通过页表索引，PagedAttention可以快速地定位到KV Cache中的向量，减少随机访存操作，从而提高推理速度。
- 更好的可扩展性： 通过页面交换机制，PagedAttention可以扩展KV Cache的容量，支持处理更长的序列。
四、代码实现（PyTorch示例）
以下是一个简化的PagedAttention的PyTorch实现示例，用于说明其核心思想。请注意，这只是一个示例，并没有包含所有的优化和细节。
import torch
class PagedAttention:
def __init__(self, page_size, num_pages, head_dim):
self.page_size = page_size # 每个页面存储的token数量
self.num_pages = num_pages # 总共的页面数量
self.head_dim = head_dim # 每个头的维度
self.kv_cache = torch.zeros(num_pages, page_size, 2, head_dim).cuda() # (num_pages, page_size, 2, head_dim)，2表示Key和Value
self.page_table = {} # 存储每个token的页面索引，token_id: [page_index, token_index_in_page]
self.free_pages = list(range(num_pages)) # 可用的页面索引列表
self.token_counter = 0 # 用于生成唯一的token ID
def allocate_page(self):
"""分配一个空闲页面"""
if not self.free_pages:
return None # 没有空闲页面了
page_index = self.free_pages.pop(0)
return page_index
def free_page(self, page_index):
"""释放一个页面"""
self.free_pages.append(page_index)
self.free_pages.sort() # 为了方便管理，保持页面索引有序
def add_token(self, key, value):
"""添加一个token的KV到缓存"""
token_id = self.token_counter
self.token_counter += 1
# 查找是否有可用的页面
if token_id not in self.page_table:
page_index = self.allocate_page()
if page_index is None:
return None, token_id # 页面分配失败
# 找到一个空闲的页面，将token的KV存储到该页面
token_index_in_page = 0
self.kv_cache[page_index, token_index_in_page, 0] = key
self.kv_cache[page_index, token_index_in_page, 1] = value
self.page_table[token_id] = [page_index, token_index_in_page]
return token_id, page_index # 返回token_id和页面索引
return token_id, None # 如果token_id已经存在，返回token_id和None
def get_kv(self, token_ids):
"""根据token ID获取KV向量"""
keys = []
values = []
for token_id in token_ids:
page_index, token_index_in_page = self.page_table[token_id]
key = self.kv_cache[page_index, token_index_in_page, 0]
keys.append(key)
value = self.kv_cache[page_index, token_index_in_page, 1]
values.append(value)
return torch.stack(keys), torch.stack(values)
def remove_token(self, token_id):
"""移除一个token"""
if token_id in self.page_table:
page_index, token_index_in_page = self.page_table[token_id]
# 将token对应的页面设为未使用
# 这里简单地将KV置零，实际应用中可能需要更复杂的操作
self.kv_cache[page_index, token_index_in_page, :] = 0
del self.page_table[token_id]
# 判断页面是否为空，如果为空，则释放页面
is_page_empty = True
for other_token_id, (other_page_index, _) in self.page_table.items():
if other_page_index == page_index:
is_page_empty = False
break
if is_page_empty:
self.free_page(page_index)
def clear(self):
"""清空KV Cache"""
self.kv_cache.zero_()
self.page_table = {}
self.free_pages = list(range(self.num_pages))
self.token_counter = 0示例用法：
# 初始化PagedAttention
page_size = 16
num_pages = 4
head_dim = 64
paged_attention = PagedAttention(page_size, num_pages, head_dim)
# 模拟添加token
key1 = torch.randn(head_dim).cuda()
value1 = torch.randn(head_dim).cuda()
token_id1, page_index1 = paged_attention.add_token(key1, value1)
print(f"Token ID 1: {token_id1}, Page Index 1: {page_index1}")
key2 = torch.randn(head_dim).cuda()
value2 = torch.randn(head_dim).cuda()
token_id2, page_index2 = paged_attention.add_token(key2, value2)
print(f"Token ID 2: {token_id2}, Page Index 2: {page_index2}")
# 获取KV向量
token_ids = [token_id1, token_id2]
keys, values = paged_attention.get_kv(token_ids)
print(f"Keys shape: {keys.shape}, Values shape: {values.shape}")
# 移除token
paged_attention.remove_token(token_id1)
# 清空缓存
paged_attention.clear()五、PagedAttention与FlashAttention的结合
PagedAttention可以与FlashAttention等其他优化技术结合使用，进一步提升推理性能。FlashAttention通过重新排序计算顺序，减少了HBM（High Bandwidth Memory）的访问次数，从而提高了计算效率。将PagedAttention与FlashAttention结合使用，可以同时优化显存利用率和计算效率，实现更快的推理速度。
六、实验数据
以下表格展示了PagedAttention在不同模型和序列长度下的性能提升（数据为示例，实际结果可能因硬件和软件配置而异）：
| 模型 | 序列长度 | 传统KV Cache | PagedAttention | 性能提升 | 显存占用降低 |
|---|---|---|---|---|---|
| LLaMA-7B | 2048 | 10GB | 6GB | 1.5x | 40% |
| LLaMA-7B | 4096 | OOM | 8GB | N/A | N/A |
| LLaMA-13B | 2048 | 18GB | 12GB | 1.4x | 33% |
| GPT-3-175B | 2048 | OOM | 80GB | N/A | N/A |
注意： OOM 表示 Out Of Memory 错误，N/A 表示 Not Applicable。
这些数据表明，PagedAttention可以显著降低显存占用，提高推理速度，并支持处理更长的序列。
七、PagedAttention的局限性
PagedAttention虽然带来了显著的性能提升，但也存在一些局限性：
- 页面大小的选择： 页面大小的选择会影响性能。如果页面太小，会导致页表过大，增加访存开销。如果页面太大，可能会浪费显存空间。
- 页面交换的开销： 页面交换需要将KV Cache从显存移动到CPU内存，这会带来额外的开销。
- 实现复杂度： PagedAttention的实现相对复杂，需要仔细地管理页表和页面分配。
八、未来发展方向
未来，PagedAttention可以朝着以下方向发展：
- 自适应页面大小： 根据序列长度和模型大小，动态地调整页面大小，以获得最佳的性能。
- 更智能的页面交换策略： 采用更智能的页面交换策略，减少页面交换的开销。
- 硬件加速： 利用硬件加速技术，例如GPU上的页表缓存，进一步提高PagedAttention的性能。
总结一下PagedAttention带来的好处
总的来说，PagedAttention是一种创新的KV Cache管理机制，通过分页管理，页面共享和页面交换等技术，解决了KV Cache显存碎片化和吞吐量瓶颈问题，显著提高了大模型的推理性能，使得长序列推理成为可能，并提升了显存的利用效率。