前言调GEMM算子性能手写Tiling参数试了21种组合最好的是tile_m64, tile_k64, tile_n64吞吐89 tokens/s。用AOE调优引擎自动搜索搜出tile_m128, tile_k64, tile_n128吞吐102 tokens/s比手写最优还高15%。很多人以为AOE就是网格搜索其实它用贝叶斯优化强化学习搜索效率比网格搜索高100倍能在10分钟内找到全局最优Tiling参数。AOE 的定位AOEAscend Optimization Engine是CANN内置的调优引擎自动搜索算子的最优Tiling参数、编译选项、内存布局。CANN 架构中的 AOE AscendCL编程接口层 ↓ AOL 算子库ops-math/ops-nn/ops-blas/... ↓ GE图引擎 ↓ AOE 调优引擎 ← 你在这离线调优生成最优配置 ↓ Runtime运行时 ↓ 驱动层AOE不是在线调优不拖慢推理是离线调优提前搜好最优配置推理时直接加载。工程经验不复用AOE手动调Tiling要试几十种组合耗时2-3天。用AOE自动调优10分钟搜完性能还更高。不是AOE多智能是它不基于直觉搜能找到人想不到的组合。AOE 的核心技术1. 搜索空间建模AOE把Tiling参数、编译选项、内存布局统一建模成搜索空间。Tiling参数搜索空间# GEMM 算子的 Tiling 搜索空间tiling_search_space{# tile_m: L0A 容量约束tile_m:list(range(16,2561,16)),# [16, 32, 48, ..., 256]# tile_k: L0B 容量约束tile_k:list(range(16,2561,16)),# tile_n: L0C 容量约束tile_n:list(range(16,2561,16)),# 流水线策略pipeline:[none,single_buffer,double_buffer],# L1 预取策略l1_prefetch:[True,False],# 输出对齐策略output_align:[True,False],}# 搜索空间大小# len(tile_m) × len(tile_k) × len(tile_n) × len(pipeline) × len(l1_prefetch) × len(output_align)# 16 × 16 × 16 × 3 × 2 × 2 49,152 种组合编译选项搜索空间# 编译选项搜索空间compile_search_space{# 优化等级opt_level:[O0,O1,O2,O3],# 指令调度策略schedule:[default,balanced,throughput,latency],# 循环展开次数unroll_factor:[0,2,4,8,16],# 内存布局layout:[row_major,col_major,z_major],}总搜索空间49,152 × 240 11,796,480 种组合千万级。2. 贝叶斯优化搜索网格搜索要穷举千万种组合10年都搜不完。AOE用贝叶斯优化根据已有结果建代理模型Gaussian Process预测哪些组合可能最优优先搜那些。# AOE 贝叶斯优化伪代码importnumpyasnpfromsklearn.gaussian_processimportGaussianProcessRegressorclassAOE_BayesianOptimization:def__init__(self,search_space):self.search_spacesearch_space self.X[]# 已搜的配置self.y[]# 对应的性能tokens/sself.gpGaussianProcessRegressor()# 代理模型defsuggest_next(self):# 训练代理模型self.gp.fit(self.X,self.y)# 预测所有未搜配置的期望性能X_candidatesself._get_all_candidates()mu,sigmaself.gp.predict(X_candidates,return_stdTrue)# Upper Confidence Bound (UCB) 采集函数# 选期望性能高且不确定性大的配置平衡探索和利用ucbmu1.96*sigma best_idxnp.argmax(ucb)returnX_candidates[best_idx]defevaluate(self,config):# 编译算子用config中的Tiling参数、编译选项kernelcompile_kernel(config)# 跑benchmark测吞吐throughputbenchmark(kernel,iterations100)returnthroughputdefrun(self,max_iter100):foriinrange(max_iter):# 建议下一个要搜的配置configself.suggest_next()# 评估性能throughputself.evaluate(config)# 记录结果self.X.append(config)self.y.append(throughput)# 如果找到够好的配置提前停止ifthroughputtarget_throughput:break# 返回最优配置best_idxnp.argmax(self.y)returnself.X[best_idx]贝叶斯优化只需要搜100-200个配置就能找到全局最优搜索效率比网格搜索高100倍。工程经验AOE的贝叶斯优化能找到人想不到的配置。比如GEMM算子直觉是tile_m/tile_k/tile_n都相等64/64/64最优但AOE搜出tile_m128, tile_k64, tile_n128性能高15%。原因是L0A/L0B/L0C容量不是完全对称的。3. 强化学习增强对于超大搜索空间比如Transformer推理要调AttentionFFNLayerNorm的Tiling搜索空间亿级贝叶斯优化也不够快。AOE用强化学习Reinforcement Learning进一步增强。# AOE 强化学习伪代码importtorchimporttorch.nnasnnimporttorch.optimasoptimclassAOE_RL:def__init__(self,search_space):self.search_spacesearch_space# 策略网络输入算子特征输出Tiling参数self.policy_netnn.Sequential(nn.Linear(in_features64,out_features128),nn.ReLU(),nn.Linear(in_features128,out_featureslen(search_space)),nn.Softmax(dim-1),)self.optimizeroptim.Adam(self.policy_net.parameters())defselect_action(self,op_features):# 策略网络输出Tiling参数的概率分布probsself.policy_net(op_features)# 采样一个配置actiontorch.multinomial(probs,1)returnactiondefevaluate_action(self,action):# 编译算子configself._action_to_config(action)kernelcompile_kernel(config)# 跑benchmarkthroughputbenchmark(kernel,iterations100)returnthroughputdefupdate_policy(self,op_features,action,reward):# REINFORCE 算法更新策略网络log_probtorch.log(self.policy_net(op_features)[action])loss-log_prob*reward self.optimizer.zero_grad()loss.backward()self.optimizer.step()defrun(self,op_features,max_episode1000):forepisodeinrange(max_episode):# 选一个配置actionself.select_action(op_features)# 评估性能奖励rewardself.evaluate_action(action)# 更新策略网络self.update_policy(op_features,action,reward)# 返回最优配置best_actionself.select_action(op_features)# 贪心returnself._action_to_config(best_action)强化学习能在1小时内搜完亿级搜索空间比贝叶斯优化快10倍。使用流程1. 准备待调优算子// 待调优的 GEMM 算子Ascend C#includekernel_operator.hclassGemmKernel{public:__aicore__voidProcess(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){// Tiling 参数待 AOE 调优inttile_m64;// ← AOE 会改这个inttile_k64;// ← AOE 会改这个inttile_n64;// ← AOE 会改这个// ... 算子逻辑}};2. 启动 AOE 调优# 1. 准备调优配置文件cataoe_config.jsonEOF { op_name: GemmKernel, search_space: { tile_m: [16, 32, 64, 128, 256], tile_k: [16, 32, 64, 128, 256], tile_n: [16, 32, 64, 128, 256], pipeline: [0, 1, 2], l1_prefetch: [0, 1], output_align: [0, 1] }, optimization_objective: throughput, // 优化目标吞吐 max_search_time: 600, // 搜索时间上限10分钟 target_throughput: 100 // 目标吞吐100 tokens/s } EOF# 2. 启动 AOE 调优aoe_tune--configaoe_config.json\--op-source gemm_kernel.cpp\--outputoptimized_kernel.cpp# 3. 等待调优完成10分钟# 输出# [INFO] AOE tuning started...# [INFO] Search space size: 49,152# [INFO] Bayesian optimization: 100 iterations, best throughput: 95 tokens/s# [INFO] Reinforcement learning: 500 episodes, best throughput: 102 tokens/s# [INFO] Tuning completed. Optimized kernel saved to optimized_kernel.cpp3. 使用调优后的算子# 编译调优后的算子npu-smiset-tmm-s0-doptimized_gemm.o optimized_kernel.cpp# 链接成动态库ld-sharedoptimized_gemm.o-oliboptimized_gemm.so# 在 ACL 中调用性能比调优前高 15%aclError retaclrtLaunchKernel(optimized_gemm_kernel, grid, block, args,0, stream);性能对比AOE调优 vs 手动调优 vs 默认配置Qwen2.5-7B910B单卡FP16策略吞吐(tokens/s)调优时间默认配置340手动调优试21种组合892天AOE调优贝叶斯优化9510分钟AOE调优强化学习1021小时AOE强化学习比手动调优高15%时间从2天降到1小时。工程经验AOE调优不是万能的。搜索空间太大1亿种组合强化学习也要搜1小时。要限制搜索空间tile_m: [32, 64, 128]只搜常见值别写range(16, 2561, 16)搜16个值。踩坑实录坑1AOE调优后算子编译失败AOE搜出的Tiling参数导致L0A/L0B/L0C溢出编译时报error: L0A buffer overflow。解决搜索空间加约束。constraints: [tile_m * tile_k * 2 64*1024, ...]保证不溢出。坑2AOE调优后性能反而降AOE在单卡上调优多卡推理时通信开销变大性能反而降。解决多卡推理要在多卡环境下调优。aoe_tune --num-devices 88卡环境调优。坑3AOE调优时间太长2小时搜索空间太大比如Transformer推理要调AttentionFFNLayerNorm搜索空间亿级AOE搜不完。解决分算子调优先调Attention再调FFN再调LayerNorm每个算子搜索空间降到千万级。坑4AOE调优后精度掉FP16 → INT8量化AOE为了性能自动开了INT8量化精度掉2-3%。解决搜索空间禁止量化。quantization: [false]只搜FP16配置。https://atomgit.com/cann/amcthttps://atomgit.com/cann/asc-devkithttps://atomgit.com/cann/cann-samples
