模型服务与推理

模型训练完成后，需要将其部署为可访问的服务。高效的模型服务需要解决显存管理、请求调度、吞吐量优化等多个挑战。本章将系统介绍模型推理的核心技术和部署实践。

推理与训练的区别

模型推理和训练在计算特性上有显著差异：

特性	训练	推理
计算模式	前向+反向传播	仅前向传播
批次大小	固定或较大	动态变化
延迟要求	相对宽松	严格（毫秒级）
显存使用	参数+梯度+优化器状态	仅参数+激活
精度要求	高（影响收敛）	可适当降低

推理服务的核心目标是：在满足延迟要求的前提下，最大化吞吐量。

推理的核心指标

指标	说明	重要性
TTFT	首个 Token 延迟	用户体验
TPS	每秒生成 Token 数	生成速度
吞吐量	每秒处理请求数	系统容量
显存利用率	GPU 显存使用率	资源效率

延迟 vs 吞吐量权衡：

提高吞吐量通常需要增大批次大小，但这会增加延迟。实际部署需要根据场景平衡两者：

交互式应用（如聊天机器人）：优先低延迟，接受较低吞吐量
批量处理（如文档摘要）：优先高吞吐量，接受较高延迟

推理优化技术

模型量化

量化是将高精度浮点数转换为低精度表示，减少模型大小和计算量。

量化类型：

类型	说明	精度损失
FP16	半精度浮点	极小
INT8	8位整数	小
INT4	4位整数	中等
GPTQ/AWQ	智能量化	小

量化方法对比：

方法	特点	适用场景
训练后量化（PTQ）	无需重训练，快速部署	资源受限场景
量化感知训练（QAT）	精度更高，需要训练	对精度要求高
GPTQ	针对 Transformer 优化	LLM 量化
AWQ	保护重要权重	LLM 高精度量化
GGUF	CPU 推理优化	边缘设备部署

示例：使用 bitsandbytes 进行量化

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    load_in_4bit=True,
    torch_dtype=torch.float16,
    device_map="auto"
)

GPTQ 量化示例：

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

# 量化配置
quantize_config = BaseQuantizeConfig(
    bits=4,              # 4位量化
    group_size=128,      # 分组大小
    desc_act=False       # 是否使用激活顺序
)

# 量化模型
model = AutoGPTQForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantize_config
)
model.quantize(calibration_data)
model.save_quantized("llama-2-7b-gptq")

KV Cache 优化

在自回归生成中，KV Cache 存储之前计算的 Key 和 Value，避免重复计算。

问题：KV Cache 随序列长度线性增长，可能占用大量显存

KV Cache 显存估算公式：

\text{KV Cache} = 2 \times \text{layers} \times \text{heads} \times \text{head\_dim} \times \text{seq\_len} \times \text{batch\_size} \times \text{bytes}

以 Llama-2-7B 为例（FP16）：

Layers: 32, Heads: 32, Head dim: 128
单个 token 的 KV Cache: $2 \times 32 \times 32 \times 128 \times 2 = 512$ KB
2048 序列长度 + 批次 32: $512 \text{KB} \times 2048 \times 32 = 32$ GB

优化方案：

PagedAttention（vLLM）：将 KV Cache 分页管理，按需分配
MQA/GQA：多查询注意力，减少 KV 数量
Sliding Window：限制注意力窗口大小

GQA（Grouped Query Attention）原理：

标准注意力：
Q: [head0, head1, head2, head3, head4, head5, head6, head7]
K: [head0, head1, head2, head3, head4, head5, head6, head7]
V: [head0, head1, head2, head3, head4, head5, head6, head7]

GQA（4组）：
Q: [head0, head1, head2, head3, head4, head5, head6, head7]
K: [head0, head0, head1, head1, head2, head2, head3, head3]
V: [head0, head0, head1, head1, head2, head2, head3, head3]

KV Cache 减少 4 倍

Flash Attention

Flash Attention 是一种高效的注意力计算方法，通过减少 HBM（高带宽内存）访问次数显著提升速度和降低显存占用。

核心原理

传统注意力计算的瓶颈在于频繁访问 HBM。以序列长度 $n$ 为例，标准注意力需要：

计算 $QK^T$ ：从 HBM 读取 $Q$ 和 $K$ ，写入 $n^2$ 大小的注意力矩阵
应用 Softmax：读取并写入 $n^2$ 大小的矩阵
与 $V$ 相乘：再次读取 $n^2$ 大小的注意力矩阵

整个过程中， $O(n^2)$ 大小的中间结果需要多次在 HBM 和计算单元之间传输，造成严重的带宽瓶颈。

Flash Attention 通过分块计算和重计算策略解决这一问题：

