graph TD
subgraph "传统 KV Cache 分配"
A["请求 A<br/>(预分配 max_len=2048)"]
B["请求 B<br/>(预分配 max_len=2048)"]
C["请求 C<br/>(预分配 max_len=2048)"]
D["未使用空间<br/>(碎片化)"]
end
subgraph "PagedAttention"
E["Block Pool<br/>(固定大小 block)"]
E --> F["Block 0 → A[0:16]"]
E --> G["Block 1 → B[0:16]"]
E --> H["Block 2 → A[16:32]"]
E --> I["Block 3 → C[0:16]"]
E --> J["Block 4 → 空闲"]
end
第29章 开源推理引擎深度解析
第29章 开源推理引擎深度解析
推理引擎是 LLM 服务化的”最后一公里”。同一个模型,用不同的推理引擎部署,性能可以差 10 倍。理解推理引擎的内部原理,不是可选项——如果你不知道 PagedAttention 怎么工作,你就不知道你的推理服务为什么慢。
29.1 章节导入
2023 年 6 月,UC Berkeley 的 SkyComputing Lab 发表了一篇论文,介绍了一个叫 PagedAttention 的技术。这个灵感来自操作系统虚拟内存管理的想法,将 LLM 推理的吞吐量提升了 2-4 倍。
这篇论文的开源实现就是 vLLM——今天最流行的开源推理引擎。
但 vLLM 不是唯一的选择。TensorRT-LLM 在 NVIDIA GPU 上提供了极致性能,SGLang 引入了 RadixAttention 来优化多轮对话,llama.cpp 让你在 MacBook 上跑 70B 模型,Triton Inference Server 提供了企业级的模型管理。
本章将深入这五大推理引擎的内部原理,帮助你理解它们各自擅长什么、为什么擅长,以及如何选择。
29.2 vLLM 架构深度拆解
29.2.1 为什么需要 PagedAttention
标准 LLM 推理中,KV Cache 的管理是一个严重的内存碎片问题:
在传统方案中,每个请求需要预分配 max_seq_len 大小的连续内存空间。如果请求只生成了 100 个 token,但预分配了 2048 的空间,就浪费了 95% 的内存。
更严重的是,当多个请求的长度差异很大时,内存碎片化会导致可用内存充足但无法分配新请求。
29.2.2 PagedAttention 的设计
PagedAttention 借鉴了操作系统的虚拟内存分页机制:
- 物理内存按固定大小的 Block 分配(通常 16 个 token 为一个 block)
- 逻辑序列通过 Block Table 映射到物理 Block
- Block 可以不连续——逻辑上连续的 token 可以存在不连续的物理 block 中
- 空闲 Block 可以被任何请求使用
# PagedAttention 的 Block Table(简化)
class BlockTable:
def __init__(self, num_blocks, block_size=16):
self.block_size = block_size
self.num_blocks = num_blocks
# 物理存储池
self.kv_pool = KVPool(num_blocks, block_size)
# 每个序列的逻辑→物理映射
self.block_tables = {} # seq_id -> [block_ids]
def allocate(self, seq_id):
"""为新序列分配第一个 block"""
block_id = self.kv_pool.alloc()
self.block_tables[seq_id] = [block_id]
def append_token(self, seq_id, k, v):
"""追加一个 token 的 KV"""
seq_len = len(self.block_tables[seq_id]) * self.block_size
# 计算在最后一个 block 中的位置
block_idx = seq_len // self.block_size
offset = seq_len % self.block_size
if offset == 0:
# 需要新 block
new_block = self.kv_pool.alloc()
self.block_tables[seq_id].append(new_block)
# 写入物理 block
physical_block = self.block_tables[seq_id][-1]
self.kv_pool.write(physical_block, offset, k, v)
def free(self, seq_id):
"""序列完成后释放所有 block"""
for block_id in self.block_tables[seq_id]:
self.kv_pool.free(block_id)
del self.block_tables[seq_id]29.2.3 Continuous Batching 的实现
vLLM 将 PagedAttention 与 iteration-level scheduling 结合,实现了高效的 Continuous Batching:
graph TD
subgraph "Step 1"
A1["Batch: [A, B, C]"]
A2["A: 正在生成 token 5"]
A3["B: 正在生成 token 3"]
A4["C: 正在生成 token 8(即将完成)"]
end
subgraph "Step 2"
B1["C 完成,D 加入"]
B2["A: 正在生成 token 6"]
B3["B: 正在生成 token 4"]
B4["D: 正在生成 token 1(新请求)"]
end
subgraph "Step 3"
C1["B 完成,E 加入"]
C2["A: 继续..."]
C3["D: 继续..."]
C4["E: 新请求"]
end
关键在于:每个 decoding step(生成一个 token),调度器都可以: 1. 加入新请求:在当前 batch 中添加新序列 2. 移除完成的请求:将完成的序列从 batch 中移除 3. 抢占低优先级请求:当内存不足时,可以 swap 出部分请求
# vLLM Scheduler 的核心逻辑(简化)
class Scheduler:
def __init__(self, max_num_seqs, max_tokens_per_seq):
self.waiting = [] # 等待队列
self.running = [] # 正在运行的序列
self.swapped = [] # 被换出的序列
self.max_num_seqs = max_num_seqs
def schedule(self):
"""每个 decoding step 的调度决策"""
scheduled = []
# Step 1: 尝试从 waiting 加入新请求
while self.waiting and len(self.running) < self.max_num_seqs:
seq = self.waiting.pop(0)
if self._can_allocate(seq):
self.running.append(seq)
scheduled.append(("prefill", seq))
else:
# 内存不够,尝试换出
if self.running:
victim = self.running.pop()
self.swapped.append(victim)
self._swap_out(victim)
break
# Step 2: 从 swapped 恢复
while self.swapped:
seq = self.swapped.pop(0)
if self._can_swap_in(seq):
self.running.append(seq)
scheduled.append(("decode", seq))
# Step 3: 对 running 的序列做 decode
for seq in self.running:
if seq.is_finished():
self._free(seq)
else:
scheduled.append(("decode", seq))
return scheduled29.2.4 Prefix Caching
vLLM 还实现了 Prefix Caching——多个请求共享相同的 KV Cache 前缀。
这在以下场景中非常有用: - 所有请求共享同一个系统提示 - 多轮对话中共享之前的对话历史 - Few-shot 示例共享
# Prefix Caching 示意
"""
请求 A: [系统提示] + [Few-shot 示例] + [问题 1]
请求 B: [系统提示] + [Few-shot 示例] + [问题 2]
→ 共享前缀的 KV Cache 只需计算一次!
"""
# vLLM 自动识别共享前缀
vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
--enable-prefix-caching # 启用前缀缓存
# 效果:
# 第一个请求: 正常计算
# 第二个请求(共享系统提示): 跳过系统提示的 KV 计算
# → TTFT 减少 50%+(如果系统提示很长)29.3 TensorRT-LLM 优化原理
29.3.1 什么是 TensorRT-LLM
TensorRT-LLM 是 NVIDIA 开源的推理优化库,在 NVIDIA GPU 上提供了最高性能的 LLM 推理。它的核心是 Kernel 融合 和 INT8/FP8 量化。
graph LR
subgraph "标准 PyTorch 推理"
A["矩阵乘"] --> B["Bias Add"]
B --> C["LayerNorm"]
C --> D["GELU"]
D --> E["矩阵乘"]
E --> F["..."]
end
subgraph "TensorRT-LLM"
G["Fused Kernel<br/>(矩阵乘 + Bias + LN + GELU)"]
G --> H["Fused Kernel<br/>(矩阵乘 + ...)"]
end
29.3.2 Kernel 融合
PyTorch 中每个操作都是一个独立的 CUDA kernel 调用。每次 kernel 调用都需要: 1. 启动开销(~5-10 μs) 2. 从 HBM 读取输入 3. 计算 4. 写回 HBM
TensorRT-LLM 将多个操作融合到一个 kernel 中,减少了 kernel 启动开销和中间结果的内存读写:
# 标准实现(未融合)
def linear_layer_standard(x, weight, bias):
y = torch.matmul(x, weight.T) # Kernel 1: 读取 x, weight → 写入 y
y = y + bias # Kernel 2: 读取 y, bias → 写入 y
y = F.gelu(y) # Kernel 3: 读取 y → 写入 y
return y
# 融合实现(TensorRT-LLM 内部)
def linear_layer_fused(x, weight, bias):
# 单个 CUDA Kernel:矩阵乘 + Bias 加 + GELU
return fused_linear_gelu(x, weight, bias)
# 只读取一次 x, weight, bias,只写入一次结果29.3.3 SmoothQuant:INT8 量化
TensorRT-LLM 支持 SmoothQuant——一种不需要精度损失的 INT8 量化方案:
核心思想:将激活值的”难量化”部分转移到权重上。
\[ \mathbf{y} = (\mathbf{x} \cdot \text{diag}(s)) \cdot (\text{diag}(s)^{-1} \cdot \mathbf{W}) = \hat{\mathbf{x}} \cdot \hat{\mathbf{W}} \]
通过选择合适的缩放因子 \(s\),让 \(\hat{\mathbf{x}}\) 和 \(\hat{\mathbf{W}}\) 都在 INT8 的表示范围内。
# SmoothQuant 的简化实现
def smooth_quant_calibration(model, calibration_data):
"""收集激活值的统计信息并计算缩放因子"""
activation_scales = {}
# Step 1: 收集每层激活值的 max
hooks = register_activation_hooks(model, activation_scales)
model(calibration_data)
remove_hooks(hooks)
# Step 2: 收集权重的 max
weight_scales = {}
for name, param in model.named_parameters():
if 'weight' in name:
weight_scales[name] = param.abs().max(dim=1)[0]
# Step 3: 计算缩放因子
# s = (max_act^α / max_wt^(1-α))^(1/2)
alpha = 0.5 # 平衡因子
for name in activation_scales:
act_max = activation_scales[name]
wt_max = weight_scales[name]
scales = (act_max.pow(alpha) / wt_max.pow(1 - alpha)).sqrt()
scales = scales.clamp(min=1e-5)
# 应用缩放
# 不修改模型逻辑,只调整数值范围
with torch.no_grad():
# 激活值乘 s
model.get_submodule(name).weight.copy_(
model.get_submodule(name).weight * scales.view(1, -1)
)
# 权重除 s(在前一层的输出端乘 s)
prev_layer = get_previous_linear(model, name)
if prev_layer:
prev_layer.weight.copy_(
prev_layer.weight / scales.view(-1, 1)
)
return model29.3.4 TensorRT-LLM 部署
# TensorRT-LLM 的部署流程(比 vLLM 复杂)
# Step 1: 转换模型为 TensorRT 引擎
python build.py \
--model_dir meta-llama/Meta-Llama-3-8B-Instruct \
--dtype float16 \
--use_gpt_attention_plugin \
--use_gemm_plugin \
--output_dir ./trt_engines/llama-8b
# Step 2: 启动 Triton 服务
# 使用 TensorRT-LLM Backend for Triton
cd TensorRT-LLM/backends/triton
python fill_template.py \
--in_dir ./all_models/inflight_batcher_llm \
--out_dir ./triton_models \
--tensorrt_llm_model_dir ./trt_engines/llama-8b \
--tokenizer_dir meta-llama/Meta-Llama-3-8B-Instruct
# Step 3: 启动 Triton Server
docker run --gpus all \
-v ./triton_models:/models \
nvcr.io/nvidia/tritonserver:24.01-py3 \
tritonserver --model-repository=/modelsTensorRT-LLM 的代价: 1. 编译时间长:每个模型需要编译为 TensorRT 引擎,耗时 10-60 分钟 2. 灵活性低:编译后的引擎绑定了特定的序列长度、batch size 范围 3. 硬件锁定:编译后的引擎只在特定 GPU 架构上运行 4. 但性能最好:通常比 vLLM 快 20-50%
29.4 SGLang:RadixAttention 和调度
29.4.1 SGLang 的核心创新
SGLang(Structured Generation Language)来自 UC Berkeley 的同一个团队,它针对 结构化生成(如 JSON、代码、多轮对话)做了专门优化。
核心创新是 RadixAttention——用 Radix Tree(基数树)管理多请求之间的 KV Cache 共享:
graph TD
R["Root<br/>(空)"]
R --> S["System Prompt"]
S --> C1["Context A"]
S --> C2["Context B"]
C1 --> Q1["Question 1"]
C1 --> Q2["Question 2"]
Q1 --> A1["Answer 1"]
style R fill:#f9f
style S fill:#bbf
style C1 fill:#bfb
style Q1 fill:#fbb
在 RadixAttention 中,多轮对话的每一轮都复用之前的 KV Cache——因为之前的对话历史是完全相同的 prefix。
29.4.2 SGLang 的使用
import sglang as sgl
# 声明式程序
@sgl.function
def multi_turn_chat(s, question: str):
s += "You are a helpful assistant."
s += "User: " + question
s += "Assistant: " + sgl.gen("answer", max_tokens=256)
# 可以根据答案做后续逻辑
if "I don't know" in s["answer"]:
s += "User: Please think harder."
s += "Assistant: " + sgl.gen("answer2", max_tokens=256)
# 批量推理(自动优化 KV Cache 共享)
states = multi_turn_chat.run_batch([
{"question": "What is 2+2?"},
{"question": "What is 3+3?"},
{"question": "What is the meaning of life?"},
])
# 结构化输出(JSON)
@sgl.function
def extract_info(s, text: str):
s += f"Extract from: {text}"
s += sgl.gen("name", regex=r'"name":\s*"[^"]+"')
s += sgl.gen("age", regex=r'"age":\s*\d+')29.5 llama.cpp / Ollama:端侧推理
29.5.1 llama.cpp
llama.cpp 是 Georgi Gerganov 创建的纯 C/C++ LLM 推理库。它的核心设计理念是:零依赖、纯 CPU、极致优化。
graph LR
subgraph "llama.cpp 的技术栈"
A["GGUF 格式<br/>(量化模型存储)"]
B["CPU 推理<br/>(AVX2/NEON/AVX-512)"]
C["Metal 加速<br/>(Apple Silicon)"]
D["CUDA 加速<br/>(NVIDIA GPU)"]
end
A --> E["轻量推理引擎"]
B --> E
C --> E
D --> E
29.5.2 GGUF 格式
GGUF(GPT-Generated Unified Format)是 llama.cpp 专用的模型文件格式:
# GGUF 文件结构(概念性)
"""
GGUF 文件包含:
├── metadata (KV pairs)
│ ├── general.name = "LLaMA 3 8B"
│ ├── general.architecture = "llama"
│ ├── llama.context_length = 8192
│ ├── llama.embedding_length = 4096
│ ├── llama.block_count = 32
│ └── tokenizer.ggml.model = "gpt2"
├── tensor info
│ ├── token_embd.weight: [4096, 128256]
│ ├── blk.0.attn_q.weight: [4096, 4096]
│ └── ...
└── tensor data (量化后的)
"""GGUF 支持多种量化方案:
| 量化类型 | bits/weight | 压缩比 (vs FP16) | 质量损失 |
|---|---|---|---|
| F16 | 16 | 1x | 0% |
| Q8_0 | 8.5 | ~1.9x | <1% |
| Q6_K | 6.6 | ~2.4x | ~1% |
| Q5_K_M | 5.7 | ~2.8x | ~2% |
| Q4_K_M | 4.8 | ~3.3x | ~2-3% |
| Q3_K_M | 3.9 | ~4.1x | ~4-5% |
| Q2_K | 2.6 | ~6.2x | ~7-10% |
29.5.3 Ollama:让所有人都能用
Ollama 在 llama.cpp 之上封装了一层易用的 CLI 和 API:
# 一行命令运行模型
ollama run llama3.1:8b
# 相当于 llama.cpp 的:
# 1. 下载 GGUF 文件
# 2. 设置模型参数
# 3. 运行推理服务器
# 4. 提供类似 OpenAI 的 API
# Ollama API
curl http://localhost:11434/api/chat -d '{
"model": "llama3.1:8b",
"messages": [
{"role": "user", "content": "Hello!"}
],
"stream": true
}'29.5.4 Apple Silicon 上的推理
Apple Silicon(M1/M2/M3/M4)的统一内存架构使得端侧 LLM 推理非常有吸引力:
# Mac Studio M2 Ultra (192GB 统一内存) 可以运行 70B 模型!
ollama run llama3.1:70b
# 使用 MLX(Apple 官方框架)获得更好性能
pip install mlx-lm
mlx_generate --model meta-llama/Meta-Llama-3-8B-Instruct \
--prompt "Explain quantum computing" \
--max-tokens 256| 硬件 | 模型 | 量化 | 速度 (tok/s) | 备注 |
|---|---|---|---|---|
| M2 Pro 16GB | LLaMA 3 8B | Q4_K_M | ~20 | 可用 |
| M2 Max 32GB | LLaMA 3 8B | F16 | ~30 | 流畅 |
| M2 Max 64GB | LLaMA 3 70B | Q3_K_M | ~5 | 可用 |
| M2 Ultra 192GB | LLaMA 3 70B | Q4_K_M | ~15 | 流畅 |
29.6 Triton Inference Server
29.6.1 Triton 的定位
Triton 是 NVIDIA 的通用推理服务器,不限于 LLM——它可以部署 TensorFlow、PyTorch、ONNX、TensorRT 等多种框架的模型:
graph TD
subgraph "Triton 架构"
A["HTTP/gRPC 客户端"] --> B["Triton Server"]
B --> C["Model 1<br/>(TensorRT-LLM)"]
B --> D["Model 2<br/>(PyTorch)"]
B --> E["Model 3<br/>(ONNX)"]
B --> F["Model 4<br/>(TensorFlow)"]
G["Model Repository<br/>(文件系统)"] --> B
H["Metrics<br/>(Prometheus)"] <--> B
end
29.6.2 企业级特性
# Triton 的模型配置文件(config.pbtxt)
name: "llama-3-8b"
platform: "tensorrt_llm"
max_batch_size: 32
input [
{
name: "input_ids"
data_type: TYPE_INT32
dims: [-1]
}
]
output [
{
name: "output_ids"
data_type: TYPE_INT32
dims: [-1]
}
]
dynamic_batching {
preferred_batch_size: [4, 8, 16]
max_queue_delay_microseconds: 100000
preserve_ordering: true
}
instance_group [
{
kind: KIND_GPU
count: 1
gpus: [0]
}
]Triton 的核心价值: 1. 多框架支持:一个服务管理所有模型 2. 动态批处理:自动聚合请求 3. 模型版本管理:支持 A/B 测试和金丝雀发布 4. Prometheus 指标:内建监控 5. 多 GPU/多节点:水平扩展
29.7 引擎选型决策
graph TD
A["场景是什么?"] --> B{"生产环境<br/>高吞吐?"}
B -->|是| C{"只有 NVIDIA GPU?"}
C -->|是| D{"需要极致性能?"}
D -->|是| E["TensorRT-LLM + Triton"]
D -->|否| F["vLLM"]
C -->|否| G["vLLM"]
B -->|否| H{"端侧 / 本地?"}
H -->|Mac| I["Ollama / MLX"]
H -->|Linux CPU| J["llama.cpp"]
H -->|开发测试| K["Ollama"]
L{"结构化输出<br/>多轮对话?"} -->|是| M["SGLang"]
L -->|否| F
29.7.1 性能对比参考
以 LLaMA 3 8B 在 A100 80GB 上的推理性能为例:
| 引擎 | 吞吐 (tok/s) | TTFT (ms) | 显存利用 | 易用性 |
|---|---|---|---|---|
| vLLM | 3,200 | 80 | 85% | ⭐⭐⭐⭐ |
| TensorRT-LLM | 4,500 | 50 | 90% | ⭐⭐ |
| SGLang | 3,500 | 75 | 85% | ⭐⭐⭐ |
| TGI | 2,800 | 90 | 80% | ⭐⭐⭐⭐ |
| llama.cpp (CUDA) | 1,200 | 150 | 70% | ⭐⭐⭐⭐⭐ |
实用建议: 1. 默认选 vLLM:除非你有特殊需求,vLLM 是 80% 场景下的最佳选择 2. 追求极致性能 → TensorRT-LLM:但要接受复杂的编译流程 3. 多轮对话 / 结构化输出 → SGLang:RadixAttention 在这些场景有显著优势 4. 个人/开发 → Ollama:不要为开发环境搭 GPU 集群 5. 企业/多模型 → Triton:如果你需要管理多种类型的模型
29.8 小结
推理引擎的选择对 LLM 服务的成本和性能有着决定性影响。vLLM 通过 PagedAttention 让”好”的推理变得容易,TensorRT-LLM 让”最好”的推理成为可能但需要更多工程投入,llama.cpp/Ollama 让推理民主化到每个人都能用。
理解这些引擎的内部原理——PagedAttention 的分页机制、Kernel 融合的优化原理、RadixAttention 的共享策略——能帮助你在遇到性能问题时快速定位瓶颈。
29.9 延伸阅读
- Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023 — vLLM 的核心论文
- NVIDIA (2024). TensorRT-LLM Best Practices. NVIDIA 官方文档
- Zheng et al. (2024). SGLang: Efficient Execution of Structured Language Model Programs. arXiv:2312.07104
- Gerganov (2023). llama.cpp. https://github.com/ggerganov/llama.cpp
- NVIDIA Triton Documentation. https://docs.nvidia.com/deeplearning/triton-inference-server/
- Xiao et al. (2023). SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models. ICML 2023