Linear approximation attention operator acceleration method based on NPU

Abstract

The invention discloses a linear approximation attention operator acceleration method based on NPU, and the method combines linear approximation attention with an Ascend C programming language, and achieves an Ascend C operator for calculating the linear approximation attention. According to hardware characteristics of the mercuration NPU, optimization of vector calculation is carried out, the number of instructions used for line-by-line summation and division is reduced, and time consumption is reduced; according to hardware characteristics of the mercuric chloride NPU, pipeline parallel optimization is carried out, large matrix multiplication is decomposed into multiple small matrix multiplication and asynchronous matrix multiplication, execution of matrix multiplication and vector calculation tasks can be coordinated more efficiently, and the utilization rate of hardware resources is maximized; cache occupation is optimized according to hardware characteristics of the mercuric chloride NPU, and the number of blocks of matrixes such as Q, K, P, Z and the like can be reduced under some input scales, so that the number of cycles of calculation and carrying is reduced, and the operation time of operators is shortened.

Description

Technical Field

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

[0002] Ascend C is a programming language launched by Ascend's self-developed heterogeneous computing architecture (CANN) for operator development scenarios. It natively supports C and C++ standards and specifications, maximizing the compatibility with user development habits. Through key technologies such as multi-layer interface abstraction, automatic parallel computing, and twin debugging, it greatly improves operator development efficiency, helping AI developers complete operator development and model tuning and deployment at low cost.

[0003] A major feature of the Ascend C programming language is its strong correlation with hardware. The core computing unit in the Ascend chip is AICore, and its abstract structure is shown in the following figure: Figure 1 shown.

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

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

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

[0007] Figure 1The 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; L0A Buffer and L0B Buffer are used to store the two matrices currently being multiplied; L0C Buffer stores the result matrix of the current matrix multiplication.

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

[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] When implementing Ascend C operators, it often involves accessing physical memory at all levels of the NPU. When managing physical memory at different levels, Ascend C uses an abstract logical location (TPosition) to represent each level of storage, replacing the concept of on-chip physical storage and concealing the hardware architecture. Common TPosition types include GM, VECIN, and VECOUT. GM corresponds to physical memory Global Memory, while VECIN and VECOUT correspond to physical memory Unified Buffer.

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

[0012] Vector calculations are performed by the Vector Computation Unit. Both the source and destination operands are stored in the Unified Buffer, and multiple iterations of calculations can be performed. During each iteration, the Vector Computation Unit retrieves eight consecutive datablocks (each datablock is 32 bytes) from the Unified Buffer, performs calculations, and writes the results back to the corresponding eight datablocks. The Vector Computation Unit supports interleaved vector calculations, which can be configured using dataBlockStride (the address stride between different datablocks in the same iteration) and repeatStride (the address stride between the same datablocks in adjacent iterations).

[0013] 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 compute 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 compute unit also controls program flow, scalar computations are converted to vector computations whenever possible to avoid synchronization blocking issues. Since the vector compute unit's cache size is fixed, for intermediate computations, its cache memory usage can be minimized, reducing the total number of transfers and thus reducing vector unit transfer time. 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 compute unit.

[0014] When performing attention calculation, Attention(Q, K, V) = softmax(QK T )V, because QK T The time and space complexity of matrix calculation is the square of the sequence length. If K can be calculated first T V, we can achieve a computational complexity that is linearly related to the sequence length. The problem is that softmax cannot be easily disassembled.

[0015] Katharopoulos et al. proposed the iconic linear attention, which shows that The function limit should be non-negative. The function uses elu(x)+1, which significantly improves memory and computational efficiency. At the same time, because it has limited ability to capture global dependencies, it is more suitable for real-time and streaming processing. However, when trying to migrate the linear approximate attention operator to the Ascend NPU, the following difficulties arise:

[0016] 1. Currently, there is no linear approximate attention operator that can run on the Ascend NPU.

[0017] 2. If the calculation method of the linear approximate attention operator is not changed and it is directly migrated to the Ascend NPU, the internal steps require more instructions to complete the calculation. For example, for the activation function elu(x), Ascend does not officially provide an API for direct calculation. If it is calculated directly, it will take a long time.

[0018] 3. The linear approximate attention operator has many dependencies between multiple steps and needs to be calculated serially, which cannot fully utilize the hardware features of the Ascend NPU.

