大模型推理优化技术实践笔记

这次项目要部署一个 7B 参数的大模型,一开始单次推理要 30 秒,用户根本接受不了。

经过一轮优化,最后压到了 2 秒。

初始状态:简单粗暴

刚开始的部署方式很简单:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("model-name")
tokenizer = AutoTokenizer.from_pretrained("model-name")

def generate(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=100)
    return tokenizer.decode(outputs[0])

代码能跑,但问题是太慢了。

打开监控一看:

  • 模型加载时间:40 秒
  • 单次推理延迟:30 秒
  • 显存占用:14GB
  • QPS:0.03

这显然没法用。

第一步:模型量化

量化是最直接有效的优化手段,把模型从 FP16 降到 INT8,显存占用减半,计算速度也能提升。

# 之前:FP16
model = AutoModelForCausalLM.from_pretrained(
    "model-name",
    torch_dtype=torch.float16
)

# 之后:INT8
model = AutoModelForCausalLM.from_pretrained(
    "model-name",
    load_in_8bit=True
)

效果立竿见影:

  • 显存占用:14GB → 7GB
  • 推理延迟:30s → 12s
  • QPS:0.03 → 0.08

但 INT8 的精度损失还是有的,有些场景下输出质量会受影响。试了 INT4,速度确实更快,但质量下降太明显,最后还是回退到 INT8。

第二步:选择合适的推理引擎

transformers 的实现虽然方便,但性能不是最优。换到专门的推理引擎试试。

对比了几个选项:

引擎优势劣势
TensorRT性能最强,NVIDIA 优化配置复杂,学习成本高
vLLM专为 LLM 优化,易用性高生态相对较新
ONNX Runtime跨平台,社区活跃LLM 优化不够深

选了 vLLM,原因很简单:专为 LLM 设计,配置简单,性能提升明显。

from vllm import LLM, SamplingParams

llm = LLM(model="model-name", quantization="awq")
sampling_params = SamplingParams(temperature=0.7, top_p=0.95)

def generate(text):
    outputs = llm.generate([text], sampling_params)
    return outputs[0].outputs[0].text

效果:

  • 推理延迟:12s → 5s
  • QPS:0.08 → 0.2
  • 显存占用:7GB → 6GB(得益于 KV Cache 优化)

第三步:KV Cache 优化

生成过程中,每个 Token 都要重新计算 Attention,KV Cache 可以缓存之前的结果,避免重复计算。

graph TB subgraph KV Cache优化技术 A[传统 KV Cache] --> B[显存占用大] B --> C[限制批量大小] D[量化 KV Cache] --> E[显存占用减少<br/>精度轻微损失] E --> F[支持更长上下文] G[PagedAttention] --> H[分页管理<br/>显存利用率提升] H --> I[支持更大的批量大小] end style B fill:#FFB6C1,stroke:#FF0000,stroke-width:2px style F fill:#90EE90 style I fill:#87CEEB

vLLM 默认就支持 PagedAttention,相当于把 KV Cache 分页管理,显存利用率更高。

实际效果:

  • 支持的批量大小:4 → 16
  • 吞吐量提升:3x

第四步:Flash Attention

传统 Attention 的实现需要存储完整的 Attention 矩阵,空间复杂度是 O(N²)。Flash Attention 通过分块计算,把复杂度降到 O(N)。

sequenceDiagram participant User as 用户查询 participant Attn as 注意力机制 participant Mem as 显存 User->>Attn: 输入序列长度 N alt 传统注意力 Attn->>Mem: 分配 N×N 矩阵<br/>显存 O(N²) Mem-->>Attn: 分配成功 Attn->>Attn: 计算注意力 else Flash Attention Attn->>Attn: 分块计算<br/>显存 O(N) Note over Attn: 无需显式存储<br/>注意力矩阵 end Attn-->>User: 返回结果

在 vLLM 中 Flash Attention 默认开启,效果体现在长文本场景下特别明显:

  • 上下文长度 2K 时,推理速度提升 2x
  • 上下文长度 8K 时,推理速度提升 4x

第五步:动态批处理

单个请求的吞吐量总是有限,如果能同时处理多个请求,资源利用率会更高。

但问题是:不同请求的输入长度、输出长度都不一样,怎么高效批处理?

vLLM 支持连续批处理(Continuous Batching),简单说就是:

  • 新请求来了,随时加入批处理
  • 某个请求结束了,立即释放资源
  • 动态调整批量大小
llm = LLM(
    model="model-name",
    max_num_batched_tokens=4096,  # 单批次最大 Token 数
    max_num_seqs=32               # 最大并行请求数
)

实际效果:

  • 吞吐量提升:5x
  • 平均延迟:5s → 3s

把各阶段延迟放在一起看,量化、换引擎和批处理每一步都在往下压,而不是某一项单独救场。

7B 模型推理优化各阶段单次延迟对比

从 30 秒到 2 秒的跨度说明:瓶颈会随优化阶段变化,需要按实测数据逐步定位,而不是一次性堆技术。

最终效果对比

指标初始状态优化后提升
单次推理延迟30s2s15x
QPS0.030.827x
显存占用14GB6GB2.3x
支持并发数11616x

踩过的坑

坑一:INT4 精度损失太大

一开始直接上 INT4 量化,速度确实快了,但输出质量下降明显。有些简单的问题都能答错。

后来才发现,不同的层对量化的敏感度不一样。简单粗暴的全局 INT4 不可行,要么用更精细的量化策略(如 AWQ),要么回退到 INT8。

我们的做法是:INT8 作为基线,有性能瓶颈的场景再考虑其他方案。

坑二:显存不足导致 OOM

批量开到 32 之后,偶尔会 OOM。查了半天发现是上下文长度的问题:

# 错误:不同请求的上下文长度差异太大
requests = [
    generate("很长的输入文本...", max_tokens=1000),
    generate("短的输入", max_tokens=100),
]

长的请求占用了大量 KV Cache,短的请求资源不够。

解决:限制最大上下文长度,或者把长短请求分开批处理。

坑三:精度和性能的权衡

优化到一定程度之后,每一步都要在精度和性能之间权衡:

  • 量化提升了性能,但损失了精度
  • 批处理提升了吞吐,但增加了延迟
  • KV Cache 节省了计算,但增加了显存占用

没有银弹,只能根据业务场景选方案。我们的场景对实时性要求不高,更在意吞吐量,所以偏向批处理。

什么时候该优化

不是所有场景都需要这么折腾:

值得优化的场景

  • 高并发、低延迟要求(如客服机器人)
  • 成本敏感(如大规模商用部署)
  • 用户体验直接受推理速度影响

不值得优化的场景

  • 低频使用(如内部工具)
  • 预算充足(直接加硬件)
  • 对延迟不敏感(如离线生成)

写在最后

大模型推理优化这东西,理论讲再多不如实际踩一次坑。

但也不是一开始就上各种优化技术。先用简单方案跑起来,测性能,找瓶颈,再针对性优化。

很多时候,瓶颈不是你想的那样。一开始我以为瓶颈在计算,后来发现其实是显存带宽。优化思路完全变了。


这次优化花了两个月,中间走了不少弯路。但回头看,从 30 秒到 2 秒的提升,确实值得。

版权声明: 本文首发于 指尖魔法屋-大模型推理优化技术实践笔记https://blog.thinkmoon.cn/post/14-llm-inference-optimization/) 转载或引用必须申明原指尖魔法屋来源及源地址!