vLLM 工作原理

1. 为什么理解推理引擎很重要
2. Tokenization
3. Prefill 阶段 vs Decoding 阶段
4. KV Cache
5. Block 结构与内存计算
6. Page KV Cache
7. Continuous Batching
8. GPU 显存分配策略
9. 代码

当用户向 LLM 发送一个 prompt（如 “Give me a brief history of chocolate”），在这个 prompt 到输出之间究竟发生了什么？

推理引擎的重要性

• OpenAI、ChatGPT、Gemini 等闭源引擎, 我们无法了解其内部机制
• 推理效率直接关联成本, GPU 资源消耗越少，成本越低
• **vLLM *：
  • state-of-the-art serving throughput
  • 开源且高效

技术创新:

• Paged Attention（分页注意力机制）
• Continuous Batching（连续批处理）
• Speculative Decoding（推测解码）
• Chunk Prefill（分块预填充）
• Prefix Caching（前缀缓存）

---

2. Tokenization（分词）

语言模型无法理解人类语言（英文），必须将文本转换为数值格式

分词过程

用户输入文本 → Token IDs (数值格式)

Prompt	Token 数量	Token IDs 示例
P1: “hi my name is”	5 tokens	15411, 12427, 56298
P2: “today is a beautiful summer day”	7 tokens	…
P3: “hello there”	2 tokens	…

• Tokenizer 由模型决定, 除非指定自定义 tokenizer，每个模型都有配套的 tokenizer
• 现代语言模型使用 Byte Pair Encoding (BPE) 子词分词方案
• 分词后，后续处理全部以 Token IDs 为单位，不再涉及原始文本

---

3. Prefill 阶段 vs Decoding 阶段

Token 到达 → 等待队列 → Prefill 阶段 → Decoding 阶段 → 输出

目的：为所有 prompt tokens 计算 Key 和 Value，并存储到 KV Cache 中

过程：

1. 将 prompt 的所有 token 转换为 token embeddings
2. 通过所有 Transformer 层
3. 计算 Q、K、V 矩阵
4. 计算 attention scores
5. 预测第一个新 token
6. 将计算结果存储到 KV Cache

特点：

• 处理所有 prompt tokens（批量处理）
• 只生成一个 token
• 必须完成才能进入 decoding

Decoding 阶段

目的：基于已缓存的 KV Cache，逐个生成新 token

过程：

1. 将新生成的 token 添加到序列
2. 计算新 token 的 K、V（复用已有的 K、V）
3. 预测下一个 token
4. 更新 KV Cache

特点：

• 每次只处理一个 token
• KV Cache 被复用，避免重复计算
• 循环直到达到终止条件

优先级

• Decoding 请求优先于 Prefill 请求
• 多个 prompt 同时处理时，部分可能处于 Prefill，部分可能处于 Decoding

等待队列中的请求类型

请求类型	状态	说明
Prefill	Waiting	KV Cache 尚未创建，正在预填充
Decoding	Running	KV Cache 已存在，正在生成新 token

---

4. KV Cache

为什么需要 KV Cache

问题：每次生成新 token 时，是否需要重新计算所有 K、V？

答案：不需要，已计算的 K、V 可以复用

KV Cache 的工作原理

初始状态：4 个 tokens [hi, my, name, is]
         ↓
计算 K1, V1（每层 Transformer）
         ↓
存储到 KV Cache（4×4 矩阵）
         ↓
预测新 token [Alice]
         ↓
新状态：5 个 tokens [hi, my, name, is, Alice]
         ↓
复用 KV Cache，只计算新 token 的 K、V
         ↓
预测下一个 token

• 避免重复计算所有 tokens 的 K、V
• 大幅减少计算量
• 需要额外的显存存储 KV Cache

KV Cache 是理解 vLLM 如何工作的基础构建块

---

5. Block 结构与内存计算

每个 Transformer 层都有 KV Blocks 用于存储 K、V 矩阵。

Block 的结构

┌─────────────────────────────────────┐
│  Key Matrix                          │
│  16 rows × (num_heads × head_dim) cols│
├─────────────────────────────────────┤
│  Value Matrix                        │
│  16 rows × (num_heads × head_dim) cols│
└─────────────────────────────────────┘

参数	默认值	说明
每块 tokens 数	16	最大容量
行数	16	等于 token 数量
列数	num_heads × head_dim	嵌入维度
矩阵数	2	Key + Value

每个 Block 内存 = 2 × num_tokens × num_heads × head_dim × bytes_per_parameter

参数来源（以 TinyLlama 为例）：

参数	值	来源
num_layers	22	Transformer 层数
num_kv_heads	4	使用 GQA（Group Query Attention）
head_dim	64	每个 head 的维度
bytes_per_element	2	16-bit 浮点数

计算：

每个 Block = 2 × 16 × 4 × 64 × 2 bytes
          = 32,768 bytes
          ≈ 0.034 MB

总 Block 数量计算

总 KV Cache 显存预算：10.86 GB

