NPU-based sparse attention operator acceleration method

Abstract

The invention discloses a sparse attention operator acceleration method based on an NPU, and the method comprises the steps: combining a Linformer with an Ascend C programming language, and achieving an Ascend C operator for calculating the Linformer based on the NPU; according to the method, an accelerated calculation mode of combining Linformer with Flash-Attention2 is adopted, so that the whole operator can be subjected to calculation in a blocking manner; according to hardware characteristics of the mercuric chloride NPU, optimization of vector calculation is carried out, the number of instructions used by multiple vector calculations such as vector subtraction, line-by-line summation and line-by-line max calculation for calculating safety softmax is reduced, and time consumption is reduced; aiming at the problem that front and back data cannot be parallel due to more dependence, the calculation efficiency of a sparse approximate attention operator is improved by using an algorithm for constructing parallel calculation by emitting Cube in advance; basic block segmentation is carried out on matrix multiplication, and the carrying time of matrix multiplication is shortened.

Description

Technical Field

[0001] The present invention relates to a code implementation and optimization method for a sparse attention operator, and in particular to an NPU-based sparse attention operator acceleration method. Background Art

[0002] Ascend C is a programming language launched by Huawei for the Ascend AI processor. One of the characteristics of this programming language is its strong correlation with hardware. The core computing unit in the Ascend chip is the AI ​​Core, and its abstract structure is as follows: Figure 1 shown.

[0003] Huawei's Ascend NPU910B4 is designed with a separate architecture, splitting the AI ​​processor into two independent cores: matrix computing (AI Cube, AIC) and vector computing (AI Vector, AIV). Each core has its own Scalar unit and can independently load its own code segment, thereby decoupling matrix and vector computing. Under the unified scheduling of the system software, they work together to achieve the effect of optimizing computing efficiency. Data is transferred between AIV and AIC through Global Memory, allowing them to perform matrix and vector calculations in parallel without data dependencies. Its core structure includes three key components: computing unit, storage unit, and transport unit. These three work together to support the overall performance and efficiency of the system.

[0004] The AI ​​processor contains three basic computing units: Cube (matrix) computing unit, Vector (vector) computing unit and Scalar (scalar) computing unit. The Vector computing unit is responsible for performing vector operations, the Cube computing unit is responsible for performing matrix operations, and the Scalar computing unit is responsible for various types of scalar data operations and program flow control.

[0005] There is an internal storage unit in the AI ​​processor, and the data in the external storage needs to be loaded into the internal storage to complete the corresponding calculation. Figure 1 The GM in the operator is generally used to store the input and output of the operator as a whole, as well as the intermediate quantities generated by the operation. The size of the internal storage unit is much smaller than the external storage unit.

[0006] Figure 1 The main storage units in the AIC architecture core in the upper and middle part are L1 Buffer: L1 buffer is used to temporarily store part of the input of matrix operation, thereby reducing the number of reads and writes from outside the core; L0ABuffer and L0BBuffer are used to store the two matrices currently in matrix multiplication operation; L0C Buffer stores the result matrix of the current matrix multiplication operation.

[0007] Figure 1 The main storage unit in the AIV architecture core in the middle and lower half is the Unified Buffer, which serves as a unified buffer to store the input and output of the AIV architecture vector calculation and scalar calculation.

[0008] To facilitate data transmission and handling in the AI ​​processor, the AI ​​processor also includes three types of handling units: the MTE1 handling unit is generally responsible for moving the matrix data that needs to be calculated from the L1 buffer to the L0A buffer and L0B buffer; the MTE2 handling unit is responsible for moving the data that needs to be calculated from the global memory to the AIV core's cache Unified Buffer, and is also responsible for moving data from the global memory to the AIC core's cache L1 buffer; the MTE3 handling unit is responsible for moving the required data from the AIV core's cache Unified Buffer to the global memory, and is also responsible for moving data from the AIC core's cache L0C buffer for temporary storage of results to the global memory.

[0009] Instructions in the AI ​​processor will be scheduled to five independent classification queues (including Vector instruction queue, Cube instruction queue and MTE1 / MTE2 / MTE3 instruction queue), and then assigned to the corresponding execution units for processing. Instructions in the same instruction queue are executed in sequence according to the order in which they enter the queue, while different instruction queues can run in parallel. The overall execution efficiency can be improved by the parallel execution of multiple instruction queues. For data dependency problems that may exist in the parallel process, the event synchronization module is used to insert synchronization instructions to coordinate the synchronization operations of the pipeline. The event synchronization module controls the execution status of each pipeline at all times and analyzes the dependencies between different pipelines to solve data dependency and synchronization problems.

[0010] Since the main computing units for kernel function calculations are vector computing units and cube computing units, which are also divided into two different cores: AIV core and AIC core, there are three different programming scenarios when implementing Ascend C operators: vector programming scenarios using only AIV cores, matrix programming scenarios using only AIC cores, and fusion operator programming scenarios using both AIV and AIC cores.

[0011] Vector calculations are completed through the Vector calculation unit. The source and destination operands of the vector calculation are stored in the UnifiedBuffer, and multiple iterative calculations can be performed. In each iteration, the Vector calculation unit takes out 8 consecutive datablocks (each datablock is 32 bytes) from the UnifiedBuffer, performs calculations, and writes the results back to the corresponding 8 datablocks. The vector calculation unit supports vector calculations with intervals, which can be configured through dataBlockStride (the address step size of different datablocks in the same iteration) and repeatStride (the address step size of the same datablock between adjacent iterations).

[0012] The runtime of the Ascend operator is closely related to the hardware architecture. The main computational time can be roughly divided into vector computation time, scalar computation time, matrix multiplication computation time, and data transfer time. For vector computation time, it is necessary to reduce vector computation instructions by leveraging the vector computation unit's ability to iterate multiple times. Controlling the address interval calculation of operands can reduce vector computation instructions and cache space usage; changing the implementation of vector computations can also reduce vector computation instructions. For scalar computation time, since the scalar computation unit also controls program flow, scalar computations are converted to vector computations whenever possible to avoid synchronization blocking issues. For matrix multiplication, the computation time is fixed when the left and right matrices are of fixed size. However, because matrix multiplication involves block-based computation, the matrix transfer time is affected when the basic block sizes vary, which can be optimized. In addition to reducing the computation time of individual components, the Ascend C operator not only implements parallel scheduling from multiple perspectives, but also from the perspective of the computation unit.