[0019] 4. Direct migration of the linear approximate attention operator generates a large number of intermediate quantities within the calculation steps, occupying the Ascend NPU cache and increasing the number of calculation cycles. Summary of the Invention

[0020] The present invention provides an NPU-based linear approximate attention operator acceleration method. This method combines linear approximate attention with the Ascend C programming language to implement an Ascend C operator for calculating linear approximate attention. 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. Vector calculations are optimized based on the hardware characteristics of the Ascend NPU, reducing the number of instructions used for row-by-row summation and division, thereby reducing time consumption. Pipeline parallelization is optimized based on the hardware characteristics of the Ascend NPU, decomposing large matrix multiplications into multiple small matrix multiplications and asynchronous matrix multiplication operations, which can more efficiently coordinate the execution of matrix multiplication and vector calculation tasks and maximize the utilization of hardware resources. Cache occupancy is optimized based on the hardware characteristics of the Ascend NPU. Under certain input scales, the number of blocks for matrices such as Q, K, P, and Z can be reduced, thereby reducing the number of calculation and transfer cycles, and thus reducing the runtime of the operator.

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

[0022] A linear approximate attention operator acceleration method based on NPU, comprising the following steps:

[0023] Step 1: K TThe matrix is ​​divided into blocks of size B according to the d dimension. r *N matrix, move it from global memory to the UnifiedBuffer of the Vector calculation unit, and split the Q matrix into blocks of size B according to the N dimension n *d matrix, move it from global memory to the UnifiedBuffer of the vector computing unit. T =elu(K T )+1 activation function update, during the update process of the activation function of Q and K, the cache memory usage can be optimized:

[0024] Step 1: Compare the Q / K vector with the 0 scalar using CompareScalar. If the element in Q / K is larger, it is recorded in bit 1, otherwise it is recorded in bit 0 to obtain the comparison result vector, saving space for the 0 vector.

[0025] Step 2: Generate an element-wise natural exponential e for the Q / K vector Q / e K vector;

[0026] Step 3: Generate the Q+1 / K+1 vector. This can directly overwrite the Q / K vector in the cache without creating a new space.

[0027] Step 4: Obtain the updated Q / K vector by comparing the result vector from the Q+1 / K+1 vector and the exponential vector;

[0028] Step 2: Perform pipeline parallel optimization: first perform asynchronous matrix multiplication operations While performing matrix multiplication, the vector calculation unit synchronously calculates the activation function update of block Q and the activation function update of K T The i-th block of the matrix The row-by-row summation process in N dimensions obtains a multi-block matrix B, where the row-by-row summation calculation can be optimized for vector calculation based on the hardware characteristics of the NPU:

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

[0030] Step 2: Perform vector addition 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 with the vectors in the blank space through vector addition.

[0031] Step 3: If K T If the column dimension of the matrix cannot be divided into 64, then the remainder is used as the tail block and the accumulation is also performed once, and K TThe row-by-row summation of the matrix is ​​converted to the row-by-row summation of the (m, 64) space that is convenient for the Vector calculation unit to calculate;

[0032] Step 3: Move the Q matrix after the activation function update to the global memory, and calculate the matrix multiplication QA in the matrix multiplication calculation unit to obtain the matrix P;

[0033] Step 4: Calculate the matrix multiplication QB in the matrix multiplication calculation unit to obtain the normalized vector Z block;

[0034] Step 5: Move P and Z from global memory to the local memory of the vector computing unit in block form. Expand the Z block in the vector computing unit and normalize the P block row by row by dividing the Z block by the vector to obtain the result matrix. The expansion of the Z block can optimize the cache memory usage:

[0035] Step 1: Utilize the vector computation unit to support vector computation with intervals. By configuring the block step size between adjacent iterations of the Z matrix to 0, the vector division of the P matrix and the Z matrix can be directly calculated. This optimization can omit the expansion process and save cache space.

[0036] Step 2: By optimizing the cache memory usage, the problem is solved by:

[0037] 11*d*N / t*size+(N+d+d / t)*32≤192kb

[0038] Convert to

[0039] 8*d*N / t*size+(N+d+d / t)*32≤192kb

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

[0041] 1. This paper organically combines linear approximate attention calculation with the Ascend C programming language to implement a completely innovative linear approximate attention operator.

