深入解析 PagedAttention：vLLM 显存优化的核心魔法

如果你了解过大模型推理优化的，一定听说过 PagedAttention —— 这个让 vLLM 声称吞吐量提升 2-4 倍的核心技术。它到底是如何工作的？

传统 KV Cache 的痛点

在 Transformer 推理过程中，每个 token 都需要计算注意力，这需要读取之前所有 token 的 Key 和 Value。于是产生了 KV Cache：

Token 1: KV₁
Token 2: KV₁, KV₂  
Token 3: KV₁, KV₂, KV₃
...

问题来了：

1. 内存碎片化：生成的序列长度不确定，导致显存分配碎片化
2. 内存预分配：需要预先分配可能需要的最大长度（如 4096）
3. 显存浪费：即使只生成了 100 个 token，也占着 4096 的空间

传统方案就像给每个客人预分配一间酒店房间，即使只住一晚，也要付整间的钱。

---

PagedAttention 的核心思想

PagedAttention 借鉴了操作系统的虚拟内存 + 分页管理思想：

核心创新

概念	操作系统	PagedAttention
内存单位	Page (4KB)	Block (16KB)
映射机制	页表	Block Table
内存分配	按需分配	按需分配
共享内存	写时复制	引用计数

工作原理

[物理块 0] → Token 1, 2, 3, 4
[物理块 1] → Token 5, 6, 7, 8
[物理块 2] → Token 9, 10, 11, 12
...

Block Table:
序列1: [0, 1, 2, ...]
序列2: [0, 1, -1, -1, ...]  # 共享前缀

关键优势：

• 显存按需分配，不预分配
• 不同序列可以共享物理块
• 碎片化问题迎刃而解

---

源码解析

vLLM 中 PagedAttention 的核心实现：

# vllm/attention.py (简化版)
class PagedAttention:
    def forward(self, query, key, value, block_table, seq_len):
        # 1. 逻辑分页 → 物理分页转换
        physical_blocks = block_table[seq_len // block_size]
        
        # 2. 分块计算注意力
        for i in range(0, seq_len, block_size):
            block_k = key[:, i:i+block_size]
            block_v = value[:, i:i+block_size]
            # 计算当前块的注意力
            attn_output = self._attend(query, block_k, block_v)
        
        return attn_output

---

显存对比实验

方案	序列长度	显存占用	碎片率
传统 contiguous	2048	16GB	35%
PagedAttention	2048	11GB	<5%

效果：显存减少 30%+，吞吐量提升 2-4 倍

---

多序列共享

PagedAttention 还支持前缀缓存：

# 多个对话共享 system prompt
Conversation 1: [System] → [User1] → [AI1]
Conversation 2: [System] → [User2] → [AI2]

# System prompt 的 KV Cache 完全共享！
# 物理块: [System blocks] [User1 blocks] [AI1 blocks]
#                ↓共享↓
#               [User2 blocks] [AI2 blocks]

引用计数机制：

• 块引用数为 0 时回收
• 写入时复制（Copy-on-Write）

---

总结

PagedAttention 的核心贡献：

1. 分页管理：将连续的 KV 序列分块存储
2. 按需分配：消除预分配带来的浪费
3. 共享优化：多序列共享前缀，减少重复计算
4. 碎片控制：固定大小块，避免内存碎片化

这就是 vLLM 能够支持更高并发、更大 batch 的秘密武器。下次当你使用 vLLM 获得高性能时，记得背后是操作系统的智慧在支撑！

---

下期预告：SGLang 的 RadixAttention 如何进一步优化多轮对话？