[0013] When performing attention calculations, the time complexity and space complexity of the calculations are O(N 2Wang et al. (Wang S, Li BZ, Khabsa M, et al. Linformer: self-attention with linear complexity [A]. arXiv, 2020.) wanted to approximate the attention matrix through low-rank matrix decomposition, compressing the original N×N attention matrix to N×k, where k is much smaller than N, and thus proposed the Linformer model. Linformer proves through the Johnson–Lindenstraus theorem that when the dimension of the low-rank matrix is ​​large enough (specifically, the logarithm of the sequence length), it can approximate the original attention matrix P with a very small error. At the same time, Linformer proves that singular value decomposition obtains most of the information of the low-rank matrix in the same dimension. However, since singular value decomposition requires decomposing the matrix into three matrices, it increases the computational complexity. Therefore, Linformer proposes to approximate the attention matrix by projecting the K and V matrices into low dimensions. The low-rank assumption is consistent with the attention distribution characteristics in many natural language processing tasks and is suitable for NLP tasks such as long document classification and summary generation.

[0014] First, the input and output of the Linformer operator are given:

[0015] The input of the operator is five: Q, K, V matrices (all of which are in the shape of R N×d ), and the E and F matrices for dimension reduction K and V matrices (both of which are in the shape of R N×k ), the output of the operator is the result matrix O (whose shape is R N×d ).

[0016] Then give its calculation formula:

[0017] Attention(Q,K,V)=softmax(Q(EK) T )(FV)(1)

[0018] When performing calculations, calculate the matrix multiplication EK to obtain the matrix K', calculate the matrix multiplication FV to obtain the matrix V', and then calculate the matrix multiplication QK' T To ensure that the softmax calculation does not cause numerical overflow, we need to find the maximum value of each row of matrix S and subtract the maximum value of each row from each row. x , and then sum each row to get the matrix P to achieve the safe softmax calculation effect. Finally, calculate the matrix multiplication PV' to get the result matrix.

[0019] With the rapid development of AI technology, the importance of AI processors in various application scenarios has become increasingly prominent. Huawei's Ascend NPU, with its decoupled architecture that decouples matrix and vector computations, and its accompanying Ascend C programming language, has opened up a new world in AI computing. In the field of natural language processing, the Linformer model effectively reduces the complexity of attention calculations through low-rank matrix decomposition, demonstrating strong application potential. However, when attempting to deploy the Linformer model on the Ascend NPU, the following challenges arise:

[0020] 1. Currently, there is no Linformer operator that can run on Ascend NPU.

[0021] 2. If the Linformer operator's calculation method is not changed and it is directly migrated to the Ascend NPU, the softmax problem of directly calculating large matrices will still occur. The Ascend NPU's vector computing module has limited cache, making direct calculation impossible.

[0022] 3. Linformer steps require many instructions to complete calculations. For example, for steps such as vector subtraction, row-by-row summation, and row-by-row max calculation for calculating safesoftmax, if these steps are calculated directly, it will take a long time due to too many instructions.

[0023] 4. Linformer has many dependencies between multiple steps, requiring serial calculations, which cannot fully utilize the hardware features of the Ascend NPU.

[0024] 5. Linformer's matrix multiplication block calculation requires a long transfer time. Summary of the Invention

[0025] The present invention provides an NPU-based sparse attention operator acceleration method, which combines Linformer with the Ascend C programming language and implements an Ascend C operator for calculating Linformer based on the NPU. This operator belongs to the hybrid programming operator in the Ascend C operator classification and implements the four stages of the overall operation to form an Ascend C hybrid programming operator with certain scalability and robustness. Linformer is combined with the accelerated calculation method of Flash-Attention2, so that the entire operator can be calculated in blocks. Vector calculation is optimized according to the hardware characteristics of Ascend NPU, reducing the vector subtraction, row-by-row summation, row-by-row max calculation, and The number of instructions used for multiple vector calculations is reduced, reducing the time consumption; to address the problem of multiple dependencies between front and back data that cannot be parallelized, an algorithm that constructs parallel calculations by emitting Cube in advance is used to improve the computational efficiency of the sparse approximate attention operator; matrix multiplication is divided into basic blocks to reduce the transfer time of matrix multiplication.

[0026] The purpose of the present invention is achieved through the following technical solutions:

[0027] A sparse attention operator acceleration method based on NPU, comprising the following steps:

[0028] Step 1: Divide the Q matrix into blocks according to N dimensions. piece The size of each block is Br×d, and the E and F matrices are divided into blocks according to the k dimension. Blocks, each block size is B c ×N;

[0029] Step 2: Name the block of the Q matrix Q i , the blocks of E and F matrices are named E j , F j ;

[0030] Step 3: The outer loop traverses the blocks according to the Q matrix, and the inner loop traverses the blocks according to the E and F matrices. In the inner loop, the Linformer operator is divided into four steps:

[0031] Step 1: Calculate the similarity matrix:

[0032] Step 1-1: Calculate the matrix multiplication E in the matrix multiplication calculation unit j K gets the matrix K' j , that is, the K matrix block after dimensionality reduction;

[0033] Step 1-2: Calculate matrix multiplication in the matrix multiplication calculation unit Get the matrix S' ij , as the preliminary similarity matrix block;

[0034] Step 2: Perform safe softmax calculation on the similarity matrix blocks:

[0035] Step 2-1, scale the similarity matrix and change S' ij Move from global memory to local memory of vector computing unit, for S' ij Perform vector multiplication

[0036] Step 2-2, safe softmax calculation:

[0037] (1) If it is the first time of the inner loop, the vector calculation unit will calculate S' ij Perform the max operation row by row to obtain the inmax vector, and the vector calculation unit calculates S' ij The element of the vector is subtracted from the maximum value of each row, that is, the inmax vector is subtracted, and the vector calculation unit calculates S' ij All elements of S' are natural exponentials, and the vector calculation unit is ij Perform row-by-row summation to obtain the insum vector;

[0038] (2) If it is not the first time of the inner loop, the vector calculation unit calculates S' ij Perform the max operation row by row to obtain the inmax vector. The vector calculation unit calculates the max vector in inmax and the max in the previous inner loop. The vector calculation unit calculates S' ij Subtract the max vector from the elements of the inner loop, take the natural exponent and sum them row by row to get the insum vector, calculate the difference between the max vector of the previous inner loop and the current temporary maximum value of each row max vector after the update of this inner loop through vector subtraction, then take the natural exponent to get the expmax vector used to update the temporary result matrix and each row and insum vector, finally update the current temporary sum vector for each row by multiplying the expmax vector and the sum vector of the previous inner loop and add the insum vector of this round, and update the latest sum and max to the insum vector and inmax vector respectively through transfer;

[0039] Step 3: Calculate the temporary result matrix blocks:

[0040] Step 3-1: Calculate the matrix multiplication F in the matrix calculation unit j V gets the matrix V' j , that is, the V matrix block after dimensionality reduction;

[0041] Step 3-2: Calculate the matrix multiplication S' in the matrix multiplication calculation unit ij V' j Get the temporary result matrix O' ij ;

[0042] Step 4: Update and accumulate the results of the previous inner loop:

[0043] If it is the first time of the inner loop, no update and accumulation process is required; if it is not the first time of the inner loop, the temporary result matrix O' of the previous round is ij-1 Move from global memory to the local memory of the vector computing unit, and calculate the expmax vector and O' by vector multiplication with intervals ij-1 Update to get the correct O' ij-1 , the temporary result matrix O' of this roundij Move from global memory to the local memory of the vector computing unit, and add the correct O' ij-1 and O' ij Add to get a new temporary result matrix O' ij ; If it is the last round of the inner loop, the result matrix O is obtained at this time i It is necessary to update the globally correct result matrix O by performing vector division with the sum vector i ;

[0044] Step 4: Optimize vector calculations based on the hardware characteristics of the NPU:

[0045] Step 4.1. Optimize the vector calculation for calculating safe softmax: Focus on the characteristics of the vector calculation unit and perform calculations in blocks. Each iteration can perform vector subtraction calculations on 8 data units. Instead of performing m iterations in one loop, only n / 128 loops are required.

[0046] Step 4.2: Optimize vector calculation for row-by-row summation and row-by-row max:

[0047] Step 1: Apply for a blank space of (m,64);

[0048] Step 2: Perform vector addition / max calculation on the first (m, 64)-shaped vector and the second block. The resulting vector is recorded in the blank space of (m, 64). The remaining blocks are also accumulated / maxed with the vectors in the blank space by vector addition / max calculation.

[0049] Step 3: If the column dimension of the matrix cannot be divided evenly by 64, then the remainder is used as the tail block and the accumulation / max is performed in the same way. The row-by-row summation / max of the matrix is ​​converted into the row-by-row summation / max of the (m, 64) space that is convenient for the vector calculation unit to calculate.

[0050] Step 4.3. O i =DIV(O i ,sum) vector calculation is optimized:

[0051] Step 1: The address step size of different datablocks in the same iteration is configured as 1, and the address step size of the same datablock between adjacent iterations is configured as d / perElementByBlk, where perElementByBlk is the number of elements contained in each datablock;

[0052] Step 2: Set the address step size of different datablocks within the same iteration of expmax / sum to 0, and the address step size of the same datablock between adjacent iterations to 1. This will reduce the number of calculation instructions used during the calculation.

[0053] Step 5: Optimize the matrix multiplication calculation by reasonable blocks according to the hardware characteristics of the NPU:

[0054] Step 51: Based on the (M, N, K) matrix multiplication and the cache size limits of the AIC core's L0ABuffer, L0B Buffer, and L0CBuffer, set size limits for basicM, basicN, and basicK. The left matrix basicM*basicK after dicing can be placed in the L0ABuffer, the right matrix basicK*basicN after dicing can be placed in the L0BBuffer, and the calculated partial results can be placed in the L0C Buffer.

[0055] Step 52: According to the size limits of basicM, basicN, and basicK, the largest basicM and basicN, as well as the basicK that matches them, are selected to divide the left and right matrices into basic blocks;

[0056] Step 53: When performing matrix multiplication, the basic block is first moved from global memory into the AIV core's L1 Buffer, then from the L1 Buffer to the L0A Buffer and L0B Buffer. The basic block matrix multiplication operation then begins, and the result is stored in the L0C Buffer. Because matrix multiplication requires cumulative summation after block calculation to obtain the correct result for a block of the large result matrix, after performing the K / basicK basic block matrix multiplication operation, the correct partial result of the matrix multiplication is moved out until all basic block operations are completed.

[0057] Step 6: Perform pipeline parallel optimization based on the hardware characteristics of the NPU:

[0058] The calculation flow of the Linformer operator is divided into Cube1: Matrix K' is obtained by EK matrix multiplication, Cube2: Matrix S is obtained by matrix multiplication of the transpose of Q and K', Cube3: Matrix V' is obtained by FV matrix multiplication, Cube4: Result matrix O is obtained by PV' matrix multiplication, Vector1: Normalize the calculation result S of Cube2 and calculate Flashsoftmax to obtain matrix P, Vector2: Update and accumulate the results of the previous inner loop;

[0059] Step 61: Allocate an additional block of memory to each computational flow in the global memory;

[0060] Step 62: After the calculation of Cube1 and Cube2 in the first inner loop is completed, do not proceed with the subsequent process. Calculate Cube1 and Cube2 in the second inner loop in another space first. The synchronous parallel calculation requires the calculation of Vector1 that depends on the results of Cube1 and Cube2 in the first inner loop.

[0061] Step 6.3: Calculate Cube3 and Cube4 in the first inner loop, and calculate Vector1 in the second inner loop in parallel. Calculate Cube1 and Cube2 in the third inner loop, and calculate Vector2 in the first inner loop in parallel. Calculate Cube3 and Cube4 in the second inner loop.

[0062] Step 64: Continue the loop as described above, and determine which of the two memory blocks should be written to based on the current flag. If there is a data dependency, the flag should also be used to determine which block the data needs to be taken from.

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

[0064] 1. This paper organically combines sparse attention calculation with the Ascend C programming language and the accelerated calculation method of Flash-Attention2, allowing the entire operator to be calculated in blocks, realizing a completely innovative sparse attention operator.

[0065] 2. This paper optimizes multiple vector calculations such as vector subtraction, row-by-row summation, and row-by-row max calculation for the safe softmax calculation of the sparse attention operator implemented on the Ascend NPU.

[0066] 3. This invention optimizes the reasonable block calculation of matrix multiplication in the sparse attention operator implemented on the Ascend NPU, and reasonably divides the basic blocks to reduce the operator's transfer time.

[0067] 4. This invention uses an algorithm for constructing parallel computing by launching Cube in advance for the sparse attention operator implemented on the Ascend NPU, thereby performing pipeline parallel optimization.

[0068] 5. The above optimization has been tested on Huawei Ascend NPU training products, improving the computational efficiency of the linear approximate attention operator. BRIEF DESCRIPTION OF THE DRAWINGS

[0069] Figure 1 This is the AI ​​Core abstract architecture diagram;

[0070] Figure 2 It is the computational logic diagram of the sparse approximate attention operator.

[0071] Figure 3 It is a block-by-block update logic diagram for softmax calculation.

[0072] Figure 4 This is a schematic diagram of vector subtraction for calculating safe softmax.

[0073] Figure 5 It is a schematic diagram of vector subtraction flow.

[0074] Figure 6 This is a schematic diagram of matrix multiplication block calculation.

[0075] Figure 7 This is a schematic diagram of the sparse approximate attention operator before and after pipeline optimization.

[0076] Figure 8 This is a schematic diagram of the parallel process of early emission pipeline.

[0077] Figure 9 Comparison of the total time (left) and transfer time (right) for matrix multiplication in different blocks.

[0078] Figure 10 is the running time of the sparse approximate attention operator vector calculation before and after optimization.

[0079] Figure 11 is the running time of the sparse approximate attention operator before and after pipeline parallel optimization. DETAILED DESCRIPTION

[0080] The technical solution of the present invention is further described below with reference to the accompanying drawings, but is not limited thereto. Any modification or equivalent replacement of the technical solution of the present invention that does not depart from the spirit and scope of the technical solution of the present invention should be included in the scope of protection of the present invention.

[0081] The present invention provides a sparse attention operator acceleration method based on NPU. When the method is implemented on Huawei Ascend NPU hardware, due to the limited cache memory of its vector computing unit, if the calculation is performed directly, R N ×k Softmax is calculated on the vector. When the input is large, a portion of the input must be moved in for calculation, then removed, and the next portion of the input must be moved in for calculation. Overall, the Q, E, and F matrices are divided into blocks and divided into two inner and outer loops. Multiple loops are performed to calculate the final sparse approximate attention result.

[0082] The input data of this operator is 3 R N×d The Q, K, V matrices, and two R shapes k×N The E and F matrices, the output data is a R N×dThe O matrix.

[0083] First, the Q matrix is ​​divided into blocks according to the N dimension piece The size of each block is Br×d. Then, the E and F matrices are divided into blocks according to the k dimension. Blocks, each block size is B c × N. The outer loop traverses the blocks of the Q matrix, and the inner loop traverses the blocks of the E and F matrices. Next, the blocks of the Q matrix are named Q i , the blocks of E and F matrices are named E j , F j .

[0084] The Linformer operator in the inner loop can be divided into four steps, which can be summarized as calculating the similarity matrix block, performing safe softmax calculation on the similarity matrix block, calculating the temporary result matrix block, and updating and accumulating the results of the previous inner loop. The specific steps are as follows:

[0085] Step 1: Calculate the similarity matrix. First, reduce the dimension of the K matrix and calculate the matrix multiplication in the matrix multiplication calculation unit. Get the matrix That is, the K matrix is ​​divided into blocks after dimensionality reduction. Next, the matrix multiplication is calculated in the matrix multiplication calculation unit. Get the matrix As a preliminary similarity matrix block.

[0086] Step 2: Calculate the safe softmax of the similarity matrix blocks. First, scale the similarity matrix and convert S' ij Move from global memory to local memory of vector computing unit, for S' ij Perform vector multiplication The safe softmax calculation is divided into two processes. If it is the first time in the inner loop, the vector calculation unit calculates S' ij Perform the max operation row by row to obtain the inmax vector. The vector calculation unit calculates S' ij The vector calculation unit calculates S' ij All elements of S' are natural exponentials. ijPerform row-by-row summation to obtain the insum vector. If this is not the first time in the inner loop, the max vector of the previous inner loop will affect the maximum value of each row, and the sum vector of the previous inner loop will affect the sum of each row. After calculating the inmax vector in the same way, it is necessary to first calculate inmax and the max vector in the max of the previous inner loop in the vector calculation unit. Next is the same subtraction of the max vector, taking the natural exponent and summing row by row to obtain the insum vector. Calculate the difference between the max vector of the previous inner loop and the current temporary maximum value of each row max vector after the update of this inner loop through vector subtraction, and then take the natural exponent to obtain the expmax vector used to update the temporary result matrix and the sum insum vector of each row. Finally, perform vector multiplication of the expmax vector and the sum vector of the previous inner loop to update and add the insum vector of this round to obtain the current temporary sum vector for each row. And update the latest sum and max to the insum vector and inmax vector respectively through transfer.

[0087] Step 3: Calculate the temporary result matrix blocks. First, reduce the dimension of the V matrix and calculate the matrix multiplication in the matrix calculation unit. Get the matrix That is, the V matrix is ​​divided into blocks after dimensionality reduction. Next, the matrix multiplication calculation unit calculates the matrix multiplication Get the temporary result matrix

[0088] Step 4: Update and accumulate the results of the previous inner loop. If it is the first time of the inner loop, no update and accumulation process is required. If it is not the first time of the inner loop, the temporary result matrix O' of the previous round is ij-1 Move from global memory to the local memory of the vector computing unit, and calculate the expmax vector and O' by vector multiplication with intervals ij-1 Update to get the correct O' ij-1 , the temporary result matrix O' of this round ij Move from global memory to the local memory of the vector computing unit, and add the correct O' ij-1 and O' ij Add to get a new temporary result matrix O' ij If this is the last round of the inner loop, the result matrix O is obtained at this time i It is necessary to update the globally correct result matrix O by performing vector division with the sum vector i .

[0089] Based on the hardware characteristics of Huawei's Ascend NPU, we first optimized vector computing. When using the vector computing unit for calculations, the smallest operable data unit is 32 bytes. During a single iteration, the vector computing unit retrieves eight data units from the vector cache, the UnifiedBuffer, for calculation. Furthermore, the vector computing unit can be configured to execute multiple iterations of control instructions.

[0090] When calculating softmax in blocks, exponential calculations (exp) are unstable in practice, such as overflow and decreased accuracy beyond a certain range. Therefore, using safe softmax, subtracting the calculated max vector from the x vector before performing the exponential calculation results in the following scenario: a (m, 32B) max vector and the original (m, n) x vector are calculated. The required operation is to subtract the current maximum value of each row from each element of x. A natural approach is to directly broadcast the (m, 32B) max vector to an (m, n) max vector and then perform the vector subtraction directly on the x vector. This seemingly simple operation incurs significant time and space overhead. First, a temporary (m, n) space must be allocated to store the broadcasted max vector. The vector calculation unit supports two types of broadcast operations: 1. Expand the max vector, expanding each element to the size of a data unit, and then performing multiple transfers to expand the vector. 2. Extract the elements of each row of the max vector and expand them by n. The first method is very time-consuming to carry out multiple transfers and will occupy a large amount of cache memory; the second method involves scalar calculations and vector calculations, and requires multiple loops. The two different pipelines need to wait for each other, which is also very time-consuming.

[0091] At this point, focusing on the characteristics of the vector calculation unit, calculations can be performed by row. Each 32-byte element in each row of the x vector can be subtracted from the element in the corresponding row of the max vector. In this case, one loop requires n / 128 iterations, requiring m loops. Calculations can be performed by column, with one data unit subtracted per iteration. This requires m / 16 iterations per loop, requiring n / 8 loops. Calculations can also be performed by block, with eight data units subtracted per iteration. This requires m iterations per loop, requiring n / 128 loops.

[0092] It can be seen that the total number of iterations of the three calculation methods is the same, because the number of elements calculated in one iteration and the total number of elements are fixed, but this does not mean that the calculation time of the three schemes is the same. Before the start of each loop's Vector operation, shift operations and other possible pre-loop operations are required. When the total number of iterations is the same, the calculation time for vector subtraction remains unchanged, but when the number of loops increases, the total calculation time will increase with the increase in the time-consuming pre-loop operations. In the sparse approximate attention operator, in addition to the vector calculation of calculating safe softmax, there are also row-by-row summation, row-by-row max, and O i =DIV(O i , sum) and other vector calculations are optimized.

[0093] Then, we optimize the matrix multiplication by reasonable block calculation. For two matrices A and B, the standard matrix multiplication AB(R M×K *R K×N ) has a computational complexity of O(M*N*K). Without changing the matrix calculation process and only performing block calculations on the matrix, the number of multiplication and addition operations for matrix multiplication will not decrease regardless of the block division method. The block division method has a greater impact on the processing time required for matrix multiplication. Matrix multiplication block division is to split (M, N, K) into (basicM, basicN, basicK). To calculate a basicM*basicN block in the result matrix C, the blocks in matrix A and matrix B must be matrix multiplied and then accumulated to obtain the result. In the matrix multiplication calculation unit, only one matrix multiplication result of matrix A block and matrix B block is calculated at a time. In such block operations, the transport time involved in matrix multiplication includes the MTE2 time for transporting the basic block from global memory to L1 Buffer, the MTE1 time for transporting the two basic blocks undergoing multiplication from L1 Buffer to L0ABuffer and L0B Buffer respectively, and the MTE3 time for transporting the result of the matrix multiplication of the basic block from L0CBuffer to global memory. When performing matrix multiplication, the basic block is first moved from global memory to the cache L1 Buffer of the AIV core, and then from L1 Buffer to the cache L0ABuffer and L0B Buffer. After that, the basic block matrix multiplication operation is started and the result is saved in L0C Buffer. Since the matrix multiplication block calculation requires cumulative summation to obtain the correct result of a certain block of the large result matrix, the MTE3 transport operation will be performed only after the K / basicK basic block matrix multiplication operation is performed. Therefore, the transport operation generally does not become a bottleneck for matrix multiplication. Observing the transport time calculation formulas of MTE1 and MTE2, it is found that the biggest difference is that the denominators are BW respectively.MTE1 , BW MTE2 , MTE1 is for the transfer between internal caches, MTE2 is for moving from global memory to cache, and the bandwidth of the two is BW MTE1 Much larger than BW MTE2 , then the most restrictive one is t MTE2 .

[0094] Assuming the current basic block size is (basicM, basicN, basicK), the corresponding MTE2 transfer time calculation formula for transferring a basic block is:

[0095]

[0096] The time calculation formula for calculating a basic block is:

[0097]

[0098] The ratio of the time it takes for MTE2 to move a basic block to the time it takes for the matrix multiplication unit to calculate a basic block is calculated using the following formula:

[0099]

[0100] From this calculation formula, it can be concluded that if the basic block sizes of basicN and basicM can be increased, the impact of the transport time on the calculation time can be reduced.

[0101] From the perspective of total transport capacity, when performing matrix multiplication block operations, the left matrix needs to be transported Second, the matrix on the right needs to be moved times, then the total transport volume is:

[0102]

[0103] Among them, M, N, and K are all fixed values ​​before blocking, so only basicN and basicM can be changed. Increasing their sizes can reduce the total handling capacity.

[0104] In the sparse approximate attention operator, the matrix multiplication operation K' j =Matmul(E j ,K) is Let's assume that N=1024, d=128, B c= 128. The matrix calculations are now executed with M = 128, N = 128, and K = 1024. Assume there are three possible partitioning schemes. Because the cache memory size of the matrix calculation unit is limited, we set basicM*basicN*basicK to constants. The three schemes (basicM, basicN, basicK) are (32, 128, 512), (64, 128, 256), and (128, 128, 128), respectively. According to the calculation formula, the transportation restriction when basicM, basicN, and basicK are (32, 128, 512) is the greatest, and the transportation restriction when (64, 128, 256) is second. The transportation restriction when basicM, basicN, and basicK are (128, 128, 128) is the least. According to the MTE2 transportation time calculation formula, it can be calculated that the MTE2 transportation time when basicM, basicN, and basicK are (32, 128, 512) is the longest, followed by (64, 128, 256), and the time used by (128, 128, 128) is the shortest.

[0105] Finally, pipeline parallel optimization is performed. The computation flow of this operator can be divided into Cube1: matrix multiplication of EK to obtain matrix K′; Cube2: matrix multiplication of the transpose of Q and K′ to obtain matrix S; Cube3: matrix multiplication of FV to obtain matrix V′; Cube4: matrix multiplication of PV′ to obtain the result matrix O; Vector1: a series of processes for normalizing the result S of Cube2 and calculating Flashsoftmax to obtain matrix P; and Vector2: a process for updating and accumulating the results of the previous inner loop. The computations involve many dependencies. Here, assuming that the computations in two loops are combined into one loop, it can be seen that the computations of Cube1 and Cube2 in the second loop only require the computations of Cube1 and Cube2 in the first loop and do not depend on the results of other computations in the previous loop. The same is true for other computation flows. Therefore, the gaps between the pipelines can be filled by pre-calculating the Cubes. By pre-calculating the Cubes, the dependencies of the subsequent loops are already calculated, thereby decoupling the multiple computation steps and enabling parallel computation of the Cube and Vector units, achieving optimal efficiency.

[0106] Example:

[0107] In this example, the sparse approximate attention operator is implemented in the software environment of CANN Community Edition 8.0.RC3.alpha001 and the hardware environment is the Ascend NPU training series product. The pseudocode for the specific sparse approximate attention operator implementation is shown in Algorithm 1.

[0108]

[0109]

[0110] After giving the pseudo code of the overall implementation process, we first introduce the optimized code example of vector calculation. The Sub interface in Ascend C is called to perform vector subtraction with intervals. The address steps between different datablocks of the single iteration of the src vector as the result, the src vector as the minuend, and the max vector as the subtrahend are configured to be 1, 1, and 0 respectively; at the same time, the address steps of the same datablock between adjacent iterations of the src vector as the result, the src vector as the minuend, and the max vector as the subtrahend are configured to be width / perElementByBlk, width / perElementByBlk, and 1 respectively. Where width is B c , that is, S' ij The width of the datablock, perElementByBlk is the number of elements contained in each datablock, each datablock is 32 bytes, for example, when the element is half type, each element occupies 2 bytes of space, then each datablock contains 16 elements. c ≤128, that is, the number of elements contained in 8 datablocks in this scenario, can be completed in one loop, and B is performed in one loop r Repeat the calculation, and calculate B in each repeat. c elements. When B c >128, which needs to be cycles, and B is performed in one cycle r The calculation is repeated 1 times, with 128 elements calculated in each repeat. The specific implementation of the optimized vector subtraction is shown in pseudo code 1.

[0111] Pseudocode 1 optimized vector reduction implementation

[0112]

[0113] Next, we'll introduce the pipelined implementation of pre-launch Cubes. If you want to pre-calculate a Cube, using the existing implementation, the second calculation of Cube2 might occur while Vec1, which depends on the results of the first Cube2 calculation, is being calculated simultaneously. This overwrites the results of the first Cube2 calculation, resulting in partially correct and partially incorrect Vec1 calculations. To avoid data trampling, an additional block of global memory is allocated for each computation flow. The current flag determines which of these two blocks should be written to. Data dependencies also use flags to determine which block to retrieve data from. Overall, the first loop performs the calculations on Cube1 and 2 independently. The second loop performs Cube1 and 2, Vector1 (depending on the results of Cube1 and 2 from the first loop), and Cube3 and 4. The third loop continues normally until the penultimate loop, which performs Vector1, Cube3 and 4, and Vector2. The final loop only performs Vector2. This requires two more loops than the normal inner and outer loops, making loopQ+2 (i.e., the outer loop plus 2). Furthermore, during the last two outer loops, only one inner loop is executed. At the same time, since the operations required for the first, middle, and final inner loops of steps 2 and 4 in the detailed steps above are different, and this looping method disrupts the loop order, it is also necessary to provide markers. To ensure the correct calculation results of the operator, it is also necessary to correctly transmit expmax between Vec1 and Vec2 across loops. This requires that two vectors be allocated for expmax and distinguished by identifiers. In addition, the correctness of the sum used to update the output O at the end of the inner loop must be determined. After the last sum calculation in the inner loop is completed, the sum is written to global memory and read back in later. The specific implementation of Cube early emission pipeline parallelism is shown in pseudocode 2.

[0114] Pseudocode 2Cube early launch pipeline parallelism

[0115]

[0116]

[0117] In this operator, a total of four matrix multiplications are performed. They are respectively to calculate the matrix multiplication Get the matrix Compute matrix multiplication Get the matrix Compute matrix multiplication Get the matrix Compute matrix multiplication Get the temporary result matrix Here E j Let's take a look at the matrix multiplication of a matrix and a -matrix. First, the matrix multiplication is So by calling Ascend C's SetOrgShape(B c The ,N,d) interface sets the shape of matrix multiplication. The left matrix's TPosition is GM, its CubeFormat is ND, its DataType is half, and the transpose enable parameter is false. The right matrix's TPosition is GM, its CubeFormat is ND, its DataType is half, and the transpose enable parameter is false. The result matrix's TPosition is GM, its CubeFormat is ND, and its DataType is half. Because Bias is not required, the isBias parameter is set to false. By calling Ascend C's SetFixSplit(128,128,128) method, we adjust basicM, basicN, and basicK in the basic block split to reduce the transfer time of the matrix multiplication operation.

[0118] Operator correctness test and performance test:

[0119] (1) Operator correctness test

[0120] First, we counted the common values ​​of d for common large language models. The results are shown in Table 1. Common values ​​of d are usually 64 and 128.

[0121] Table 1 Common values ​​of single-head dimensions of common large language models

[0122]

[0123] The correctness test uses the random method of the numpy library in python to first generate three random number matrices Q, K, and V with a shape of N×d and a data type of half as input data. Then, the random method of the numpy library in python is used to generate two random number matrices E and F with a shape of k×N and a data type of half as the dimensionality reduction matrices in the input data. First, the K matrix is ​​reduced in dimensionality, and the matrix multiplication result K' matrix of the E matrix and the K matrix is ​​calculated by numpy.matmul. Next, the similarity matrix is ​​calculated, and the matrix multiplication result S matrix of the transposed matrix of the Q matrix and the K' matrix is ​​calculated by numpy.matmul. The similarity matrix is ​​scaled, and each element of the similarity matrix is ​​divided by numpy.divide. Next, the safe softmax of the scaled similarity matrix is ​​calculated: numpy.max is used to find the maximum value of each row of the similarity matrix, and the maximum value of each row is subtracted from the elements of each row of the similarity matrix; the exponential value of the similarity matrix is ​​calculated using numpy.exp and numpy.sum and divided by the sum of the row, and the elements of the similarity matrix are converted into a probability distribution form to obtain the weight matrix. Finally, numpy.matmul is used to calculate the matrix multiplication result V' matrix of the E matrix and the K matrix to reduce the dimension of the V matrix, and numpy.matmul is used to calculate the matrix multiplication of the weight matrix and the V' matrix to obtain the resulting golden matrix. Taking golden as the true value result, the true value result calculated by python and the calculation result calculated by the development operator are compared element by element. The comparison method is:

[0124] |ab|≤(atol+rtol*|b|)(6)

[0125] Where a is the result calculated by the development operator, b is the result calculated by Python, and atol=10 -6 , rtol=10 -9 , if formula (6) holds, the calculation of this element is correct, otherwise it is wrong.

[0126]

[0127] In formula (7), error_num is the number of element calculation errors, all_num is the total number of elements in the result, and etol is the element calculation error rate. In the verification indicator, etol≤10 -4 The operator is considered to be accurate.

[0128] Table 2 Accuracy of basic scheme and optimized scheme

[0129]

[0130] (2) Operator performance test

[0131] First, the matrix multiplication operation is separated to show the optimization effect of the block. c =128. Set basicM to 32, 64, and 128 respectively, keep basicN at 128, and set basicK to 512, 256, and 128.

[0132] Depend on Figure 9 It can be seen that as basicM continues to increase, the MTE2 transfer time and total time of different blocks of matrix multiplication are decreasing, which is basically consistent with the above analysis results.

[0133] Next, we compare the operator execution time before and after vector optimization, where N = 512. The x-axis represents the size of the computational dimension d, starting from a small 16 and doubling each time from 16 to 128.

[0134] Depend on Figure 10 According to the statistical results in

[15] , although vector calculation optimization optimizes the calculation process and running time of the operator, the optimization range is small under inputs of various scales, with the maximum optimization range being 3.9%.

[0135] Finally, we compare the operator execution time before and after implementing Cube early emission pipeline parallelism. N = 512, and the x-axis represents the size of the computational dimension d, starting with a small 16 and doubling each time from 16 to 128.

[0136] Depend on Figure 11 As can be seen from the running time, the operation of launching the Cube in advance to create a parallel environment has a significant effect on the optimization of the operator. The overall optimization efficiency exceeds 10%, and the highest optimization efficiency reaches 23.73%.

Claims

1. A sparse attention operator acceleration method based on NPU, characterized by The method comprises the following steps: Step 1: Divide the Q matrix into blocks according to N dimensions. piece The size of each block is Br×d, and the E and F matrices are divided into blocks according to the k dimension. Blocks, each block size is B c ×N; Step 2: Name the block of the Q matrix Q i , the blocks of E and F matrices are named E j , F j ; Step 3: The outer loop traverses the blocks according to the Q matrix, and the inner loop traverses the blocks according to the E and F matrices. In the inner loop, the Linformer operator is divided into four steps: Step 1: Calculate the similarity matrix: Step 1-1: Calculate the matrix multiplication E in the matrix multiplication calculation unit j K gets the matrix K' j , that is, the K matrix block after dimensionality reduction; Step 1-2: Calculate the matrix multiplication Q in the matrix multiplication calculation unit i K' j T Get the matrix S' ij , as the preliminary similarity matrix block; Step 2: Perform safesoftmax calculation on the similarity matrix blocks: Step 2-1, scale the similarity matrix and change S' ij Move from global memory to local memory of vector computing unit, for S' ij Perform vector multiplication Step 2-2, safe softmax calculation; Step 3: Calculate the temporary result matrix blocks: Step 3-1: Calculate the matrix multiplication F in the matrix calculation unit j V gets the matrix V' j , that is, the V matrix block after dimensionality reduction; Step 3-2: Calculate the matrix multiplication S' in the matrix multiplication calculation unit ij V' j Get the temporary result matrix O' ij ; Step 4: Update and accumulate the results of the previous inner loop: If it is the first time of the inner loop, no update and accumulation process is required; if it is not the first time of the inner loop, the temporary result matrix O' of the previous round is ij-1 Move from global memory to the local memory of the vector computing unit, and calculate the expmax vector and O' by vector multiplication with intervals ij-1 Update to get the correct O' ij-1 , the temporary result matrix O' of this round ij Move from global memory to the local memory of the vector computing unit, and add the correct O' ij-1 and O' ij Add to get a new temporary result matrix O' ij ; If it is the last round of the inner loop, the result matrix O is obtained at this time i It is necessary to update the globally correct result matrix O by performing vector division with the sum vector i ; Step 4: Optimize vector calculations based on the hardware characteristics of the NPU: Step 4.

1. Optimize the vector calculation for calculating safe softmax: Focus on the characteristics of the vector calculation unit, perform calculations in blocks, and perform vector subtraction calculations on 8 data units in each iteration. Instead of performing m iterations in one loop, only n / 128 loops are required. Step 42: Optimize vector calculation for row-by-row summation and row-by-row max; Step 4.

3. O i =DIV(O i ,sum) vector calculation is optimized; Step 5: Optimize the matrix multiplication calculation by reasonable blocks according to the hardware characteristics of the NPU: Step 51: Based on the (M, N, K) matrix multiplication and the cache size limits of the AIC core's L0ABuffer, L0B Buffer, and L0C Buffer, set size limits for basicM, basicN, and basicK. The left matrix basicM*basicK after dicing can be placed in the L0ABuffer, the right matrix basicK*basicN after dicing can be placed in the L0B Buffer, and the partial calculated results can be placed in the L0C Buffer. Step 52: According to the size limits of basicM, basicN, and basicK, the largest basicM and basicN, as well as the basicK that matches them, are selected to divide the left and right matrices into basic blocks; Step 53: When performing matrix multiplication, the basic block is first moved from global memory to the AIV core's L1 Buffer, then from the L1 Buffer to the L0A Buffer and L0B Buffer. The basic block matrix multiplication operation then begins, and the result is stored in the L0C Buffer. After performing the K / basicK basic block matrix multiplication operation, the partially correct result of the matrix multiplication is moved out until all basic block operations are completed. Step 6: Perform pipeline parallel optimization based on the hardware characteristics of the NPU: The calculation flow of the Linformer operator is divided into Cube1: Matrix K' is obtained by EK matrix multiplication, Cube2: Matrix S is obtained by matrix multiplication of the transpose of Q and K', Cube3: Matrix V' is obtained by FV matrix multiplication, Cube4: Result matrix O is obtained by PV' matrix multiplication, Vector1: Normalize the calculation result S of Cube2 and calculate Flashsoftmax to obtain matrix P, Vector2: Update and accumulate the results of the previous inner loop; Step 61: Allocate an additional block of memory to each computational flow in the global memory; Step 62: After the calculation of Cube1 and Cube2 in the first inner loop is completed, do not proceed with the subsequent process. Calculate Cube1 and Cube2 in the second inner loop in another space first. The synchronous parallel calculation requires the calculation of Vector1 that depends on the results of Cube1 and Cube2 in the first inner loop. Step 6.3: Calculate Cube3 and Cube4 in the first inner loop, and calculate Vector1 in the second inner loop in parallel. Calculate Cube1 and Cube2 in the third inner loop, and calculate Vector2 in the first inner loop in parallel. Calculate Cube3 and Cube4 in the second inner loop. Step 64: Continue the loop as described above, and determine which of the two memory blocks should be written to based on the current flag. If there is a data dependency, the flag should also be used to determine which block the data needs to be taken from.

2. The NPU-based sparse attention operator acceleration method according to claim 1 is characterized in that The specific steps of step 2-2 are as follows: (1) If it is the first time of the inner loop, the vector calculation unit will calculate S' ij Perform the max operation row by row to obtain the inmax vector, and the vector calculation unit calculates S' ij The element of the vector is subtracted from the maximum value of each row, that is, the inmax vector is subtracted, and the vector calculation unit calculates S' ij All elements of S' are natural exponentials, and the vector calculation unit is ij Perform row-by-row summation to obtain the insum vector; (2) If it is not the first time of the inner loop, the vector calculation unit calculates S' ij Perform the max operation row by row to obtain the inmax vector. The vector calculation unit calculates the max vector in inmax and the max in the previous inner loop. The vector calculation unit calculates S' ij Subtract the max vector from the elements of the inner loop, take the natural exponent and sum them row by row to get the insum vector, calculate the difference between the max vector of the previous inner loop and the current temporary maximum value of each row max vector updated in this inner loop through vector subtraction, and then take the natural exponent to get the expmax vector used to update the temporary result matrix and each row and insum vector, and finally update the current temporary sum vector for each row by multiplying the expmax vector and the sum vector of the previous inner loop and adding the insum vector of this round, and update the latest sum and max to the insum vector and inmax vector respectively through transfer.

3. The NPU-based sparse attention operator acceleration method according to claim 1 is characterized in that The specific steps of step 42 are as follows: Step 1: Apply for a blank space of (m,64); Step 2: Perform vector addition / max calculation on the first (m, 64)-shaped vector and the second block. The resulting vector is recorded in the blank space of (m, 64). The remaining blocks are also accumulated / maxed with the vectors in the blank space by vector addition / max calculation. Step 3: If the column dimension of the matrix cannot be divided evenly by 64, then the remainder is used as the tail block and the accumulation / max is performed in the same way. The row-by-row summation / max of the matrix is ​​converted into the row-by-row summation / max of the (m, 64) space that is convenient for the vector calculation unit to calculate.

4. The NPU-based sparse attention operator acceleration method according to claim 1 is characterized in that The specific steps of step 43 are as follows: Step 1: The address step size of different datablocks in the same iteration is configured as 1, and the address step size of the same datablock between adjacent iterations is configured as d / perElementByBlk, where perElementByBlk is the number of elements contained in each datablock; Step 2: Set the address step of different datablocks within the same iteration of expmax / sum to 0, and the address step of the same datablock between adjacent iterations to 1.

Images