前沿方向:AI Infra / 推理优化
基本信息
| 属性 | 内容 |
|---|---|
| 方向 | AI Infra / 推理优化 |
| 热度 | ★★★★☆ |
| 薪资 | 最高(40-80K+) |
| 难度 | ★★★★★ |
| 人才稀缺度 | 极高 |
核心技术
推理优化技术栈
│
├── 模型压缩
│ ├── 量化 (INT8/FP16/GPTQ/AWQ)
│ ├── 剪枝 (结构化/非结构化)
│ └── 知识蒸馏
│
├── 注意力优化
│ ├── Flash Attention
│ ├── PagedAttention (vLLM)
│ └── KV Cache优化
│
├── 推理框架
│ ├── TensorRT-LLM
│ ├── vLLM
│ ├── SGLang
│ └── lmdeploy
│
└── 系统优化
├── CUDA编程
├── Triton Kernel
├── 批处理优化
└── 分布式推理量化原理
量化方法对比
| 方法 | 精度损失 | 加速比 | 适用场景 |
|---|---|---|---|
| FP16 | 无 | 1.5-2x | 通用 |
| INT8 | 小 | 2-4x | 通用 |
| INT4 | 中 | 4-8x | 极致压缩 |
| GPTQ | 小 | 4-8x | LLM |
| AWQ | 小 | 4-8x | LLM |
量化代码示例
python
# PyTorch动态量化
model = model.eval()
quantized_model = torch.quantization.quantize_dynamic(
model, {nn.Linear}, dtype=torch.qint8
)
# HuggingFace GPTQ量化
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"model_name",
load_in_4bit=True, # bitsandbytes
torch_dtype=torch.float16
)
# AWQ量化
from awq import AutoAWQForCausalLM
model, tokenizer = AutoAWQForCausalLM.from_pretrained(
"model_name", quant_type="awq"
)Flash Attention
核心思想
标准Attention: O(N²) 显存
FlashAttention: O(N) 显存
通过分块计算 + online softmax,
避免materialize整个N×N attention matrix实现原理
python
def flash_attention(Q, K, V, block_size=64):
"""
Flash Attention简化实现
核心:分块计算 + 数值稳定
"""
B, H, N, D = Q.shape
scale = D ** -0.5
# 输出和归一化因子
O = torch.zeros_like(Q)
l = torch.zeros((B, H, N, 1))
m = torch.full((B, H, N, 1), -float('inf'))
# 分块计算
for i in range(0, N, block_size):
Q_i = Q[..., i:i+block_size, :]
m_i = torch.maximum(m[..., i:i+block_size, :], torch.sum(Q_i * K[..., :i+block_size, :], dim=-1) * scale)
# online softmax trick
exp_mi = torch.exp(m[..., i:i+block_size, :] - m_i)
exp_QK = torch.exp(Q_i * scale @ K.transpose(-2, -1) - m_i.unsqueeze(-1))
p = exp_QK / (exp_QK.sum(dim=-1, keepdim=True) + 1e-10)
# 更新输出
O[..., i:i+block_size, :] = torch.exp(m[..., i:i+block_size, :] - m_i) * O[..., i:i+block_size, :] + p @ V
return OPagedAttention (vLLM)
设计动机
问题:传统KV Cache是连续显存,碎片化严重
解决:类似操作系统的分页管理,按需分配显存核心概念
python
# vLLM的PagedAttention
from vllm import LLM, SamplingParam
llm = LLM(model="meta-llama/Llama-2-7b-chat-hf")
# 自动管理KV Cache
# 显存利用率提升2-4倍
# 吞吐量提升数倍TensorRT-LLM
优化技术
1. Kernel Fusion:合并多个kernel为一个
2. FP8量化:使用FP8计算
3. In-flight批处理:动态批处理
4. 张量并行:多卡并行使用示例
python
# TensorRT-LLM
from tensorrt_llm import LLM
model = LLM(model="llama-7b")
# 自动优化推理
# 相比HuggingFace加速5-10倍推理优化路线图
入门(2-4周):
1. 了解量化原理(INT8/FP16)
2. 使用vLLM部署模型
3. 了解KV Cache机制
进阶(1-2月):
4. 学习CUDA编程基础
5. 阅读vLLM源码
6. 实现简单量化
深入(3-6月):
7. Flash Attention原理
8. TensorRT开发
9. 自定义Kernel (Triton)就业方向
| 岗位 | 技能要求 | 薪资范围 |
|---|---|---|
| 推理优化工程师 | CUDA/TensorRT/vLLM | 45-80K |
| ML Infra工程师 | 分布式/云原生/GPU | 40-70K |
| 模型部署工程师 | 量化/优化/服务化 | 35-60K |
学习资源
开源项目
- vLLM - PagedAttention实现
- TensorRT-LLM - NVIDIA推理优化
- SGLang - 高效LLM推理框架
- lmdeploy - 书生浦语推理优化
必读资料
- Flash Attention论文 ⭐⭐⭐⭐⭐
- vLLM论文 ⭐⭐⭐⭐⭐
- CUDA编程文档
推理优化是AI Infra的核心,高端人才极度稀缺,薪资最高