SuperAPI AI 工具集

triton-operator-performance-optim

编程与开发

查看项目

更新于 10-16 05:00

技能介绍

该技能用于优化昇腾 NPU 亲和的 Triton 算子性能,覆盖 UB 溢出、Cube 利用率与 Tiling 策略设计。底线要求:优化后精度需与 PyTorch-NPU 对齐(rtol=1e-3, atol=1e-3),不通过则回退;保持原有输入形状与 dtype 的泛化性,不得硬编码;优先级为正确性 > 泛化性 > 性能。性能比定义 Ratio = torch_npu 耗时 / Triton 耗时,大于 1.0 表示 Triton 更快。

工作流:先做算法审视;再用 time.time() 快速筛选中/大/小规模端到端,必要时用 msprof 精确定位 kernel 瓶颈;按瓶颈类型优化,Memory-Bound 侧重向量化访存、UB 缓存复用与算子融合,Compute-Bound 侧重 Cube 适配与 Block 调优,Latency-Bound 侧重增大并行度与减少同步,基础顺序为 BLOCK/GIRD size→连续访存→UB 复用→编译时常量;随后做硬件特化,如 BLOCK_M/N/K 取 16 倍数、累加器 FP32、UB 总大小 <192KB、单值 32B 对齐、1D Grid 不超过物理核数、大矩阵用对角线调度、post-dot 用 tl.parallel(bind_sub_block=True);按需做算子融合与 Double Buffer(tl.multibuffer 或 tl.compile_hint);最后必须做精度、泛化性、性能与端到端回归验证。

限制与反模式包括:不以单一规模数据做决策,不跳过 msprof 直接改 kernel,不为性能牺牲精度或破坏泛化性,避免 FP16 直接归约、非 16 倍数 BLOCK 做矩阵乘、BLOCK 超 UB、热路径用 tensor.item()、循环内 if 分支、2D Tiling 只算数据 buffer、预计算 offset 做 2D broadcasting、kernel 内用 broadcast stride 访问辅助张量、归约类双 pass、msprof 对比 kernel 同名、大矩阵不用对角线调度等。相关加载文档有 optimization-patterns.md、ascend-terminology.md、triton-ascend-api.md、tiling-strategies.md、triton-api-reference.md。