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 函数定义
• 支持自动微分
• 支持多种数据类型

接口示例：

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 数据类型支持

阶段 3：规格约束

3.1 输入 Tensor 约束

3.2 输出 Tensor 约束

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 - 完整的文档模板
• ascend-terminology.md - Ascend 术语表
• tiling-strategies.md - Tiling 策略详解

官方文档：

• Triton 编程文档
• Triton-ascend 编程优化指南
• Triton-ascend 性能优化指南