9.0 KiB
9.0 KiB
Paged Attention
背景与动机
Paged Attention是vLLM（Virtual Machine for LLM Serving）提出的注意力机制改进，旨在解决LLM推理中的KV Cache内存管理问题。
传统KV Cache的问题
问题1: 预分配固定大小，连续内存占用
问题2: 不同序列长度差异大，内存浪费
问题3: 无法动态调整，序列长度受预分配限制
传统方法内存分配:
Sequence 1: [Block 0: 512 tokens] [Block 1: 512 tokens] ...
Sequence 2: [Block 0: 512 tokens] [Block 1: 512 tokens] ...
问题：短的序列仍然占用完整Block，导致内存碎片化
Paged Attention原理
核心思想
借鉴操作系统虚拟内存的页调度思想：将KV Cache分页管理，每页固定大小，按需分配，不要求连续物理内存。
物理内存布局:
Page 0: [Physical Block 0] ← Sequence 1 使用
Page 1: [Physical Block 1] ← Sequence 1 使用
Page 2: [Physical Block 2] ← Sequence 2 使用
Page 3: [Physical Block 3] ← Sequence 2 使用
Page 4: [Physical Block 0] ← Sequence 3 使用（共享/回收）
逻辑页与物理块映射
class PagedAttentionCache:
"""KV Cache的页式管理"""
def __init__(self, block_size=16, num_blocks=1000):
self.block_size = block_size # 每块token数
self.num_blocks = num_blocks
# 逻辑页→物理块映射表
# 对于序列1: logical_pages = [0, 1] → physical_blocks = [5, 8]
# 对于序列2: logical_pages = [0, 1] → physical_blocks = [3, 12]
self.block_mapping = {} # seq_id -> list of physical_block_ids
# 物理块内存池
self.memory_pool = [None] * num_blocks
def get_physical_blocks(self, seq_id):
"""获取序列的物理块列表"""
return self.block_mapping.get(seq_id, [])
def allocate_sequence(self, seq_id, num_tokens):
"""为新序列分配物理块"""
num_blocks_needed = ceil(num_tokens / self.block_size)
physical_blocks = []
for _ in range(num_blocks_needed):
block_id = self._find_free_block()
physical_blocks.append(block_id)
self.block_mapping[seq_id] = physical_blocks
return physical_blocks
def append_tokens(self, seq_id, num_new_tokens):
"""扩展序列的KV Cache"""
current_blocks = self.block_mapping[seq_id]
current_tokens = len(current_blocks) * self.block_size
if current_tokens + num_new_tokens > len(current_blocks) * self.block_size:
# 需要分配新块
new_blocks_needed = ceil((current_tokens + num_new_tokens) % self.block_size / self.block_size)
for _ in range(new_blocks_needed):
block_id = self._find_free_block()
current_blocks.append(block_id)
物理块管理
块分配策略
class BlockManager:
"""物理块管理器"""
def __init__(self, num_blocks=10000):
self.num_blocks = num_blocks
# 空闲块链表
self.free_blocks = list(range(num_blocks))
# 块引用计数（用于共享）
self.ref_count = [0] * num_blocks
# 块状态追踪
self.block_states = ['free'] * num_blocks
def allocate(self, num_blocks=1):
"""分配连续物理块"""
if len(self.free_blocks) < num_blocks:
raise OutOfMemoryError()
allocated = []
for _ in range(num_blocks):
block_id = self.free_blocks.pop()
self.block_states[block_id] = 'allocated'
allocated.append(block_id)
return allocated
def free(self, block_ids):
"""释放物理块"""
for block_id in block_ids:
self.block_states[block_id] = 'free'
self.free_blocks.append(block_id)
内存布局
GPU Memory Layout (示例):
Block 0: [KV Cache for tokens 0-15] ← Sequence A
Block 1: [KV Cache for tokens 16-31] ← Sequence A
Block 2: [KV Cache for tokens 0-15] ← Sequence B
Block 3: [KV Cache for tokens 0-15] ← Sequence C
Block 4: [KV Cache for tokens 16-31] ← Sequence C
...
Block 999: [KV Cache for tokens 496-511] ← Sequence Z
Prefix Caching机制
共享前缀
多个请求往往共享系统提示（system prompt）或用户模板的前缀：
Prompt 1: "System: You are helpful. User: What is AI?"
Prompt 2: "System: You are helpful. User: Tell me about ML"
共享前缀: "System: You are helpful. User: "
专有后缀: "What is AI?" / "Tell me about ML"
前缀哈希与缓存
import hashlib
class PrefixCache:
"""前缀缓存管理"""
def __init__(self):
self.cache = {} # hash -> (block_ids, reference_count)
self.hash_prefix = 64 # 前缀长度用于哈希
def compute_prefix_hash(self, token_ids):
"""计算前缀哈希"""
prefix = token_ids[:self.hash_prefix]
return hashlib.sha256(bytes(prefix)).hexdigest()
def lookup_prefix(self, token_ids):
"""查找匹配的前缀缓存"""
prefix_hash = self.compute_prefix_hash(token_ids)
if prefix_hash in self.cache:
return self.cache[prefix_hash]
return None
def store_prefix(self, token_ids, block_ids):
"""存储前缀的KV Cache块"""
prefix_hash = self.compute_prefix_hash(token_ids)
self.cache[prefix_hash] = (block_ids, ref_count=1)
缓存命中检测
def check_prefix_hit(prompt_tokens):
"""检查前缀缓存是否命中"""
# 1. 计算系统前缀哈希
system_prefix = extract_system_prefix(prompt_tokens)
system_hash = hash_prefix(system_prefix)
# 2. 检查缓存
cached_blocks = prefix_cache.lookup(system_hash)
if cached_blocks:
# 3. 验证完整前缀匹配
if verify_prefix_match(prompt_tokens, cached_blocks):
return cached_blocks
return None
缓存收益分析
| 场景 | 无Prefix Caching | 有Prefix Caching |
|---|---|---|
| 100个相同系统前缀的请求 | 100 × prefix_compute | 1 × prefix_compute + 99 × cache_hit |
| 前缀长度=1024 tokens | 100 × O(1024²) | 1 × O(1024²) + 99 × O(0) |
| 节省计算量 | - | ~99% |
vLLM实现
核心API
from vllm import LLM, SamplingParams
# 初始化引擎
llm = LLM(
model="meta-llama/Llama-2-7b-hf",
tensor_parallel_size=2, # 张量并行
block_size=16, # 每块16个token
max_num_seqs=256, # 最大并发序列数
max_num_batched_tokens=8192 # 最大批处理token数
)
# 推理请求
outputs = llm.generate(
prompts=["Hello, world!", "How are you?"],
sampling_params=SamplingParams(temperature=0.8, top_p=0.95)
)
配置参数
| 参数 | 说明 | 典型值 |
|---|---|---|
| block_size | 每物理块的token数 | 16 |
| max_num_seqs | 最大并发序列数 | 256 |
| max_num_batched_tokens | 每批最大token数 | 8192 |
| gpu_memory_utilization | GPU显存利用比例 | 0.9 |
内部实现流程
1. 请求入队（Request Enqueue）
└→ 解析prompt，构建Sequence对象
2. 调度（Scheduling）
└→ Continuous Batching决定当前batch
└→ 检查prefix cache命中
3. 前向传播（Forward Pass）
└→ PagedAttention计算
└→ 块管理器更新映射
4. 采样（Sampling）
└→ 取出新生成的token
└→ 检查序列是否完成
5. 输出（Output）
└→ 序列完成则输出结果
└→ 释放物理块
Paged Attention与传统Attention对比
| 特性 | 传统Attention | Paged Attention |
|---|---|---|
| 内存管理 | 连续分配 | 分页管理 |
| 内存碎片 | 高 | 低 |
| 最大序列长度 | 预分配上限 | 受总显存限制 |
| 序列扩展 | 不可变 | 动态可扩展 |
| 共享前缀 | 不支持 | 支持 |
| 实现复杂度 | 低 | 中 |
性能收益
吞吐提升
配置: LLaMA-7B, A100 80GB
| 方法 | Throughput (req/s) | GPU Memory Utilization |
|------|-------------------|------------------------|
| HuggingFace (naive) | 5.2 | 45% |
| vLLM (Paged Attention) | 24.8 | 92% |
| 提升 | ~4.8x | 2x |
内存效率
序列长度分布: [128, 256, 512, 1024, 2048]
传统方法（预分配2048）:
平均内存浪费 = 1 - 平均长度 / 2048
典型值 = 1 - 512 / 2048 = 75%
Paged Attention:
内存浪费 < 5%
与其他优化技术的关系
- Continuous Batching：Paged Attention是Continuous Batching的底层支持
- KV Cache：Paged Attention是KV Cache的物理块实现
- Prefix Caching：依赖Paged Attention的块映射机制
- Flash Attention：可与Paged Attention结合使用
注意事项
- 块大小选择：过大导致内存浪费，过小导致管理开销
- 块回收：序列完成后需及时回收物理块
- 共享块引用计数：共享前缀时需正确维护引用计数
- Prefill阶段：新请求的首次前向传播需要特殊处理