SuperAPI AI 工具集

triton-operator-code-gen

编程与开发

查看项目

更新于 10-16 05:00

技能介绍

用途:把 Ascend NPU 算子设计文档或直接需求转换为可执行的 Triton kernel 代码,并附一个基本正确性(smoke)测试。

核心原则:按“计算逻辑 → Tiling 策略 → 代码实现”的顺序推进,顺序不可颠倒。

主要能力与用法:

  1. 需求确认:提取数学公式、输入输出规格与约束,用伪代码描述计算过程并与用户确认。
  2. Tiling 策略:已有设计文档则按其实现,否则在此设计。核间切分遵循 grid = 物理核数(get_npu_aicore_num() 或 get_npu_vectorcore_num()),核内循环处理多任务并自行负载均衡;超过物理核数的 grid 只会增加开销。UB 空间安全 BLOCK_SIZE = (196608 - 32) / (缓冲区数 × dtype大小) × 0.8。
  3. 生成 kernel:按算子类型选模板——归约类(vector core)、GEMM/注意力(AI core)、激活/损失/索引/MoE/后处理(vector core)、卷积(AI core)。Triton API 优先级为 tl.tanh/tl.erf/tl.sqrt,其次 tl.math.*,最后 libdevice。
  4. 生成测试:单 shape × 单 dtype 的 smoke test,确认可编译运行且结果正确;全面精度验证由 precision-eval skill 负责。

使用方式:需按阶段加载参考资源——设计 Tiling 时读 hardware-architecture.md,生成 Kernel 时读 templates.md,查 API 签名时完整阅读 triton-api-reference.md。

限制与常见陷阱:不使用 PyTorch 替代核心计算(matmul、conv、norm、scan 等),kernel 中不用第三方库或逐元素计算;归约需升 FP32;避免 grid > 65535 或 grid ≠ 物理核数;NPU 上不使用 GPU 专用参数(num_stages/num_warps/num_ctas);循环内不可 return/break,tensor 不支持 [i] 下标,改用 tl.where、tl.gather、tl.extract_slice;大矩阵用对角线调度;归约类用单 pass。资料未提供具体安装命令与价格信息。