A method for optimizing the attention output correction and rescaling stage of an NPU-based FlashAttention operator in parallel

Abstract

This invention discloses a pipelined parallel optimization method for the attention output correction and rescaling stage of the FlashAttention operator based on an NPU. The method is as follows: 1. Decompose the computation process of the attention output correction and rescaling stage into: a multiplication operation for correcting historical attention outputs and an addition operation for accumulating the attention output of the current block; 2. Move the historical attention output correction multiplication operation, originally located in the next round, to the end of the attention weight calculation stage of the current round; 3. Decentralize the addition accumulation operation in the attention output correction and rescaling stage to the write-back process of the attention output calculation stage; 4. Prune the entire attention operator pipeline. This invention achieves efficient parallel issuance of transport instructions and computation instructions by decoupling the computation logic and deeply rearranging the assembly instructions, thereby improving the utilization rate of hardware units.

Description

Technical Field

[0001] This invention belongs to the field of code optimization technology for core operator algorithms of large models, and relates to a pipelined parallel optimization method. Specifically, it relates to a pipelined parallel optimization method for the attention output correction and rescaling stages of the FlashAttention forward operator (FA operator) based on NPU. Background Technology

[0002] Self-attention, a core component of the Transformer architecture, empowers deep learning models with the crucial ability to capture long-range dependencies in input sequences. However, the time and space complexity of classic self-attention operators increases quadratically with sequence length. This characteristic leads to a sharp increase in memory usage and computational overhead, posing a severe computational bottleneck for large models when processing long input sequences, significantly limiting the scale of model pre-training and the efficiency of inference deployment. Therefore, efficient acceleration optimization of self-attention operators has become a research hotspot in both academia and industry.

[0003] Represented by the FlashAttention series, these methods do not alter the model structure but delve into the storage hierarchy and instruction pipeline characteristics of the accelerator, reconstructing the computation order through block computation and operator fusion. Building upon the foundations laid by FlashAttention-1 / 2 in reducing global memory access and fully utilizing on-chip cache, FlashAttention-3 further introduces a producer-consumer asynchronous parallel model. Utilizing the asynchronous memory copying and thread bundle scheduling mechanisms of the underlying hardware, it achieves extreme latency hiding for matrix multiplication and Softmax operations. The latest FlashAttention-4 is deeply adapted to next-generation hardware architectures, proposing a "lazy" online Softmax rescaling and exponential operation simulation approximation technique. By eliminating up to 90% of redundant rescaling computations at the algorithm level and using high-throughput multiply-accumulate instructions to mask complex vector operation latency, it brings computational throughput close to the physical limits of the hardware. Because this precise acceleration scheme, evolving from memory-constrained to computationally pipelined, can significantly improve training and inference speeds without sacrificing model performance, it has become a core evolutionary direction for large model operator optimization.

[0004] In line with the trend of hardware-aware optimization, the domestic Ascend computing platform has also actively introduced and deepened the engineering adaptation of the FlashAttention series of algorithms. In the Ascend 910B AI processor and its supporting CANN heterogeneous computing architecture, the official high-performance fusion operator library is not limited to basic algorithm implementation, but deeply integrates the block computing and operator fusion architecture of FlashAttention-1 / 2, and further draws on the extreme optimization ideas of FlashAttention-3 / 4 for the underlying hardware pipeline.

[0005] Specifically, the existing Ascend FA operator adaptation scheme cleverly achieves parallel queueing of transport instructions and computation instructions by designing a Ping-Pong buffer strategy on the vector side, and uses an efficient memory handover mechanism to mask data transport latency. In terms of the collaboration between matrix and vector computation units, it introduces an asynchronous multi-stage pipelined parallel mechanism similar to FlashAttention-3 / 4. Through early firing and Cube dual firing strategies, the single-round serial computation link is deconstructed into four highly overlapping stages (attention score calculation, attention weight calculation, attention output calculation, attention output correction and rescaling).

