An NPU-based sparse attention operator acceleration method

Abstract

The application discloses a sparse attention operator acceleration method based on an NPU, which combines Linformer with an Ascend C programming language, realizes an Ascend C operator for calculating Linformer based on the NPU, combines the acceleration calculation mode of Flash-Attention2 with Linformer, so that the whole operator can be calculated in blocks, optimizes vector calculation according to the hardware characteristics of the Ascend NPU, reduces the number of instructions for multiple vector calculations such as vector subtraction, row-by-row summation and row-by-row max calculation of safe softmax, and reduces time consumption, uses the algorithm for constructing parallel calculation by emitting a Cube in advance to solve the problem that front and rear data dependence cannot be parallel, improves the calculation efficiency of the sparse approximate attention operator, and divides a basic block for matrix multiplication to reduce the carrying time of matrix multiplication.

Description

Technical Field

[0001] This invention relates to a code implementation and optimization method for sparse attention operators, specifically to an NPU-based method for accelerating sparse attention operators. Background Technology

[0002] Ascend C is a programming language developed by Huawei for its Ascend AI processor. A key characteristic of this language is its strong hardware dependence. The core computing unit in the Ascend chip is the AI ​​Core, whose abstract structure is shown in the diagram below. Figure 1 As shown.

[0003] Huawei's Ascend NPU910B4 features a discrete architecture, splitting the AI ​​processor into two independent cores: an AI Cube (AIC) core for matrix computation and an AI Vector (AIV) core for vector computation. Each core has its own Scalar unit, capable of independently loading its own code segments, thus decoupling matrix and vector computations. Under unified system software scheduling, they cooperate to optimize computational efficiency. AIV and AIC exchange data via Global Memory, enabling parallel matrix and vector computations that do not depend on each other. Its core components include three key elements: a computing unit, a storage unit, and a data transfer 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 computing unit, Vector computing unit, and 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] AI processors contain internal storage units. Data needs to be loaded from external storage into the internal storage to complete the necessary calculations. The external storage units are... Figure 1 In general, the GM in the memory is used to store the input, output and intermediate data generated by the operation of the operator as a whole. The size of the internal storage unit is much smaller than that of the external storage unit.

[0006] Figure 1 The main storage units within the upper half of the AIC architecture core include: L1 Buffer: L1 buffer is used to temporarily store part of the input for matrix operations, thereby reducing the number of reads and writes from outside the core; L0A Buffer and L0B Buffer are used to store the two matrices currently being multiplied; L0C Buffer stores the result matrix obtained from the current matrix multiplication.

[0007] Figure 1 The main storage unit within the lower half of the AIV architecture core is the Unified Buffer, which serves as a unified buffer to store the inputs and outputs of vector and scalar computations within the AIV architecture.

[0008] To facilitate data transfer and movement within the AI ​​processor, three types of transfer units are included: the MTE1 transfer unit, which is primarily responsible for moving matrix data to be processed from the L1 Buffer to the L0A Buffer and L0B Buffer; the MTE2 transfer unit, which is responsible for moving data to be processed from global memory into the AIV core's Unified Buffer, and also for moving data from global memory into the AIC core's L1 Buffer; and the MTE3 transfer unit, which is responsible for moving the required data from the AIV core's Unified Buffer to global memory, and also for moving data from the AIC core's L0C Buffer, where temporary results are stored, to global memory.

[0009] In the AI ​​processor, instructions are scheduled into five independent classification queues (including the Vector instruction queue, the Cube instruction queue, and the MTE1 / MTE2 / MTE3 instruction queues), and then distributed to the corresponding execution units for processing. Instructions within the same queue are executed sequentially in the order they entered the queue, while different instruction queues can run in parallel. Parallel execution of multiple instruction queues improves overall execution efficiency. To address potential data dependencies during parallel processing, an event synchronization module inserts synchronization instructions to coordinate pipeline synchronization. The event synchronization module continuously monitors the execution status of each pipeline and analyzes the dependencies between different pipelines, thereby resolving data dependency and synchronization issues.

[0010] Since the main computational units of the kernel function are the Vector computational unit and the Cube computational unit, which are also divided into two different kernels, the AIV kernel and the AIC kernel, there are three different programming scenarios when implementing the Ascend C operator: vector programming scenario using only the AIV kernel, matrix programming scenario using only the AIC kernel, and fusion operator programming scenario using both the AIV kernel and the AIC kernel.

[0011] Vector computation is performed using the Vector computation unit. Both the source and destination operands are stored in a UnifiedBuffer, and multiple iterations are possible. During each iteration, the Vector computation unit retrieves eight consecutive data blocks (each 32 bytes) from the UnifiedBuffer, performs the computation, and writes the result back to the corresponding eight data blocks. The Vector computation unit supports spaced vector computations, configurable via `dataBlockStride` (the address stride between different data blocks in the same iteration) and `repeatStride` (the address stride between identical data blocks in adjacent iterations).

