XLA编译流程：算子融合Fusion

可视化演示

三步展示XLA算子融合：多算子(Mul/Add/Act)合并为单一Fusion核、融合类型分类(elementwise/producer/reduction)、以及融合后访存减少与内核启动开销降低的性能收益。

融合原理

想象你经营一家中餐厅的后厨。按照传统流程，切菜师傅把土豆切成丝后，不能直接把半成品递给旁边的炒锅师傅——必须装进保鲜盒、放进冷库；等炒锅师傅要用时，再从冷库里取出来。一份酸辣土豆丝，中间要经历好几轮“放回冷库—取出”。这些步骤不增加任何味道，却占了师傅们大半的力气和时间。

深度学习模型的推理和训练也是如此。GPU 上执行一个算子，就像后厨的一道工序：算完的结果必须写回显存（相当于冷库），下一个算子再把它读回来。让人头疼的是，现代 GPU 的计算速度增长得飞快，显存带宽却远远跟不上——这就是学术界说的内存墙（Memory Wall）。很多算子其实几微秒就算完了，数据搬运反而花了几十微秒。工程师们用一个指标来衡量这种失衡——计算访存比（Arithmetic Intensity），即每个字节的数据搬运能支撑多少次计算；这个值越低，说明算子越是被“空跑搬运”拖累。

更隐蔽的开销来自内核启动（Kernel Launch）。每执行一个算子，CPU 都要向 GPU 发起一次“开工指令”，从下发到真正执行通常要消耗微秒级的时间。如果一个模型有几百个算子，光启动开销就相当可观。把单个算子的耗时拆开看：

$$T_{\text{op}} = T_{\text{compute}} + T_{\text{memory}} + T_{\text{launch}}$$

1. T_op 总耗时：一个算子从开始到结束的墙钟时间（Wall-clock Time），是我们在性能分析里直接观测到的值。

$$T_{\text{op}}$$

2. T_compute 纯计算：GPU 真正做数学运算（如矩阵乘、逐元素加法）耗费的时间。

$$T_{\text{compute}}$$

3. T_memory 访存时间：数据从显存读入寄存器、以及把结果写回显存的时间，通常被内存墙放大。

$$T_{\text{memory}}$$

4. T_launch 启动时间：CPU 向 GPU 下发内核指令、GPU 排队调度所消耗的时间。

$$T_{\text{launch}}$$

所以算子融合的出发点很简单：既然“放回冷库再取出来”和“呼叫厨师”都这么贵，那能不能让一位厨师一口气把几道工序做完，把菜直接留在手边？这正是 XLA 要解决的第一个问题——找到那些适合“传递接力棒”的工序。

融合策略

厨师长拿到今天的菜单，开始排兵布阵。有的工序之间关系紧密——前一道刚出锅、后一道马上要用，比如“焯水”和“爆炒”，那就让同一位师傅趁热完成；有的工序天生适合“边洗边切边扔”——洗菜时顺手把烂叶扔掉，只留下能用的部分；还有的工序规矩地排成流水线，可以串成一条长线。

XLA 的融合策略也遵循类似的直觉。最经典的是生产者-消费者融合：当一个算子的输出恰好是下一个算子的输入（而且这个输出没有被其他算子使用）时，就把它们合并成一个内核，中间结果直接留在寄存器或片上缓存里，不再写回显存。这就像切好的菜直接递给炒锅师傅，冷库完全被跳过了。

另一种重要的是 reductions 融合：某些算子会把数据沿某个维度“约减”掉一大块，比如求和、求最大值。融合时把“先算出全部结果、再约减”变成“边算边约减”，只把最终的小结果写回显存，中间的大块数据根本不需要落地。此外还有循环融合，把多个逐元素算子（如 Add、Mul、激活函数）打包成一个对数据只遍历一次的循环。无论哪种策略，XLA 最终都会用启发式规则做决策，核心判据就是“计算访存比”：

$$AI = \frac{W}{Q} = \frac{\text{总浮点运算量}}{\text{总显存数据访问量}}$$

1. W 总运算量：整个算子簇需要完成的浮点运算次数（FLOPs），是融合后内核的“体力活”总量。

$$W$$

2. Q 总访存量：从显存读入加上写出的总字节数，是融合后内核的“搬运活”总量。

$$Q$$

3. AI 算术强度：两者的比值，代表每字节搬运能支撑多少次计算，是判断“该不该融合”的标尺。

$$AI = \frac{W}{Q}$$

厨师长不需要管每一根葱怎么切，只需要判断“这道菜值不值得趁热做”“洗完的菜要不要保留全部”。XLA 也一样：启发式规则扫过整张计算图，把值得合并的算子簇圈出来，再交给后端的代码生成器去“掌勺”。

融合收益

客人最直观的感受是上菜变快了。以前一道菜要经过六位师傅、六次冷库存取、六次吆喝开工；现在一位全能师傅一气呵成，冷库几乎没派上用场。出餐时间缩短，餐厅在同一时段能接待的客人自然更多——用行话说，就是“吞吐量上去了”。