分块计算：将 Q、K、V 划分为小块，每次只处理能放入 SRAM 的数据
在线 Softmax：在分块过程中逐步计算 Softmax，避免存储完整的注意力矩阵
重计算：反向传播时不存储注意力权重，而是重新计算，以计算换内存

这使内存复杂度从 $O(n^2)$ 降低到 $O(n)$ ，同时减少了 90% 以上的 HBM 访问。

版本演进

版本	主要改进	硬件支持
Flash Attention 1	分块计算、在线 Softmax	Ampere (A100)、Ada (RTX 4090)
Flash Attention 2	更好的并行性、工作分区优化	Ampere、Ada、Hopper
Flash Attention 3	异步计算、FP8 支持、Warp 特化	Hopper (H100/H200)
Flash Attention 4	CuTeDSL、Blackwell 支持	Hopper、Blackwell (B200)

Flash Attention 3 的关键技术：

Flash Attention 3 专门针对 NVIDIA Hopper 架构优化，利用了三个关键硬件特性：

异步计算：Hopper 架构的 Tensor Core 和 TMA（张量内存加速器）可以异步执行，FA3 通过 warp-specialization 实现计算和内存传输的完全重叠
交错执行：将块级矩阵乘法和 Softmax 操作交错执行，最大化硬件利用率
FP8 低精度支持：利用 Hopper 的原生 FP8 支持，实现接近 1.2 PFLOPs/s 的吞吐量

性能对比（H100 GPU）：

精度	Flash Attention 2	Flash Attention 3	提升
FP16	~35% 利用率	75% 利用率	2.0×
FP8	不支持	1.2 PFLOPs/s	-

FA3 的 FP8 实现不仅更快，而且数值误差比标准 FP8 attention 降低 2.6 倍。

使用方法

PyTorch 内置支持（推荐）：

PyTorch 2.0+ 内置了 Flash Attention 支持，通过 scaled_dot_product_attention 自动选择最优实现：

import torch
import torch.nn.functional as F

# 创建输入张量
batch_size, num_heads, seq_len, head_dim = 4, 8, 1024, 64
query = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
key = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
value = torch.randn(batch_size, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)

# 自动选择最优实现（包括 Flash Attention）
output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=None,
    dropout_p=0.0,
    is_causal=True,  # 因果注意力，适用于自回归生成
    scale=None       # 默认为 1/sqrt(head_dim)
)

flash-attn 库（高级功能）：

安装：

pip install flash-attn --no-build-isolation

使用示例：

from flash_attn import flash_attn_func, flash_attn_qkvpacked_func

# 方式一：分离的 Q、K、V
# 支持 MQA/GQA（多查询注意力/分组查询注意力）
output = flash_attn_func(
    q, k, v,
    dropout_p=0.0,
    softmax_scale=None,
    causal=True,
    window_size=(-1, -1),  # 滑动窗口注意力，(-1, -1) 表示不限制
    alibi_slopes=None      # ALiBi 位置编码
)

# 方式二：打包的 QKV（更高效）
# qkv 形状：(batch, seq_len, 3, num_heads, head_dim)
output = flash_attn_qkvpacked_func(
    qkv,
    dropout_p=0.0,
    causal=True
)

KV Cache 推理优化：

Flash Attention 提供了专门的 KV Cache 接口，用于推理加速：

from flash_attn import flash_attn_with_kvcache

# 增量解码：更新 KV Cache 并计算注意力
output = flash_attn_with_kvcache(
    q,                    # 当前查询：(batch, 1, num_heads, head_dim)
    k_cache,              # 缓存的 Key
    v_cache,              # 缓存的 Value
    k=new_k,              # 新的 Key（可选，用于更新缓存）
    v=new_v,              # 新的 Value
    cache_seqlens=seq_len,  # 当前缓存长度
    causal=True,
    softmax_scale=None
)

性能对比

显存占用：

序列长度	标准注意力	Flash Attention	节省
1K	4 GB	0.4 GB	90%
2K	16 GB	0.8 GB	95%
4K	64 GB	1.6 GB	97%
8K	256 GB	3.2 GB	98%

速度提升：

GPU	序列长度 2K	序列长度 4K	序列长度 8K
A100	2.4×	3.0×	4.0×
H100	2.8×	3.5×	4.5×

最佳实践

优先使用 PyTorch 内置接口：除非需要高级功能（如滑动窗口、ALiBi），否则推荐使用 F.scaled_dot_product_attention
确保硬件兼容：Flash Attention 2 需要 Ampere 及以上架构（A100、RTX 3090/4090、H100）
使用 BF16：BF16 比 FP16 数值稳定性更好，推荐使用：

