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 和量化开始(收益最大、成本最低),然后根据实际瓶颈逐步引入更高级的优化技术。记住:过早优化是万恶之源,先测量再优化。