用途:在昇腾(Ascend)NPU 上对指定 Triton 向量类算子做单算子深度性能优化,目标是把算子性能提升到用户要求的至少 x 倍,未达标则持续迭代。明确禁止用入图方式提性能,因为图优化由模型侧整网入图或 Piecewise 完成。
主要能力与优化思路:优先保证正确性,每次修改后做正确性验证与性能测量,改前备份源码,追求手术级精准修改。核心手段包括按 Unified Buffer 容量规划单次循环可处理的最大 Token 数 N(单 Kernel UB 总容量 192 KB,留余量后按 170 KB 的 50% 即 85 KB 估算,N×S_token ≤ 85×1024,用整数除法而非 tl.cdiv 计算最大处理量,以启用 Double Buffering);批量并行处理多 Token 减少访存;用 mask 处理尾块并保持各核 tensor 形状一致;把与 pid、循环变量无关的 Scalar 计算外提或合并;index_select 等非连续访问逐行读取;加载与计算交织以隐藏访存延迟;多写入流边算边写;用 tl.arange 生成二维索引避免读离散行;避免 tl.where、避免对同一 tensor 多次 insert_slice;规约优先选最大维度;循环期间不变的参数用 tl.constexpr,可能变化的(如 batch_size、seq_len)用普通动态参数。
使用方式:先配置环境 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) 为基线并与优化后对比,达标签出标准化优化结果报告(算子信息、性能对比表、优化技术清单、关键修改说明)。可按需查阅 references/hardware_constraints.md 与 references/troubleshooting.md。
限制与约束:仅限单算子模式;kernel 内 if-else 分支同名变量 Shape 必须一致;tl.load 应加载连续多行,离散数据逐行加载;传入 tensor 须内存连续,必要时用 .contiguous();避免在 tl.load 中用 other 参数(内部触发 tl.where 妨碍并行),可先无掩码 load 再用 tl.where 加 mask,或在访问连续时用 tl.insert_slice;避免复用 tl.load 与 tl.store 的变量名。其他平台支持、收费与价格信息暂未提供。