[0042] 2. This invention optimizes the vector calculations for the row-by-row summation operation of the linear approximate attention operator implemented on the Ascend NPU, and also optimizes cache memory usage, enabling the ingestion of more data and reducing the number of computation and transfer cycles. It also decomposes large matrix multiplications into smaller ones, subdividing the pipeline into smaller areas for pipeline parallelization. This method has been tested on the Atlas A2 training series and has improved the computational efficiency of the linear approximate attention operator. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] Figure 1 This is the AI ​​Core abstract architecture diagram.

[0044] Figure 2 This is a schematic diagram of the basic implementation steps of the linear approximate attention operator.

[0045] Figure 3 is the linear approximate attention operator K T Schematic diagram of the basic implementation of matrix row-by-row summation.

[0046] Figure 4 It is the basic scheme of linear approximate attention operator K T Schematic diagram of the matrix row-by-row summation pipeline.

[0047] Figure 5 is the linear approximate attention operator K T Schematic diagram of the matrix row-by-row summation optimization implementation.

[0048] Figure 6 is the linear approximate attention operator K T Schematic diagram of pipeline implementation for row-by-row matrix summation optimization.

[0049] Figure 7 This is a pipeline comparison diagram before and after pipeline parallel optimization of the linear approximate attention operator.

[0050] Figure 8 This is the pipeline diagram of the vector calculation unit before and after the linear approximate attention operator is optimized in parallel.

[0051] Figure 9 This is a time comparison chart of cache memory optimization solutions.

[0052] Figure 10 This is a time comparison chart of the basic solution and the optimized solution (N=256).

[0053] Figure 11 This is a time comparison chart of the basic solution and the optimized solution (N=512).

[0054] Figure 12 This is a time comparison chart between the optimization solution and the CPU implementation solution. DETAILED DESCRIPTION

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

[0056] The present invention provides a linear approximate attention operator acceleration method based on NPU, wherein the input data of the linear approximate attention operator is 3 R N×d The Q, K, V matrices, the output data is a size of R N×dO matrix. Figure 2 As shown, the basic implementation steps of the linear approximate attention operator are as follows:

[0057] Step 1: Q and K given by input T The matrix is ​​used to update the activation function. T The matrix is ​​divided into blocks of size B according to the d dimension. r *N matrix, move it from global memory to the Unified Buffer of the Vector calculation unit. T =elu(K T )+1 update, then perform ReduceSum on N dimensions to obtain the multi-block matrix B The reason why B needs 8 columns is that the vector calculation unit requires that the first address and the operation length are both aligned to 32Byte. Originally, the B matrix is ​​the calculation of the denominator in the expression, but it is also included in step 1 because if the B matrix calculation is combined with K T The block activation function of the matrix is ​​updated separately. If these two steps are separated, K needs to be moved in twice. T The matrix is ​​transferred to the UnifiedBuffer of the Vector calculation unit, and the two are merged to reduce the time consumption. Next, the Q matrix is ​​divided into blocks of size B according to the N dimension. n *d matrix, move it from global memory to the UnifiedBuffer of the vector calculation unit, and update Q=elu(Q)+1 after the vector calculation unit divides the blocks. In the process of updating the activation function, here, take the Q matrix as an example, the operation of the K matrix is ​​the same. If each unit element of Q is assigned a value after traversing each value to determine whether it is positive or negative, the performance will be greatly reduced due to too many loops. Therefore, the implementation here is to generate a Q+1 vector by summing the vector element by element and the scalar, and generate an element-by-element natural exponential e for Q. Q Next, the sign of each element of the entire vector is determined in a vectorized manner. First, an all-zero vector of the same size as the Q block is generated. This is then compared element-by-element with the zero vector. If the element in Q is larger, the bit is set to 1; otherwise, the bit is set to 0, resulting in a comparison result vector. Finally, the comparison result vector is used to select the updated Q vector from the Q+1 vector and the exponential vector. However, the vectorized activation function process generates a large number of intermediate quantities.

[0058] Step 2: Calculate the numerator of the linear approximate attention operator mathematical expression. T The matrix is ​​moved out to global memory and the matrix multiplication K is calculated in the matrix multiplication calculation unit. T V(R d×N *R N×d ) to get the matrix A(Rd ×d ), at the same time, the K after the activation function is updated T The matrix is ​​summed in N dimensions to obtain the B matrix, and the obtained B is aggregated in the local memory (R d×8 ), prepare for the next matrix multiplication with Q. Then, the Q matrix after the activation function update is moved out to the global memory, and the matrix multiplication QA(R N×d *R d×d ) to obtain the matrix P(R N×d ).

