Triton Ascend 算法草图设计 Skill
你的任务是基于以下固定配置设计高质量的算法草图(sketch):
- 目标 DSL: triton_ascend
- 目标框架: torch
- 目标后端: ascend
- 目标架构: {{ arch }}
⚠️ 你仅生成算法草图,不生成可执行代码。草图用于指导后续的代码生成(triton-op-coding)。
输入信息
你将获得以下信息:
- 任务描述和规格说明 — 算子任务格式的算子需求(包含
Model类) - 相关的知识和示例 — UnifiedSketch DSL 规范和设计模式(见下方知识加载规则)
- 执行历史 — 之前的设计反馈和优化建议(迭代设计时)
知识加载规则
必选知识(每次设计都加载)
@references/sketch-design.md— UnifiedSketch DSL 语法规范、核心操作、设计模式、最佳实践硬件规格 详细硬件规格参考:
@../npu-arch/references/npu-arch-guide-triton.md和@../npu-arch/references/npu-hardware-params.md使用
read工具读取对应架构的硬件规格文件。
手写优化案例(根据任务选择最相关的 2 个)
根据任务描述中的算子类型,从以下案例中选择最相关的 2 个加载。选择依据:算子类型匹配 > 数据规模接近 > 优化模式相似。
| 类别 | 案例文件 | 核心优化 |
|---|---|---|
| Elementwise | @references/cases/elemwise-broadcast-2d.md |
2D 广播:小维不切分、循环外加载 |
@references/cases/elemwise-broadcast-3d.md |
跨轴 3D 广播:两阶段 kernel | |
@references/cases/elemwise-cast.md |
int8→fp16:二次切分 + 用满 UB | |
@references/cases/elemwise-concat.md |
Slice+Concat 融合:精确切片 load | |
@references/cases/elemwise-zeros.md |
小 shape:少核、减调度开销 | |
| Index | @references/cases/index-histogram.md |
直方图:预排序 + 二分查找 |
@references/cases/index-put.md |
批量 load 索引到 UB、get_element 复用 | |
| MatMul | @references/cases/matmul-swizzle2d.md |
固定核心数 grid、Swizzle2D 块重排 |
| Reduction | @references/cases/reduction-amax-large.md |
M≪N:reduce 轴多核 + 原子 + 二次切分 |
@references/cases/reduction-amax-medium.md |
中等规模:矩阵累加再归约 | |
@references/cases/reduction-amax-small.md |
极小 shape:grid=1 最优 | |
@references/cases/reduction-amin-atomic.md |
原子 amin:两种原子方案对比 | |
@references/cases/reduction-amin-large.md |
超大 1D:二次切分 + 重组 | |
@references/cases/reduction-amin-medium.md |
大 N 维 amin:矩阵 min 再轴归约 | |
@references/cases/reduction-amin-small.md |
1D amin:并行度平衡 | |
@references/cases/reduction-mean-large.md |
mean 行二次切分 | |
@references/cases/reduction-mean-medium.md |
mean reduce 第一轴:重组 | |
@references/cases/reduction-prod-small.md |
prod:tl.reduce + 自定义 mul | |
@references/cases/reduction-sum-fused.md |
elemwise + sum 融合 | |
@references/cases/reduction-sum-large.md |
大规模 sum:重组 | |
@references/cases/reduction-weighted-swiglu.md |
3D SwiGLU backward:reshape + 行二次切分 |
按需加载的知识
| 条件 | 加载文档 |
|---|---|
任务描述中包含 hint 标记(@hint:, @range_hint 等) |
@references/hint-mode.md |
设计模式
- 仔细阅读
task_desc中Model.forward()的参考实现 - 理解算子的数学逻辑和计算模式
- 判断算子类型(elementwise / reduce / matmul / attention / 复合)
- 根据目标硬件架构,选择合适的并行化策略和内存访问模式
- 使用 UnifiedSketch DSL 设计算法草图
输出要求
直接输出 sketch op_name { ... } 格式的算法草图,如果任务描述中包含 hint 标记,在草图末尾附上"设计适用范围"注释(格式见 hint-mode.md)。
设计原则
- 设计清晰的、可理解的算法流程
- 遵循 Ascend NPU 硬件特性的最佳实践(core 级别并行、内存层次)
- 考虑目标硬件架构的优化机会(并行度、内存访问模式、数据对齐)
- 标注优化点和权衡决策(使用
@llm_hint注解) - 数值正确性优先,性能次之
草图特点
算法草图应该:
- 高层抽象: 关注算法逻辑和优化策略,而非实现细节
- 易于理解: 便于 triton-op-coding 转换为可执行的 Triton Ascend 代码
- 包含优化提示: 标注并行化、内存优化、循环展开等机会
思考要求
重要:思考过程中请只做框架级别的分析和决策,例如:
- 算子类型判断(elementwise / reduce / matmul 等)
- 选择什么并行策略(core 级并行、数据切分方式)
- Tile 大小选择(考虑 NPU UB 容量和对齐要求)
- 数据类型如何处理
不要在思考过程中写出完整的草图,完整草图只在最终输出中给出。