317 lines
9.0 KiB
Markdown
317 lines
9.0 KiB
Markdown
# 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 使用（共享/回收）
|
||
```
|
||
|
||
### 逻辑页与物理块映射
|
||
|
||
```python
|
||
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)
|
||
```
|
||
|
||
## 物理块管理
|
||
|
||
### 块分配策略
|
||
|
||
```python
|
||
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"
|
||
```
|
||
|
||
### 前缀哈希与缓存
|
||
|
||
```python
|
||
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)
|
||
```
|
||
|
||
### 缓存命中检测
|
||
|
||
```python
|
||
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
|
||
|
||
```python
|
||
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结合使用
|
||
|
||
## 注意事项
|
||
|
||
1. **块大小选择**：过大导致内存浪费，过小导致管理开销
|
||
2. **块回收**：序列完成后需及时回收物理块
|
||
3. **共享块引用计数**：共享前缀时需正确维护引用计数
|
||
4. **Prefill阶段**：新请求的首次前向传播需要特殊处理 |