Auto-tuning

Auto-tuning自动调优 / 自动自动性能优化是现代 AI 编译系统如 TVM、Triton、Ansor、Halide、XLA 等和高性能计算HPC中的核心技术。它的主要目标是针对特定的硬件架构如 GPU、NPU、CPU自动探索极度庞大的编译参数空间Search Space为高频算子如 Gemm、Conv2D、FlashAttention寻找最优的代码生成策略从而榨干硬件的峰值性能。为什么需要 Auto-tuning传统的高性能算子开发依赖顶级性能专家手写 Assembly/CUDA 代码如 cuBLAS、cuDNN。但面对如今的大模型时代这种模式遇到了巨大瓶颈硬件架构高度异构与复杂NVIDIAHopper/Blackwell、AMDCDNA3、华为升腾Ascend、Google TPU 等不同架构的 Register 数量、Shared Memory/SRAM 容量、Memory Bandwidth 以及 Tensor Core 架构截然不同。算子与 Shape 爆炸LLM 中不仅有各种 Transformer 变体而且动态 Shape如不同 Prompt 长度、Batch Size、KV Cache 状态使得“一套代码调优所有场景”变得不再可能。组合爆炸的优化空间一个看似简单的矩阵乘法CA×BC A \times BCA×B涉及到的编译参数组合可能多达上百万种Tile SizesThread Block / Warp / Thread 级的数据块切分大小。Loop Transformations循环展开Unrolling、循环重排Reordering、循环向量化Vectorization。Memory ManagementShared Memory double buffering双缓冲、Register Tiling、Swizzling消除 Bank Conflict。Pipeline/Parallelism流水线深度、线程块映射逻辑。人力根本无法枚举并找出最优解Auto-tuning 相当于用“计算”代替“人力”。Auto-tuning 的基本工作流程一个典型的 Auto-tuning 系统包含以下 4 个核心模块构成的闭环┌────────────────┐ │ Search Space │ │ (Program/IR) │ └───────┬────────┘ │ ▼ ┌─────────────────┐ Predict ┌───────────────┐ Select ┌─────────────────┐ │ Machine Learning│ ─────── │ Optimization │ ─────── │ Code Generator │ │ Cost Model │ │ Search Engine │ │ (Triton/TVM IR) │ └────────┬────────┘ └───────────────┘ └────────┬────────┘ ▲ ▲ │ │ │ ▼ │ Feedback │ Feedback ┌─────────────────┐ └───────────────────────────┴───────────────── │ Hardware Test │ │ (Kernel Timing) │ └─────────────────┘搜索空间定义Search Space Definition算子编译器将算法逻辑抽象为程序模板Schedule Template / IR并将待调优的参数定义为搜索维度如block_size_x ∈ [16, 32, 64, 128]。搜索算法Search Engine在巨大的组合空间中寻找最优解。常用搜索策略包括随机采样Random Search/ 格点搜索Grid Search遗传算法Genetic Algorithm, GA/ 模拟退火Simulated Annealing贝叶斯优化Bayesian Optimization代价模型Cost Model如果在真实 GPU 上测试每一个候选配置调优可能需要数天甚至数周。Cost Model通常基于 XGBoost、GNN 或 MLP的作用是在不开辟真实 GPU 编译运行的前提下快速预测某个编译参数组合的性能过滤掉 99% 的劣质配置。真实硬件测速与反馈Benchmark Feedback挑选出 Cost Model 预测排名前KKK的候选 Kernel在物理显卡上实际编译执行并精准测量耗时Kernel Duration。真实耗时会反馈给 Cost Model 进行在线学习Online Fine-tuning使其预测越来越准。主流 Auto-tuning 技术路线的演进从历史演变来看Auto-tuning 经历了从手写模板调优到全自动程序生成再到结合 AI 大模型指导的三个阶段阶段 1基于模板的调优Template-based Tuning代表TVM AutoTVM、Halide。特点性能专家为特定算子手写一个带有占位符Parameters的 Schedule 模板Auto-tuning 只负责在给定的取值范围内“填空”。优缺点性能上限极高依赖专家经验但开发成本极高。如果不手写模板系统就无能为力。阶段 2无模板的全自动搜索Template-free / Ansor-style Tuning代表TVM Ansor (AutoScheduler)。特点摒弃手写模板。系统直接输入算子的数学表达式Compute Definition通过层次化程序生成器Hierarchical Program Generator自动生成极其庞大的搜索空间并使用梯度提升树XGBoost或图神经网络预测性能。优缺点极大地解放了人力通用性极强但搜索耗时依然较长。阶段 3JIT 与动态/运行时调优Just-In-Time Dynamic Tuning代表Triton Auto-tuner、PyTorchtorch.compile(Inductor)。特点在 Triton 中用户可以使用装饰器triton.autotune为 Kernel 指定多种属性组合如num_warps、num_stages、BLOCK_SIZE_M。在程序运行初期系统进行JIT 编译与 Warmup 测速选择当前 Shape 下性能最好的配置文件挂载执行并将结果Cache在磁盘中。Auto-tuning 的核心挑战与未来趋势编译与搜索耗时Tuning Overhead问题对一个大模型的全套算子进行全量 Auto-tuning 可能会耗费数小时甚至数天。解法离线调优 Kernel Database将调优好的结果持久化到数据库引入迁移学习Transfer Learning将卡 A 上的调优经验迁移到卡 B。动态 Shape 难题Dynamic Shapes问题大模型推理时Prompt 长度SeqLen是变化的。针对 SeqLen512 调优出的最优 Kernel在 SeqLen1 时性能可能极其糟糕。解法基于分段Bucketing的多 Kernel 预调优 运行时根据 Shape 动态路由结合 Symbolic Execution符号化执行。LLM 辅助编译与代码生成AI for Compiler近年来利用 LLM / RL强化学习来引导编译器策略如 Seed Code 提取、死锁避让、自动配置参数成为热点方向能够极大地缩短 Search Engine 的收敛时间。