# 启用 BF16 自动混合精度
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Head Dimension 优化：Head dimension 为 64 或 128 时性能最佳，大于 192 时需要 A100/H100

连续批处理

传统批处理需要等待所有请求完成才能处理下一批，造成资源浪费。连续批处理动态调整批次：

传统批处理：
请求 A: [生成中...] [等待其他请求]
请求 B: [生成中...完成]
请求 C: [生成中...] [等待其他请求]

连续批处理：
请求 A: [生成中...]
请求 B: [生成中...完成]
请求 C: [生成中...]
        ↑ 新请求立即加入

连续批处理的优势：

指标	传统批处理	连续批处理
GPU 利用率	30-50%	80-90%
平均延迟	高	低
吞吐量	基准	2-3× 提升

投机解码

使用小模型预测多个 token，大模型验证：

小模型预测: [token1, token2, token3, token4]
大模型验证: [✓, ✓, ✓, ✗]
最终接受: [token1, token2, token3]

可以显著提升生成速度，特别是对于大模型。

投机解码的实现原理：

# 伪代码示意
def speculative_decode(draft_model, target_model, prompt, max_tokens):
    tokens = tokenize(prompt)
    
    while len(tokens) < max_tokens:
        # 小模型生成 K 个候选 token
        draft_tokens = draft_model.generate(tokens, num_tokens=K)
        
        # 大模型验证
        accepted = 0
        for i, token in enumerate(draft_tokens):
            target_prob = target_model.probability(tokens + draft_tokens[:i+1])
            draft_prob = draft_model.probability(tokens + draft_tokens[:i+1])
            
            if random() < min(1, target_prob / draft_prob):
                accepted += 1
            else:
                break
        
        tokens.extend(draft_tokens[:accepted])
    
    return tokens

投机解码的性能提升：

场景	加速比	说明
简单文本	2-3×	小模型预测准确率高
复杂推理	1.2-1.5×	预测准确率较低
代码生成	1.5-2×	代码结构可预测

vLLM

vLLM 是目前最流行的开源 LLM 推理引擎，以其高效的显存管理和高吞吐量著称。

核心技术

PagedAttention：受操作系统虚拟内存启发，将 KV Cache 分页管理：

传统方式：为每个请求预分配最大长度的连续内存
┌────────────────────────────────────────────────────┐
│ 请求1: [████████████████████████████████████████] │
│ 请求2: [████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░] │
│ 请求3: [████████████████████░░░░░░░░░░░░░░░░░░░░] │
└────────────────────────────────────────────────────┘
显存浪费严重

PagedAttention：按需分配内存块
┌────┬────┬────┬────┬────┬────┬────┬────┐
│ P1 │ P2 │ P3 │ P1 │ P4 │ P2 │ P5 │ P3 │
└────┴────┴────┴────┴────┴────┴────┴────┘
高效利用显存

安装和使用

pip install vllm

离线推理：

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-hf")
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=128)

prompts = [
    "什么是机器学习？",
    "请解释深度学习的概念。",
    "自然语言处理有哪些应用？"
]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(output.outputs[0].text)

启动 API 服务：

vllm serve meta-llama/Llama-2-7b-hf \
    --host 0.0.0.0 \
    --port 8000 \
    --tensor-parallel-size 2

调用 API：

import openai

client = openai.OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="dummy"
)

response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[{"role": "user", "content": "你好！"}],
    max_tokens=100
)
print(response.choices[0].message.content)

关键参数

参数	说明	建议值
`tensor-parallel-size`	张量并行数	GPU 数量
`gpu-memory-utilization`	GPU 显存利用率	0.9
`max-model-len`	最大序列长度	根据需求
`block-size`	内存块大小	16

TensorRT-LLM

TensorRT-LLM 是 NVIDIA 官方的高性能推理引擎，针对 NVIDIA GPU 进行了深度优化。

核心优势

极致性能：充分利用 NVIDIA 硬件特性
量化支持：INT8、INT4、FP8 等多种精度
优化内核：高度优化的 CUDA 内核

使用流程

构建引擎：将模型转换为 TensorRT 引擎

python convert_checkpoint.py \
    --model_dir /path/to/llama \
    --output_dir /path/to/checkpoint \
    --dtype float16

trtllm-build \
    --checkpoint_dir /path/to/checkpoint \
    --output_dir /path/to/engine \
    --gemm_plugin float16

运行推理：

from tensorrt_llm.runtime import ModelRunner

runner = ModelRunner.from_dir("/path/to/engine")
output = runner.generate("你好，请介绍一下自己。")
print(output)

与 Triton 集成

Triton Inference Server 提供生产级的服务能力：

# 启动 Triton Server
tritonserver --model-repository=/path/to/models

Ray Serve