把收益量化：假设一条计算链有 n 个算子，未融合时要启动 n 个内核，产生 n-1 份中间结果，每份中间结果都要经历一次“写回显存 + 再读入”的往返。融合成一个内核后，这些开销几乎全部消失。

融合到底能省多少时间？可以用这个式子估算：

$$T_{\text{saved}} = (n-1) \times T_{\text{launch}} + \sum_{i=1}^{n-1} \frac{Q_{\text{intermediate},i}}{B_{\text{mem}}}$$

其中 $Q_{\text{intermediate},i}$ 是第 i 份中间结果的体积，$B_{\text{mem}}$ 是显存带宽。前半部分省下的是内核启动次数，后半部分省下的是中间结果搬运时间。

1. 启动节省：n 个算子合成 1 个内核，启动次数从 n 次直接降到 1 次。

$$(n-1) \times T_{\text{launch}}$$

2. 访存节省：每一份中间结果都免掉一次“写显存 + 读显存”的往返，除以带宽即节省的时间。

$$\sum_{i=1}^{n-1} \frac{Q_{\text{intermediate},i}}{B_{\text{mem}}}$$

3. 总节省时间：把启动开销与访存开销两部分相加，就是融合带来的真实加速。

$$T_{\text{saved}} = \text{启动节省} + \text{访存节省}$$

更妙的是，这种加速是“免费午餐”——厨师并没有多花力气，只是少走了冤枉路。在真实 GPU 上，逐元素算子链（比如 add→relu→mul）融合后访存量常常能下降 50%~90%，内核数量从几十个降到几个，GPU 利用率显著提升，端到端推理延迟随之下降。这正是融合成为 XLA 核心优化手段的原因。

XLA实现

如果说上面的策略是“厨师长的经验”，那 XLA 就是把经验变成了一套自动化后厨管理系统。你不需要告诉它“这道菜要趁热加工”，只要把原料丢进系统，它自己会读菜谱、排工序、合并操作，最后端出一盘热菜。

XLA 的编译对象是一张计算图，叫做 HLO 图，图中的每个节点就是一个算子。系统里有一个专门的融合 pass（一个优化步骤），在图上做一遍深度遍历，把符合策略的相邻算子标记为同一个融合簇，再交给后端的代码生成器（GPU 后端通常生成 CUDA 内核）。

融合决策本身是一个权衡问题。一个融合内核会把多个算子的逻辑全部写进同一个 kernel，中间值要尽量留在寄存器里，因此非常依赖寄存器数量和共享内存容量。启发式判定的通用形式可以写成：

$$\text{Fuse}(A, B) \iff R_{\text{fused}} < R_{\text{limit}} \;\wedge\; \frac{Q_{\text{unfused}} - Q_{\text{fused}}}{Q_{\text{unfused}}} > \delta$$

1. 资源约束：融合后内核需要占用的寄存器数量 R_fused，必须低于当前硬件给出的上限 R_limit。

$$R_{\text{fused}} < R_{\text{limit}}$$

2. 访存收益：融合前后访存量之差，占未融合访存量的比例，必须大于阈值 δ 才有融合价值。

$$\frac{Q_{\text{unfused}} - Q_{\text{fused}}}{Q_{\text{unfused}}} > \delta$$

3. 组合判定：两个条件必须同时满足，系统才执行融合操作。

$$\text{Fuse}(A, B) \iff \text{条件1} \wedge \text{条件2}$$

这套系统还配了“监控摄像头”：XLA 提供 HLO 转储、图形可视化等调试选项，工程师可以直观地看到哪些算子被合并了、哪些没有、为什么没有；配合性能分析工具统计内核数量和访存量，就能验证融合到底带来了多少收益。整条链路高度自动化，这也是它成为生产级编译器的重要原因。

PyTorch集成

普通食客不需要进后厨点菜——他们在前台下单，后厨系统自动把订单分解成一道道工序。PyTorch 用户就是食客，写下模型、调用一次接口，编译栈自动把 Python 层面的算子翻译成高效的 GPU 内核。

PyTorch 2.0 引入的 torch.compile 是当前最主流的入口。它内部由三大部分组成：Dynamo 负责“捕获计算图”——把 Python 代码的执行轨迹还原成一张静态图；AOTAutograd 负责把反向传播也展开成图；最后交给后端编译器生成内核。默认后端是 Inductor（基于 Triton），但在 PyTorch/XLA 这类框架中，图会直接转换成 HLO 交给 XLA 编译，从而完整享受算子融合。

无论是 Inductor 的 Triton 融合还是 XLA 的 HLO Fusion，思想同源。用一个加速比公式概括编译带来的收益：

$$\text{Speedup} = \frac{T_{\text{eager}}}{T_{\text{compiled}}} = \frac{\sum_{i=1}^{n} T_{\text{op},i}}{\sum_{j=1}^{m} T_{\text{fused},j}}, \quad m \ll n$$