[0059] Mathematical expression of linear approximate attention operator:

[0060]

[0061] Step 3: Calculate the denominator in the mathematical expression of the linear approximate attention operator and calculate the remaining part. In the matrix multiplication calculation unit, calculate the matrix multiplication QB(R N×d *R d×8 ) to obtain the normalized vector Z block (R N×8 ).

[0062] Step 4: Move P and Z from the global memory to the local memory of the vector computing unit in block form, expand the Z block in the vector computing unit, and normalize the P block row by row by dividing the Z block by the vector to obtain the result matrix.

[0063] At the same time, according to the hardware characteristics of Huawei Ascend NPU, vector calculation optimization is first performed. T The process of generating matrix B by summing row by row is specifically expanded as follows: for each row of the matrix, the sum operation is first performed to obtain a scalar result b. However, due to the characteristics of the vector calculation unit, if these results are needed in subsequent calculations, b must be expanded to a vector format that conforms to the datablock unit. The schematic diagram is as follows Figure 3 As shown. This process involves the operation of the Scalar computing unit. The Scalar computing unit is not only responsible for scalar calculations, but also undertakes the process control task of the program. In this process, multi-layer loops will frequently call the Scalar computing unit, which may cause other computing pipelines to be blocked. In addition, the synchronization wait between the Scalar and Vector computing units will also introduce additional delays. Considering the computing characteristics of the Vector computing unit that can perform multiple iterative calculations, it is very suitable for large-block vector calculations. First, apply for a blank space of (m,64) and place K TThe matrix is ​​considered as a vector composed of multiple blocks of vectors in the shape of (m, 64). First, the first block and the second block are vector-added. The resulting vector is recorded in the blank space of (m, 64). The remaining blocks are also accumulated with the vectors in the blank space through vector addition. If K T If the column dimension of the matrix cannot be divided evenly by 64, then the remainder is used as the tail block and is also accumulated once. The schematic diagram is as follows Figure 5 As shown. At this time, through such a scheme, K T The row-by-row summation of the matrix is ​​converted to the row-by-row summation of the (m, 64) space. This conversion is mainly to adapt to the characteristics of the Vector calculation unit. In each iteration, the Vector calculation unit will take out 8 datablocks from the UnifiedBuffer, corresponding to the size of 64*sizeof(float). While iterating the taken-out vector, the elements of the vector can be accumulated and summed. This calculation method avoids the need to execute multiple times through the loop, and makes full use of the ability of the Vector calculation unit to iterate multiple times. Through multiple iterations, more efficient calculations are achieved, thereby significantly improving the execution speed of the summation process. The pipeline diagram is shown above. Figure 6 As shown in the figure, the basic implementation calls 2n vector calculation instructions and n scalar calculation instructions. The optimized solution, however, only calls m / 64 vector calculation instructions. The efficiency improvement of the optimized solution is particularly significant when m is small and n is large.

[0064] Then we proceed with pipeline parallel optimization. After analyzing the calculation flow, we find that the matrix multiplication K in step 2 T V does not need to be in all blocks The activation functions of the blocks can be updated before they can be calculated. After the activation function of is updated, the matrix multiplication of the block is calculated, which can be calculated in parallel with some Vector operations in step 1. Specifically, After the activation function is updated, the asynchronous matrix multiplication operation is performed first. While performing matrix multiplication, the vector calculation unit synchronously calculates the activation function update of block Q and the activation function update of K T The i-th block of the matrix The increased communication time in the row-by-row summation process is much smaller than the total computing time reduced by the vector computing unit through pipeline parallelization. The pipeline diagram before and after parallelization is shown in the figure. Figure 7 As shown, the parallel front and back vector calculation unit pipeline is as follows Figure 8 Through this decomposition strategy and asynchronous matrix multiplication operations, the execution of matrix multiplication and vector calculation tasks can be coordinated more efficiently, maximizing the utilization of hardware resources.

