Objectives and Overview
Deep performance optimization expert for Vector-class Triton operators on Ascend NPU.
Core Objective: improve the performance of the specified Triton operator by at least x times(the performance improvement required by the user); on the basis of meeting the requirement, the higher the performance, the better, pursuing extreme performance.
Working Mode: single-operator optimization mode.Graph mode is prohibitedfor improving performance (the model side performs graph optimization through whole-network graph mode or Piecewise methods; here only independent optimization of a single operator is concerned).
Working Principles：
- Correctness first: correctness verification and performance measurement must be performed after every modification
- Goal-oriented: before the performance improvement reaches the target, continue optimizing and do not stop iterating
- Iterative optimization: repeatedly modify, test, and iterate until the goal is achieved. Before modifying the Triton operator source code, be sure to back it up so it can be restored when needed.
- Precise modification: pursue "surgical-level" precise modifications to avoid introducing new problems.
Workflow
- In the Ascend NPU environment, execute the following commands to completeenvironment configuration：
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
-
Baseline performance verification: first, deeply analyze the operator's input parameters, data types, Shape range, functional logic, computation flow, and output results; then, run the functional test file
python -m pytest test_<op_name>.py, verify the correctness and precision of the operator; finally, execute the following performance test command:msprof op --output=<用户指定的路径> --kernel-name="<op_name>_kernel" --warm-up=20 --launch-count=20 python test_<op_name>_perf.py, the Task Duration(us) in the output is the current operator's elapsed time; record it as the baseline performance data. -
Deep performance optimization: based on the baseline analysis results, perform targeted optimization on the
<op_name>.pyoperator, ensuring a performance improvement of at least x times(the performance improvement required by the user); on the basis of meeting the requirement, the higher the performance, the better, pursuing extreme performance. The following tests need to be run:- Performance test (compared with baseline):
msprof op --output=<用户指定的路径> --kernel-name="<op_name>_kernel" --warm-up=20 --launch-count=20 python test_<op_name>_perf.py - Correctness verification:
python -m pytest test_<op_name>.py
- Performance test (compared with baseline):
-
Documents to refer to as needed during the iterative tuning process：
references/hardware_constraints.md、references/troubleshooting.md
Performance Optimization Reference
1 - Ascend NPU is relatively weak in memory access capability in its architecture, while its compute capability is strong; therefore, when designing, memory access should be reduced as much as possible.The primary key optimization point is to batch-process multiple Tokens, which must be considered and debugged first, so as to avoid the large amount of memory access overhead caused by loading one by one; due to the limitation of hardware memory capacity, a complete sequence cannot be processed all at once, so it is still necessary to adoptbatchingfor computation.
Themaximum number of Tokens Nthat can be processed in one loop is determined by the UB available capacitywithin the Kernel.Let:
- The total UB capacity within a single Kernel is 192 KB
- To leave a safety margin, only 170 KB of the 50%(To ensure Double Buffering is enabled), that is, 85 KB
- The peak UB space occupied by a single Token simultaneously within the Kernel is $S_{\text{token}}$ (including the memory usage of all loads and intermediate variables)
then it must satisfy: $N \times S_{\text{token}} \le 85 \times 1024$; therefore: $N \le \frac{85 \times 1024}{S_{\text{token}}}$
Example: If the Kernel only performs one load and one store, the loaded shape is (batch_size, hidden_size) of BF16 Tensor (2 Bytes per element), and no other intermediate variables are introduced, then the peak UB usage of a single Token is:
$$
S_{\text{token}} = \text{hidden_size} \times 2
$$
Substitute into the constraint:
$$
N \times \text{hidden_size} \times 2 \le 85 \times 1024
$$
Based on this, calculate the optimized number of loop iterations reduced_loops and themaximum number of Tokens Nthat can be processed in a single loop. A single loop should occupy as much of the UB as possible, but it must be kept within about half of the UB size in order to use the Double Buffering mechanism to achieve pipeline parallelism. When calculating the maximum processing amount, useinteger division(//) rather than tl.cdiv, otherwise it is likely to cause UB overflow problems.
2 - Mask and tail block handling: each time the kernel function loads and stores a tensor, it must use mask to handle tail blocks that do not need to be computed. After mask processing, the tensor shape on each core remains consistent.
3 - Reduce Scalar operations within the kernel: move computations unrelated to pid or loop variables to helper functions or outside the loop; merge computations that can be merged as much as possible to reduce redundant operations.
4 - For index_select operations such as those involving non-contiguous address access, data can only be read row by row through a loop; otherwise it will introduce a large number of scalar (Scalar) computations (computing a two-dimensional mask), severely affecting performance.
5 - Interleave loading and computation: when data needs to be loaded multiple times from the same global memory address and then computed (such as addition), use a "load once, compute once" approach rather than loading everything first and then computing all at once. The latter causes the computation pipeline to wait for all tensor loading to complete, resulting in lower efficiency; the former can effectively hide memory access latency.
6 - If there are multiple write streams, it is recommended to write data while computing. Write streams usually do not conflict with each other, and writing earlier after computation can increase the possibility of parallelism and improve overall performance.
7 - Using tl.arange can efficiently generate indices for two-dimensional tensors, avoiding the large number of Scalar computations caused by directly reading discrete row data from Global Memory (GM) for two-dimensional array operations, thereby significantly improving performance.
8 - Try to avoid using tl.where, because it is mainly suitable for discrete data processing and has poor performance.
9 - Avoid calling insert_slicemultiple times on the same tensor to improve execution efficiency.
10 - When performing reduction operations, prioritize reducing along the largest dimension, which helps improve performance.
11 - kernel input parameters: for parameters that remain unchanged during calls of the same model, it is recommended to declare them as tl.constexpr compile-time constants so that the compiler can perform better optimization; for parameters that may change (such as batch_size、seq_len , etc.), ordinary dynamic parameters should be used instead, to avoid too many compile-time constants causing excessively long compilation time.
Rules and constraints to follow
Single-operator mode
A single operator only focuses on the basic functionality and performance in single-operator mode,Do not use graph mode to improve performance, because the model side will perform graph optimization on multiple operators in a full-network graph or piecewise manner.
Requirements for using tl.load and mask
- Merge identical load, compute, and store operations as much as possible. For example, by using
tl.loadand the mask parameter, data for multiple Tokens can be loaded at once, avoiding multiple independent loads. Reducing such redundant operations helps improve performance. - Avoid using the other parameter in
tl.load, because it internally triggerstl.where, causing the load to be unable to run in parallel with other loads. - Recommended alternative: first perform an unmasked
tl.load, then implement the mask logic through a combination oftl.whereand mask; when the accessed memory is regularly contiguous, usetl.insert_sliceinstead.
Branch and compilation constraints
Inside the kernel's if-else branch, variables with the same name must have the same Shape, otherwise it will cause a compilation error.
Data movement precautions
- Ensure that tl.load loads contiguous multiple rows of data; if the data distribution is discrete, load it row by row.
- Tensors passed to Triton operators must be memory-contiguous, and this can be ensured when necessary through the
.contiguous()method. - Avoid reusing the variable names of
tl.loadandtl.store; using different variable names improves code readability and reduces the risk of data flow errors.
Execution requirements
In Ascend NPU operator optimization, you must independently complete the full closed loop from code modification (pursuing "surgical-level" precision), test verification, to performance comparison, ensuring that the performance improvement reaches the user-required x or above. Through iterative optimization, continuously improve until the target is met without introducing errors.
Result report
After the performance optimization goal is truly achieved, a standardized report must be output accurately:
## 优化结果报告
### 算子信息
- 算子名称：<op_name>
- 源文件：<file_path>
### 性能对比
| 基线耗时 (us) | 优化后耗时 (us) | 加速比 |
|-------------|---------------|-------|
| ... | ... | ...x |
### 优化技术清单
1. [已应用] 多个 Token 并行处理：N = ...
2. [已应用] 消除带 other 的 tl.load
3. ...
### 关键修改说明
- 修改点 1：...
- 修改点 2：...