1. Eager 总耗时：未编译时每个算子单独执行，总耗时是所有算子耗时的累加。

$$\sum_{i=1}^{n} T_{\text{op},i}$$

2. 编译后总耗时：融合内核逐一执行，内核数量 m 远小于算子数量 n。

$$\sum_{j=1}^{m} T_{\text{fused},j}$$

3. 加速比：两者相除，衡量整体性能提升的倍数。

$$\text{Speedup} = \frac{T_{\text{eager}}}{T_{\text{compiled}}}$$

对用户来说，这就像点餐后不用关心后厨怎么安排——torch.compile(model) 一行代码，编译器自动完成图捕获、融合、代码生成。同一套模型，Eager 模式逐个算子跑、编译模式融合后跑，基准测试一对比，融合的威力立刻可见。

优化案例

再好的理论也要拿招牌菜验证。Transformer 就是 AI 界的招牌菜，几乎每家大模型餐厅都在做。这道菜工序极多：多头注意力、Softmax、层归一化、MLP……每一道工序之间都藏着大量“冷库往返”。

以 Attention 融合为例，Q、K、V 三个线性投影本来是三次独立的矩阵乘法，每次都要把结果写回显存。融合后，可以把三个投影拼成一个大矩阵，一次矩阵乘法同时算出 Q、K、V，省掉两次内核启动和两次中间结果落地。更极致的做法参考 FlashAttention 的思路：把 QK^T 和 PV 的计算放进片上缓存，Softmax 的中间矩阵根本不落显存。

Attention 融合前后的访存量差异，可以用一个式子说明：

$$Q_{\text{unfused}} = Q_{QK} + Q_{S} + Q_{PV}, \qquad Q_{\text{fused}} = Q_{QK} + Q_{PV}$$

其中 $Q_S$ 是 Softmax 的注意力分数矩阵（大小为序列长度的平方）。序列越长，$Q_S$ 越大，不落显存的收益就越夸张。

1. 未融合访存量：QK 矩阵乘的结果、Softmax 分数矩阵、PV 矩阵乘的结果全部要写显存。

$$Q_{\text{unfused}} = Q_{QK} + Q_{S} + Q_{PV}$$

2. 融合后访存量：Softmax 分数矩阵 S 留在片上不再写回显存，只剩两个矩阵乘的读写。

$$Q_{\text{fused}} = Q_{QK} + Q_{PV}$$

3. 节省幅度：节省的部分正是规模为序列长度平方的 S 矩阵。

$$Q_{\text{saved}} = Q_{S} = O(N^{2})$$

Transformer 之外，MLP 融合也是重灾区。把“线性层 + GELU + Dropout”这些逐元素算子揉进相邻的矩阵乘内核后，一个 Decoder Layer 的内核数量往往能减少一半以上。放到整张网上看，融合后模型吞吐量提升 30%~50% 是常见结果——冷库没人去了，菜自然上得快。

---

AI 应用场景

Python 演示

"""算子融合效果对比: Eager 模式 vs torch.compile 编译模式"""import timeimport torch  def build_mlp():    """一个典型 MLP 块: Linear -> GELU -> Linear -> GELU"""    return torch.nn.Sequential(        torch.nn.Linear(256, 256),        torch.nn.GELU(),        torch.nn.Linear(256, 256),        torch.nn.GELU(),    )  def bench(model, x, iters=50):    """预热后返回平均耗时(毫秒)"""    with torch.no_grad():        for _ in range(5):          # 预热: 触发编译与缓存            model(x)        if x.is_cuda:            torch.cuda.synchronize()        t0 = time.perf_counter()        for _ in range(iters):            model(x)        if x.is_cuda:            torch.cuda.synchronize()    return (time.perf_counter() - t0) / iters * 1000  def main():    dev = "cuda" if torch.cuda.is_available() else "cpu"    x = torch.randn(2048, 256, device=dev)     eager = build_mlp().to(dev).eval()    compiled = torch.compile(build_mlp().to(dev).eval())     te = bench(eager, x)    tc = bench(compiled, x)     print(f"Eager    : {te:.3f} ms")    print(f"Compiled : {tc:.3f} ms")    print(f"Speedup  : {te / tc:.2f}x (得益于算子融合)")  if __name__ == "__main__":    main()

1. torch.nn.Sequential(...)定义 4 个连续算子的 MLP 链，正是一串适合融合的“工序链”。
2. build_mlp().to(dev).eval()把模型搬到 GPU（或 CPU）并切到推理模式，保证计时公平。
3. torch.compile(build_mlp().to(dev).eval())一行开启编译模式，后端自动做图捕获与算子融合。
4. bench(eager, x)` / `bench(compiled, x)同一套预热+计时逻辑分别测量两种模式，避免编译延迟污染结果。
5. torch.cuda.synchronize()阻塞等待 GPU 执行完毕，保证计时不包含异步排队的误差。
6. te / tcEager 耗时除以编译耗时，得到融合带来的加速比，通常大于 1。

常见误区

面试题