SuperAPI AI 工具集

external-gitcode-ascend-triton-operator-code-gen

编程与开发

查看项目

更新于 10-15 17:19

技能介绍

用途:把 Ascend NPU 算子设计文档或直接需求转为可执行的 Triton kernel 代码,并产出单 shape、单 dtype 的基本正确性测试,全面精度验证由 precision-eval skill 负责。

流程与原则:顺序为计算逻辑→Tiling 策略→代码实现,不可颠倒;先用伪代码描述计算并与用户确认。已有设计文档则按其 Tiling 实现,否则自行设计。核间切分要求 grid 等于物理核数(get_npu_aicore_num() 或 get_npu_vectorcore_num()),核内循环处理多任务并自行计算负载;超出物理核数的 grid 只会增加开销。UB 空间参考 196KB,安全 BLOCK_SIZE 按 (196608-32)/(缓冲区数×dtype大小)×0.8 估算。

API 与模板:优先 tl.tanh / tl.erf / tl.sqrt,其次 tl.math.*,最后 triton.language.extra.ascend.libdevice 或 triton.language.extra.libdevice。按算子类型选模板:归约类用 vector core;GEMM/注意力用 AI core 并考虑对角线调度;激活、损失、索引、MoE、后处理用 vector core;卷积用 AI core。可调用 tensor 分配、形状操作与数据搬运等 torch 接口,核心计算必须由 Triton kernel 完成。

反模式与限制:不得用 torch.matmul/mm/bmm/einsum、F.conv*、nn.Conv*/Norm* 或 F.layer_norm 等替代核心 kernel,也不得写了 @triton.jit 却不调用。避免忽略 UB 大小、归约不升 FP32、grid>65535 或 grid≠物理核数、kernel 内用第三方库或逐元素计算、矩阵乘法退化为 vector core、在 NPU 使用 GPU 专用参数、归约用双 pass、for/while 内 return/break、tensor[i] 下标。常见问题包括 UB 溢出(减小 BLOCK_SIZE)、coreDim 超限(增大 BLOCK_SIZE 或设 TRITON_ALL_BLOCKS_PARALLEL=1)、精度损失(升 FP32)、超大 shape 需 int64、归约慢与 L2 颠簸。