Ray Serve 是 Ray 生态中的模型服务框架，支持复杂的推理流水线。

基本使用

from ray import serve
from starlette.requests import Request

@serve.deployment
class MyModelDeployment:
    def __init__(self):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.model = AutoModelForCausalLM.from_pretrained("gpt2")
        self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
    
    async def __call__(self, request: Request):
        data = await request.json()
        inputs = self.tokenizer(data["text"], return_tensors="pt")
        outputs = self.model.generate(**inputs, max_length=50)
        return self.tokenizer.decode(outputs[0])

deployment = MyModelDeployment.bind()
serve.run(deployment)

多模型组合

@serve.deployment
class Preprocessor:
    def __call__(self, text):
        return text.lower().strip()

@serve.deployment
class Model:
    def __init__(self, preprocessor):
        self.preprocessor = preprocessor
        self.model = load_model()
    
    async def __call__(self, request):
        text = await request.json()
        processed = await self.preprocessor(text)
        return self.model(processed)

deployment = Model.bind(Preprocessor.bind())
serve.run(deployment)

自动扩缩容

@serve.deployment(
    autoscaling_config={
        "min_replicas": 1,
        "max_replicas": 10,
        "target_num_ongoing_requests_per_replica": 5
    }
)
class MyModel:
    pass

Triton Inference Server

Triton 是 NVIDIA 开源的高性能推理服务器。

核心特性

多框架支持：TensorFlow、PyTorch、ONNX、TensorRT
动态批处理：自动合并请求提高吞吐
模型编排：支持多模型流水线

配置示例

config.pbtxt：

name: "llama"
backend: "tensorrtllm"
max_batch_size: 32
dynamic_batching {
    max_queue_delay_microseconds: 100
}
instance_group [
    {
        count: 2
        kind: KIND_GPU
    }
]

性能指标与监控

关键指标

指标	说明	优化方向
TTFT	首个 Token 延迟	减少模型加载、优化调度
TPS	每秒生成 Token 数	提高计算效率
吞吐量	每秒处理请求数	批处理、并行
显存利用率	GPU 显存使用率	KV Cache 管理

监控工具

Prometheus + Grafana：通用监控方案
vLLM 内置指标：/metrics 端点
NVIDIA DCGM：GPU 级别监控

实践建议

选择推理引擎

场景	推荐方案
快速部署、高吞吐	vLLM
NVIDIA GPU、极致性能	TensorRT-LLM
复杂流水线	Ray Serve
多模型、企业级	Triton

优化策略

量化：根据精度要求选择 FP16、INT8 或 INT4
批处理：使用连续批处理提高吞吐
并行：多 GPU 张量并行
缓存：利用 KV Cache 避免重复计算

成本优化

Spot 实例：使用云厂商的竞价实例降低成本
自动扩缩容：根据负载动态调整资源
模型压缩：量化、蒸馏减小模型体积

小结

模型服务是 AI 应用落地的关键环节。本章介绍了：

推理特性：理解推理与训练的区别，把握核心性能指标
优化技术：量化、KV Cache 优化、连续批处理、投机解码
推理引擎：vLLM、TensorRT-LLM 的使用方法
服务框架：Ray Serve、Triton 的部署方案

选择合适的推理引擎、优化推理性能、构建可靠的服务架构，是提供高质量 AI 服务的基础。vLLM 以其易用性和高性能成为当前最流行的选择，而 TensorRT-LLM 则在 NVIDIA GPU 上提供极致性能。

参考资料

官方文档

vLLM 官方文档 - PagedAttention 和高性能推理
TensorRT-LLM GitHub - NVIDIA 官方推理引擎
Ray Serve 文档 - 复杂推理流水线
Triton Inference Server - 企业级模型服务

技术论文

PagedAttention 论文 - vLLM 核心技术
Flash Attention 论文 - 高效注意力计算
Speculative Decoding 论文 - 投机解码原理

量化技术

GPTQ 论文 - LLM 训练后量化
AWQ 论文 - 激活感知权重量化

推理与训练的区别​

推理的核心指标​

推理优化技术​

模型量化​

KV Cache 优化​

Flash Attention​

核心原理​

版本演进​

使用方法​

性能对比​

最佳实践​

连续批处理​

投机解码​

vLLM​

核心技术​

安装和使用​

关键参数​

TensorRT-LLM​

核心优势​

使用流程​

与 Triton 集成​

Ray Serve​

基本使用​

多模型组合​

自动扩缩容​

Triton Inference Server​

核心特性​

配置示例​

性能指标与监控​

关键指标​

监控工具​

实践建议​

选择推理引擎​

优化策略​

成本优化​

小结​

参考资料​

官方文档​

技术论文​

量化技术​