用途:在昇腾 NPU 上对 Vector 类 Triton 算子做单算子深度性能优化,目标是将指定算子性能提升至少用户要求的 x 倍,达标后越高越好。禁止使用入图方式提升性能,只关注单算子独立优化。
主要能力与优化要点:优先批量处理多个 Token 以减少访存;按 UB 容量规划单次循环最大 Token 数 N(UB 总容量 192 KB,取 170 KB 的 50% 即 85 KB 以启用 Double Buffering,满足 N×S_token≤85×1024,用整数除法 // 计算);用 mask 处理尾块;减少 kernel 内 Scalar 运算;index_select 等非连续访问逐行读取;加载与计算交织;多写入流边算边写;用 tl.arange 生成二维索引;避免 tl.where、避免重复 insert_slice;规约优先选最大维度;不变参数用 tl.constexpr,可变参数用动态参数。
使用方式:先配置环境 export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:$LD_LIBRARY_PATH && source /usr/local/Ascend/ascend-toolkit/set_env.sh;用 python -m pytest test_<op_name>.py 验证正确性;用 msprof op --output=<用户指定的路径> --kernel-name="<op_name>kernel" --warm-up=20 --launch-count=20 python test<op_name>_perf.py 测性能,Task Duration(us) 为基线。修改前备份源码,每次修改后都需正确性验证与性能测量,迭代至达标,并按模板输出优化结果报告。
限制:仅适用于昇腾 NPU 与 Triton 单算子场景;tensor 须内存连续;kernel 内 if-else 同名变量 Shape 必须一致;避免 tl.load 的 other 参数;具体性能倍数、支持算子范围与版本要求资料未提供。