[0065] Finally, cache usage is optimized. Based on the architectural characteristics of the Huawei Ascend NPU 910B4, its vector computing unit is equipped with a 192KB Unified Buffer. During the block segmentation process, the cache capacity must be fully utilized to reduce loop counts, thereby minimizing header overhead and improving computational efficiency. When requesting memory space in the Unified Buffer, the logical locations of the requested data structures must be specified: VECIN and VECOUT. When moving data from global memory, data can only be moved to the data structure in the VECIN logic. When moving data from global memory, data can only be moved from the data structure in the VECOUT logic. In the basic implementation, a large number of intermediate quantities are generated during the activation function update of the Q and K matrices. Furthermore, during the row-by-row normalization of the P matrix and the Z vector, Z is expanded. Both the intermediate quantities and the expansion of Z need to be stored in the Unified Buffer. In fact, during the update process of the activation function, there is no need to open new space for the Q+1 / K+1 vector. As long as the Q / K vector is compared with the 0 vector first, and then the Q / K vector in the cache is directly overwritten, there will be no impact on the correctness. At the same time, for the 0 vector compared with the Q / K vector, it can be compared with the 0 scalar through CompareScalar to achieve the same effect. For the expansion of Z, the vector calculation unit can be used to support vector calculations with intervals. By configuring the block step size between adjacent iterations of the Z matrix to 0, the vector division of the P matrix and the Z matrix can be directly calculated. These two cache memory optimization measures reduce the cache memory usage. By optimizing the cache memory usage, the problem is reduced to:

[0066] 11*d*N / t*size+(N+d+d / t)*32≤192kb(2)

[0067] Convert to

[0068] 8*d*N / t*size+(N+d+d / t)*32≤192kb(3)

[0069] For certain input sizes, the number of blocks used to partition matrices such as Q, K, P, and Z can be reduced, thereby reducing the number of calculation and transfer cycles and, in turn, the runtime of the operator. This optimization strategy not only reduces the frequency of data transfers but also improves cache hit rates, effectively increasing the resource utilization of the computing unit.

[0070] Example:

[0071] In this embodiment, the software environment for implementing the linear attention operator is CANN Community Edition 8.0.RC3.alpha001, and the hardware environment is the AtlasA2 training series product.

[0072] 1. Class definition and main function implementation

[0073] The class of this operator is named LinearAttention, which mainly includes 2 main functions, namely the initialization function Init() and the logic execution function Process(); and 8 functions, namely Compute(), MatmulComputeP(), MatmulComputeZ(), CopyInP(), ComputeResult(), CopyOut(), ComputeReduceSum(), EluCompute(). Considering the time consumed in calculation, it is necessary to calculate the numerator and denominator in parallel with the activation function update. Compute(), MatmulComputeP(), MatmulComputeZ(), ComputeResult() correspond to the four steps of the present invention one by one, EluCompute() is the function used to update the activation function, and ComputeReduceSum corresponds to K T Sum row by row, CopyInP() is used to move the intermediate quantity P, and CopyOut() is used to move the result out.

[0074] The rest of the text mainly includes the definition of the matrix multiplication class, which is named mm1, mm2, and mm3 respectively; the setting of the matrix multiplication execution information structure, which is named tiling1, tiling2, and tiling3; the definition of the output of the intermediate quantity of the off-board input, that is, the definition of the input as qGlobal, kGlobal, and vGlobal, corresponding to the input Q, K, and V matrices respectively. The intermediate quantities are defined as aGlobal, pGlobal, and qOutGlobal, corresponding to the intermediate quantities A, P matrices, and the Q matrix after the activation function is updated. The output is defined as resultGlobal. For the declaration of the TQue variable on the board, first set Tposition to VECOUT, including zQueue, which is used to store the calculation result Z matrix of the denominator; including qOutQueue and kOutQueue, which are used to store the blocks Q and K after the activation function is updated. T Matrix; including bAllQueue, used to store the complete B matrix; including resultQueue, used to store the calculated result matrix. Then set Tposition to VECIN, including qQueue and kQueue, used to store the moved-in block input Q and K T ; Includes pQueue, used to store the intermediate amount P of the moved blocks; includes bQueue, used to store the blocks K TThe row-by-row summation expands the result. Declare the TBuf variable on the board. TBuf variables are primarily used for calculations and not for data transfer. Set expBuf and compareBuf for exponential operations and comparison bit storage in the activation function, respectively. Set sumBuf and sumlBuf for temporary storage and result storage in the row-by-row summation, respectively. Enter N, d, and t as the number of tiles, and tileLength as the tile matrix size.

[0075] The operator's initialization function, Init(), performs three main operations. The first is to load the split information for the three matrix multiplications. The second is to locate the external input and output memory locations using the SetGlobalBuffer() method. The third is to initialize all requested Tque and Tbuf instances, which requires explicit memory allocation.