[0012] The runtime of the Ascend operator is closely related to the hardware architecture. The main computation times can be broadly categorized into vector computation time, scalar computation time, matrix multiplication computation time, and data transport time. For vector computation time, the ability of the vector computation unit to iterate multiple times can be leveraged to reduce vector computation instructions: controlling the address interval of operands can reduce vector computation instructions and cache space usage; changing the implementation of vector computation can also reduce instructions. For scalar computation time, since the scalar computation unit is also responsible for program flow control, scalar computation should be converted into vector computation as much as possible to avoid affecting program synchronization and blocking issues. For matrix multiplication computation time, when the sizes of the left and right matrices are fixed, the computation time is fixed. However, because it involves block-based matrix multiplication, different basic block sizes affect matrix transport time, which can be optimized. Besides reducing the computation time of individual parts, the Ascend C operator can be parallelized not only from multiple perspectives but also from the perspective of computation units.

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

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

[0015] The operator takes five inputs: Q, K, and V matrices (all of which have shape R). N×d ), and the E and F matrices (both of which have shape R) used for dimensionality reduction of the K and V matrices. N×k The operator outputs a result matrix O (of shape R). N×d ).

[0016] Then, the calculation formula is given:

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

[0018] During the calculation, matrix multiplication EK is performed to obtain matrix K', matrix multiplication FV is performed to obtain matrix V', and then matrix multiplication QK' is performed. T Given a matrix S, to prevent softmax calculation from causing numerical overflow, we find the maximum value for each row of elements in matrix S and then subtract that maximum value from each subsequent row element. Next, we calculate e for each element x in the updated matrix S. x Then, by summing each row, we obtain matrix P to achieve the computational effect of safe softmax. Finally, we calculate matrix multiplication PV' to obtain the result matrix.

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

[0020] 1. There are currently no Linformer operators that can run on Ascend NPU.

[0021] 2. If the computation method of the Linformer operator is not changed and it is directly migrated to the Ascend NPU, the problem of needing to directly calculate the softmax of large matrices will still exist. The Ascend NPU vector computation module has limited cache and cannot perform the calculation directly.

[0022] 3. The Linformer steps require a lot of instructions to complete the calculations. For example, the vector subtraction, row-by-row summation, and row-by-row max calculation for calculating safesoftmax would take a long time if calculated directly due to the large number of instructions.

[0023] 4. The Linformer has many dependencies between its multiple steps, requiring sequential computation, which makes it difficult to make good use of the Ascend NPU's hardware features.

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

[0025] This invention provides an NPU-based method for accelerating sparse attention operators. This method combines Linformer with the Ascend C programming language, implementing an Ascend C operator for Linformer computation on the NPU. This operator belongs to the hybrid programming operator category of Ascend C operators, implementing all four stages of the overall computation to form a scalable and robust Ascend C hybrid programming operator. The method combines Linformer with Flash-Attention2 acceleration, allowing the entire operator to be computed in blocks. Vector computation is optimized based on the hardware characteristics of the Ascend NPU, reducing the vector subtraction, row-by-row summation, and row-by-row max calculation required for safe softmax. The number of instructions used for multiple vector calculations is reduced, thus reducing the time consumption; to address the problem of multiple data dependencies that cannot be parallelized, a parallel computing algorithm for constructing Cubes in advance is used, which improves the computational efficiency of the sparse approximate attention operator; matrix multiplication is divided into basic blocks to reduce the transport time of matrix multiplication.

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

[0027] An NPU-based sparse attention operator acceleration method includes the following steps:

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

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

[0030] Step 3: The outer loop iterates through blocks based on the Q matrix, and the inner loop iterates through blocks based on the E and F matrices. Within the inner loop, the Linformer operator is divided into four steps:

[0031] Step 1: Calculate the similarity matrix:

[0032] Step 1-1: Calculate matrix multiplication E in the matrix multiplication calculation unit. j K obtains matrix K' j That is, the K matrix is ​​divided into blocks after dimensionality reduction;

[0033] Steps 1-2: Calculate matrix multiplication in the matrix multiplication calculation unit. Obtain matrix S' ij This serves as the initial block division of the similarity matrix;

[0034] Step 2: Perform safe softmax calculation for block-based similarity matrix segmentation:

[0035] Step 2-1: Scale the similarity matrix, reducing S' ij The data is moved from global memory to the local memory of the vector computation unit, for S' ij Perform vector multiplication

[0036] Step 2-2, Calculation of safe softmax:

[0037] (1) If it is the first iteration of the inner loop, the vector calculation unit is set to S' ij Perform a row-by-row max operation to obtain the inmax vector, and then use the vector computation unit to calculate S'. ij Subtracting the maximum value of each row from the elements of S', i.e., subtracting the inmax vector, the vector computation unit calculates S'. ij All elements are taken as natural exponents, and the vector computation unit is S' ij Perform a row-by-row summation operation to obtain the insum vector;

[0038] (2) If it is not the first iteration of the inner loop, the vector calculation unit for S' ij Perform a row-by-row max operation to obtain the inmax vector. The vector calculation unit calculates the max vector in inmax and the max vector in the previous inner loop. The vector calculation unit calculates S'. ij Subtract the max vector from the elements, take the natural exponent and sum row by row to get the insum vector. Calculate the difference between the max vector from the previous inner loop and the current temporary maximum value max vector after the current inner loop update by subtracting the vector, then take the natural exponent to get the expmax vector used to update the temporary result matrix and the sum insum vector of each row. Finally, update the current temporary sum sum vector by multiplying the expmax vector and the sum vector from the previous inner loop, and add the current insum vector. Then, by moving the latest sum and max, update the insum vector and inmax vector respectively.

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

