GPU 推理性能瓶颈
大模型推理的性能瓶颈主要来自三个方面:计算瓶颈——矩阵乘法的浮点运算;显存带宽瓶颈——从显存读取数据的速率;通信瓶颈——多GPU之间的数据传输。不同的优化技术针对不同的瓶颈。
FlashAttention:注意力机制的革命
FlashAttention 是近年最重要的推理优化技术之一。它通过分块计算和重计算,将注意力机制的内存访问从 O(N²) 降低到 O(N),大幅减少显存读写:
# 使用 FlashAttention(现代框架已默认启用)
from vllm import LLM
llm = LLM(
model="deepseek-ai/deepseek-llm-7b-chat",
enforce_eager=False, # 启用 CUDA Graph
enable_flash_attention=True # 启用 FlashAttention
)
# FlashAttention 效果:
# - 显存占用降低 30-50%
# - 推理速度提升 2-4x(长序列场景)
# - 支持更长的上下文量化推理:INT8/INT4 加速
量化将模型权重从 FP16 降低到 INT8 或 INT4,在几乎不损失精度的情况下大幅提升推理速度:
# AWQ 量化(推荐)
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
model_path = "deepseek-ai/deepseek-llm-7b-chat"
quant_path = "./deepseek-7b-awq"
# 量化模型
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.quantize(
tokenizer,
quant_config={
"zero_point": True,
"q_group_size": 128,
"w_bit": 4
}
)
model.save_quantized(quant_path)
# INT4 推理:速度提升 3-4x,显存降低 75%TensorRT-LLM 加速
NVIDIA TensorRT-LLM 是最高性能的推理方案之一:
# 构建 TensorRT 引擎
python TensorRT-LLM/examples/llama/build.py \
--model_dir ./deepseek-7b \
--dtype float16 \
--use_gpt_attention_plugin float16 \
--use_gemm_plugin float16 \
--max_batch_size 8 \
--max_input_len 2048 \
--max_output_len 512 \
--output_dir ./trt_engine
# 运行推理
python TensorRT-LLM/examples/run.py \
--engine_dir ./trt_engine \
--tokenizer_dir ./deepseek-7b \
--input_text "What is AI?"TensorRT-LLM 相比 vLLM 可以额外提升 20-40% 的吞吐量。
连续批处理(Continuous Batching)
传统批处理需要等所有请求完成后才能开始下一批。连续批处理允许动态添加和移除请求,大幅提升 GPU 利用率:
# vLLM 默认启用连续批处理
llm = LLM(
model="deepseek-ai/deepseek-llm-7b-chat",
max_num_seqs=256, # 最大并发请求数
max_num_batched_tokens=8192 # 最大批处理 token 数
)
# 连续批处理的好处:
# - GPU 利用率提升 50-80%
# - 请求排队时间降低 70%
# - 支持动态并发多维优化组合策略
| 场景 | 推荐组合 |
|---|---|
| 低延迟实时对话 | FlashAttention + 连续批处理 |
| 高吞吐批量处理 | AWQ量化 + TensorRT-LLM |
| 长上下文处理 | FlashAttention + PagedAttention |
| 多GPU推理 | 张量并行 + Pipeline并行 |
| 低显存设备 | INT4量化 + CPU卸载 |
监控与调优
# 监控 GPU 使用情况
nvidia-smi dmon -s pucv -d 1
# 使用 PyTorch Profiler
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
output = model.generate(input_ids)
print(prof.key_averages().table(sort_by="cuda_time_total"))总结
GPU 推理优化是一个持续的过程。建议从 FlashAttention 和量化开始(收益最大、成本最低),然后根据实际瓶颈逐步引入更高级的优化技术。记住:过早优化是万恶之源,先测量再优化。