[0076] The Process() function of this operator is primarily responsible for executing the computational logic. It first partitions the inputs Q and K into blocks, then loops through the blocks multiple times. Within the loop, the Ascend C DataCopy interface is used to move the blocks of Q and K from global memory into the Vector computation unit. The Compute() function is then called, which updates the activation function and performs parallel computations on some of the numerator and denominator. After the loop, the MatmulComputeP() function is called to calculate the numerator, and the MatmulComputeZ() function is called to calculate the denominator. Finally, because the P matrix is ​​large, the intermediate quantities P and Z are partitioned into blocks. Multiple loops are then repeated, moving the blocks of the P matrix from global memory into the Vector computation unit. The ComputeResult() function is then called to normalize the blocks of the computation results and move them from the Vector computation unit to global memory. The specific implementation is shown in pseudocode 1.

[0077] Pseudocode 1Process() function implementation

[0078]

[0079] 2. Specific implementation of four-stage functions

[0080] (1) Compute() function implementation

[0081] In the Compute() function, K is first calculated by calling the EluCompute() function. T The activation function is updated. Then the IterateAll interface in Ascend C is called to perform K TThe matrix multiplication of V is performed by setting the parameter in <> to false, using asynchronous mode for matrix multiplication. While waiting for the result, the EluCompute() function is called to update the activation function of Q. The updated Q block is moved to the intermediate quantity outside the board for storage by calling the DataCopy interface in Ascend C, and the ComputeReduceSum() function is called to update the updated K T The sum is performed in blocks. The specific implementation of the Compute() function is shown in pseudo code 2.

[0082] Pseudocode 2Compute() function implementation

[0083]

[0084]

[0085] In the function, EluCompute() is used to calculate the activation function. For the function's argument, LocalTensor x, it uses the CompareScalar interface in Ascend C to compare each element of x with the scalar 0. If the comparison result is less than the scalar, the corresponding bit of the output result is set to 1; otherwise, it is set to 0. This saves d*N / t*sizeof(float) buffer space compared to generating a LocalTensor of all 0s and then performing an element-by-element comparison between the two LocalTensors. Next, the Exp interface in Ascend C is used to obtain the intermediate result, expTensor, by taking the element-by-element exponential of x. At this point, the Adds interface in Ascend C is used to obtain x+1, and the result is stored in x. At this point, the original value of x does not affect the calculation result, saving d*N / t*sizeof(float) buffer space. Finally, the Select interface in Ascend C is used to select elements based on the bit values ​​of the previously obtained comparison result, compareTensor. If the value is less than 0, x is exponentially converted; if it is greater than or equal to 0, x is converted to x+1. The final result is stored in y. The specific implementation of the EluCompute() function is shown in pseudo code 3.

[0086] Pseudocode 3EluCompute() function implementation

[0087]

[0088] In the function, ComputeReduceSum() performs row-by-row summation. When the input N is less than 64, the WholeReduceSum interface in Ascend C can be used to sum all data within each repeat. When N is less than 64, the portion to be summed can be limited to less than 8 datablocks at a time. The WholeReduceSum interface in Ascend C can be directly used to perform the calculation d / t times, summing N data points in each repeat. When the input N is greater than or equal to 64, the results are first accumulated into a blank space of (d / t, 64) by calling Ascend C's Add interface N / 64 times. The WholeReduceSum interface in Ascend C is then used to perform the calculation d / t times, summing 64 data points in each repeat. The specific implementation of the ComputeReduceSum() function is shown in pseudo-code 4.

[0089] Pseudocode 4ComputeReduceSum() function implementation

[0090]

[0091] (2) MatmulComputeP() function implementation

[0092] In the MatmulComputeP() function, the main function is to call Ascend C's IterateAll to perform matrix multiplication of the Q matrix and the A matrix to obtain the P matrix.

[0093] (3) MatmulComputeZ() function implementation

[0094] In the MatmulComputeZ() function, the main function is to call Ascend C's IterateAll to perform matrix multiplication of the Q matrix and the B matrix to obtain the Z matrix.

[0095] (4) ComputeResult() function implementation