GPU Block 数量 = 10.86 GB ÷ 0.034 MB
              = 10.86 × 1024 ÷ 0.034
              ≈ 32,357 blocks

CPU Blocks

• 默认 CPU 交换空间：4 GB per GPU
• 用于存储某些需要 offload 到 CPU 的参数（如梯度计算时的激活状态）

CPU Block 数量 = 4 GB ÷ 0.034 MB
              ≈ 11,915 blocks

Block 分配示例，对于三个 prompts：

Prompt	Token 数	所需 Block	Block 编号
P1: “hi my name is”	5	1	Block 1
P2: “today is a beautiful summer day”	7	1	Block 2
P3: “hello there”	2	1	Block 3

• 每个 Block 可容纳 16 个 tokens
• P1、P2、P3 都 < 16 tokens，各自只需 1 个 Block
• 每个 Transformer 层都维护独立的 Block 分配
• 如果有 4 层，则总 Block 数 = 3 prompts × 4 layers = 12 blocks

不同提示词长度

如果 P1 有 18 个 tokens：

Block 1：16 tokens（填满）
Block 2：2 tokens（P1 剩余）+ P2 的一部分
Block 3：P2 剩余/P3
Block 4：P3 剩余

为什么不同 Prompts 不能共享 Block

保持不同 prompts 之间的 attention 机制完全独立，防止相互干扰。

---

6. Page KV Cache（分页注意力机制）

为什么需要 Page KV Cache

问题：如果 KV Cache 只是随机占用内存块，当一个 prompt 处理完成后，新 prompt 应该发送到哪个位置？

Page Table（页表）

类似于操作系统虚拟内存中的概念：

虚拟地址	物理地址（Block 编号）
Sequence 1	Block 1
Sequence 2	Block 2
Sequence 3	Block 3

• Page Table：存储在 CPU 上，只包含指针列表，不占用 GPU 显存
• 物理地址：实际占用 GPU 显存的 Block 编号

Free Block Queue

Layer 1: [Block 4]
Layer 2: [Block 4]
Layer 3: [Block 4]
Layer 4: [Block 4]

工作流程

1. Prompt 到达
        ↓
2. 查看 Free Block Queue
        ↓
3. 分配空闲 Block 给 Prompt
        ↓
4. 更新 Page Table
        ↓
5. Prompt 处理完成
        ↓
6. Block 归还到 Free Block Queue
        ↓
7. 立即可用于下一个 Prompt

优势

• Block 复用：处理完成的 Block 立即可用
• 高效分配：新 prompt 可以快速获得可用 Block
• 内存管理：避免内存碎片化

把 GPU 显存想象成一片土地，上面有多个"摊位"。Page Table 是土地的"登记簿"，记录每个摊位（Block）被谁使用。当一个用户离开，他的摊位立即可以被下一个用户使用。

---

7. Continuous Batching（连续批处理）

静态批处理的问题

┌─────────────┬─────────────┬─────────────┬─────────────┐
│ Sequence 1  │ Sequence 2  │ Sequence 3  │ Sequence 4  │
│ ━━━■ (停止) │ ━━━━━━━━━━━ │ ━━━━━━━━━━━ │ ━━━━━━━━━━━ │
│ 空余空间    │             │             │             │
│ 不能被其    │             │             │             │
│ 他使用      │             │             │             │
└─────────────┴─────────────┴─────────────┴─────────────┘

问题：一个 Sequence 完成后，其占用的空间闲置，无法被其他 Sequence 使用。

连续批处理

┌────────────────────────────────────────────────────┐
│              连续 Batch (Continuous Batch)           │
│  ┌────┬────┬────┐                                    │
│  │ P1 │ P2 │ P3 │  → 合并为一个 Batch 处理           │
│  └────┴────┴────┘                                    │
│   5t    7t    2t                                     │
└────────────────────────────────────────────────────┘

优势：

• 第一个 Sequence 完成后，立即释放其 Block
• 其他 Sequence 可以立即使用释放的空间
• 零闲置内存：没有未被使用的内存

连续批处理的工作原理

1. 检查总 token 数是否小于 token budget（默认 2048）
2. 将所有 prompts 的 tokens 合并为一个连续 Batch
3. 记录每个 prompt 的位置范围
4. 批量处理所有 tokens
5. 为每个 prompt 预测新 token
6. 完成后立即释放 Block

Token Budget

• 默认值：2048 tokens
• 表示一次迭代中最多处理 2048 个 tokens
• 本例中 14 tokens << 2048，所有 prompts 可同时处理

Decoding 阶段的迭代过程

Iteration 1:
  P1: [5 tokens] + [1 new] = 6 tokens
  P2: [7 tokens] + [1 new] = 8 tokens
  P3: [2 tokens] + [1 new] = 3 tokens
  
Iteration 2:
  P1: [6 tokens] + [1 new] = 7 tokens
  P2: [7 tokens] + [1 new] = 8 tokens  → 如遇到 EOS 终止
  P3: [3 tokens] + [1 new] = 4 tokens
  
Iteration N:
  达到 max_tokens(50) 或 EOS → Sequence 终止