[0040] Step 3-1: Calculate matrix multiplication F in the matrix calculation unit. j V obtains matrix V' j That is, the V matrix after dimensionality reduction is divided into blocks;

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

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

[0043] If this is the first iteration of the inner loop, no update and accumulation process is needed; otherwise, the temporary result matrix O' from the previous iteration will be used. ij-1 The vector is moved from global memory to the local memory of the vector computation unit, and the expmax vector and O' are calculated through spaced vector multiplication. ij-1 The update yielded the correct O' ij-1 The temporary result matrix O' of this roundij The data is moved from global memory to the local memory of the vector computation unit, and the correct O' is calculated using vector addition. ij-1 and O' ij The addition yields a new temporary result matrix O' ij If this is the last iteration of the inner loop, the resulting matrix O is... i The globally correct result matrix O needs to be updated by performing a vector division with the sum vector. i ;

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

[0045] Step 4. Optimize the vector calculation for safe softmax: Focus on the characteristics of the vector calculation unit, calculate in blocks, and perform vector subtraction calculation for 8 data units in each iteration. One loop requires m iterations, but only n / 128 loops are needed.

[0046] Step 4.2: Optimize the row-by-row summation and row-by-row max calculation using vector calculations:

[0047] Step 1: Allocate blank space (m, 64);

[0048] Step 2: Add vectors to the first (m,64) shape vector and calculate the maximum vector from the second block. Record the resulting vector in the empty space of (m,64). For the remaining blocks, add vectors to the vector in the empty space and calculate the maximum vector.

[0049] Step 3: If the column dimension of the matrix cannot be divided by 64, then take the remainder as the tail block and perform the same accumulation / max calculation, converting the row-by-row summation / max calculation of the matrix into the row-by-row summation / max calculation of the (m,64) space for the Vector computing unit to calculate.

