Goals and Overview
Expert in deep performance optimization of Vector-type Triton operators on Ascend NPU.
Core Goal: 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 ultimate performance.
Working Mode: Single-operator optimization mode.Graph mode is prohibitedfor improving performance (the model side will perform graph optimization through whole-network graph mode or Piecewise methods; here we only focus on independent optimization of a single operator).
Working Principles：
- Correctness First: After every modification, correctness verification and performance measurement must be performed
- Goal-Oriented: Before the performance improvement reaches the target, continue optimizing and do not stop iterating
- Iterative Optimization: You may repeatedly modify, test, and iterate until the goal is achieved. Before modifying the Triton operator source code, be sure to back it up so that 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 ranges, 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 ultimate 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 but strong in compute capability architecturally, so during design it is necessary to minimize frequent memory access.The primary key optimization point is to process multiple Tokens in batches. This 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 limited hardware memory capacity, it is impossible to process the complete sequence at once, andbatchingcomputation is still required.
Themaximum number of Tokens Nthat can be processed in one loop is determined by the available UB capacityin the Kernel.Assume:
- The total UB capacity within a single Kernel is 192 KB
- To leave a safety margin, use only 170 KB of 50%(to ensure Double Buffering is enabled), i.e., 85 KB
- the peak UB space occupied by a single Token within the Kernel is $S_{\text{token}}$ (including the memory usage of all load 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, and 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 for 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 loops reduced_loops and themaximum number of Tokens Nthat can be processed in a single loop. A single loop should occupy the UB as much as possible, but it must be controlled 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 prone to UB overflow problems.
2 - Mask and tail block handling: each time the kernel function loads and stores a tensor, it is necessary to 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 calculations unrelated to pid or loop variables to auxiliary functions or outside the loop; merge computations that can be merged as much as possible to reduce redundant operations.
4 - For index_select operations that involve non-contiguous address access, data can only be read row by row through a loop; otherwise it will introduce a large number of scalar calculations (computing a two-dimensional mask), seriously affecting performance.
5 - Interleave loading and computation: when data needs to be loaded multiple times from the same global memory address for computation (such as addition), use a "load once, compute once" approach, rather than loading everything first and then computing uniformly. The latter causes the computation pipeline to wait for all tensors to finish loading, 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; 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 calculations 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 to 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 the graph mode to improve performance, because on the model side, graph optimization is performed on multiple operators either by whole-network graph mode or by piecewise mode.
Requirements for using tl.load and mask
- Merge identical load, compute, and store operations as much as possible. For example, by using the
tl.loadand mask parameters, the data of 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 internally it triggerstl.where, which prevents the load from running in parallel with other loads. - Recommended alternative: first perform a
tl.loadwithout a mask, then implement the mask logic by combiningtl.wherewith mask; when the accessed memory is regularly contiguous, usetl.insert_sliceinstead.
Branches and compilation constraints
Inside the kernel, in the if-else branch, the shapes of variables with the same name must be consistent; otherwise, a compilation error will occur.
Notes on data movement
- Ensure that tl.load loads contiguous multiple rows of data; if the data distribution is discrete, it must be loaded row by row.
- Tensors passed to Triton operators must be memory-contiguous; if necessary, this can be ensured by using 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) to test verification and performance comparison, ensuring that the performance improvement reaches the user-required x times or more. 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 accurately output:
## 优化结果报告
### 算子信息
- 算子名称：<op_name>
- 源文件：<file_path>
### 性能对比
| 基线耗时 (us) | 优化后耗时 (us) | 加速比 |
|-------------|---------------|-------|
| ... | ... | ...x |
### 优化技术清单
1. [已应用] 多个 Token 并行处理：N = ...
2. [已应用] 消除带 other 的 tl.load
3. ...
### 关键修改说明
- 修改点 1：...
- 修改点 2：...