# 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阶段**：新请求的首次前向传播需要特殊处理