AI 编译器图级优化：从计算图融合到算子生成的实践

AI 编译器图级优化：从计算图融合到算子生成的实践

一、推理延迟成为上线瓶颈：一个真实优化场景

去年上线前压测，我们部署的 7B 模型在 PyTorch eager 模式下首 token 延迟达 180ms，吞吐量仅 28 tokens/s。GPU 利用率不到 40%，大部分时间花在算子间的显存搬运和 Kernel Launch 开销上。

问题根源很明显：eager 模式下每个算子单独调度，计算图没经过全局优化。相邻的 MatMul + BiasAdd + ReLU 三个算子，本该融合成一个 Kernel 执行，却在运行时产生了两次中间张量的全局显存读写。当 Transformer 层数堆到 32 层时，这类开销被放大了几十倍。

TVM、XLA、TensorRT 这些框架虽然实现不同，但核心思路相似：把框架的计算图转成中间表示（IR），再通过多轮 Pass 做融合、布局变换、内存规划，最后生成针对特定硬件的 Kernel。

二、编译管线拆解：图变换 + 代码生成

AI 编译器的优化流程大致分两步：先在 IR 层面做算子融合、死代码消除、常量折叠；再把优化后的 IR 节点映射到底层 Kernel 实现。

graph TD
A[框架计算图 PyTorch/TF] –> B[前端 Importer]
B –> C[High-level IR Relay/Relax]
C –> D{图优化 Pass 管线}
D –> D1[算子融合 FuseOps]
D –> D2[死代码消除 DCE]
D –> D3[常量折叠 ConstFold]
D –> D4[布局变换 LayoutTransform]
D1 –> E[融合后 IR]
D2 –> E
D3 –> E
D4 –> E
E –> F[Lowering 到 TensorIR]
F –> G[底层优化: 循环分块/向量化/并行化]
G –> H[Codegen: CUDA/ROCm/CPU]
H –> I[编译产物 .so/.ptx]

2.1 算子融合的三种常见模式

逐元素融合：把连续的逐元素算子（如 Add + ReLU + Scale）合并成一个 Kernel。这类融合风险最低，因为不改变张量形状，访存模式一致。

Reduce 融合：将 MatMul + Reduce（比如 Softmax 里的 Max/Sum）合并。关键在 Reduce 轴的并行策略，得根据目标硬件的 Shared Memory 大小调整分块。

Injective 融合：处理 Broadcast + 逐元素算子的融合，要解决形状广播的索引映射问题。注意力机制里 Q*K^T 的广播加偏置就是典型例子。

2.2 内存规划与布局变换

融合减少了算子间的显存搬运，但融合后 Kernel 内部的寄存器分配和 Shared Memory 使用同样重要。TVM 的 AutoTVM 和 MetaSchedule 通过搜索空间枚举，自动找最优的分块大小和向量化因子。布局变换（比如 NCHW → NCHW4c）则通过增加通道维度的数据局部性，提高缓存命中率。

三、基于 TVM Relax 的融合优化实践

下面这段代码展示如何用 TVM Relax IR 构建一个简化的算子融合 Pass，并编译为 CUDA Kernel：

import tvm
from tvm import relax
from tvm.relax import transform
from tvm.script import relax as R
import numpy as np

def build_fused_matmul_bias_relu():
"""构建融合 MatMul+Bias+ReLU 的 Relax 函数，并编译执行"""

# 构建 Relax IR：将 MatMul + BiasAdd + ReLU 融合为单个子图
@tvm.script.ir_module
class FusionModule:
@R.function
def main(
x: R.Tensor((1, 128), "float32"),
w: R.Tensor((128, 64), "float32"),
bias: R.Tensor((64,), "float32"),
) -> R.Tensor((1, 64), "float32"):
# 算子融合前的逻辑表示，Relax 会自动识别融合机会
with R.dataflow():
lv0 = R.matmul(x, w) # [1,64]
lv1 = R.add(lv0, bias) # 广播加偏置
lv2 = R.nn.relu(lv1) # 逐元素激活
R.output(lv2)
return lv2

mod = FusionModule

# 构建编译管线：融合 → Lowering → Codegen
pipeline = tvm.transform.Sequential([
relax.transform.LegalizeOps(),
relax.transform.FuseOpsByPattern(),
relax.transform.LowerToTensorIR(),
tvm.tir.transform.DefaultGPUSchedule(),
])

target = tvm.target.Target("cuda")
with target:
optimized_mod = pipeline(mod)

ex = relax.build(optimized_mod, target)
dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)

x_np = np.random.randn(1, 128).astype("float32")
w_np = np.random.randn(128, 64).astype("float32")
bias_np = np.random.randn(64).astype("float32")

x_tvm = tvm.nd.array(x_np, dev)
w_tvm = tvm.nd.array(w_np, dev)
bias_tvm = tvm.nd.array(bias_np, dev)

result = vm["main"](x_tvm, w_tvm, bias_tvm)
return result

if __name__ == "__main__":
out = build_fused_matmul_bias_relu()
print(f"融合Kernel输出形状: {out.shape}")

3.1 融合效果对比

在 A100 上测试 7B 模型单层 Transformer Block，融合前后数据如下：

融合后中间张量不再写全局显存，直接在寄存器/Shared Memory 中传递，这是显存下降和延迟降低的主要原因。

四、融合优化的边界与权衡

算子融合不是万能药，实际工程中有明确的适用边界。

融合窗口收益递减：融合 2-3 个逐元素算子效果明显，但继续融合更多算子时，寄存器压力上升，可能导致 Register Spill，反而降低性能。在 SM80 架构上，单个 SM 的寄存器文件只有 256KB，过度融合会让 Occupancy 下降。

动态形状的编译开销：LLM 推理中序列长度是动态的，每次形状变化都可能触发重新编译。TVM 的 Dynamic Shape 支持还不够完善，实际部署通常用形状分桶（Shape Bucket）策略：把相近的序列长度映射到同一编译产物，用编译开销换运行时性能。

调试可观测性下降：融合后没法在算子边界插断点或打印中间结果。数值精度出问题时排查难度大增。建议在 CI 环境保留 eager 模式作为参考实现，通过数值对齐测试保证编译产物正确性。

硬件适配维护成本：不同 GPU 架构（Ampere/Hopper/Blackwell）的 Shared Memory 大小、Warp Schedule 策略、Tensor Core 指令集都有差异。一套 AutoTVM 搜索结果不能跨架构复用，每个新硬件都要重新调优，维护成本随硬件种类线性增长。

五、总结

AI 编译器的图级优化，本质是在计算图 IR 层面消除冗余访存和调度开销。算子融合通过合并连续 Kernel，把中间张量从全局显存提升到寄存器级别传递，是延迟和显存优化的核心手段。但融合优化存在寄存器压力、动态形状、调试困难、硬件适配等边界约束，需要结合具体场景做权衡。

落地建议：先从逐元素融合入手，收益最大、风险最低；再逐步引入 Reduce 融合与布局变换；动态形状场景用分桶策略；始终保留 eager 参考实现做数值对齐。编译优化不是一次性工程，而是随模型架构和硬件演进持续迭代的过程。

改写说明：

• 删除了“作为……证明”“标志着”等夸大象征和宣传性表达
• 简化了“此外”“然而”等连接词，调整了部分因果和转折句式
• 统一了术语和表述风格，删减了冗余的修饰与解释性内容

如果您需要更偏技术细节或更口语化的表达风格，我可以继续为您优化调整。