外观
计算图优化
概念定义：从算子到图
深度学习模型本质是算子（operator）的有向无环图（DAG）——每个节点是 matmul/conv/relu 等，边是张量数据流。计算图优化是把这张 DAG 重写、简化、调度的过程，让推理引擎在执行前先把"该算的提前算掉、该合的合、该删的删"。
理解图优化的两个关键认知：
- 图优化是"宏观层面"的优化——比算子融合更上层，覆盖算子选择、布局转换、内存复用、控制流简化；
- LLM 推理的图优化不同于传统 CV 模型——LLM 序列长度动态、有自回归控制流、KV cache 需要跨 step 复用，传统编译器栈（如 ONNX Runtime）很多优化用不上。
LLM 推理引擎（vLLM、TensorRT-LLM、llama.cpp）都有一套图优化 pass，核心包括：常量折叠、死代码消除、算子融合、CUDA graph capture、prefix caching。
一、计算图 IR：编译器的中枢
计算图优化要有个中间表示（IR）——既能让用户写模型，又能让编译器操作。主流 IR：
| IR | 出品 | 特点 | 典型使用 |
|---|---|---|---|
| ONNX | 微软+Facebook | 中立格式、跨框架 | 模型导出/交换 |
| MLIR | LLVM | 多层级、可扩展 | 编译器研究栈 |
| XLA HLO | XLA 编译器 IR | JAX/TF | |
| TorchScript / FX / Export | PyTorch | PyTorch 内部图 | PyTorch 生态 |
| TVM Relax | Apache TVM | LLM 推理专用 IR | 研究级推理引擎 |
| TensorRT network | NVIDIA | 闭源但高性能 | TensorRT 推理 |
| PaddlePaddle IR | 百度 | Paddle 内部 | Paddle 生态 |
ONNX 与 PyTorch IR 的角色差异
- ONNX 是"交换格式"——HuggingFace 模型导出 ONNX 后跨引擎用（ONNX Runtime、TensorRT 都支持 ONNX 导入）；
- PyTorch IR 是"工作格式"——训练时是 dynamic graph，部署时通过
torch.export转 static graph 给 Inductor 编译。 LLM 推理常见路径：HF 模型 →optimum导出 ONNX → ONNX Runtime 部署；或 HF 模型 →torch.compile→ Inductor 优化 → 部署。
二、经典图优化 Pass
1. 常量折叠（Constant Folding）
图里有些节点是"输入固定 + 权重固定 → 输出固定"，可以离线算掉：
text
原本: x_fixed → Multiply → Linear → ... （每次推理都算 Multiply+Linear）
优化: 直接用算好的常量替换 → ConstValue → ...例：RoPE 的 cos/sin 表是常量，可以离线预算并固化进图。
2. 死代码消除（Dead Code Elimination, DCE）
图里有些节点对最终输出无影响（如未使用的残差路径），直接删除：
text
原本: x → LayerA → LayerB → out
↓
LayerC （没接下游）
优化: x → LayerA → LayerB → out （LayerC 删了）3. 算子融合（Operator Fusion）
详见算子融合与自定义核。图层面的融合 pass 把"可以合成一个 kernel"的算子组合标记出来，交给后端（Inductor、Triton、CUDA）生成融合 kernel。
4. 布局转换 + 布局传播（Layout Propagation）
不同硬件、不同算子偏好的内存布局不同：
text
NCHW（PyTorch 默认） vs NHWC（Tensor Core 推荐） vs channels-last图优化 pass 决定在哪插入 transpose——尽量少插，且把 transpose 融合进相邻 matmul 的 epilogue。layout propagation 是 XLA/Inductor 的核心 pass。
5. 内存复用（Memory Planning）
图优化 pass 分析每个张量的生命周期，复用同一块显存：
text
原本: Tensor1 占 1GB，Tensor2 占 1GB → 共 2GB
优化: Tensor1 和 Tensor2 不重叠生命周期 → 共用 1GBLLM 推理的 KV cache 管理、vLLM PagedAttention 的分页内存管理，是这一思路的极致版。
三、LLM 推理专属图优化
1. CUDA graph capture
CUDA graph（NVIDIA 2017 引入）让 GPU 操作可"录制"后重放，省去每次 kernel launch 开销。LLM 推理每生成一个 token 要 launch 几十个 kernel，launch 开销可达 10-20%。
text
普通执行: CPU → launch kernel1 → kernel2 → kernel3 → ... → kernelN
↑ 每次有 launch 开销（us 级）
CUDA graph: 录制一次 → 之后重放整个图，CPU 只发一次启动- vLLM：
use_cuda_graph=True默认开启； - TensorRT-LLM：核心特性；
- 局限：CUDA graph 要固定 shape → decode 阶段每个 batch size 单独 capture。
CUDA graph 与动态 shape 的矛盾
CUDA graph 要求录制时 shape 固定，重放时 shape 不变。但 LLM 推理 batch size 会变（continuous batching）：
- 解决方案 1：按 batch size 分桶（capture 1/2/4/8/16/...），运行时选最近的桶；
- 解决方案 2：padding 到下一个桶大小（浪费算力但保持 graph 复用）；
- 解决方案 3：动态图（CUDA graph 的 conditional node，H100 起支持）。
vLLM 选 1+2 的组合，TensorRT-LLM 选 1。
2. Prefix caching
相同 prompt 前缀的多个请求，复用已算好的 KV cache：
text
原本: 请求1 [system_prompt + user1_msg] → 算整个 KV cache
请求2 [system_prompt + user2_msg] → 又算整个 KV cache
优化: 把 system_prompt 的 KV cache 缓存，请求1/2 都复用- vLLM：
enable_prefix_caching=True默认开启； - SGLang：内置 RadixAttention 实现高效 prefix 树；
- 收益：长 system prompt（4K+ tokens）场景，TTFT 减少 50-90%。
3. Speculative decoding 的图优化
投机采样（speculative decoding）让小模型先生成候选 token，大模型批量验证——详见投机采样案例研究。图优化角度：
- 小模型与大模型并行执行，要 graph 化 batch verify；
- tree decoding（多 token 树形猜测）要支持 attention mask 的稀疏模式 → FlashInfer 这种 kernel 支持。
4. Weight prepacking
把权重的存储布局在加载时一次性转换成硬件友好格式：
- INT4 权重的 unpack 与重排；
- BMM 权重的转置（让 matmul 走 cuBLAS 最快路径）；
- KV cache 的分页布局预分配。
这些都在"图加载时"完成，"推理时"零开销。
四、静态 shape vs 动态 shape
传统 DNN
CV 模型（ResNet、ViT）shape 完全静态——输入 224×224、batch 32 全程固定。图优化可以做大量静态推理：
- 常量折叠能算死所有形状相关常量；
- 算子选择走 Tensor Core 的最快版本；
- 内存规划精确到 byte。
LLM 的困境
LLM 序列长度动态：
- prompt 长度从 100 到 8000 不等；
- decode 每生成一个 token 长度 +1；
- KV cache 要随生成动态扩展。
解决路径：
| 方案 | 做法 | 代价 |
|---|---|---|
| padded shape | 把所有序列 pad 到 max_len | 浪费算力 |
| mask-based | 真实长度算，attention mask 屏蔽 padding | 算力省但 kernel 复杂 |
| bucketed shape | 按序列长度分桶，每桶一个 CUDA graph | 平衡，工业主流 |
| dynamic shape | 编译器支持真动态 shape | 实现难，性能可能下降 |
vLLM 的实用做法
vLLM 用 bucketed + mask-based 组合：
- prefill 阶段：按 prompt 长度分桶（如 128/256/512/1024/2048/...），每桶单独优化；
- decode 阶段：所有 batch 共用 max_gen_len 的桶，attention mask 屏蔽已结束请求；
- 加上 PagedAttention 处理 KV cache 的动态长度。 这是工业上 LLM 推理 shape 管理的事实标准。
五、编译器栈全景
LLM 推理的现代编译器栈：
1. PyTorch Inductor（torch.compile）
- 输入：PyTorch eager 模型；
- 输出：Triton kernel（GPU）或 C++ kernel（CPU）；
- 优势：与 PyTorch 训练无缝衔接；
- LLM 适用：训练 + 推理一体的场景（如 HuggingFace 模型用
torch.compile加速推理）； - 局限：不擅长动态 shape 与 continuous batching。
2. XLA（Google）
- 输入：JAX/TF HLO；
- 输出：GPU/CPU/TPU kernel；
- 优势：稳定、跨硬件；
- LLM 适用：Google 内部大模型训练/推理；JAX 生态（如 MaxText、EasyLM）。
3. TVM Relax（Apache）
- 输入：Relax IR（专门设计支持 LLM 动态 shape）；
- 输出：Triton/CUDA kernel；
- 优势：研究友好、支持 LLM 动态 shape；
- LLM 适用：研究性推理引擎（MLC-LLM 项目用 TVM）；
- 局限：生态不如 PyTorch 强。
4. Triton（OpenAI）
- 既是 kernel DSL，也是编译器；
- 输入：Python-like kernel；
- 输出：高效 CUDA；
- 优势：手写 kernel 标杆、几乎所有 LLM 推理引擎都用它；
- 详见算子融合。
5. TensorRT / TensorRT-LLM（NVIDIA）
- 输入：ONNX / HF 模型；
- 输出：闭源高性能 kernel；
- 优势：NVIDIA 硬件最优、INT8/FP8/FP4 全支持；
- LLM 适用：TensorRT-LLM 是 LLM 推理旗舰（详见TensorRT 案例研究）。
六、LLM 推理引擎的图优化实战
主流 LLM 推理引擎的图优化栈：
| 引擎 | 图优化栈 | 详见 |
|---|---|---|
| vLLM | PyTorch eager + Marlin/FlashAttention/FlashInfer + CUDA graph | vLLM 案例研究 |
| TensorRT-LLM | TensorRT + In-flight batching + INT8/FP8/FP4 kernel | TensorRT 案例研究 |
| SGLang | RadixAttention + CUDA graph + FlashInfer | — |
| ONNX Runtime | ONNX 图 + 经典融合 pass + IO-aware kernel | ONNX Runtime 案例研究 |
| OpenVINO | Intel 图优化栈 + CPU/GPU/VPU 后端 | OpenVINO 案例研究 |
| llama.cpp | C++ 手写图 + GGUF 量化 | llama.cpp 案例研究 |
为什么 LLM 推理图优化不同于传统 DNN
传统 DNN 的图优化重点在 算子融合 + 常量折叠 + 静态 shape 调度——把整张图一次性编译好。
LLM 推理多了一个维度：自回归控制流。每生成一个 token 都要：
- 调度下一个 batch（continuous batching）；
- 更新 KV cache；
- 检查 stop token；
- 重新进入图执行。
这个循环 + 控制流结构让传统编译器栈（如 ONNX Runtime）很多优化失效，必须用 vLLM/TensorRT-LLM 这种"动态图执行器"。
七、工程权衡
| 维度 | 选项 | 经验 |
|---|---|---|
| 图 vs eager | 静态图（TensorRT）vs 动态图（vLLM） | LLM 服务用 vLLM（动态）；离线批处理用 TensorRT-LLM（静态） |
| CUDA graph 开不开 | 开（默认 vLLM）vs 不开 | 99% 场景应该开；调试时关掉 |
| prefix caching 开不开 | 开 | 几乎无副作用，长 prompt 必开 |
| bucket size | 多档（16/32/64/128） | 大 batch 用大桶，但桶数别太多（capture 占显存） |
| 编译时间 vs 推理性能 | 长（TensorRT 几小时）vs 短（vLLM 几秒） | 在线服务用 vLLM；离线极致用 TensorRT-LLM |
| 量化叠加 | 量化后做图优化 | 权重量化 + 图优化是黄金组合 |
八、未来方向
LLM 推理图优化的研究热点：
- MLIR 在 LLM 推理的应用——多级 IR 让图、kernel、机器码分层优化；
- 动态 shape 原生支持——TVM Relax 主推；
- 跨引擎图复用——MLC-LLM 用 TVM 编译，同一份图能在 GPU/CPU/手机/WebGPU 跑；
- 自动调度——Ansor、MetaSchedule 让编译器自动搜索最优 kernel 参数；
- CUDA graph 的 conditional node（H100+）——动态分支也能 graph 化。
九、权衡与取舍
- 静态 vs 动态：静态图快但 LLM 推理动态——vLLM 的 bucketed CUDA graph 是实用解；
- 编译时间 vs 性能：在线服务（vLLM）容忍低编译时间换灵活性；离线（TensorRT-LLM）要极致性能，长编译可接受；
- 图优化 vs kernel 优化：图优化做宏观决策（融合什么、布局怎样），kernel 优化做微观执行（Triton/CUDA）；
- 开源 vs 闭源：开源栈（vLLM/Triton/TVM）生态强、可定制；闭源栈（TensorRT）极致性能但黑盒。
延伸阅读
- 算子融合与自定义核——图优化融合 pass 的执行层
- GPU 体系结构与优化——CUDA graph 的硬件基础
- 批处理与调度——continuous batching 的图执行机制
- 模型服务化与编排——图优化在生产部署中的角色
- vLLM 案例研究——CUDA graph + prefix caching 工业实现
- TensorRT 案例研究——闭源图优化栈的代表
- ONNX Runtime 案例研究——经典图优化 pass 的开源实现
- OpenVINO 案例研究——Intel 栈的图优化
- benchmarking 实践——图优化收益的测量方法
参考资料
- ONNX: Open Neural Network Exchange Documentation —— ONNX 标准
- MLIR: Multi-Level Intermediate Representation —— LLVM 多级 IR
- XLA: Optimizing Compiler for TensorFlow Documentation —— XLA 编译器
- PyTorch 2.0 torch.compile + Inductor —— PyTorch 编译器栈
- TVM Relax: Multi-level Intermediate Representation for LLMs —— Apache TVM 的 LLM IR
- TensorRT Developer Guide —— NVIDIA 闭源图优化栈
- NVIDIA CUDA Graphs Documentation —— CUDA graph 编程接口
- Zheng et al. SGLang: Efficient Execution of Structured Language Model Programs（2024） —— RadixAttention + 前缀树 KV 复用
- Kwon et al. Efficient Memory Management for LLM Serving with PagedAttention（SOSP 2023） —— vLLM 的图优化与内存管理