终止条件

• EOS Token（End of Sequence）
• max_tokens 达到上限（本例设为 50）
• 到达终止条件后，Block 立即释放回 Free Block Queue

---

8. GPU 显存分配策略

GPU 显存分区

┌────────────────────────────────────────────────────────────────────┐
│                        GPU VRAM (总显存)                             │
├──────────────┬──────────────┬───────────────┬───────────────────────┤
│ Model Weights│ PiTorch Mem │  Activation  │   KV Cache            │
│   (绿色)      │  (橙色)      │  Peak Mem    │   (绿色主体)           │
│   2 GB       │   0.05 GB   │   0.3 GB     │   10.86 GB            │
└──────────────┴──────────────┴───────────────┴───────────────────────┘

各组成部分

组成部分	典型大小	说明
Model Weights	~2 GB	模型参数（以 TinyLlama 为例）
PiTorch Mem	0.05 GB	非 torch 内存（优化器状态等）
Activation Peak Mem	0.3 GB	反向传播时存储的激活状态
KV Cache Reserved	10.86 GB	Paged Attention 使用的最大空间

GPU 显存的重要性

• GPU VRAM 是最珍贵的资源
• vLLM 将大部分空间分配给 KV Cache
• 越大的 VRAM → 越多的并发请求 → 越低的单位成本

CPU-GPU 内存交换:

• 什么时候需要交换：

  • 梯度计算时
  • 反向传播时
  • 激活状态存储

• Swap Space：

  • 默认 4 GB per GPU
  • 用于临时存储需要 offload 到 CPU 的参数

---

9. 代码

运行 vLLM

# 安装 vLLM
!pip install vllm

# 定义三个 prompts
prompts = [
    "hi my name is",           # P1
    "today is a beautiful summer day",  # P2
    "hello there"              # P3
]

# 加载模型（使用 TinyLlama - 约 10 亿参数）
llm = LLM(model=" NousResearch/TinyLlama-1.1B-Chat-v1.0")

# 生成配置
output = llm.generate(
    prompts,
    SamplingParams(
        temperature=0.8,
        top_p=0.95,
        max_tokens=50
    )
)

# 打印结果
for o in output:
    print(f"Prompt: {o.prompt}")
    print(f"Completion: {o.outputs[0].text}")
    print("---")

运行输出

# 模型加载时显示的信息
# kveda blocks: 32357  → GPU 允许的最大 Block 数量
# cpu blocks: 11915    → CPU 交换空间允许的 Block 数量

# 处理时间
# Graph capture finished in 36 seconds
# Creating KV cache took around 42 seconds

---

10. 总结

Prompt 的完整旅程:

1. 到达 (Arrival)
   用户输入 prompt → 文本形式
   
2. 分词 (Tokenization)
   文本 → Token IDs (数值格式)
   
3. 等待队列 (Waiting Queue)
   进入 Prefill 或 Decoding 等待区域
   
4. Prefill 阶段
   ├─ 检查总 token 数 < token budget (2048)
   ├─ 合并为连续 Batch
   ├─ 分配 Block（每 Block 16 tokens）
   ├─ 更新 Page Table 和 Free Block Queue
   ├─ 计算 K、V 并存储到 KV Cache
   └─ 生成第一个新 token
   
5. Decoding 阶段
   ├─ 将新 token 添加到序列
   ├─ 复用已有 KV Cache
   ├─ 计算新 token 的 K、V
   ├─ 预测下一个 token
   └─ 重复直到终止条件
   
6. 释放 (Cleanup)
   达到终止条件 → Block 释放 → 加入 Free Block Queue

附录

概念	简明定义
Prefill	为所有 prompt tokens 计算 K、V 并缓存
Decoding	基于 KV Cache 逐个生成新 token
KV Cache	存储已计算的 Key/Value 矩阵，避免重复计算
Block	存储 K/V 矩阵的单元，每块 16 tokens
Page Table	CPU 上的映射表，记录哪个 Block 属于哪个 Sequence
Free Block Queue	CPU 上维护的空闲 Block 列表
Continuous Batching	动态合并多个 prompts 到一个 Batch 处理
Paged Attention	基于页表的高效 KV Cache 管理机制

vLLM 高效的原因:

1. Paged Attention：通过页表机制实现 Block 的即时复用
2. Continuous Batching：避免内存闲置，最大化 GPU 利用率
3. 合理的显存分配：将大部分显存预留给 KV Cache
4. Decoding 优先：确保用户体验（首 token 延迟低）

显存与成本的关系

更高效的显存利用 → 同时处理更多请求 → 单请求成本降低

这就是为什么近年来 AI 服务价格不断下降的原因之一。

---

GPU 附录

GPU 型号	VRAM	适用场景
T4	16 GB	小规模推理（TinyLlama 级别）
RTX A4000	16 GB	小规模推理
A100	40/80 GB	生产级推理
H100	80 GB	大规模生产部署

本例中 TinyLlama 只需 ~12 GB VRAM，T4 级别的 GPU 完全够用。