[0050] Step 43, to O i =DIV(O i Optimize vector calculations (sum):

[0051] Step 1, The address step size of different datablocks within 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: Configure the address step size of different data blocks within the same iteration of expmax / sum to 0, and configure the address step size of the same data block between adjacent iterations to 1. This will reduce the number of calculation instructions used during the calculation.

[0053] Step 5: Optimize matrix multiplication by performing reasonable block-based computation based on the NPU's hardware characteristics:

[0054] Step 51: Based on the matrix multiplication (M, N, K) and the buffer size limits of the AIC core's L0ABuffer, L0BBuffer, and L0CBuffer, size limits are given for basicM, basicN, and basicK. The left matrix basicM*basicK after being divided into blocks can be placed in L0ABuffer, the right matrix basicK*basicN after being divided into blocks can be placed in L0BBuffer, and the calculated partial results can be placed in L0CBuffer.

[0055] Step 52: Based on the size constraints of basicM, basicN, and basicK, select the largest basicM and basicN, and the matching basicK, and divide the left and right matrices into basic blocks;

[0056] Step 53: During matrix multiplication, the basic blocks are first moved from global memory into the AIV core's L1 buffer, then from the L1 buffer to the L0A buffer and L0B buffer. After that, the basic block matrix multiplication operation begins, and the result is stored in the L0C buffer. Since the matrix multiplication is performed in blocks and then summed to obtain the correct result of a certain part of the large result matrix, after the K / basicK basic block matrix multiplication operation, the correct result of the matrix multiplication is moved out until all basic block operations are completed.

[0057] Step Six: Optimize pipelined parallelism based on the NPU's hardware characteristics:

[0058] The computation flow of the Linformer operator is divided into: Cube1: calculates matrix K′ through EK matrix multiplication; Cube2: calculates matrix S by matrix multiplication of the transposes of Q and K′; Cube3: calculates matrix V′ through FV matrix multiplication; Cube4: calculates the result matrix O through PV′ matrix multiplication; Vector1: normalizes the calculation result S of Cube2 and calculates matrix P using Flashsoftmax; Vector2: updates and accumulates the result of the previous inner loop.

[0059] Step 61: Allocate an additional block of memory in global memory for each computation stream;

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

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

[0062] Step 64: Continue the loop as described above, allocating which of the two memory blocks to write to based on the current flag. Also, for data dependencies, determine which block to retrieve the data from based on the flag.

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

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

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

[0066] 3. This invention optimizes the matrix multiplication calculation in the sparse attention operator implemented on Ascend NPU by reasonable block division and reduces the operator's transport time by reasonable basic block partitioning.

[0067] 4. This invention uses an algorithm for constructing parallel computation by pre-emitting Cubes for sparse attention operators implemented on Ascend NPU, thereby performing pipelined parallel optimization.

[0068] 5. The above optimizations have been tested in actual operation on Huawei Ascend NPU training series products, improving the computational efficiency of the linear approximate attention operator. Attached Figure Description

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

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

[0071] Figure 3 This is the block-based update softmax calculation logic diagram.

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

[0073] Figure 5 This is a schematic diagram of vector flow reduction.

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

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

[0076] Figure 8 This is a schematic diagram of the pre-launch parallel process.

[0077] Figure 9 The comparison shows the total time (left) and transport time (right) for different blocks of matrix multiplication.

[0078] Figure 10 It represents the runtime before and after optimization of the sparse approximate attention operator vector computation.

[0079] Figure 11 It represents the runtime before and after pipelined parallel optimization of the sparse approximate attention operator. Detailed Implementation

[0080] 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.

[0081] This invention provides an acceleration method for sparse attention operators based on an NPU. When this method is implemented on the Huawei Ascend NPU hardware, due to the limited cache memory of its vector computation unit, direct computation requires processing R... N ×k The softmax operation is performed on the vector. When the input size is large, it is necessary to move a portion of the input for calculation and then move it out, and then move the next portion of the input for calculation. Overall, the Q, E, and F matrices are divided into blocks, with inner and outer loops, and the final sparse approximate attention result is calculated in multiple loops.

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

[0083] First, the Q matrix is ​​divided into blocks according to the N-dimensional plane. piece Each block is Br×d in size. Then, the E and F matrices are divided into blocks according to the k-dimensional dimension. Blocks, each block is B in size c ×N. The outer loop iterates through blocks based on the Q matrix, and the inner loop iterates through blocks based on the E and F matrices. Next, the blocks of the Q matrix will be named Q... i The blocks of matrices E and F 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 blocks, performing safe softmax calculation on the similarity matrix blocks, calculating the temporary result matrix blocks, and updating and accumulating the results from the previous inner loop. The specific steps are as follows:

[0085] Step 1: Calculate the similarity matrix. First, reduce the dimensionality of the K matrix, then perform matrix multiplication in the matrix multiplication unit. Obtain the matrix This involves partitioning the dimensionality-reduced K matrix into blocks. Next, matrix multiplication is performed in the matrix multiplication computation unit. Obtain the matrix This serves as an initial block division of the similarity matrix.

[0086] Step 2: Perform safe softmax calculation for block-based similarity matrix segmentation. First, scale the similarity matrix to S' ij The data is moved from global memory to the local memory of the vector computation unit, for S' ij Perform vector multiplication The safe softmax calculation consists of two processes. If it's the first iteration of the inner loop, the vector calculation unit calculates S'. ij Perform a row-by-row max operation to obtain the inmax vector. The vector computation unit calculates S'. ij Subtract the maximum value of each row from the elements of S', which is equivalent to subtracting the inmax vector. The vector computation unit calculates S'. ij All elements are taken as natural exponents. The vector computation unit for S' ijThe process involves summing row by row to obtain the `insum` vector. If this isn't the first iteration of the inner loop, the `max` vector from the previous inner loop influences the maximum value of each row, and the `sum` vector from the previous inner loop influences the sum of each row. After calculating the `inmax` vector, the `max` vector from both `inmax` and the `max` vector from the previous inner loop is calculated in the vector calculation unit. Next, the `max` vector is subtracted, and the natural exponent is taken, resulting in the `insum` vector. The difference between the `max` vector from the previous inner loop and the updated temporary maximum value `max` vector for each row in the current inner loop is calculated by subtracting the vector from the previous inner loop, and then the natural exponent is taken to obtain the `expmax` vector used to update the temporary result matrix and the `insum` vector for each row. Finally, the `expmax` vector is multiplied by the `sum` vector from the previous inner loop to update and add the current `insum` vector to obtain the current temporary `sum` vector for each row. The latest `sum` and `max` values ​​are then moved to update the `insum` and `inmax` vectors respectively.

[0087] Step 3: Calculate the temporary result matrix in blocks. First, reduce the dimensionality of matrix V, then perform matrix multiplication in the matrix calculation unit. Obtain the matrix That is, the V matrix after dimensionality reduction is divided into blocks. Next, matrix multiplication is performed in the matrix multiplication computation unit. Obtain the temporary result matrix

[0088] Step 4: Update and accumulate the results of the previous inner loop. If this is the first iteration of the inner loop, no update and accumulation process is needed. If it is not the first iteration of the inner loop, then the temporary result matrix O' from the previous round is... ij-1 The vector is moved from global memory to the local memory of the vector computation unit, and the expmax vector and O' are calculated through spaced vector multiplication. ij-1 The update yielded the correct O' ij-1 The temporary result matrix O' of this round ij The data is moved from global memory to the local memory of the vector computation unit, and the correct O' is calculated using vector addition. ij-1 and O' ij The addition yields a new temporary result matrix O' ij If this is the last iteration of the inner loop, the resulting matrix O is... i The globally correct result matrix O needs to be updated by performing a vector division with the sum vector. i .

[0089] Based on the hardware characteristics of Huawei Ascend NPU, vector computation is optimized first. When using the vector computation unit for computation, the smallest operable data unit is 32 bytes. In one iteration, the vector computation unit will retrieve 8 data units from the vector buffer unit UnifiedBuffer for computation. At the same time, the vector computation unit can be executed multiple times through configuration control instructions.

[0090] When performing block-based softmax calculations, the exponential calculation (exp) suffers from instability in practice, such as numerical overflow and decreased accuracy beyond a certain range. Therefore, a safe softmax approach is used, subtracting the calculated max vector from the x vector before performing the exponential calculation. This leads to the following scenario: a (m, 32B) max vector is calculated, along with the original (m, n) x vector. The required operation is to subtract the current maximum value of each row's x element. A natural approach would be to directly broadcast the (m, 32B) max vector to a (m, n) max vector and then perform the vector subtraction operation directly from the x vector. While seemingly simple, this approach incurs significant time and space overhead. First, a temporary (m, n) space is needed to store the broadcasted max vector. Second, the vector calculation unit supports two broadcast operations: 1. Expanding the max vector by one data unit per element, then expanding it further through multiple transfers. 2. Extracting the elements of each row of the max vector and expanding it by n. The first method involves multiple transfers, which is time-consuming and consumes a lot of cache memory. The second method involves scalar and vector calculations and requires multiple loops. The two different pipelines need to wait for each other, which is also quite time-consuming.

[0091] At this point, focusing on the characteristics of the vector calculation unit, we can perform calculations row-by-row. Each 32-byte element in each row of vector x can be subtracted from the corresponding element in the row of vector max. In this case, one loop requires n / 128 iterations, and m loops are needed. Alternatively, we can perform calculations column-by-column, with each iteration performing a vector subtraction of one data unit. One loop requires m / 16 iterations, and n / 8 loops are needed. Finally, we can perform calculations block-by-block, with each iteration performing a vector subtraction of eight data units. One loop requires m iterations, and n / 128 loops are needed.

[0092] It can be seen that the total number of iterations is the same for all three calculation methods because the number of elements calculated in one iteration and the total number of elements are fixed. However, this does not mean that the computation time of the three schemes is the same. Each loop's Vector operation requires a shift operation and other possible pre-loop operations. With the same total number of iterations, the computation time for vector subtraction remains constant, but as the number of loops increases, the total computation time increases with the time spent on pre-loop operations. In sparse approximate attention operators, besides optimizing the vector computation for safe softmax, there are also row-by-row summation, row-by-row max calculation, and... O i =DIV(O i Optimize multiple vector calculations such as sum).

[0093] Next, we optimize the matrix multiplication by dividing it into reasonable blocks. For two matrices A and B, the standard matrix multiplication AB(R) is performed. M×K *R K×N The computational complexity of matrix multiplication is O(M*N*K). Without changing the matrix computation flow, and only performing block-based computations, the number of multiplication and addition operations in matrix multiplication is not reduced by the block-based approach. The block-based approach has a greater impact on the data transfer time required for matrix multiplication. Block-based matrix multiplication divides (M, N, K) into (basicM, basicN, basicK). For the operation of a basicM*basicN block in the result matrix C, matrix multiplication of the blocks in matrix A and matrix B is performed, followed by summation to obtain the result. In a single matrix multiplication unit, only one matrix multiplication result between a block of matrix A and a block of matrix B is calculated at a time. In this block-based operation, the time involved in matrix multiplication includes MTE2 (moving basic blocks from global memory into the L1 buffer), MTE1 (moving the two multiplied basic blocks from the L1 buffer to the L0A and L0B buffers respectively), and MTE3 (moving the result of the matrix multiplication from the L0C buffer to global memory). During matrix multiplication, basic blocks are first moved from global memory into the AIV core's L1 buffer, then from the L1 buffer to the L0A and L0B buffers. Afterward, the basic block matrix multiplication operation begins, and the result is stored in the L0C buffer. Because the matrix multiplication block calculation requires summation to obtain the correct result for a specific part of the large result matrix, the MTE3 move-out operation only occurs after the K / basicK basic block matrix multiplication operation. Therefore, the move-out operation generally does not become a bottleneck in matrix multiplication. Observing the formulas for calculating the time of MTE1 and MTE2, a significant difference is found in the denominators BW.MTE1 BW MTE2 MTE1 handles data movement between internal caches, while MTE2 handles data movement from global memory into the cache. The bandwidth (BW) of both is... MTE1 Much larger than BW MTE2 Then, the most restrictive value at this time is t. MTE2 .

[0094] Assuming the current basic block size is (basicM, basicN, basicK), the formula for calculating the MTE2 transport time for moving one basic block is:

[0095]

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

[0097]

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

[0099]

[0100] This calculation formula shows that increasing the size of basicN and basicM blocks can reduce the impact of transport time on computation time.

[0101] From the perspective of total transport volume, when performing block operations for matrix multiplication, the left-side matrix needs to be transported. Next, the right-hand matrix needs to be moved. Then, the total amount of material handled is:

[0102]

[0103] M, N, and K are all fixed values ​​before partitioning, so only basicN and basicM can be changed. Increasing their size can reduce the total amount of material handled.

[0104] In the sparse approximate attention operator, the matrix multiplication operation K' j =Matmul(E j K) is Let's assume N = 1024, d = 128, and B... c=128. At this point, we input M=128, N=128, K=1024 for matrix calculation. Assume there are three block partitioning schemes. Because the cache memory size of the matrix calculation unit is limited, we set basicM*basicN*basicK as constants. In 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 restrictions are greatest when basicM, basicN, and basicK are (32, 128, 512), followed by (64, 128, 256), and least restrictive when basicM, basicN, and basicK are (128, 128, 128). Using the MTE2 transportation time calculation formula, it can also be calculated that the MTE2 transportation time is longest when basicM, basicN, and basicK are (32, 128, 512), followed by (64, 128, 256), and shortest when (128, 128, 128).

[0105] Finally, pipelined parallel optimization is performed. The computational flow of this operator can be divided into: Cube1: matrix multiplication of EK to obtain matrix K′; Cube2: matrix multiplication of Q and the transpose of 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 including normalization and Flashsoftmax calculation of the result S calculated by Cube2 to obtain matrix P; Vector2: updating and accumulating the result of the previous inner loop. There are many dependencies during the computation. Assuming that the computations of two loops are placed in one loop, it can be observed that the computations of Cube1 and Cube2 in the second loop only need to be performed after 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 applies to other computational flows. Therefore, the gaps between pipelines can be filled by pre-compiling Cubes. By pre-compiling Cubes, the dependencies of subsequent loop processes have been pre-calculated, thereby decoupling the multi-step computations and allowing the Cube and Vector computation units to compute in parallel, achieving the goal of optimizing efficiency.

[0106] Example:

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

[0108]

[0109]

[0110] After providing the pseudocode for the overall implementation, we first introduce an optimized code example for vector computation. The Sub interface in Ascend C is called to perform intervaled vector subtraction. The address steps between different datablocks in a single iteration are configured as follows: 1, 1, 0 for the result's src vector, the minuend's src vector, and the subtrahend's max vector. Simultaneously, the address steps between adjacent iterations of the same datablock in the result's src vector, the minuend's src vector, and the subtrahend's max vector are configured as width / perElementByBlk, width / perElementByBlk, 1, respectively. Where width is B... c That is, S' ij The width of `perElementByBlk` is the number of elements contained in each datablock. Each datablock is 32 bytes. For example, when the element is of type half, each element occupies 2 bytes of space, so each datablock contains 16 elements. When B... c ≤128, meaning the number of elements contained in 8 datablocks in this scenario, can be completed in one loop. One loop performs B... r The calculation is performed for each repeat, and B is calculated in each repeat. c 1 element. When B c >128, which requires... One loop, B is performed in one loop. r The calculation is performed in 128 repeats, with each repeat calculating 128 elements. The specific optimized vector subtraction implementation is shown in pseudocode 1.

[0111] Optimized vector subtraction implementation in pseudocode 1

[0112]

[0113] Next, we will introduce the implementation method of pre-launching Cube in a pipelined parallel manner. With the original implementation, if we want to pre-calculate Cube once, it's possible that when calculating Cube2 for the second time, Vec1, which depends on the result of the first Cube2 calculation, is simultaneously calculating. In this case, the result of the first Cube2 calculation is overwritten, resulting in a partially correct and partially incorrect Vec1 calculation. To avoid data corruption, we need to allocate an extra block of memory in global memory for each computation stream, and allocate which of these two memory blocks to write to based on the current flag. Simultaneously, for data dependencies, we also need to determine which block to retrieve data from based on the flag. In the overall flow, the first loop performs Cube1 and Cube2 calculations separately. The second loop performs Cube1, Cube2, Vector1 (depending on the result of Cube1, Cube2 from the first loop), and Cube3, Cube4. The third loop proceeds normally until the penultimate loop, performing Vector1, Cube3, Cube4, and Vector2. The final loop only performs Vector2. This requires two more loops than normal inner and outer loops, increasing loopQ by 2 (i.e., adding 2 to the outer loop), and during the last two outer loops, only one inner loop is executed. Meanwhile, since the operations required for the first, middle, and last iterations of the inner loops in steps two and four above differ, and this looping pattern disrupts the loop order, a marker is also needed. To ensure the correctness of the operator calculation results, it is also necessary to correctly transfer expmax, which operates between Vec1 and Vec2 across loops. This requires allocating two vector spaces for expmax and using identifiers to distinguish them. Furthermore, it is necessary to ensure the correctness of sum, used to update output O at the end of the inner loop. After the last calculation of sum in the inner loop, sum is written to global memory for subsequent reading. The specific implementation of the Cube pre-emission pipelined parallelism is shown in pseudocode 2.

[0114] Pseudocode 2Cube Early Launch Pipeline Parallelism

[0115]

[0116]

[0117] This operator performs a total of four matrix multiplications. These are the calculations of matrix multiplication. Obtain the matrix Calculate matrix multiplication Obtain the matrix Calculate matrix multiplication Obtain the matrix Calculate matrix multiplication Obtain the temporary result matrix Here, E is used. j Example of matrix multiplication of matrices and matrices. First, the matrix multiplication is... Therefore, by calling Ascend C's SetOrgShape(B) method... c The `,N,d)` interface sets the shape of the matrix multiplication. The left matrix's `TPosition` is `GM`, `CubeFormat` is `ND`, `DataType` is `half`, and the transpose enable parameter is `false`. The right matrix's `TPosition` is `GM`, `CubeFormat` is `ND`, `DataType` is `half`, and the transpose enable parameter is `false`. The result matrix's `TPosition` is `GM`, `CubeFormat` is `ND`, and `DataType` is `half`. Because `Bias` is not needed, the `isBias` parameter is set to `false`. The `SetFixSplit(128,128,128)` method in Ascend C is called to adjust `basicM`, `basicN`, and `basicK` in the basic block split, thereby reducing the transport time of matrix multiplication operations.

[0118] Operator correctness testing and performance testing:

[0119] (1) Operator correctness test

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

[0121] Table 1 Common values ​​for single-head dimension in 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 of shape N×d and data type half as input data. Next, the `random` method of the NumPy library generates two random number matrices E and F of shape k×N and data type half as dimensionality reduction matrices in the input data. First, the K matrix is ​​reduced in dimensionality by calculating the matrix multiplication result K' of matrices E and K using `numpy.matmul`. Then, the similarity matrix is ​​calculated by multiplying the transposes of matrices Q and K' using `numpy.matmul` to obtain the matrix S. Finally, the similarity matrix is ​​scaled by dividing each element of the similarity matrix by numpy.divide. Next, the safe softmax of the scaled similarity matrix is ​​calculated: The maximum value of each row of the similarity matrix is ​​found using `numpy.max`, and the maximum value of that row is subtracted from each element of the similarity matrix. The exponent of the similarity matrix is ​​calculated using `numpy.exp` and `numpy.sum` and divided by the sum of the rows, converting the elements of the similarity matrix into a probability distribution form to obtain the weight matrix. Finally, the matrix multiplication result V' of the E and K matrices is calculated using `numpy.matmul`. The dimension of the V matrix is ​​reduced, and the matrix multiplication of the weight matrix and V' matrix is ​​calculated using `numpy.matmul` to obtain the golden matrix. The golden matrix is ​​used as the ground truth, and the ground truth result calculated by Python is compared element-wise with the result calculated by the developed operator. The comparison method is as follows:

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

[0125] Where 'a' represents the calculation result obtained by the developed operator, 'b' represents the calculation result obtained by Python, and 'atol=10'. -6 rtol=10 -9 If formula (6) is true, the element is calculated correctly; otherwise, it is incorrect.

[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. The validation metric is etol≤10. -4 Then the operator is considered to be calculated accurately.

[0128] Table 2 Accuracy of Basic and Optimized Schemes

[0129]

[0130] (2) Operator performance test

[0131] First, separating matrix multiplication operations makes it easier to demonstrate the optimization effect of block partitioning. Let N = 1024, d = 128, and B... c =128. Take values ​​of 32, 64, and 128 for basicM respectively, keep basicN unchanged at 128, and take values ​​of 512, 256, and 128 for basicK.

[0132] Depend on Figure 9 It can be seen that as basicM increases, the MTE2 transport time and total time for different blocks of matrix multiplication decrease, which is basically consistent with the analysis results above.

[0133] Next, we compare the time taken to run the operator before and after vector optimization, where N = 512, and the x-axis represents the size of d in the computation dimension, starting from a smaller 16 and doubling each time, growing from 16 to 128.

[0134] Depend on Figure 10 The statistical results show that although vector computation optimization has improved the computation process and running time of operators, the optimization magnitude is small under various input sizes, with the largest magnitude being 3.9%.

[0135] Finally, the time taken to run the operator was compared before and after the Cube was pre-launched in pipelined parallelism. Here, N = 512, and the x-axis represents the size of d in the computational dimension, starting from a smaller 16 and doubling each time, increasing 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, with an overall optimization efficiency of over 10% and the highest optimization efficiency reaching 23.73%.

Claims

1. A method for accelerating sparse attention operators based on NPU, characterized in that... The method includes the following steps: Step 1: Divide the Q matrix into blocks according to the N-dimensional plane, into... piece ... Each piece is Br d. Divide the E and F matrices into blocks according to the k-dimensional plane, resulting in... Blocks, each block is [size missing] N; Step 2: Name the blocks of the Q matrix as The blocks of matrices E and F are named , ; Step 3: The outer loop iterates through blocks based on the Q matrix, and the inner loop iterates through blocks based on the E and F matrices; the Linformer operator in the inner loop is divided into four steps: Step 3-1: Calculate the similarity matrix: Step 3-1-1: Calculate matrix multiplication in the matrix multiplication calculation unit. Obtain the matrix That is, the K matrix is ​​divided into blocks after dimensionality reduction; Step 3-1-2: Calculate matrix multiplication in the matrix multiplication calculation unit. Obtain the matrix This serves as the initial block division of the similarity matrix; Step 3-2: Perform safe softmax calculation for block-based similarity matrix segmentation: Step 3-2-1: Scale the similarity matrix. Moving data from global memory to the local memory of the vector computation unit, for Perform vector multiplication ; Step 3-2-2: Calculate safe softmax; Step 3-3: Calculate the temporary result matrix in blocks: Step 3-3-1: Calculate matrix multiplication in the matrix calculation unit. Obtain the matrix That is, after dimensionality reduction Matrix partitioning; Step 3-3-2: Calculate matrix multiplication in the matrix multiplication calculation unit. Obtain the temporary result matrix ; Steps 3-4: Update and accumulate the results of the previous inner loop: If this is the first iteration of the inner loop, no update and accumulation process is needed; otherwise, the temporary result matrix from the previous iteration will be used. The vector sum of expmax is calculated by moving data from global memory to the local memory of the vector computation unit and performing intervaled vector multiplication. The update was correct. The temporary result matrix for this round The data is moved from global memory to the local memory of the vector computation unit, and the correct vector is added. and The addition yields a new temporary result matrix. If it's the last iteration of the inner loop, the resulting matrix at this point... The globally correct result matrix needs to be updated by performing a vector division with the sum vector. ; Step 4: Optimize vector computation based on the hardware characteristics of the NPU: Step 4-1: Optimize the vector calculation for safe softmax: Focus on the characteristics of the vector calculation unit, calculate in blocks, perform vector subtraction calculation for 8 data units in each iteration, and only need to perform n / 128 iterations instead of m iterations in one loop. Step 4-2: Optimize the vector calculation for row-by-row summation and row-by-row max calculation; Step 4-3, for = , = Optimize vector computation; Step 5: Optimize matrix multiplication by performing reasonable block-based computation based on the NPU's hardware characteristics. Step 5-1: Based on the matrix multiplication (M, N, K) and the buffer size limits of the AIC core's L0A Buffer, L0B Buffer, and L0C Buffer, give the size limits for basicM, basicN, and basicK, and the left matrix after block division. The right matrix after being sliced ​​and placed in the L0A Buffer The results can be placed in the L0BBuffer, and some of the calculated results can be placed in the L0C Buffer; Step 5-2: Based on the size constraints of basicM, basicN, and basicK, select the largest basicM and basicN, and the matching basicK, and divide the left and right matrices into basic blocks. Step 5-3: During matrix multiplication, the basic blocks are first moved from global memory into the AIV core's L1 buffer, then from the L1 buffer to the L0A and L0B buffers. After that, the basic block matrix multiplication operation begins, and the result is stored in the L0C buffer. After the K / basicK basic block matrix multiplication operation, the partially correct result of the matrix multiplication is retrieved, until all basic block operations are completed. Step 6: Optimize pipelined parallelism based on the NPU's hardware characteristics: The computational flow of the Linformer operator is divided into Cube1: obtaining the matrix through EK matrix multiplication. Cube2: Combine Q and Cube3: Matrix S is obtained by matrix multiplication using the transpose of Cube3. Cube4: via Matrix multiplication yields matrix O; Vector1: normalizes the result S from Cube2 and performs Flashsoftmax calculation to obtain matrix P; Vector2: updates and accumulates the result from the previous inner loop. Step 6-1: Allocate an additional block of memory in global memory for each computation stream; Step 6-2: After the calculation of Cube1 and Cube2 in the first inner loop is completed, the subsequent process is not carried out. Instead, the calculation of Cube1 and Cube2 in the second inner loop is performed in another space. The synchronous parallel calculation of Vector1 depends on the results of Cube1 and Cube2 in the first inner loop. Step 6-3: Calculate Cube3 and Cube4 of the first inner loop, calculate Vector1 of the second inner loop in parallel, calculate Cube1 and Cube2 of the third inner loop, calculate Vector2 of the first inner loop in parallel, and calculate Cube3 and Cube4 of the second inner loop in parallel. Step 6-4: Continue the loop as described above, allocating which of the two memory blocks to write to based on the current flag. Also, for data with data dependencies, determine which block to retrieve the data from based on the flag.

2. The NPU-based sparse attention operator acceleration method according to claim 1, characterized in that... The specific steps of step 3-2-2 are as follows: (1) If it is the first iteration of the inner loop, the vector calculation unit is... Perform a row-by-row max operation to obtain the inmax vector, and then calculate the vector computation unit. Subtracting the maximum value of each row from the elements, i.e., subtracting the inmax vector, is calculated by the vector computation unit. All elements are taken as natural indices, and the vector computation unit is... Perform a row-by-row summation operation to obtain the insum vector; (2) If it is not the first iteration of the inner loop, the vector calculation unit is... Perform a row-by-row max operation to obtain the inmax vector. The vector calculation unit calculates inmax and the max vector from the previous inner loop. The vector calculation unit then calculates... Subtract the max vector from the elements, take the natural exponent and sum row by row to get the insum vector. Calculate the difference between the max vector from the previous inner loop and the current temporary maximum value max vector after the current inner loop update by subtracting the vector, then take the natural exponent to get the expmax vector used to update the temporary result matrix and the sum insum vector of each row. Finally, update the current temporary sum sum vector by multiplying the expmax vector and the sum vector from the previous inner loop, and add the current insum vector. Then, update the latest sum and max vectors to the insum and inmax vectors respectively by moving them.

3. The NPU-based sparse attention operator acceleration method according to claim 1, characterized in that... The specific steps of step 4-2 are as follows: Step 4-2-1: Allocate blank space for (m, 64); Step 4-2-2: Add vectors to the first (m,64) shape vector and calculate the maximum vector from the second block. Record the resulting vector in the empty space of (m,64). For the remaining blocks, add vectors to the vector in the empty space and calculate the maximum vector. Step 4-2-3: If the column dimension of the matrix is ​​not divisible by 64, then take the remainder as the tail block and perform the same accumulation / max calculation, converting the row-by-row summation / max calculation of the matrix into the row-by-row summation / max calculation of the (m,64) space for convenient calculation by the Vector computing unit.

4. The NPU-based sparse attention operator acceleration method according to claim 1, characterized in that... The specific steps of step 4-3 are as follows: Step 4-3-1, / The address step size of different datablocks within 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 4-3-2, / The address step size of different datablocks within the same iteration is configured to be 0, and the address step size of the same datablock between adjacent iterations is configured to be 1.

Images