当搜索空间足够大时优化效率本身就是性能工程的一部分。背景与问题FlashAttention 是大模型推理和训练中不可或缺的核心算子。在昇腾平台上torch_npu提供了高度优化的融合实现作为性能基线。但当我们需要通过 PTOParallel Tile Operation[l1] 来定制 FlashAttention 的行为——比如适配不同的 attention 变体、调整精度策略、或者做更细粒度的算子融合——就需要从 PTO 层面重新实现并优化。问题在于PTO 版本的 FlashAttention 初始性能与融合基线差距很大。在我们的测试场景中seq_len8192的 full-sequence dense attention初始实现的运行时间为1760.489 ms而torch_npu融合基线只要16.948 ms。差距接近两个数量级。面对这个 gap我们选择让 Agent 自主迭代优化。但很快我们发现真正难的不是Agent 能不能想到优化点而是每个优化点都要花真金白银去远端验证。这篇文章分享的就是如何通过引入 PTO Cost Model改变 Agent 的搜索策略在 10 轮迭代内从平台期走向持续下降的完整过程。为什么 PTO 适合 Agent 来优化在讲 Cost Model 之前先说一个前提这轮优化之所以适合让 Agent 参与核心原因是PTO 的 tile 级语义天然对 Agent 友好。传统底层 kernel 框架的语义非常碎片化——一大段底层指令、寄存器搬运、隐式同步混在一起Agent 很难稳定地抓住优化杠杆。PTO 不一样。它把 FlashAttention 自然拆成三类 tile 操作阶段操作类型说明QK 阶段cube matmulQuery 与 Key 的矩阵乘Softmax 阶段vector row op行级归一化PV 阶段cube matmulScore 与 Value 的矩阵乘这种表达不仅对人更好读对 Agent 来说更重要——它更容易形成这里的瓶颈大概是什么下一步最该改哪个 tile的判断。更关键的是PTO 把tile 级的乱序执行和流水线掩盖空间也暴露了出来double buffer、prefetch、load/compute/store overlap、reduction loop 的软件流水线——这些都可以被重排、被比较、也可以被 cost model 估价。简单来说PTO 不只是让人能写 kernel它还让Agent 更容易理解 kernel。PTO 将 FlashAttention 拆解为 QK / Softmax / PV 三类 tile 操作每类操作的计算类型和数据搬运路径边界清晰Agent 自主优化柿子先挑软的捏没有 Cost Model 的第一阶段Agent 表现可以说还行。它按照一套严格的迭代规则每轮只验证一个假设、正确性先于性能、记录每轮结果把大 shape 上的性能从1760.489 ms一路压到了46.111 ms。主要动作包括开启 cube block利用多核并行将 reduction loop 流水化Softmax row batch 从 1 推到 2实现 head overlap 两路并发但是相比于torch_npu融合基线性能还有1倍多的差距我们的目标显然不止于此。山重水复疑无路搜索进入迷宫当性能来到 46 ms 左右后Agent 仍然能源源不断地提出看上去合理的假设调 softmax block、改 QK launch cap、换 stream 分发方式、把两路 overlap 扩到三路或四路……问题是这些方向不跑到远端机器上编译、benchmark几乎无法提前判断值不值得。于是搜索进入了一个尴尬的阶段性能在4654 ms 区间来回横跳。看到的结果是差一点、退一点、直接 crash的循环。在大搜索空间里继续优化这句话本质上信息量很低。迭代 25-35 轮性能在 4654 ms 区间反复波动缺乏明确下降趋势见蓝色区域柳暗花明又一村引入 PTO Cost Model转折不是某个 tile 参数而是搜索策略的改变先别急着上机先给候选一个低成本排序。我们给 Agent 补了一层 Cost Model。它很轻——不依赖 PTO runtime也不试图预测绝对 wall-clock 时间。它做的事情很朴素把三阶段 FlashAttention 的 QK、softmax、PV 拆开用 PTO ISA 的 cycle 常数和带宽参数估算谁更值得先测的优先级。这改的不是答案而是搜索顺序环节没有 Cost Model有了 Cost Modelprompt 策略继续优化、再试一个参数先沿 QK → PV → runtime 排优先级第一层过滤器远端 benchmark本地 estimatorbenchmark 的角色既筛候选也确认结果只确认结果错误方向代价一次完整的远端编译运行校验本地几毫秒先淘汰低价值候选搜索轨迹4654 ms 区间横跳46.1 → 39.8 → 32.5 → 25.8 → 23.8 → 21.3一路往下Cost Model 作为本地预筛选层在远端 benchmark 之前过滤低价值候选表中第三行最关键Cost Model 把 benchmark 从第一层过滤器降回了最后一层裁判。以前提一个候选就得去远端交学费现在先本地做便宜筛选只把靠前的候选送去上机确认。量化对比同样进入平台期后差距有多大公平的比较不是看有了 cost model 之后每一步多漂亮而是同样已经进入平台期的前提下搜索效率差了多少对比维度无 Cost Model平台期有 Cost Modelguided sweep搜索方向runtime 微调、stream remap、dispatch cache大 tiling runtime overlap 高价值方向正式候选数11 轮r25-r3512 个正式候选keep 数量5 次6 次crash 数量2 次3 次best 变化46.568 → 46.111 ms46.111 → 21.287 ms净收益-0.457 ms-24.824 ms相对改善-0.98%-53.84%没有 Cost Model 时11 轮只抠掉不到 1% 的 runtime overhead有了之后bench 时间重新集中到大 tiling 和四路 grouped streams 这些真正能往下掉的方向。Cost Model 具体押中了什么在 Cost-Model-Guided 阶段Agent 沿着一条清晰的路径连续命中了 6 个关键 keep优先验证方向具体动作median 变化单步改善大 tilingQKQK row tile 推到 128×6446.111 → 39.767 ms-13.76%大 tilingQKQK score tile 推到 128×12839.767 → 32.516 ms-18.23%大 tilingQKQK score tile 推到 128×25632.516 → 31.875 ms-1.97%大 tilingPVPV row tile 推到 128×12831.875 → 25.797 ms-19.07%大 tilingPVPV score tile 推到 128×25625.797 → 23.833 ms-7.62%多流优化四路 grouped streams23.833 → 21.287 ms-10.68%大 tiling 不是撞出来的是先被筛出来的这条线的逻辑非常清楚对于seq_len8192的 full-sequence attention大量代价都和 tile 数量、数据搬运次数、cube 重复次数绑在一起。Cost Model 不需要精准预测到底快多少只需要把更少的 tile、更少的往返搬运、更少的 stage 开销排到前面搜索路径就变得清楚了。当然Cost Model 也不是万能的。更大的 tile 不一定总是更快——把 QK 硬推到更极端的组合出现了回退再往极端方向推编译器甚至会直接报错。Cost Model 做的是筛选不是拍板。最终结论仍然靠真实 benchmark 来定。四路 grouped streams 的时机问题多流优化这件事更能说明 Cost Model 的价值。在更早阶段四路 overlap 并不是一个好选择——kernel 自身的 tile geometry 还没收敛launch 开销和计算本体的平衡还没达到合适的比例。那时试过四路并发结果并不理想。直到大 tiling 把 QK 和 PV 两侧的 kernel 代价压下来之后问题才变成单个 head 的 kernel 已经够轻了runtime 层是不是值得重新回来做多流Cost Model 对多流的帮助不是直接算出四路最快而是更早收敛 kernel geometry让注意力更早转到 runtime 层的问题上。在大 tiling 稳定后四路 grouped streams 带来了 10.68% 的进一步改善。但继续推到八路结果又回退了——这恰好说明了整件事的边界。三种角色的分工回头看这条优化曲线整个闭环里存在三种角色角色擅长的事在优化中的定位Agent扩空间——持续生成候选、快速落代码、快速回传结果提供搜索广度Cost Model压空间——低成本筛选候选、排优先级降低搜索成本人定义问题——把模糊的继续优化改写成高信息密度的方向引导搜索方向典型的好问题长这样现在到底还该不该继续扫 tile如果扫 tile应该先压 QK 还是先压 PVtile geometry 基本收敛以后runtime overlap 是不是终于值得重测三种角色一旦分清楚优化就不再是看谁能想到更多点子而是看谁能更便宜地决定下一轮该测什么。Agent扩空间、Cost Model 压空间、人定义问题的协作闭环最终结果从整体来看初始性能1760.489 msAgent 自主优化后46.111 ms提升 97.4%Cost Model 引导后21.287 ms在平台期基础上再降 53.84%最终达到 Ascend C 基线的约 80%从 1760 ms 到 21 ms 的完整优化路径Cost Model guided 阶段贡献了平台期后的主要收益写在最后如果只把 Cost Model 理解成一个辅助脚本那就把它看轻了。在搜索空间大、单次验证成本高的问题里Cost Model 本质上是一层搜索经济学——它不替代 benchmark但让 benchmark 只为真正值得的候选买单。一句话总结这轮实践Agent扩空间Cost Model 压空间人负责定义问题。PTO FlashAttention 这次从 1760 ms 走到 21 ms当然是性能优化。但更重要的是它验证了一个方法论当搜索空间足够大时优化效率本身就是性能工程的一部分。本文使用Skill链接https://gitcode.com/cann/pto-isa/tree/master/agents/skills/pto-costmodel-agentic-kernel-optimizationCostModel代码仓链接https://gitcode.com/cann/pto-isa/tree/master/include/pto/costmodel欢迎加入CANN开源社区PTO-ISA开源仓参与讨论https://gitcode.com/cann/pto-isa