第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 的管理是一个严重的内存碎片问题:

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

在传统方案中,每个请求需要预分配 max_seq_len 大小的连续内存空间。如果请求只生成了 100 个 token,但预分配了 2048 的空间,就浪费了 95% 的内存。

更严重的是,当多个请求的长度差异很大时,内存碎片化会导致可用内存充足但无法分配新请求。

29.2.2 PagedAttention 的设计

PagedAttention 借鉴了操作系统的虚拟内存分页机制

  1. 物理内存按固定大小的 Block 分配(通常 16 个 token 为一个 block)
  2. 逻辑序列通过 Block Table 映射到物理 Block
  3. Block 可以不连续——逻辑上连续的 token 可以存在不连续的物理 block 中
  4. 空闲 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 scheduled

29.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 model

29.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=/models
Warning

TensorRT-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% ⭐⭐⭐⭐⭐
Tip

实用建议: 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