[0006] The calculation process of the FA operator implemented based on Ascend hardware is as follows: Figure 1 As shown, it is mainly divided into four parts: attention score calculation, attention weight calculation, attention output calculation, and attention output correction and rescaling. Attention score calculation means performing Matmul calculation on the Query and the transposed Key. Attention weight calculation means performing position encoding, scaling, Mask, Softmax, and DropOut on the attention score. Attention output calculation means performing Matmul calculation on the result of the attention weight and the Value. Attention output correction and rescaling means multiplying the attention output of the previous round with the intermediate result exp value of the current round of Softmax, completing the correction, and then adding it to the attention output of the current round. The last inner loop iteration is divided by the sum value (the sum of rows) in the Softmax result to complete the rescaling operation of the attention output. The formula for calculating the exp value is shown in formula (1).

[0007]

[0008] To improve the utilization of the Ascend 910B computing unit, existing best practices have adopted a variety of parallel optimization strategies, including V-side pipelined parallelism, CV-side pipelined parallelism, and the IBShare strategy on the C-side.

[0009] The core computing unit in the Ascend chip is the AI ​​Core, which includes three basic computing units: Cube (matrix) computing unit, Vector (vector) computing unit, and Scalar (scalar) computing unit. It also includes storage units (including hardware storage and data transfer units) and data transfer units. The hardware architecture is divided into two types based on whether the Cube and Vector computing units are deployed on the same core: coupled architecture and discrete architecture. The Ascend 910B uses a discrete architecture, and its hardware structure diagram is shown below. Figure 2 As shown.

[0010] The Ascend 910B splits the AI ​​Core into two independent cores: Matrix Calculation (AI Cube, AIC) and Vector Calculation (AI Vector, AIV). Each core has its own Scalar unit, capable of independently loading its own code segments. This decouples matrix and vector calculations, allowing them to work together under unified system software scheduling to optimize computational efficiency. Data is transferred between AIV and AIC via off-chip memory (L2 Cache and GM).

[0011] In the matrix operation module, the SLU controls the scalar instruction flow, the L1 buffer register stores the input matrix data, the L0A and L0B registers store the two input matrix data after data format preprocessing, the CUBE (the unit that actually performs matrix operations) performs various matrix operations, and the L0C register stores the result matrix. The matrix calculation module is primarily responsible for performing matrix multiplication-related operations. In the vector operation module, there is only one general-purpose register, UB (Unified Buffer), used to store all vector data, and the VEC (the unit that actually performs vector calculations) is used to perform all vector operations. The vector calculation module is primarily responsible for performing all vector operations.

[0012] In addition to storage and computing units, AI Core also contains data transfer units. For example... Figure 2 As shown, MTE2 represents a load instruction, with its load path from GM to the L1 Buffer, L0A / L0B of the AIC, and UB of the AIV. MTE1 represents a load instruction within the AIC, with its load path from the L1 Buffer to L0A / L0B. MTE3 represents a load instruction within the AIV, with its load path from UB to GM. Fixp represents a load instruction within the AIC, with its load path from L0C to the L1 Buffer and GM.

[0013] After performing performance analysis on the Ascend 910B AI processor on the FA operator (the corresponding operator in this fusion operator library is aclnnFlashAttentionScore) in the Ascend official fusion operator library (v0.8-8.0RC3), it was found that there is a pipeline blockage phenomenon in the computational pipeline of the attention output correction and rescaling stages.

[0014] The specific bottleneck analysis is as follows:

[0015] By analyzing the pipeline diagram of the attention output correction and rescaling stages, the pipeline diagram is as follows: Figure 3 As shown, in the baseline implementation, the loop body of the attention output correction and rescaling stages mainly consists of basic multiplication and addition instructions, resulting in low arithmetic strength. This sparse computation instruction flow cannot provide enough clock cycles to mask the high latency of data loading and unloading, leading to a serial execution mode of data loading-computation-data loading. The hardware unit spends most of its time in an idling state waiting for data loading. This inefficient instruction fill rate not only wastes the computing power of the vector computation unit but also blocks the pipeline advancement of the entire FA operator. Summary of the Invention