[0096] The ComputeResult() function primarily leverages the hardware feature of the Ascend NPU's vector compute unit, which supports gapped vector computations. It calls the Div interface in Ascend C to perform gapped vector division. The address steps between different data blocks within a single iteration of the result, p, and z vectors are configured to be d / 8, d / 8, and 1, respectively. Furthermore, the address steps within the same data block between adjacent iterations of the result, p, and z vectors are configured to be 1, 1, and 0, respectively, to calculate the normalized result. The coordinates of the result and p vectors are moved by the same d*8 coordinates between each repeat. Due to its different shape, the z vector moves by 64 coordinates each time. Due to the different computation blocks, the p vector only moves that block, while the z vector remains intact. Therefore, the z vector's coordinates need to be corrected by i*N / t*8 for block i. A total of N / t / 8 vector divisions are performed, with d / 8 vector divisions performed within each repeat. The specific implementation of the ComputeResult() function is shown in pseudocode 5.

[0097] Pseudocode 5ComputeResult() function implementation

[0098]

[0099] 3. Operator kernel function body and matrix multiplication information setting

[0100] The partial implementation of the overall operator kernel function is shown in pseudo code 6.

[0101] Pseudocode 6: Partial implementation of operator kernel function

[0102]

[0103]

[0104] In the above code, line 1 applies for the TPipe type, Ascend C's unique on-chip memory management system. This is a common step for all Ascend C operators. Lines 2-7 call CopyTiling to load the matrix tile information. Line 8 declares the LinearAttention class variable, linearAttention. Line 9 executes the Init function. Line 10 registers information for the matrix multiplication calculation. Line 11 executes the Process function. This completes the kernel function for the Ascend_CoppersmithWinograd operator. Using the advanced matrix multiplication API requires writing host-side tiling code, which includes matrix multiplication operation registration information and operation scale settings.

[0105] In this operator, a total of three matrix multiplications are performed, namely K T The matrix multiplication of the block matrix and the V matrix, the matrix multiplication of the Q matrix and the A matrix, and the matrix multiplication of the Q matrix and the B matrix. Here we take the matrix multiplication of the Q matrix and the A matrix as an example. First, the matrix multiplication is R N×d *R d×d Therefore, the shape of the matrix multiplication is set by calling the SetOrgShape(N,d,d) interface of Ascend C. The TPosition of the left matrix is ​​GM, the CubeFormat is of type ND, the DataType is of type float, and the transpose enable parameter is false; the TPosition of the right matrix is ​​GM, the CubeFormat is of type ND, the DataType is of type float, and the transpose enable parameter is false; the TPosition of the result matrix is ​​GM, the CubeFormat is of type ND, and the DataType is float; because Bias is not needed, the isBias parameter is set to false.

[0106] 4. Operator correctness test and performance test

[0107] (1) Operator correctness test

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

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

[0110]

[0111] The correctness test uses the random method of the numpy library in Python to generate three random number matrices Q, K, and V with a shape of N×d and a data type of float as input data. T The matrix uses the numpy.where method to update the activation function first. If x is less than 0, it is equal to e x , if x is greater than or equal to 0, then it is equal to x+1. Get the updated Q and K T After the matrix is ​​formed, the updated K is calculated by numpy.matmul T The matrix multiplication result of the matrix A and the matrix V is calculated by numpy.matmul, and the matrix multiplication result of the updated matrix Q and the matrix A is calculated by numpy.sum. TThe matrix is ​​summed in N dimensions to obtain the matrix B. Then, the matrix Z is calculated by multiplying the updated matrix Q and the matrix B using numpy.matmul. Finally, the P matrix and the Z matrix are divided row by row using numpy.divide to obtain the golden matrix. Taking the golden matrix as the true value, the true value calculated by Python and the result calculated by the development operator are compared element by element using the following comparison method:

[0112] |ab|≤(atol+rtol*|b|) (4)

[0113] 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 (4) holds true, the calculation of this element is correct, otherwise it is wrong.

[0114]

[0115] The error in formula (5) num Count the number of errors for the element, all num is the total number of elements in the result, etol is the element calculation error rate, and etol in the verification indicator is ≤10 -4 The operator is considered to be accurate.

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

[0117]

[0118]

[0119] (2) Operator performance test

[0120] While ensuring correctness, we compared the runtime of the linear approximate attention operator implemented before and after optimization to reduce cache usage. N is 256, and the x-axis is the size of the computational dimension d, starting from a small 16 and doubling each time from 16 to 128. Each group is calculated 10 times, and the y-axis is the average runtime of the algorithm.

[0121] Depend on Figure 9 It can be seen that when d is relatively small, that is, when the input scale is small, the extent of block optimization is not obvious or even non-existent. This is because when the input scale is small, even if the memory usage in the cache is not optimized, there is no need for blocking and calculation can be performed directly. When d gradually increases, the optimization effect reaches a maximum of 23.30%.

