# Triton Operator Design

> 生成适用于 Ascend NPU 的 Triton 算子需求文档。当用户需要设计新的 Triton 算子、编写算子需求文档、进行算子性能优化设计时使用。

- Skill: `small-cat/triton-operator-design` (Agent Skill, multi-file: 4 files)
- Install (CLI): `npx skillmds@latest add small-cat/triton-operator-design`
- Raw SKILL.md: https://api.skillmd.com/api/skills/small-cat/triton-operator-design/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Coding & Dev Tools
- Author: small-cat (https://skillmd.com/u/small-cat)
- Updated: 2026-09-22
- Page: https://skillmd.com/skills/small-cat/triton-operator-design

---


# Triton 算子需求文档生成

## 工作流概览

生成 Triton 算子需求文档分为以下阶段：

1. **需求分析** → 产出：功能定义、竞品对比
2. **原型设计** → 产出：API 接口定义
3. **规格约束** → 产出：输入输出约束、硬件限制
4. **特性实现** → 产出：Tiling 策略、Kernel 实现方案

## 阶段 1：需求分析

### 1.1 功能分析

**必须包含**：

- **算子功能说明**：清晰描述算子的作用和应用场景
- **数学公式**：给出核心计算公式，使用标准数学符号
- **变量说明表**：

| 变量 | 类型 | 含义 | 约束条件 |
|------|------|------|----------|
| [变量名] | [类型] | [含义] | [约束] |

**关键术语**（必须准确使用）：

- **GM (Global Memory)**：全局内存，DDR 上的大容量存储
- **UB (Unified Buffer)**：统一缓冲区，AI Core 中 Vector Core 内部的高速缓存
- **L1 Buffer**：一级缓冲区，AI Core 中 Cube Core 内部的缓存
- **AI Core**：昇腾处理器的计算核心，A2/A3 通常有 24 个，每个包含 1 个 Cube 计算核和 2 个 Vector 计算核
- **Tiling**：数据切分策略，将大任务分解为小块
- **归约操作**：sum、mean、max 等降维计算
- **升精度/降精度**：FP16→FP32 或 FP32→FP16 的类型转换

### 1.2 竞品方案分析

**必须包含**：

- **竞品算子列表**：

| 竞品名称 | 来源框架 | 接口定义 | 实现功能 | 约束限制 |
|----------|----------|----------|----------|----------|
| [名称] | [框架] | [接口] | [功能] | [约束] |

- **对比分析**：
  - 功能对比：各框架支持的功能差异
  - 性能对比：不同硬件平台的性能表现
  - 设计借鉴：可参考的优秀设计

## 阶段 2：原型设计

### 2.1 接口定义

**Triton 接口特点**：

- 使用 Python 函数定义
- 支持自动微分
- 支持多种数据类型

**接口示例**：

```python
def triton_operator(
    input: torch.Tensor,
    param1: torch.Tensor,
    param2: Optional[torch.Tensor] = None,
    eps: float = 1e-6
) -> torch.Tensor:
    """
    [算子功能描述]
    
    Args:
        input: 输入张量，形状为 [..., D]
        param1: 参数1，形状为 [D]
        param2: 参数2（可选），形状为 [D]
        eps: 微小常数
    
    Returns:
        输出张量，形状与 input 相同
    """
    pass
```

### 2.2 接口说明表

| 参数名称 | 类型 | 输入/输出 | 说明 | 约束条件 |
|----------|------|-----------|------|----------|
| [参数名] | [类型] | [输入/输出] | [说明] | [约束] |

### 2.3 数据类型支持

| 接口类型 | 支持的数据类型 | 数据格式 |
|----------|----------------|----------|
| Triton | FLOAT16, BF16, FLOAT | ND |

## 阶段 3：规格约束

### 3.1 输入 Tensor 约束

| 约束项 | 约束条件 | 说明 |
|--------|----------|------|
| Shape | [具体约束] | [说明] |
| 数据类型 | [支持的类型] | [说明] |
| 数据格式 | ND | 统一使用 ND 格式 |
| 内存对齐 | 16字节或32字节 | 硬件要求 |

### 3.2 输出 Tensor 约束

| 约束项 | 约束条件 | 说明 |
|--------|----------|------|
| Shape | [具体约束] | [说明] |
| 数据类型 | [具体约束] | [说明] |

### 3.3 硬件约束

**必须考虑的硬件限制**：

- **AI Core 架构**：
  - A2/A3 通常有 24 个 AI Core
  - 每个 AI Core 包含 1 个 Cube 计算核和 2 个 Vector 计算核
  - Cube Core 专用于矩阵计算，Vector Core 专用于向量计算
- **UB 缓冲区大小**：192KB（A2/A3），Vector Core 专用
- **L1 Buffer 大小**：通常为 1MB（A2/A3），Cube Core 专用
- **内存对齐要求**：
  - UB 缓冲区必须 32 字节对齐
  - 单值缓冲区（如均值）需要 32B 空间（即使逻辑上只需 4B）
- **数据类型大小**：FP16=2B, BF16=2B, FP32=4B

## 阶段 4：特性实现方案

### 4.1 Tiling 切分

**这是最关键的部分，必须详细说明**。

#### 4.1.1 核间切分策略

**必须包含**：

1. **切分原则**：
   - 如何划分任务到多个 AI Core
   - 为什么选择这种切分方式
   - 如何保证负载均衡

2. **计算方法**：
   ```
   输入: x[B, D]
   
   // 步骤1: 计算每个 Core 处理的数据量
   data_per_core = ceil(total_size / num_cores)
   
   // 步骤2: 计算当前 Core 的数据范围
   core_start = core_id * data_per_core
   core_end = min((core_id + 1) * data_per_core, total_size)
   ```

3. **示例**：
   - 给出具体的输入形状
   - 展示切分结果
   - 说明每个 Core 处理的数据范围

#### 4.1.2 核内循环策略

**必须包含**：

1. **UB 空间计算**：

   ```
   UB 总大小: 192KB
   数据类型大小: FP16=2B, FP32=4B
   
   单次循环需要的缓冲区:
   - 输入缓冲区: [大小] × [类型大小]
   - 中间缓冲区: [大小] × [类型大小]
   - 输出缓冲区: [大小] × [类型大小]
   
   单次循环可处理的数据量 = UB总大小 / 单次循环总空间
   ```

2. **缓冲区分配策略**：
   - 列出所有需要的缓冲区
   - 说明每个缓冲区的大小和用途
   - 考虑对齐要求

3. **精度处理策略**：
   - 是否需要升精度计算（FP16→FP32）
   - 在哪个阶段升精度
   - 在哪个阶段降精度

### 4.2 Kernel 实现

#### 4.2.1 计算流图

**必须绘制数据流图**：

```
输入张量 (GM)    参数张量 (GM)
    │                │
    ▼                ▼
[加载到UB]       [加载到UB]
    │                │
    ▼                ▼
[计算步骤1]      [预处理]
    │                │
    ▼                │
[计算步骤2]      ───┘
    │
    ▼
[最终计算]
    │
    ▼
输出张量 (GM)
```

**关键点**：

- 标注每个步骤的数据类型
- 标注 GM↔UB 的数据传输
- 标注精度转换的位置

#### 4.2.2 核心实现逻辑

**按输入数据类型分别说明**：

**FP32 输入类型**：
1. 核间任务分配
2. UB 缓冲区管理（列出所有缓冲区）
3. 计算流程（按步骤详细说明）

**FP16/BF16 输入类型**：
1. 核间任务分配
2. UB 缓冲区管理（包括升/降精度缓冲区）
3. 计算流程（包括精度转换步骤）

**硬件优化要点**：
- 向量化计算
- 数据复用
- 内存访问优化
- 对齐处理

## 绝对不要做的事

- ❌ 使用模糊的术语（如"适当切分"、"合理分配"）
- ❌ 忽略硬件约束（UB 大小、对齐要求）
- ❌ 不说明 Tiling 的具体计算方法
- ❌ 不区分不同数据类型的处理策略
- ❌ 不标注数据流图中的数据类型
- ❌ 忽略归约操作的对齐要求（必须 32B）
- ❌ 混淆 Vector Core 和 Cube Core 的用途（Vector Core 用于向量计算，Cube Core 用于矩阵计算）
- ❌ 忽略 UB 和 L1 的区别（UB 是 Vector Core 专用，L1 是 Cube Core 专用）

## 常见陷阱

### 陷阱 1：忽略 UB 大小限制

**症状**：设计的方案超出 UB 容量

**解决**：
1. 计算所有缓冲区的总大小
2. 确保总大小 < UB 总大小
3. 如果超出，调整单次循环处理的数据量

### 陷阱 2：忽略内存对齐

**症状**：硬件报错或性能下降

**解决**：
1. UB 缓冲区按 32 字节对齐
2. 单值缓冲区（均值、方差等）分配 32B 空间
3. 使用 `ceil(实际大小, 32)` 计算分配空间

### 陷阱 3：精度损失

**症状**：FP16 输入时计算结果不准确

**解决**：
1. 归约操作前升精度到 FP32
2. 在 FP32 精度下完成所有计算
3. 最后再降精度到输出类型

### 陷阱 4：Tiling 策略不合理

**症状**：性能不佳或无法处理大 shape

**解决**：
1. 根据算子特点选择切分维度
2. 确保每个 Core 独立完成计算
3. 避免跨 Core 的数据依赖

## 质量检查清单

在完成文档后，检查以下内容：

### 需求分析
- [ ] 算子功能说明清晰
- [ ] 数学公式正确且完整
- [ ] 变量说明表包含所有关键变量
- [ ] 竞品分析涵盖主流框架
- [ ] 术语使用准确（GM、UB、AI Core 等）

### 原型设计
- [ ] 接口定义完整
- [ ] 参数说明详细
- [ ] 数据类型支持明确

### 规格约束
- [ ] 输入输出约束完整
- [ ] 硬件约束明确（UB 大小、L1 大小、对齐要求）
- [ ] 区分 Vector Core 和 Cube Core 的用途
- [ ] 边界条件说明清楚

### Tiling 切分
- [ ] 核间切分策略有具体计算方法
- [ ] 核内循环策略有 UB 空间计算
- [ ] 缓冲区分配详细
- [ ] 有具体示例

### Kernel 实现
- [ ] 计算流图清晰
- [ ] 数据类型标注完整
- [ ] 不同输入类型分别说明
- [ ] 硬件优化要点明确

## 参考资源

详细的设计指南和示例，请参考：
- [triton-operator-template.md](references/triton-operator-template.md) - 完整的文档模板
- [ascend-terminology.md](references/ascend-terminology.md) - Ascend 术语表
- [tiling-strategies.md](references/tiling-strategies.md) - Tiling 策略详解

**官方文档**：
- [Triton 编程文档](https://triton-lang.org/main/index.html)
- [Triton-ascend 编程优化指南](https://gitcode.com/Ascend/triton-ascend/tree/main/docs)
- [Triton-ascend 性能优化指南](https://gitcode.com/Ascend/triton-ascend/blob/main/docs/zh/migration_guide/performance_guidelines.md)