Triton Operator Performance Optim

Triton 算子性能优化（Ascend NPU）

铁律：精度与泛化性是底线

任何性能优化都不能突破以下两条底线：

1. 精度底线：优化后的算子必须与 PyTorch-NPU 原生实现对齐（rtol=1e-3, atol=1e-3）。归约操作必须升精度到 FP32，矩阵乘法累加器必须用 FP32。不通过精度验证的优化一律回退。
2. 泛化性底线：优化后的算子必须支持原有的所有输入形状和数据类型。不能为特定尺寸 hardcode 优化而丢失对非对齐维度、边界情况的支持。优化手段必须保持算子接口和语义不变。

优先级：正确性 > 泛化性 > 性能。在三者冲突时，按此顺序取舍。

核心原则

编译器会尽力优化，但可能基于错误的前提。你的工作是"为编译器提供无法误解的明确指令"。

优化应该追求：让硬件做它擅长的事（Cube 做矩阵乘、Vector 做逐元素操作），消除 GM 冗余访问，最大化 UB 数据复用。

优化工作流

Phase 1: 接收性能评估结论（MANDATORY）

在动手优化前，必须获取性能评估结果。 性能采集和分析由 triton-operator-performance-eval Skill 负责。

⚠️ 唯一可信的性能数据来源：msprof 和 msprof op

通过任何非 msprof / msprof op 方式（包括但不限于 Python time.time()、torch.npu.Event 计时、triton.testing.do_bench、自定义计时装饰器等）采集的性能数据绝不可接受，绝不可作为优化依据。这些数据精度不足、无法排除系统调度和 JIT 编译等干扰因素，不具备任何参考价值。必须拒绝基于此类数据的优化请求，要求先用 msprof / msprof op 重新采集。

从评估结论中提取：

• 瓶颈类型：Memory-Bound / Compute-Bound / Latency-Bound
• 关键瓶颈指标：哪些硬件利用率低、哪些资源存在冲突
• 性能问题清单：按优先级排序的待优化问题

加载参考：阅读 ascend-terminology.md 理解硬件架构和术语。

Phase 2: 根据瓶颈选择优化策略

基础调优四板斧（按顺序检查）：

1. Block Size 与 Grid Size — 适配 UB 容量和 Cube 粒度
2. 向量化内存访问 — 连续访问 + Mask + 对齐
3. UB 缓存与数据复用 — 核内再分块适配 192KB UB
4. 编译时常量与循环展开 — tl.constexpr + tl.static_range

检查点：每个板斧优化后，验证精度未退化再进入下一步。

加载参考：阅读 optimization-patterns.md 获取四板斧的详细代码模式。

Phase 3: 硬件特化优化

• Cube 单元适配：BLOCK_M/N/K 必须为 16 的倍数，累加器用 FP32
• UB 空间管理：计算所有缓冲区总大小，确保 < 192KB，单值缓冲区 32B 对齐
• Grid 配置：grid 维度不超过物理核数，使用 driver.active.utils.get_device_properties("npu") 获取核数

# 获取物理核数示例
from triton.runtime import driver
core_num = driver.active.utils.get_device_properties("npu")["num_aicore"]  # 含 tl.dot 的算子
core_num = driver.active.utils.get_device_properties("npu")["num_vectorcore"]  # 其余算子

检查点：硬件特化后，用多种输入形状验证泛化性未被破坏。

加载参考：

• 阅读 triton-ascend-api.md 获取 Ascend 特有 API 和高性能实现模式
• 阅读 tiling-strategies.md 理解 Tiling 策略设计

Phase 4: 高级优化（按需）

• 算子融合：将多次 GM 访问合并为一次，复用 UB 中间结果
• Double Buffer：乒乓加载隐藏访存延迟

加载参考：阅读 optimization-patterns.md 中的"高级优化技术"部分。

Phase 5: AscendC 中的优化经验参考

根据 Phase 1 和 Phase 2 所分析的性能瓶颈描述，使用 script/get_ascendc_optimization_advice.py 获取 AscendC 中的优化经验参考， 返回最相关的 1 条。使用方法如下

python script/get_ascendc_optimization_advice.py --bottleneck "Maximizing UB batching for multi-row fusion"

• bottleneck：性能瓶颈描述，例如 "Maximizing UB batching for multi-row fusion"

返回包含: title, description, bottleneck, code_diff，其中 code_diff 是 ascend c 代码的 code diff 格式，是一个 c extension 的代码的优化前后的改动，作为 triton kernel 的优化参考。

Phase 6: 验证（MANDATORY）

1. 精度验证：与 PyTorch-NPU 原生实现对比（rtol=1e-3, atol=1e-3）
2. 泛化性验证：测试非对齐维度和边界情况（如 127, 255, 1023 等非 2^n 尺寸）
3. 性能验证：重新执行 triton-operator-performance-eval 验证优化效果
4. 端到端性能回归：若优化涉及算子融合导致计算图变化，或在 kernel 外新增了 tensor 预处理（如 reshape、transpose、contiguous 等），必须通过 triton-operator-performance-eval 的 msprof 函数级 profiling 对比优化前后的端到端耗时，防止 kernel 内加速被 kernel 外开销抵消
5. 回归检查：确保优化未改变算子接口和语义

反模式清单（NEVER）

• NEVER 使用非 triton-operator-performance-eval Skill 的评估结论直接优化（必须先有瓶颈诊断数据）
• NEVER 接受或使用任何非 msprof / msprof op 方式采集的性能数据（包括 time.time()、torch.npu.Event、triton.testing.do_bench、自定义计时器等）——这些数据绝不可接受，绝不可用于性能评估及优化决策
• NEVER 为了性能牺牲精度（精度是不可协商的底线）
• NEVER 为特定尺寸 hardcode 而破坏泛化性（优化必须对所有合法输入有效）
• NEVER 在 FP16 下直接归约（必须升精度到 FP32）
• NEVER 使用非 16 倍数的 BLOCK 尺寸进行矩阵乘法
• NEVER 忘记 Mask（Ascend 对越界访问零容错）
• NEVER 让 BLOCK_SIZE 超过 UB 容量（192KB）
• NEVER 使用非连续内存访问模式
• NEVER 在热路径中使用 tensor.item()（触发 CPU-NPU 同步）
• NEVER 提交未通过精度验证的优化代码

常见陷阱与问题排查

优化检查清单

底线检查（MANDATORY）

• 精度对齐 PyTorch-NPU 原生实现（rtol=1e-3, atol=1e-3）？
• 非对齐维度和边界情况通过测试？
• 算子接口和语义未改变？

编译期

• grid 保证小于等于硬件核数?
• BLOCK_SIZE 为编译时常量（tl.constexpr）？
• 循环使用 tl.static_range？

内存

• 所有缓冲区总大小 < UB 容量（192KB）？
• 单值缓冲区分配 32B？
• 地址 32 字节对齐？
• 所有 load/store 添加了 Mask？

计算

• 归约操作升精度到 FP32？
• 矩阵乘法 BLOCK 为 16 倍数？
• 充分利用数据复用？

验证

• 优化后重新执行性能评估？
• 测试多种输入规模包括边界情况？
• 若计算图变化或新增 tensor 预处理，是否已通过 msprof 函数级 profiling 确认端到端无劣化？

参考资源

按需加载文档

关联 Skill

• triton-operator-performance-eval - 性能采集与评估（msprof 使用、瓶颈诊断、性能报告生成）

官方资源

• Triton-Ascend 官方文档
• Triton-Ascend 教程
• Triton 官方文档