[0122] Next, the average running time of the basic solution and the final optimized solution were compared under the same input scale changes.

[0123] like Figure 10 As shown in the figure, the optimized solution runs significantly faster than the basic solution at any input scale tested. As the input scale increases, the optimization efficiency gradually increases. The lowest optimization efficiency is 20.7%, and the highest optimization efficiency can reach 46%.

[0124] After increasing the N value in the test, due to the limited cache size, the block division of the basic solution and the optimized solution is basically the same, and the block optimization has no significant effect. Figure 11 This demonstrates the effectiveness of the optimization scheme's pipeline parallelization and vector computation optimization. In the N = 512 test, especially at the larger input size of d = 128, the optimization efficiency dropped to 23.1% due to the inability to reduce the number of blocks. When d = 32, the optimization scheme's efficiency reached a peak of 33.9% in this set of experiments.

[0125] At the same time, the efficiency of the CPU implementation scheme and the tuning scheme are compared, such as Figure 12 As shown in the figure, it can be seen that the calculation time of the CPU solution is more than 15 times or even more than that of the tuning solution. This shows that the hardware-friendly implementation solution of this operator in this article accelerates the calculation.

Claims

1. A linear approximate attention operator acceleration method based on NPU, wherein the method comprises the following steps: Step 1: K T The matrix is ​​divided into blocks of size B according to the d dimension. r *N matrix, move it from global memory to the UnifiedBuffer of the Vector calculation unit, and split the Q matrix into blocks of size B according to the N dimension n *d matrix, move it from global memory to the UnifiedBuffer of the vector computing unit. T =elu(K T )+1 activation function update; Step 2: Perform pipeline parallel optimization: first perform asynchronous matrix multiplication operations While performing matrix multiplication, the vector calculation unit synchronously calculates the activation function update of block Q and the activation function update of K T The i-th block of the matrix The row-by-row summation process in N dimensions obtains a multi-block matrix B; Step 3: Move the Q matrix after the activation function update to the global memory, and calculate the matrix multiplication QA in the matrix multiplication calculation unit to obtain the matrix P; Step 4: Calculate the matrix multiplication QB in the matrix multiplication calculation unit to obtain the normalized vector Z block; Step 5: Move P and Z from the global memory to the local memory of the vector computing unit in block form, expand the Z block in the vector computing unit, and normalize the P block row by row by dividing the Z block by the vector to obtain the result matrix.

2. The NPU-based linear approximate attention operator acceleration method according to claim 1 is characterized in that in step 1, during the update process of the activation functions of Q and K, cache memory occupancy is optimized: Step 1: Compare the Q / K vector with the scalar 0 using CompareScalar. If the element in Q / K is larger, it is recorded as bit 1, otherwise it is recorded as bit 0 to obtain the comparison result vector; Step 2: Generate an element-wise natural exponential e for the Q / K vector Q / e K vector; Step 3: Generate Q+1 / K+1 vectors and directly overwrite the Q / K vectors in the cache; Step 4: Obtain the updated Q / K vector by comparing the result vector from the Q+1 / K+1 vector and the exponential vector.

3. The NPU-based linear approximate attention operator acceleration method according to claim 1 is characterized in that In step 2, the row-by-row summation is optimized for vector calculation based on the hardware characteristics of the NPU: Step 1: Apply for a blank space of (m,64); Step 2: Perform vector addition 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 with the vector in the blank space through vector addition. Step 3: If K T If the column dimension of the matrix cannot be divided into 64, then the remainder is used as the tail block and the accumulation is also performed once, and K T The row-by-row summation of the matrix is ​​converted to the row-by-row summation of the (m, 64) space that is convenient for the Vector calculation unit to calculate.

4. The NPU-based linear approximate attention operator acceleration method according to claim 1 is characterized in that In step 5, the cache memory usage is optimized for the expansion of the Z block: Step 1: Use the vector calculation unit to support vector calculation with intervals. Configure the block step size between adjacent iterations of the Z matrix to be 0 to directly calculate the vector division of the P matrix and the Z matrix. Step 2: By optimizing the cache memory usage, the problem is solved by: 11*d*N / t*size+(N+d+d / t)*32≤192kb Convert to 8*d*N / t*size+(N+d+d / t)*32≤192kb.

Images