算子与运行时账：roofline、算子融合与 FlashAttention 的 IO 账

同一个模型，有人跑得快有人跑得慢，差距常在"数据搬了几次"。本文用三本账讲清算子层的优化逻辑：roofline 判定瓶颈、融合减少搬移、FlashAttention 把 IO 复杂度降一档。方法框架参考开源书稿《深入理解 AI Infra》（Apache-2.0），公式与算例为本站自行推导整理。

1. roofline：先判定瓶颈再优化

任何算子的执行时间 ≥ 数据搬移时间（带宽墙）与计算时间（算力墙）的较大者：

$$T = \max\left(\frac{FLOPs}{Peak}, \frac{Bytes}{BW}\right)$$

判定准则：算术强度 AI > 拐点（拐点 = Peak/BW）则算力受限，否则带宽受限。L40S 拐点 ≈ 419 FLOPs/Byte，绝大多数 LLM 算子（AI 常在 10~100）落在带宽受限区——优化方向是减少搬移，而不是堆算子。

判断：先算 AI 再选优化手段，能省掉一大半"瞎优化"。带宽受限的算子，做 CUDA 算子重写收益远大于调线程数。

2. 算子融合：省掉两次落盘

两个连续算子各自读写中间张量：

未融合: A → 写中间张量(显存) → 读回 → B      搬移 = 2×(写+读) × size
融合后: A→B 在片上(寄存器/SMEM)完成           搬移 = 单次输入/输出

融合省掉的搬移量 = 中间张量 size × 2 次写 + 2 次读。对带宽受限的算子（大多数），融合 1 个中间张量就能省 30-50% 的执行时间。这也是 FlashAttention、fused GELU、fused RMSNorm 的原理：把 IO 账算进算子边界，而不是只看 FLOPs。

3. FlashAttention：IO 复杂度降一档

标准 attention 的问题不在 FLOPs，在反复读写 S=P×Q 矩阵：

实现	显存 IO 复杂度	说明
标准 attention	O(N²)	每个注意力分数矩阵都落显存再读回
FlashAttention	O(N)	分块 + 在线 softmax，分数矩阵不落盘

序列越长，差距越大：序列 32K 时，标准 attention 的中间矩阵是 32K×32K×4 字节 ≈ 4 GB/层，FlashAttention 完全避免这笔搬移。长上下文推理与训练里，IO 复杂度比 FLOPs 复杂度更致命。

4. CUDA Graph：减少"搬指令"的开销

小算子的瓶颈可能是启动开销（launch ~微秒级 × 上千个算子）。CUDA Graph 把整个计算图固化成一次提交：

$$T_{graph} \approx T_{compute} + T_{launch\_once}, \quad vs \quad T_{loop} \approx T_{compute} + N \times T_{launch}$$

对 decode 这类"小算子多、每步重复"的负载，CUDA Graph 可省 20-50% 的步间开销（见 推理引擎优化 与 vLLM 默认启用 CUDA Graph 的原因）。

5. 可复算 Python 函数

def roofline(peak_tflops, bw_gbps, flops, bytes_moved):
    """判定算力/带宽瓶颈，返回预测耗时（ms）"""
    t_compute = flops / (peak_tflops * 1e12)          # s
    t_mem = bytes_moved / (bw_gbps * 1e9)             # s
    return max(t_compute, t_mem), "compute" if t_compute > t_mem else "bandwidth"

roofline(peak_tflops=362, bw_gbps=864, flops=2e12, bytes_moved=8e9)
# (9.3e-3, 'bandwidth') → 带宽受限，优先减少搬移

6. 交付清单

1. 优化前先算 AI 与拐点，确认瓶颈类型；
2. 带宽受限 → 找可融合的中间张量（每融合一个省 2 写 2 读）；
3. 长序列负载检查 attention 是否走 FlashAttention 类 IO 优化；
4. decode 类小算子负载启用 CUDA Graph；
5. 每项优化用 profiler 实测前后对比，存档（见 验收基线）。

---

本文为本站自行推导整理，方法框架参考开源书稿《深入理解 AI Infra》（Apache-2.0）；具体版本号与实测数据以现场为准。