SuperAPI AI 工具集

external-gitcode-ascend-triton-operator-performance-optim

设计与创作

查看项目

更新于 10-15 17:19

技能介绍

用途:优化 Ascend NPU 亲和的 Triton 算子性能,在保持正确性与泛化性的前提下提升相对 torch_npu 的执行效率。性能比定义为 Ratio = torch_npu 耗时 / Triton 耗时,Ratio > 1.0 表示 Triton 更快;优先级为正确性 > 泛化性 > 性能。

主要能力:按工作流推进。Phase 0 先审视算法本身;Phase 1 分层评估,先用 time.time() 做端到端快速筛选,不达标再用 msprof 定位 kernel 侧瓶颈;Phase 2 按 Memory-Bound、Compute-Bound、Latency-Bound 分类优化,基础顺序为 Block/Grid Size、连续访存、UB 复用、编译时常量;Phase 3 做硬件特化,如 Cube 的 BLOCK_M/N/K 取 16 倍数且累加器用 FP32、UB 缓冲区总大小小于 192KB 且单值缓冲 32B 对齐、1D Grid 不超过物理核数、大矩阵用对角线调度、post-dot 用 tl.parallel 分到 2 个 vector cores;Phase 4 按需做算子融合与双缓冲。

必要用法:优化后须通过强制验证,精度对齐 PyTorch-NPU(rtol=1e-3, atol=1e-3)不通过则回退,且不得 hardcode 特定尺寸或牺牲泛化性。可加载 optimization-patterns.md、ascend-terminology.md、triton-ascend-api.md、tiling-strategies.md、triton-api-reference.md 等资料。

限制与反模式:不做单一规模决策、不在未确认瓶颈时直接优化 kernel、热路径避免 tensor.item()、循环内避免 if 分支、2D Tiling 需计入 offset/mask/index 数组的 UB 占用、归约类算子避免双 pass、msprof 对比时避免两个 kernel 同名;FP16 直接归约、非 16 倍数 BLOCK 做矩阵乘、BLOCK_SIZE 超 UB、非连续访存均属禁止项。

适用场景:Ascend NPU 上 Triton 算子开发与调优。其余价格、平台等信息暂未提供。