8 月 4 日，2026 夏季 InfiniTensor 训练营 AI 编译器方向第三课《AI 编译器中的后端优化》开讲。
前置课程《AI 编译器中的前端优化》讲解了四类核心等价图变换、子图替换工程实践与 CUDA Graph 底层加速方案等内容。

本节课将结合计算图算子实现与代码生成流程，系统拆解后端优化的四大代表性方法。

内存管理优化

优化目标与方向

1. 降低运行时开销：调整分配/释放时机，减少 cudaMalloc/cudaFree 等高频调用。
2. 降低显存占用：优化内存规划，减少峰值显存

Naive 版本 vs 优化版本

1.Naive 版本痛点：

• 推理前一次性为所有 Tensor 分配内存，直到结束才释放。
• 大量非活跃 Tensor 持续占用显存，造成浪费。

2.优化版本：

• 生命周期分析：基于拓扑排序模拟执行过程，按 Tensor 实际使用周期规划空间。
• 引用计数机制：维护后续 OP 引用计数，归零时立即回收空间供复用。
• Lazy Allocator 实现：
  • 模拟阶段：使用红黑树维护空闲块，计算 Offset 和 Peak，不调用真实设备接口。
  • 执行阶段：仅申请一次峰值大小的连续内存，通过 Base + Offset 获取实际地址。

3.实战收益：以 Inception-v3 为例，优化后支持 Batch Size 从 256 提升至 2048。

💡 核心思想：静态分析确定生命周期 + 统一内存池实现空间复用。InfiniTensor 与 Refactor Graph 等项目均采用此思路。

算子层面优化

性能瓶颈判断

1.Roofline 模型：综合峰值算力与内存带宽，判断任务是 访存密集型 还是 计算密集型。

2.优化侧重：

• 访存密集：减少数据搬运，提高复用率（如 FlashAttention）。
• 计算密集：提高计算单元利用率（如 Tensor Core、向量化）。

循环优化

1.循环分块 (Loop Blocking)：切分迭代空间，使工作集驻留 Cache。

2.循环展开 (Unrolling)：减少控制开销，提高指令级并行度。

3.循环重排 (Reordering)：匹配内存布局（行/列优先），变跳跃访问为连续访问。

• 案例：矩阵乘法 IJK → IKJ，提升 B 矩阵缓存命中率。

4.循环融合 (Fusion)：合并同范围循环，减少开销并促进数据复用。

5.循环拆分 (Splitting)：消除依赖与条件分支，为并行化创造条件。

Online Softmax 与 FlashAttention

• Safe Softmax：减去最大值避免指数上溢。
• Online Softmax：流式更新最大值与指数和，支持分块计算。
• 核心价值：避免存储完整 N×N 注意力矩阵，显著降低长序列推理显存。

指令优化

1.CPU 向量化：SSE/AVX/Neon 指令，配合数据对齐与连续访存。

2.GPU 张量计算：Tensor Core 针对矩阵乘累加设计。

3.PTX 指令解析：理解 mma.sync.align 等指令的形状、精度与寄存器约束。

内存优化

1.预取 (Prefetch)：提前搬运数据至高速缓存。

2.双缓冲 (Double Buffering)：计算与数据传输流水线化，隐藏 IO 延迟。

Auto Tuning

为什么需要 Auto Tuning？

• 同一算子存在多种算法实现（如 cuDNN 卷积算法），最优解依赖于输入形状、数据类型及硬件型号。

• 调度参数空间巨大（Tile 大小、线程数、流水线级数等），人工试错成本高。

1. 查询缓存：命中性能数据库则直接使用最优配置。
2. 参数搜索：未命中时，通过网格搜索/贝叶斯优化/RL 等策略生成候选。
3. 实测评估：在目标硬件上运行并测量性能。
4. 反馈入库：将最优结果缓存，避免重复调优。

代码生成

Triton 语言与编译器

1.定位：：Triton 是一种基于 Python 语法设计的 GPU Kernel DSL，通过编译器自动生成高性能 GPU 代码，降低 CUDA Kernel 开发复杂度。

2.核心抽象 - Layout：描述元素到线程/Warp/内存的映射关系，处理不同硬件布局转换（如 NCHW 与 NHWC）。

3.自动优化 Pass：

• Coalescing：合并连续内存访问。
• Pipeline & Prefetch：组织计算与搬运。
• Swizzling：减少 Shared Memory 访问 Bank-conflict。

编译下降流程

Python AST → Triton IR → Triton GPU IR → LLVM IR → PTX → GPU Execution

九齿框架

• 进一步封装线程块索引、Program ID 等底层细节。
• 让开发者聚焦算子主体逻辑，降低 CUDA 编程门槛。

加入训练营

📺 完整课程内容，请观看直播或查看课程回放：

InfiniTensor 官网： https://www.infinitensor.com

B 站直播： InfiniTensor 官方直播间

视频号直播： InfiniTensor 视频号预约

答疑交流： 训练营官方社群 ⬇️

关注与交流