[0016] To address the pipeline congestion problem in baseline algorithms, this invention provides a pipelined parallel optimization method for the attention output correction and rescaling stages of the FlashAttention operator based on an NPU. The baseline algorithm uses the aclnnFlashAttentionScore implemented in the existing cann-ops-adv (v0.8-8.0RC3) repository (Huawei's official open-source repository for fusion operators), which is a forward implementation of FlashAttention. For the attention output correction and rescaling stages, this invention achieves efficient parallel issuance of transport and computation instructions through decoupling and decomposition of the computation logic and deep rearrangement of assembly instructions, thereby improving the utilization of hardware units.

[0017] The objective of this invention is achieved through the following technical solution:

[0018] A pipelined parallel optimization method for the attention output correction and rescaling stages of the FlashAttention operator based on NPU includes the following steps:

[0019] Step 1: The calculation process of the attention output correction and rescaling stage is broken down into: a multiplication operation for correcting the historical attention output, and an addition operation for accumulating the current block attention output. By analyzing their data dependencies, it can be seen that: the multiplication operation depends on the historical attention output and the exponential intermediate factor generated by the attention weight calculation stage; the addition operation depends on the current block attention output generated by the attention output calculation stage and the historical attention output result after the multiplication operation correction.

[0020] Step 2: Utilizing the asymmetry in computational complexity between the attention weight calculation phase and the attention output correction and rescaling phase (e.g., a computational ratio of 8:1), the historical attention output correction multiplication operation, originally scheduled for the next round, is moved to the end of the attention weight calculation phase in the current round. Specifically, in the last few calculation rounds of the attention weight calculation phase, after completing the Softmax calculation, an attention output correction multiplication instruction is inserted. By executing the correction multiplication operation in advance, the vector calculation unit can pre-compile the correction of the historical attention output within the dense computational gaps of the attention weight calculation phase.

[0021] Step 3: Utilizing the hardware characteristics of the matrix computation unit, the addition and accumulation operations in the attention output correction and rescaling stages are devolved to the write-back process of the attention output computation stage. Specifically, by configuring the atomic accumulation optimization option in the matrix multiplication interface of the attention output computation stage, the underlying atomic accumulation mechanism of the hardware is activated. At this time, the current block attention output generated by the attention output computation stage is directly atomically added to the original data in the target storage address (i.e., the storage address of the historical attention output correction multiplication result in Step 2) during the process of writing to global memory through the write-back pipeline.

[0022] Step 4: After completing the forward shift of the historical attention output correction multiplication operation and the hardware reconstruction of the attention output accumulation addition operation, the pipeline of the entire attention operator is pruned. Specifically, by removing all the transport and calculation instructions in the attention output correction and rescaling stages in the intermediate task cycle, only the necessary normalization scaling closing operation in the last iteration of the inner loop is retained.

[0023] Compared with the prior art, the present invention has the following advantages:

[0024] To address the issue of serialization in the computational pipeline during the attention output correction and rescaling stages, and because most loop iterations in these stages only involve multiplication (Mul) and addition (Add), efficient parallelization of computation and transport instructions is difficult to achieve even with pipeline readjustment. This invention decomposes the attention output correction and rescaling computation, retaining only the scaling computation in the final round, while implementing the Mul and Add operations in other ways, thereby achieving more efficient parallelization of transport and computation instructions. Attached Figure Description

[0025] Figure 1 This is a flowchart of the FA operator calculation based on Ascend hardware.

[0026] Figure 2 It uses the Ascend 910B hardware architecture.

[0027] Figure 3 This is the computational pipeline for the attention output correction and rescaling stage of the FA operator.

[0028] Figure 4 This is an abstract diagram of the computational pipeline of the FA operator under the Cube dual-launch strategy.

[0029] Figure 5 It is the overall computational process of the attention output correction and rescaling stage.

[0030] Figure 6 It is the computational pipeline of the attention output correction and rescaling stage after the FA operator is optimized based on the computational logic decoupling and instruction fusion optimization scheme.

[0031] Figure 7 This is an operator correctness verification.

[0032] Figure 8 This is a diagram showing the pipeline optimization verification results for the attention output correction and rescaling stages. Detailed Implementation

[0033] The technical solution of the present invention will be further described below with reference to the accompanying drawings, but it is not limited thereto. Any modifications or equivalent substitutions to the technical solution of the present invention that do not depart from the spirit and scope of the technical solution of the present invention should be covered within the protection scope of the present invention.

[0034] Given the pipeline blocking problem in the attention output correction and rescaling stages of existing FA operators, this invention provides a pipelined parallel optimization method for the attention output correction and rescaling stages of the FlashAttention operator based on an NPU. The optimization scheme of this invention is described in detail below.

[0035] To address the pipeline congestion problem in the attention output correction and rescaling stages, this invention abandons the approach of simply optimizing the pipeline scheduling of the attention output correction and rescaling stages. Instead, it reconstructs the logical topology of the entire FA operator computation graph, merging the originally sparse computation instructions into other dense vector execution stages, thereby achieving instruction flow pruning.

[0036] In the existing implementation scheme, due to the use of a dual-launch mechanism for matrix computation units, the overall computational pipeline abstraction diagram of the attention operator is as follows: Figure 4As shown in the diagram, the attention weight calculation phase and the attention output calculation phase, belonging to the same task, are executed within the same round and have a strong dependency relationship; that is, the attention output calculation phase can only start after the attention weight calculation phase is completed. However, the attention output correction and rescaling phase, belonging to the same task, is shifted one round in execution sequence. This is because the system cannot determine in advance whether the previous round of attention output calculation has been completely completed until the next round of attention output calculation begins. Since the calculation process of the attention output correction and rescaling phase strictly depends on the current block attention output produced by the previous attention output calculation phase, shifting the execution of the attention output correction and rescaling phases by one round is a prudent way to ensure the correctness of data dependencies.

[0037] Furthermore, with Figure 4 For example, the attention output correction and rescaling stage corresponding to the first task does not perform any vector calculations. Instead, it moves the block attention outputs produced by the task in the attention output calculation stage from the temporary space to another specified storage space to complete the initialization process.

[0038] This storage space is denoted as Vec2ResBuf. In the subsequent attention output correction and rescaling stages, this space is used to store the result of the attention output accumulation operation performed in the previous task. Each subsequent task's attention output correction and rescaling stage requires specific vector computations. The overall computation flow for the attention output correction and rescaling stage is as follows: Figure 5 As shown, the dashed line represents the data transmission process of the previous task, and the solid line represents the data transmission process of the current task.

[0039] The attention output correction and rescaling phase comprises two core operations: a historical attention output correction multiplication operation and an attention output accumulation addition operation. The historical attention output correction multiplication operation corrects the accumulated attention output from the previous task. As shown in the diagram, this operation has two operands: the accumulated attention output from the previous task and the exponential intermediate factor generated by the current task during the attention weight calculation phase. The attention output accumulation addition operation accumulates the output results. As shown in the diagram, this operation also has two operands: the current block attention output produced by the current task during the attention output calculation phase and the output result generated by the current task after performing the historical attention output correction multiplication operation during the attention output correction and rescaling phase.

[0040] It is important to note that the block attention output produced by the attention output computation stage, as well as the exponential intermediate factor that serves as an intermediate result in the process, are each configured with two temporary storage spaces for storing the results, corresponding to odd-numbered tasks and even-numbered tasks, respectively.

[0041] For example, the intermediate results generated by the first task in the attention weight calculation stage and the attention output calculation stage need to be delayed until the attention output correction and rescaling stage of the next round before they can be consumed; before that, the attention weight calculation stage and the attention output calculation stage of the second task have started and produced new intermediate data. Therefore, both the block attention output and the exponential intermediate factor need to be configured with two temporary storage spaces.

[0042] The two temporary storage spaces corresponding to the block attention output are denoted as MM2ResBuf[0] and MM2ResBuf[1], respectively; the two temporary storage spaces corresponding to the exponential intermediate factor are denoted as ExpBuf[0] and ExpBuf[1], respectively.

[0043] In order to break down the attention output correction and rescaling stage to achieve instruction rearrangement, this invention analyzes and breaks down each operation in this stage.

[0044] First, regarding the historical attention output correction multiplication operation, in the attention output correction and rescaling stage corresponding to the first task, the main task is to move the results in MM2ResBuf[0] to Vec2ResBuf to complete the initialization. At this time, the attention weight calculation stage of the second task has been completed, that is, ExpBuf[1] has stored the correction terms required by the second task when performing the historical attention output correction multiplication operation.

[0045] It can be observed that during the attention output correction and rescaling phase of the first task, all the data conditions required to complete the multiplication operations in the second task are already available. Therefore, this invention can advance the historical attention output correction multiplication operations in all tasks except the first one by one round.

[0046] After adjustment, the intermediate exponential factors required for the multiplication operation in the attention output correction and rescaling phase of the previous task are generated after the attention weight calculation phase of the current task (within the same round, but with the tasks staggered) is completed. Therefore, historical attention output correction multiplication operations have the potential to be performed earlier in the attention weight calculation phase of the current task.

[0047] The main obstacle in this process is that the system cannot determine in advance whether the attention output calculation phase of the previous task has been completed before the attention output calculation phase of the current task begins. Therefore, to achieve this forwarding, the instruction waiting for the end of the attention output calculation phase of the previous task needs to be executed in the attention weight calculation phase of the current task. However, since the attention output calculation phase of the previous task may cover all or part of the lifecycle of the attention weight calculation phase of the current task, this forwarding operation may cause the execution of the current task in the attention weight calculation phase to be blocked.

[0048] Further observation of existing implementations reveals a significant difference in computational scale between the attention weight calculation stage and the attention output correction and rescaling stage. Specifically, the computational cost of the attention output correction and rescaling stage is approximately 1:8 that of the attention weight calculation stage.

[0049] Since the attention weight calculation phase typically involves 8 or 16 iterations to complete all data calculations, this invention can move the historical attention output correction multiplication operation, originally belonging to the previous task, to the last one or two iterations of the attention weight calculation phase in the current task, after the Softmax calculation. This minimizes the idle time of vector computation units caused by waiting for the attention output calculation phase of the previous task to finish. Furthermore, this invention implements necessary synchronization control mechanisms to ensure the correctness of data dependencies after the instruction logic is moved forward.

[0050] The attention output accumulation operation also needs to be analyzed and broken down. The key to breaking down this addition operation is that the attention weight calculation stage and the attention output calculation stage, which belong to the same task, are executed in the same round, and the two are strongly dependent on each other. That is, the attention output calculation stage can only start after the attention weight calculation stage has completed its calculation.

[0051] By implementing the aforementioned forward-shifting scheme for the historical attention output correction multiplication operation, the temporal dependency between the correction multiplication operation and the attention output calculation stage is effectively strengthened. That is, the attention output calculation stage can only be started after the correction multiplication operation is completed. Therefore, it is only necessary to pre-output the calculation result of the correction multiplication operation to the temporary storage space corresponding to the current task's attention output calculation stage, and utilize the hardware characteristics of the matrix computation unit to enable the atomic accumulation optimization option in the matrix multiplication interface of the attention output calculation stage. This option activates the underlying hardware atomic accumulation mechanism, enabling the current block attention output produced by the attention output calculation stage to directly perform atomic addition with the original data in the target storage address when written to global memory through the write-back pipeline. This mechanism folds the three-step operation of reading the old value → accumulating the new value → writing back, which originally needed to be executed in the vector computation unit, into a single atomic instruction in the matrix computation unit's write-back stage, thus eliminating the redundant process of moving intermediate results into the vector computation unit.

[0052] Finally, redundant instruction stream pruning is performed. Thanks to the aforementioned logic decoupling and hardware integration strategy, the core computational tasks within the loop body of the attention output correction and rescaling stages have been completely stripped away. Therefore, except for the last iteration of the inner loop which needs to retain the necessary normalization scaling operation, this invention removes redundant data transport and computation instructions from all intermediate rounds in the attention output correction and rescaling stages. This refactoring compresses the originally lengthy serial pipeline into a minimally simplistic finishing operation, significantly reducing the total amount of data transport and the overall execution latency of the operator. The abstract diagram of the optimized attention operator's overall computational pipeline is as follows: Figure 6 As shown, for the sake of simplification, the attention score calculation stage is omitted from this diagram.

[0053] Example:

[0054] This experiment was conducted on Huawei Cloud AI Development Platform ModelArts, and the instance specifications selected for the cloud environment were... When compiling and deploying the FA operator, the instance specification needs to be temporarily changed to... To speed up compilation (estimated to take about 30 minutes). ascend-snt9b3 corresponds to the Ascend 910B2 AI processor, which has 24 AICs and 48 AIVs. The kernel function execution mode of the FA operator is KERNEL_TYPE_MIX_AIC_1_2, meaning that 2 AIVs share one AIC.

[0055] The baseline FA operator uses the aclnnFlash-AttentionScoreV2 implementation from the cann-ops-adv (v0.8-8.0RC3) repository (although this repository is iterating, the core computational functionality of aclnnFlashAttentionScoreV2 remains unchanged). The corresponding software environment is cann-toolkit_8.0.0, cann-kernels-910b__8.0.0, torch 2.4.0, and torch-npu 2.4.0.post2. The operator is developed in AscendC, while the test code and model are developed in Python.

[0056] Before starting the experiment, the compilation and deployment of the FA operator need to be completed. The FA operator should be deployed to the CANN package, and then the eager mode should be registered using the cppextension provided by torch. This allows the baseline or optimized FA operator to be directly called within torch-npu for subsequent testing. Assume the corresponding operator name in torch-npu is npu_flash_attention_v2.

[0057] 1. Attention output correction and rescaling baseline implementation

[0058] The baseline implementation of attention output correction and rescaling is relatively simple. For the FlashAttention2 algorithm, it mainly consists of two large loops. The outer loop iterates through the sequence length of the query, denoted as outerLoopLimits, and the inner loop iterates through the sequence length of the key / value pair, denoted as innerLoopLimits. Therefore, the total number of tasks to be executed is outerLoopLimits × innerLoopLimits, and taskIdx grows from 0 to outerLoopLimits.

[0059] ×innerLoopLimits. The calculation process for attention output correction and rescaling is completed only in the inner loop; the outer loop only affects the number of times the entire attention output correction and rescaling calculation process is executed.

[0060] The two temporary spaces for the attention output are denoted as MM2ResBuf[0] and MM2ResBuf[1], the two temporary spaces for the exponential intermediate factor are denoted as ExpBuf[0] and ExpBuf[1], and the buffer space used to store the calculation results after attention output correction is denoted as Vec2ResBuf. The selection of the temporary space mainly depends on the task ID (taskIdx). Even-numbered tasks use temporary space 0, and odd-numbered tasks use temporary space 1. The pseudocode for the attention output correction and rescaling baseline implementation is shown in pseudocode 1.

[0061] Pseudocode 1: Implementation of the ProcessVec2() function

[0062]

[0063] As can be seen from the code, the calculation process of attention output correction and rescaling exhibits a serialized characteristic. First, the attention output is loaded into MM2ResBuf[taskIdxmod2]. If it is the first round, no correction operation is needed at this time, so it is directly loaded into GM to complete the initialization. If it is not the first round, the attention output of the previous round is further loaded into Vec2ResBuf. After the loading is completed, the attention output of the previous round is corrected. The correction term is the exponential intermediate factor (refer to formula (1)). Then it is accumulated with the attention output of the current round. After the accumulation is completed, if it is the last round, the accumulated result also needs to be scaled (this scaling operation comes from the division by the sum of rows in Softmax). Finally, the result is loaded into GM, and the next round of loading can only be started after the loading is completed.

[0064] 2. Implementation of Attention Output Correction and Rescaling Pipeline Optimization

[0065] In the implementation code of attention output correction and rescaling baseline, this invention mainly provides the implementation of the inner loop. However, the implementation of a single round cannot be completed in one step. A nested loop is required within the inner loop. This is because UB cannot store all the data involved in the calculation at once, so further looping is required. The direction of the loop is the sequence length of the query. The direction of the loop for attention weight calculation is also the sequence length of the query. The difference between the two is that the length of the single step execution is different (the load of the two is different). Assuming that the step size of the loop in the attention output correction and rescaling stage is vec2s1Size, and the step size of the loop in the attention weight calculation stage is vec1s1Size, vec2s1Size is usually larger than vec1s1Size and is in a multiple relationship.

[0066] Next, let's shift our focus back to the inner loop. Based on the previous analysis, the attention weight calculation and attention output calculation in the baseline implementation are strongly dependent on being executed within the same loop, while the attention output correction and rescaling are executed in a staggered loop. The baseline implementation uses extroInfo to pre-record all the information required for all calculations under each taskIdx, allowing different calculations to be executed in staggered loops. The overall calculation flow is shown in pseudocode 2.

[0067] Pseudocode 2: Implementation of the Process() function

[0068]

[0069] Based on the extroInfo mechanism, once extroInfo is obtained, any operation corresponding to that extroInfo can be performed, such as early execution or delayed execution. Further, it is necessary to determine the timing of each operation in attention output correction and rescaling. Firstly, during the first execution of attention weight calculation, since attention output calculation has not yet been executed, no operations related to attention output correction and rescaling need to be performed. This time point is the first iteration of the inner loop. Starting from the second execution of attention weight calculation, the first execution of attention output calculation has begun, and at this time, partial attention output correction and rescaling calculations can be performed. Therefore, this invention defines a new attention weight calculation interface, denoted as ProcessVec1withVec2(extroInfoPre, extroInfo). extroInfoPre represents information from the previous task, and extroInfo represents information from the current task. Using the information from the previous task recorded in extroInfoPre, the calculation of the attention output correction and rescaling portion of the previous task can be correctly calculated. Attention output calculation also needs to perform partial attention output correction and rescaling operations starting from the second execution; this new interface is denoted as Cube2withVec2().

[0070] As for the Vector2() interface, since most of the calculations are completed through the ProcessVec1-withVec2(extroInfoPre,extroInfo) and Cube2withVec2() interfaces, Vector2 only needs to be responsible for the scaling operation in the last round. The adjusted Process() code is shown in pseudocode 3.

[0071] Pseudocode 3: Implementation of the ProcessCus() function

[0072]

[0073] As can be observed from the code, the WaitCube2Res() interface has been removed compared to the baseline implementation. This is because some operations in attention output correction and rescaling are moved to the attention weight calculation. However, in the attention weight calculation, it is impossible to determine whether the attention output calculation of the previous task has been completed. Therefore, WaitCube2Res() needs to be inserted in the attention weight calculation to explicitly wait for the attention output of the previous task to be completed and to move the data out. By completing the synchronization of the attention output calculation result of the previous task in advance in the attention weight calculation of the current task, it is not necessary to synchronize again outside.

[0074] Next, we turn our attention back to the attention weight calculation, specifically the `ProcessVec1withVec2-(extroInfoPre, extroInfo)` interface. Based on previous analysis, both the attention weight calculation phase and the attention output correction and rescaling phase loop along the sequence length of the query, but with different step sizes. For example, the attention weight calculation phase loops 16 times, while the attention output correction and rescaling phase might only loop 4 times. To reduce the premature blocking of attention weight calculation by the `WaitCube2Res()` synchronous wait instruction, and also to ensure that the necessary intermediate factors of the correction term have been calculated, the correction of historical attention outputs is placed in the last few rounds of the attention weight calculation. This allows us to easily obtain `vec2loopstartIdx` and `vec2loopLimits` through calculation.

[0075] The code for ProcessVec1withVec2(extroIn-foPre,extroInfo) is shown in pseudocode 4.

[0076] Pseudocode 4: Implementation of the ProcessVec1withVec2(extroInfoPre,extroInfo) function

[0077]

[0078] As can be seen from the code, this invention inserts the historical attention output before the Softmax calculation, and masks the input time by using the calculation time of softmax and dropout. After dropout is completed and the historical attention output is input, the correction multiplication operation of the historical attention output is started.

[0079] It's important to note that the input of historical attention outputs relies on information recorded in `extroInfoPre`, while the source of the intermediate factor for the correction term index depends on information recorded in `extroInfo`. This aligns with the previous analysis of operands in the correction operation. After completing the correction calculation, the correction result needs to be output. The output buffer space is the output space for the current task's attention output calculation. As seen in pseudocode 3, a synchronization instruction `Mte3toMte2` is inserted between the current task's attention weight calculation and attention output calculation. This means that the execution of the attention output calculation needs to wait for the output to be output during the attention weight calculation phase. This indicates that before the attention output calculation is executed, its corresponding output buffer space has already been filled with the correction result of the historical attention output. At this point, the `enAtomic` atomic addition function is enabled for the `Matmul()` high-level API in the attention output calculation. When the result is output during the attention output calculation phase, it is automatically added to the value at the output address. This method completes the accumulation calculation in the attention output correction and rescaling phases.

[0080] Finally, the Vector2() interface is pruned. The new pruned interface, simpleVector2(), is executed only in the last round of the inner loop, and after performing only the scaling operation, the result is moved out to GM. The adjusted Vector2 code is shown in pseudocode 5.

[0081] Pseudocode 5: Implementation of the simpleVector2() function

[0082]

[0083] 6. Operator correctness testing and performance testing

[0084] (1) Operator correctness test

[0085] The correct test of the operator is to compare the execution result of the optimized FA operator with the result of the FlashAttention forward process calculated by pure torch_npu according to the same calculation process. The comparison interface is the assertRtolEqual() interface in the TestCase class of torch_npu, and prec is set to 1e-3.

[0086] Operator correctness test results are as follows Figure 7 As shown, the operator's execution result is correct.

[0087] (2) Operator performance test

[0088] While ensuring correctness, this invention verifies that three pipeline optimizations improve single-operator performance compared to the baseline. The experimental verification test cases uniformly adopted a batch size of 4, an attention head size of 8, a hidden layer dimension size of 128, and the same length for Query and Kev / Valued sequences, selected from the set [2048, 4096, 6144, 8192].

[0089] from Figure 8 As can be seen, the pipelining optimization scheme in the attention output correction and rescaling stages has a significant performance gain, with a speedup ranging from 1.045x to 1.070x. This confirms the bottleneck analysis above—the attention output correction and rescaling stages originally exhibited memory access-dominated serialization characteristics due to the extremely sparse computation instructions (containing only simple multiplication and accumulation) which could not mask the MTE transport latency. By using logical forward migration and hardware atomic accumulation, redundant instruction streams were successfully eliminated, thereby freeing up a huge performance margin.

Claims

1. A pipelined parallel optimization method for the attention output correction and rescaling stages of the FlashAttention operator based on NPU, characterized in that... The method includes the following steps: Step 1: Decompose the calculation process of the attention output correction and rescaling stage into: multiplication operation for correcting historical attention output, and addition operation for accumulating the attention output of the current block. Step 2: Taking advantage of the asymmetry between the computational complexity of the attention weight calculation stage and the attention output correction and rescaling stage, the historical attention output correction multiplication operation, which was originally scheduled for the next round, is moved to the end of the attention weight calculation stage of the current round. Step 3: Utilize the hardware characteristics of the matrix calculation unit to decentralize the addition accumulation operation in the attention output correction and rescaling stage to the write-back process in the attention output calculation stage. Step 4: After completing the forward shift of the historical attention output correction multiplication operation and the hardware reconstruction of the attention output accumulation addition operation, prune the entire attention operator pipeline.

2. The pipelined parallel optimization method for the attention output correction and rescaling stage of the NPU-based FlashAttention operator according to claim 1, characterized in that... The specific steps of step two are as follows: In the last few calculation rounds of the attention weight calculation stage, after completing the Softmax calculation, an attention output correction multiplication instruction is inserted. By executing the correction multiplication operation in advance, the vector calculation unit can complete the correction of the historical attention output in advance during the dense calculation gaps of the attention weight calculation stage.

3. The pipelined parallel optimization method for the attention output correction and rescaling stage of the NPU-based FlashAttention operator according to claim 1, characterized in that... The specific steps of step three are as follows: By configuring the atomic accumulation optimization option in the matrix multiplication interface of the attention output calculation stage, the underlying atomic accumulation mechanism of the hardware is activated. At this time, the current block attention output generated by the attention output calculation stage is directly atomically added to the original data in the storage address of the historical attention output correction multiplication result in step two during the process of writing it to global memory through the write-back pipeline.

4. The pipelined parallel optimization method for the attention output correction and rescaling stage of the NPU-based FlashAttention operator according to claim 1, characterized in that... In step four, the specific steps for pruning the pipeline of the entire attention operator are as follows: remove all the handling and calculation instructions of the attention output correction and rescaling stages in the intermediate task cycle, and only retain the necessary normalization scaling and closing operations in the last iteration of the inner loop.

Images