05 - 核心架构解析
05 - 核心架构解析
深入 vLLM 引擎内部，理解其高性能推理背后的核心设计。
5.1 整体架构
vLLM 采用分层架构设计，各组件职责清晰：
┌────────────────────────────────────────────────────────────┐
│ API Layer (FastAPI) │
│ /v1/chat/completions, /v1/completions │
└────────────────────────────┬───────────────────────────────┘
│
┌────────────────────────────▼───────────────────────────────┐
│ AsyncLLMEngine │
│ (异步 LLM 引擎入口) │
└────────────────────────────┬───────────────────────────────┘
│
┌────────────────────────────▼───────────────────────────────┐
│ Scheduler │
│ 请求调度 / 批处理策略 / 抢占决策 │
│ ┌─────────────┬──────────────┐ │
│ │ Waiting Queue│ Running Queue│ │
│ └─────────────┴──────────────┘ │
└────────────────────────────┬───────────────────────────────┘
│
┌────────────────────────────▼───────────────────────────────┐
│ CacheEngine │
│ KV Cache 管理 / Block 分配 / 驱逐策略 │
│ ┌──────────────────────┐ │
│ │ GPU Block Allocator │ │
│ │ CPU Block Allocator │ │
│ └──────────────────────┘ │
└────────────────────────────┬───────────────────────────────┘
│
┌────────────────────────────▼───────────────────────────────┐
│ Worker (×N) │
│ 模型前向推理 / Attention 计算 / Token 采样 │
│ ┌────────────┬──────────────┬───────────────┐ │
│ │ Model Runner│ FlashAttention│ Sampler │ │
│ └────────────┴──────────────┴───────────────┘ │
└────────────────────────────────────────────────────────────┘
5.2 PagedAttention 深入原理
5.2.1 标准 Attention 的回顾
在标准 Transformer 的自注意力机制中，对于序列中的每个查询 token，需要计算它与所有前序 token 的注意力分数：
Attention(Q, K, V) = softmax(QK^T / √d_k) · V
其中：
Q: [seq_len, d_model] — 查询矩阵
K: [seq_len, d_model] — 键矩阵
V: [seq_len, d_model] — 值矩阵
5.2.2 KV Cache 机制
在自回归生成过程中，每生成一个新 token，只需要计算该 token 的 Query 与所有前序 token 的 Key/Value 的注意力。为避免重复计算，将前序的 K、V 缓存下来：
生成第 t 个 token 时：
Q_t: [1, d_model] — 仅当前 token 的 Query
K_1..t: [t, d_model] — 所有前序 token 的 Key（从 Cache 读取 + 新 token）
V_1..t: [t, d_model] — 所有前序 token 的 Value
注意力输出 = softmax(Q_t · K_1..t^T / √d_k) · V_1..t
5.2.3 PagedAttention 的分页机制
PagedAttention 将 KV Cache 分成固定大小的块（Block），每个块存储固定数量 token 的 K、V 向量：
Block 大小 = BLOCK_SIZE 个 token（通常 16）
每个 Block 存储：
key_cache: [num_heads, head_dim, BLOCK_SIZE]
value_cache: [num_heads, head_dim, BLOCK_SIZE]
核心数据结构：
# 简化的 Block 表设计
class BlockTable:
"""逻辑块到物理块的映射表"""
# block_table[seq_id] = [物理块ID列表]
# 例如：block_table[42] = [3, 7, 1]
# 表示序列 42 的逻辑块 0→物理块3，逻辑块1→物理块7，逻辑块2→物理块1
def __init__(self):
self.table: dict[int, list[int]] = {}
self.block_size: int = 16 # 每块存储 16 个 token
5.2.4 块分配流程
新请求到达（需要 50 个 token 的空间）：
Step 1: 计算所需块数
num_blocks = ceil(50 / 16) = 4 块
Step 2: 从空闲池分配物理块
空闲池: [0, 1, 2, 3, 4, 5, 6, 7, ...]
分配: [0, 1, 2, 3]
空闲池: [4, 5, 6, 7, ...]
Step 3: 建立映射
block_table[seq_id] = [0, 1, 2, 3]
Step 4: 逐块写入 KV Cache
写入物理块 0: token 0-15 的 K, V
写入物理块 1: token 16-31 的 K, V
写入物理块 2: token 32-47 的 K, V
写入物理块 3: token 48-49 的 K, V（部分填充）
5.2.5 注意力计算中的块寻址
# PagedAttention 核心计算过程（伪代码）
def paged_attention(
query: Tensor, # [num_seqs, num_heads, head_dim]
key_cache: Tensor, # [num_blocks, num_heads, head_dim, block_size]
value_cache: Tensor, # [num_blocks, num_heads, head_dim, block_size]
block_table: Tensor, # [num_seqs, max_blocks_per_seq]
context_lens: Tensor, # [num_seqs] — 每个序列的当前长度
) -> Tensor:
"""PagedAttention 前向计算"""
output = torch.zeros_like(query)
for i in range(num_seqs):
# 通过 block_table 获取该序列的物理块列表
blocks = block_table[i] # [物理块0, 物理块1, ...]
seq_len = context_lens[i]
q = query[i] # [num_heads, head_dim]
# 遍历每个物理块
attn_weights = []
for block_idx in range(ceil(seq_len / BLOCK_SIZE)):
physical_block = blocks[block_idx]
# 从物理块中读取 K, V
k = key_cache[physical_block] # [num_heads, head_dim, block_size]
v = value_cache[physical_block] # [num_heads, head_dim, block_size]
# 计算注意力分数
w = torch.einsum("hd,hds->hs", q, k) / sqrt(head_dim)
attn_weights.append(w)
# 拼接并 softmax
attn_weights = torch.cat(attn_weights, dim=-1)
attn_weights = softmax(attn_weights[:, :seq_len])
# 与 V 相乘（类似地遍历物理块）
# ...
return output
注意：实际实现中，PagedAttention 通过自定义 CUDA kernel 在 GPU 上高效执行，避免了逐块循环的开销。
5.3 KV Cache 内存管理
5.3.1 Block Allocator 设计
┌─────────────────────────────────────────────┐
│ Block Allocator │
│ │
│ ┌──────────────────┐ ┌─────────────────┐ │
│ │ GPU Block Pool │ │ CPU Block Pool │ │
│ │ (显存块池) │ │ (CPU 内存块池) │ │
│ │ │ │ │ │
│ │ 空闲: [5,6,7,8] │ │ 空闲: [0,1,2] │ │
│ │ 已用: [0,1,2,3,4]│ │ 已用: [] │ │
│ └──────────────────┘ └─────────────────┘ │
│ │ │ │
│ │ ┌──────────────┐ │ │
│ └──┤ Swap Space ├──────┘ │
│ │ (换入/换出) │ │
│ └──────────────┘ │
└─────────────────────────────────────────────┘
5.3.2 内存预算计算
def calculate_kv_cache_blocks(
gpu_memory_gb: float,
model_size_gb: float,
block_size: int = 16,
num_layers: int = 32,
num_kv_heads: int = 8,
head_dim: int = 128,
dtype_size: int = 2, # FP16 = 2 bytes
gpu_memory_utilization: float = 0.9,
) -> int:
"""计算可分配的 KV Cache 块数"""
# 可用 GPU 内存
total_available = gpu_memory_gb * gpu_memory_utilization
available_for_kv = total_available - model_size_gb
# 每个 block 的内存大小
# 每层需要存 key 和 value
block_mem_bytes = (
2 * # key + value
num_layers *
num_kv_heads *
head_dim *
block_size *
dtype_size
)
block_mem_gb = block_mem_bytes / (1024 ** 3)
# 可分配的 block 数量
num_blocks = int(available_for_kv / block_mem_gb)
return num_blocks
# 示例：A100 80GB，LLaMA-2 7B
blocks = calculate_kv_cache_blocks(
gpu_memory_gb=80,
model_size_gb=14, # 7B 模型 FP16 约 14GB
num_layers=32,
num_kv_heads=32,
head_dim=128,
)
print(f"可分配 KV Cache 块数: {blocks}")
print(f"可存储总 token 数: {blocks * 16}")
5.3.3 块驱逐策略
当 GPU 显存不足时，vLLM 会驱逐（eviction）部分块：
| 策略 | 说明 | 适用场景 |
|---|---|---|
| All-or-Nothing | 要么全部保留，要么全部驱逐 | 默认策略 |
| LRU | 驱逐最近最少使用的块 | 前缀缓存 |
| Swap | 将 GPU 块换出到 CPU 内存 | 内存压力大时 |
Swap 流程：
GPU 内存不足 → 选择低优先级序列 → 复制 KV Cache 到 CPU → 释放 GPU 块 → 分配给新请求
新请求完成 → 将换出的序列恢复 → 从 CPU 复制回 GPU → 继续生成
5.4 调度器（Scheduler）
5.4.1 核心职责
调度器是 vLLM 的"大脑"，负责：
- 请求排队：管理等待队列和运行队列
- 批处理组织：将多个请求组成一个前向推理批次
- 抢占决策：在资源不足时决定哪些序列暂停
- Token 预算：控制每个批次的总 token 数
5.4.2 状态机
每个请求在其生命周期中经历以下状态：
┌─────────┐
│ WAITING │ ← 请求进入等待队列
└────┬────┘
│ 调度器分配资源
▼
┌─────────┐
┌─────│ RUNNING │ ← 在 GPU 上执行推理
│ └────┬────┘
│ │
资源不足 完成/停止
│ │
▼ ▼
┌──────────┐ ┌─────────┐
│ SWAPPED │ │ FINISHED│ ← 返回结果
└──────────┘ └─────────┘
│
│ 资源恢复
▼
回到 WAITING
5.4.3 调度算法
# scheduler_simplified.py
"""简化的调度器逻辑"""
class Scheduler:
def __init__(self, max_num_seqs, max_num_batched_tokens, num_gpu_blocks):
self.waiting_queue: list[Sequence] = []
self.running_queue: list[Sequence] = []
self.swapped_queue: list[Sequence] = []
self.max_num_seqs = max_num_seqs # 最大并发序列数
self.max_num_batched_tokens = max_num_batched_tokens # 最大批 token 数
self.block_allocator = BlockAllocator(num_gpu_blocks)
def schedule(self) -> SchedulerOutput:
"""调度一批请求进行推理"""
scheduled_seqs = []
total_tokens = 0
# 1. 优先处理已在运行的序列（续生成）
for seq in self.running_queue:
if total_tokens + 1 <= self.max_num_batched_tokens:
scheduled_seqs.append(seq)
total_tokens += 1 # 每个序列生成 1 个新 token
# 2. 尝试从等待队列中加入新序列
while (self.waiting_queue and
len(scheduled_seqs) < self.max_num_seqs and
total_tokens + self._prefill_tokens(self.waiting_queue[0]) <= self.max_num_batched_tokens):
new_seq = self.waiting_queue.pop(0)
# 检查是否有足够的 GPU 块
needed_blocks = new_seq.num_required_blocks()
if not self.block_allocator.can_allocate(needed_blocks):
# 尝试抢占
if self._can_preempt(needed_blocks):
self._preempt(needed_blocks)
else:
self.waiting_queue.insert(0, new_seq)
break
self.block_allocator.allocate(new_seq)
self.running_queue.append(new_seq)
scheduled_seqs.append(new_seq)
total_tokens += new_seq.prompt_len # prefill 阶段处理所有 prompt tokens
return SchedulerOutput(scheduled_seqs=scheduled_seqs)
5.4.4 抢占策略
当 GPU 显存不足时，调度器需要抢占（preempt）低优先级的序列：
| 抢占方式 | 操作 | 恢复方式 | 性能影响 |
|---|---|---|---|
| Swap | KV Cache → CPU | Swap 回 GPU | 较小 |
| Recompute | 丢弃 KV Cache | 重新计算 | 较大（重做 prefill） |
选择抢占目标：
1. 最后加入的序列（后进先出）
2. 生成进度最少的序列
3. 优先级最低的序列
5.5 Worker 架构
5.5.1 单 Worker 流程
┌──────────────────────────────────────────┐
│ Worker │
│ │
│ ┌───────────┐ ┌───────────────────┐ │
│ │ Input Mgr │───→│ Model Runner │ │
│ │ │ │ │ │
│ │ 准备输入 │ │ 1. Embedding │ │
│ │ 整理 batch │ │ 2. Transformer │ │
│ │ │ │ 3. PagedAttention │ │
│ └───────────┘ │ 4. Sampling │ │
│ └───────┬───────────┘ │
│ │ │
│ ┌───────▼───────────┐ │
│ │ Output Mgr │ │
│ │ │ │
│ │ 收集输出 token │ │
│ │ 检查停止条件 │ │
│ │ 更新序列状态 │ │
│ └───────────────────┘ │
└──────────────────────────────────────────┘
5.5.2 多 Worker 并行
在张量并行（Tensor Parallelism）模式下，多个 Worker 协同工作：
┌──────────────┐
│ Scheduler │
└──────┬───────┘
│ 广播输入和调度决策
┌───────────────┼───────────────┐
▼ ▼ ▼
┌────────────┐ ┌────────────┐ ┌────────────┐
│ Worker 0 │ │ Worker 1 │ │ Worker 2 │
│ (GPU 0) │ │ (GPU 1) │ │ (GPU 2) │
│ │ │ │ │ │
│ 模型层 0-9 │ │ 模型层10-19│ │ 模型层20-29│
└──────┬─────┘ └──────┬─────┘ └──────┬─────┘
│ │ │
└──── AllReduce 通信 ────────────┘
│
┌──────▼───────┐
│ 最终输出 │
└──────────────┘
5.5.3 CUDA Graph 优化
vLLM 使用 CUDA Graph 捕获和重放 GPU 计算图，减少 kernel launch 开销：
# CUDA Graph 工作原理
"""
1. Warmup 阶段：
- 执行一次完整的前向推理
- 记录所有 CUDA kernel 的调用序列到 Graph 中
2. 重放阶段：
- 对于相同形状的输入，直接重放 CUDA Graph
- 跳过 Python → CUDA 的调用开销
- 显著减少小 batch 的推理延迟
3. 多 Graph 缓存：
- 不同 batch size 对应不同的 CUDA Graph
- 选择最接近的 Graph 进行重放
"""
5.6 前缀缓存（Automatic Prefix Caching）
5.6.1 工作原理
请求1: [系统提示词: 500 tokens] + [用户问题A: 50 tokens]
请求2: [系统提示词: 500 tokens] + [用户问题B: 30 tokens]
请求3: [系统提示词: 500 tokens] + [用户问题C: 80 tokens]
无前缀缓存：
KV Cache 计算量 = (500+50) + (500+30) + (500+80) = 1660 tokens
有前缀缓存：
首次计算: 500 tokens（系统提示词）→ 缓存
后续请求: 直接复用缓存
KV Cache 计算量 = 500 + 50 + 30 + 80 = 660 tokens
节省: 60%
5.6.2 哈希机制
vLLM 使用内容哈希来识别可共享的前缀：
def compute_prefix_hash(token_ids: list[int]) -> int:
"""
计算 token 序列的哈希值
相同内容的前缀会产生相同的哈希
"""
# vLLM 使用递增的哈希链
hash_value = 0
for token_id in token_ids:
hash_value = hash((hash_value, token_id))
return hash_value
5.6.3 启用前缀缓存
# 命令行启用
vllm serve model --enable-prefix-caching
# Python 启用
llm = LLM(model="model", enable_prefix_caching=True)
注意：前缀缓存对共享系统提示词的场景效果显著。如果每个请求的提示词完全不同，则无收益。
5.7 请求处理完整流程
1. 客户端发送请求
POST /v1/chat/completions
│
2. API Layer 解析请求
提取 messages, max_tokens, temperature 等
│
3. Tokenize
将文本转换为 token IDs
│
4. 创建 Sequence 对象
封装 prompt tokens, 采样参数等
│
5. 加入 Scheduler 等待队列
│
6. Scheduler 调度
检查资源，决定哪些序列参与本批次
│
7. Prefill 阶段
并行处理 prompt 的所有 token
建立初始 KV Cache
│
8. Decode 阶段（循环）
├── 每次生成 1 个 token
├── 更新 KV Cache
├── 检查停止条件（EOS, max_tokens）
└── 返回中间结果（流式）或继续
│
9. 序列完成
释放 KV Cache 块
返回最终结果
5.8 注意事项
块大小（Block Size）：默认 16。增大可减少块表大小但增加内存浪费；减小反之。一般不需调整。
内存碎片：虽然 PagedAttention 极大减少了碎片，但在长时间运行的服务中，仍需关注内存使用趋势。
CUDA Graph 兼容性：某些自定义模型可能不兼容 CUDA Graph。遇到问题时使用
--enforce-eager禁用。
多模态模型：多模态输入（如图片）的 prefill 会消耗更多 token 预算，需相应调整
max_num_batched_tokens。
5.9 扩展阅读
- Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023)
- vLLM 设计文档
- FlashAttention 论文
- Continuous Batching 论文
- vLLM 源码：scheduler.py
上一章：04 - OpenAI 兼容 API 服务 | 下一章：06 - 模型量化