An NPU-based linear approximation attention operator acceleration method

Abstract

The application discloses an NPU-based linear approximation attention operator acceleration method, which combines linear approximation attention with an Ascend C programming language, realizes an Ascend C operator for calculating linear approximation attention, optimizes vector calculation according to hardware characteristics of an Ascend NPU, reduces the number of instructions for row-by-row summation and division, and reduces time consumption, optimizes flow water parallelism according to the hardware characteristics of the Ascend NPU, decomposes large matrix multiplication into multiple small matrix multiplication and asynchronous matrix multiplication operation, can more efficiently coordinate the execution of matrix multiplication and vector calculation tasks, and maximizes the utilization rate of hardware resources, and optimizes cache occupation according to the hardware characteristics of the Ascend NPU, can reduce the block number of Q, K, P and Z matrices under some input scales, thereby reducing the cycle number of calculation and transportation, and further reducing the running time of the operator.

Description

Technical Field

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

[0002] Ascend C is a programming language developed by Ascend's self-developed heterogeneous computing architecture (CANN) for operator development scenarios. It natively supports C and C++ standards and specifications to maximize 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 and helps AI developers complete operator development and model optimization and deployment at low cost.

[0003] A key characteristic of the Ascend C programming language is its strong hardware dependence. The core computing unit in the Ascend chip is the AICore, whose abstract structure diagram is shown below. Figure 1 As shown.

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

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

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

[0007] Figure 1The main storage units within the upper-middle part of the AIC architecture core are: 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.

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

[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] When implementing operators in Ascend C, it is often necessary to call upon various levels of physical memory in the NPU. Ascend C manages different levels of physical memory using an abstract logical location (TPosition) to represent each level of storage, replacing the concept of on-chip physical memory and thus hiding the hardware architecture. Commonly used TPosition types include: GM, VECIN, and VECOUT. GM corresponds to Global Memory, while VECIN and VECOUT correspond to UnifiedBuffer.

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

[0012] Vector computation is performed using the Vector computation unit. Both the source and destination operands are stored in a Unified Buffer, and multiple iterations are possible. During each iteration, the Vector computation unit retrieves eight consecutive data blocks (each 32 bytes) from the Unified Buffer, 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).

[0013] The runtime of the Ascend operator is closely related to the hardware architecture. The main computation time can be roughly divided 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 needs to be utilized 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 vector computation 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. Since the cache size of the vector computation unit is fixed, for intermediate quantity computations, its cache memory usage can be minimized to reduce the total number of data transports, thereby reducing the data transport time of the vector unit. In addition to 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.

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

[0015] Katharopoulos et al. proposed the landmark concept of linear attention, explaining... The function should be restricted to non-negativity, and in its experiments... The function uses elu(x)+1, which significantly improves memory and computational efficiency. However, due to its limited ability to capture global dependencies, it is well-suited for real-time and streaming processing. However, when attempting to migrate the linear approximate attention operator to the Ascend NPU, the following challenges arise:

[0016] 1. There is currently no linear approximate attention operator that can run on 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 a lot of instructions to complete the calculation. For example, Ascend does not provide an API for directly calculating the activation function elu(x), so it will take a long time to calculate it directly.

[0018] 3. The linear approximate attention operator has many dependencies between multiple steps, requiring sequential computation, which makes it difficult to fully utilize the hardware features of the Ascend NPU.

[0019] 4. Direct transfer of linear approximate attention operators generates a large number of intermediate quantities within the computation steps, occupying the Ascend NPU cache and increasing the number of computation loops. Summary of the Invention

[0020] This invention provides an acceleration method for linear approximate attention operators based on an NPU. 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 category of Ascend C operators, implementing the four stages of the overall operation to form an Ascend C hybrid programming operator with certain scalability and robustness. Vector computation is optimized based on the hardware characteristics of the Ascend NPU, reducing the number of instructions used for row-by-row summation and division, thus reducing time consumption. Pipeline parallelism 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 computation tasks, maximizing the utilization of hardware resources. Cache usage is optimized based on the hardware characteristics of the Ascend NPU, which can reduce the number of blocks for Q, K, P, Z matrices under certain input scales, thereby reducing the number of computation and transport loops, and thus reducing the operator's runtime.

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

[0022] An NPU-based linear approximate attention operator acceleration method includes the following steps:

[0023] Step 1: K TThe matrix is ​​divided into blocks of size B along dimension d. r The *N matrix is ​​moved from global memory into the UnifiedBuffer of the Vector computation unit, and the Q matrix is ​​divided into blocks of size B along the N dimensions. n The matrix *d is moved from global memory into the UnifiedBuffer of the vector computation unit. After being divided into blocks in the vector computation unit, K... T =elu(K T The activation function is updated by 1+1. During the update process of the activation functions for Q and K, cache memory usage can be optimized.

[0024] Step 1: Compare the Q / K vector with the 0 scalar using CompareScalar. If an element in Q / K is larger, record it with a 1 bit; otherwise, record it with a 0 bit. This will give you the comparison result vector and save space on the 0 vector.

[0025] Step 2: Generate an element-wise natural exponent e from the Q / K vector. Q / e K vector;

[0026] Step 3: Generate Q+1 / K+1 vectors. You don't need to allocate new space; you can directly overwrite the Q / K vectors in the cache.

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

[0028] Step 2: Perform pipelined parallel optimization: First, perform asynchronous matrix multiplication operations. Simultaneously with matrix multiplication, the vector computation unit synchronously calculates the activation function update of block Q, and updates K. T The i-th block of the matrix The row-by-row summation process in the N-dimensional dimension yields multiple matrices B, where the row-by-row summation calculation can be optimized for vector computation based on the hardware characteristics of the NPU:

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

[0030] Step 2: Add the vector of the first block (m,64) and the vector of the second block, and record the resulting vector in the blank space of (m,64). The remaining blocks are also added to the vector in the blank space by vector addition.

[0031] Step 3, if K T If the column dimensions of the matrix are not divisible by 64, then the remainder is used as the tail block and accumulated again, with K... TThe row-by-row summation of the matrix is ​​transformed into the row-by-row summation of the (m,64) space, which is convenient for the Vector computation unit;

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

[0033] Step 4: Calculate 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 computation unit in blocks. Expand the Z block in the vector computation unit, and normalize the P block row-by-row by dividing the vector by the Z block to obtain the result matrix. The expansion of the Z block can be optimized by reducing cache memory usage.

[0035] Step 1: Utilize the vector computation unit to support spaced vector computation. 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 cache memory usage, the problem can be addressed 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 invention organically combines linear approximate attention computation with the Ascend C programming language to realize a completely innovative linear approximate attention operator.

[0042] 2. This invention optimizes the vector computation of the row-by-row summation operation of the linear approximate attention operator implemented on the Ascend NPU, and also optimizes cache memory usage, thereby enabling the loading of more data and reducing the number of computation and data transfer loops. Simultaneously, it decomposes large matrix multiplications into smaller matrix multiplications, subdividing the pipeline into smaller segments for pipeline parallel optimization. Actual testing on the Atlas A2 training series products has shown that this improves the computational efficiency of the linear approximate attention operator. Attached Figure Description

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

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

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

[0046] Figure 4 K is the basic scheme of the linear approximate attention operator. T A schematic diagram of the matrix row-by-row summation process.

[0047] Figure 5 K is a linear approximate attention operator T A schematic diagram illustrating the optimized implementation of row-by-row summation of a matrix.

[0048] Figure 6 K is a linear approximate attention operator T A schematic diagram of the pipeline implementation achieved by optimizing the row-by-row summation of a matrix.

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

[0050] Figure 8 This is a pipeline diagram of vector computation units before and after pipelined parallel optimization using a linear approximate attention operator.

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

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

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

[0054] Figure 12 This is a time comparison chart between the optimized solution and the CPU implementation solution. Detailed Implementation

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

[0056] This invention provides an NPU-based linear approximate attention operator acceleration method, wherein the input data of the linear approximate attention operator is three data points of shape R. N×d The Q, K, V matrices output data is a single matrix of size R. N×dThe O matrix. For example... Figure 2 As shown, the basic implementation steps of the linear approximate attention operator are as follows:

[0057] Step 1: Given the inputs Q and K T The matrix is ​​updated using the activation function. K is then... T The matrix is ​​divided into blocks of size B along dimension d. r The matrix *N is moved from global memory into the Unified Buffer of the Vector computation unit. After being divided into blocks in the Vector computation unit, K... T =elu(K T +1 update, then perform ReduceSum in N dimensions to obtain multiple matrix B The reason B needs 8 columns is that the vector computation unit requires both the starting address and the operation length to be 32-byte aligned. Originally, the B matrix is ​​part of the denominator calculation in the expression, but it's also included in step one because if the calculation of the B matrix and K were combined... T If the block activation function update of the matrix is ​​separated into these two steps, then K needs to be moved in twice. T The matrix is ​​converted to a Vector computation unit using a UnifiedBuffer, which merges the two to reduce data transfer time. Next, the Q matrix is ​​divided into blocks of size B along N dimensions. n The matrix *d is moved from global memory into the UnifiedBuffer of the vector computation unit, and then updated in the vector computation unit after being divided into blocks: Q = elu(Q) + 1. During the update of the activation function, taking the Q matrix as an example, the operation on the K matrix is ​​the same. If each element of Q were assigned a value after iterating through each value and determining its sign, the performance would be significantly degraded due to the large number of loops. Therefore, the implementation here generates a Q+1 vector by summing the vector element-wise with a scalar, and then generates an element-wise exponentiation e of Q. Q The next step is to vectorize and determine the sign of each element in the entire vector. First, a vector of all zeros of the same size as the Q-block is generated. Then, an element-wise comparison is performed between the Q-block and the zero vector. If an element in Q is larger, it is recorded with a 1 bit; otherwise, it is recorded with a 0 bit, resulting in a comparison result vector. Finally, the updated Q-vector is selected from the Q+1 vector and the exponent vector using this comparison result vector. However, the vectorization activation function process generates a large number of intermediate variables.

[0058] Step 2: Calculate the numerator of the mathematical expression for the linear approximate attention operator. Then, update K using the activation function. T The matrix is ​​moved to global memory, and matrix multiplication K is calculated in the matrix multiplication unit. T V(R d×N *R N×d ) to obtain matrix A(Rd ×d Meanwhile, K is updated with the activation function. T The matrix is ​​summed along N dimensions to obtain matrix B, and the resulting B is aggregated in local memory (R). d×8 This prepares for the subsequent matrix multiplication with Q. Then, the Q matrix, updated with the activation function, is moved to global memory, and the matrix multiplication QA(R) is calculated in the matrix multiplication unit. N×d *R d×d ) to obtain matrix P(R N×d ).

[0059] Mathematical expression for the linear approximate attention operator:

[0060]

[0061] Step 3: Calculate the denominator in the mathematical expression of the linear approximate attention operator and calculate the remainder. Calculate the matrix multiplication QB(R) in the matrix multiplication unit. N×d *R d×8 The normalized vector Z-block (R) is obtained. N×8 ).

[0062] Step 4: Move P and Z from global memory to the local memory of the vector computation unit in the form of blocks. Expand the Z block in the vector computation unit and normalize the P block row by row according to the vector divided by the Z block to obtain the result matrix.

[0063] Simultaneously, based on the hardware characteristics of Huawei Ascend NPU, optimization of vector computation is performed first. For matrix K... T The process of generating matrix B through row-by-row summation is explained in detail as follows: For each row of the matrix, a summation operation is first performed to obtain a scalar result b. However, due to the characteristics of vector computation units, if subsequent calculations require using these results, b must be expanded into a vector format conforming to datablock units, as illustrated in the diagram below. Figure 3 As shown. This process involves operations on Scalar computation units. Scalar computation units are responsible not only for scalar computations but also for program flow control. During this process, multiple loops frequently call Scalar computation units, which can potentially block other computation pipelines. Furthermore, synchronization waits between Scalar and Vector computation units introduce additional latency. Considering the computational characteristic of Vector computation units, which can perform multiple iterations, they are well-suited for computing large vectors. First, allocate a blank space of (m, 64), and then... TThe matrix can be viewed as a combination of multiple vectors of shape (m, 64). First, the first and second blocks are added together, and the resulting vector is recorded in the empty space of (m, 64). Subsequent remaining blocks are also added together with the vectors in the empty space. If K... T If the column dimensions of the matrix are not divisible by 64, then the remainder is used as the tail block and also accumulated once, as illustrated in the diagram below. Figure 5 As shown. At this point, using this method, K... T The row-by-row summation of the matrix is ​​transformed into row-by-row summation of the (m, 64) space. This transformation is primarily to adapt to the characteristics of the Vector computation unit. In each iteration, the Vector computation unit retrieves 8 data blocks from the UnifiedBuffer, corresponding to a size of 64 * sizeof(float). While iterating over the retrieved vector, the elements of that vector are summed. This computation method avoids the need for multiple iterations through loops, fully utilizing the Vector computation unit's ability to perform multiple iterations. Multiple iterations achieve more efficient computation, significantly improving the execution speed of the summation process. The pipeline diagram is shown above. Figure 6 As shown, the basic implementation calls 2n Vector calculation instructions and n Scalar calculation instructions. In the optimized implementation, only m / 64 Vector calculation instructions are called. The efficiency improvement of the optimized implementation is particularly significant when m is small and n is large.

[0064] Next, pipelined parallel optimization was performed. After analyzing the computation flow, it was found that the matrix multiplication K in step two... T V does not need to be in all blocks The activation functions can only be calculated after they have all been updated, which can be done in each block. After updating the activation function, the matrix multiplication for that block is calculated, allowing it to be computed in parallel with some vector operations in step one. Specifically, in After the activation function is updated, asynchronous matrix multiplication is performed first. Simultaneously with matrix multiplication, the vector computation unit synchronously calculates the activation function update of block Q, and updates K. T The i-th block of the matrix The increased communication time from the row-by-row summation process is far less than the total computation time reduction achieved by the vector computation unit through pipelined parallelism. The pipeline diagram before and after parallelism is shown below. Figure 7 As shown, the pipeline diagram of the parallel vector computation units is as follows: Figure 8 As shown. Through this decomposition strategy and asynchronous matrix multiplication, the execution of matrix multiplication and vector computation tasks can be coordinated more efficiently, maximizing the utilization of hardware resources.

[0065] Finally, cache optimization is performed. Considering the architectural characteristics of the Huawei Ascend NPU 910B4, its vector computation unit is equipped with a 192KB Unified Buffer cache. During the block partitioning process, the cache capacity needs to be fully utilized to reduce the number of loops, thereby minimizing header overhead and improving computational efficiency. When allocating memory space in the Unified Buffer, the logical location of the requested data structure needs to be specified: VECIN and VECOUT. When moving data from global memory, it can only be moved to the VECIN logical data structure. When moving data from global memory, it can only be moved from the VECOUT logical data structure. In the basic implementation, a large number of intermediate values ​​are generated during the update of the activation functions of the Q and K matrices, and Z is expanded during the row-by-row normalization of the P matrix and Z vector. Both the intermediate values ​​and the expanded Z need to be stored in the Unified Buffer. In fact, during the activation function update process, for the Q+1 / K+1 vector, it's not necessary to allocate new memory. As long as the Q / K vector and the 0 vector are compared first, directly overwriting the Q / K vector in the cache won't affect correctness. Similarly, for the 0 vector compared with the Q / K vector, a CompareScalar comparison with the 0 scalar achieves the same effect. For the expansion of Z, the vector computation unit can support interval-based vector computation. By configuring the block step size between adjacent iterations of the Z matrix to 0, the vector division between the P and Z matrices can be directly calculated. These two cache memory optimization measures reduce cache memory usage. By optimizing cache memory usage, the problem is solved by:

[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] Under certain input sizes, reducing the number of blocks required to partition matrices such as Q, K, P, and Z can decrease the number of computation and data transfer loops, thereby reducing operator runtime. This optimization strategy not only reduces data transfer frequency but also improves cache hit rate, effectively enhancing the resource utilization of computing units.

[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 products.

[0072] 1. Class definition and main function implementation

[0073] The class of this operator is named LinearAttention, which mainly includes two main functions: the initialization function Init() and the logical execution function Process(); and eight functions: Compute(), MatmulComputeP(), MatmulComputeZ(), CopyInP(), ComputeResult(), CopyOut(), ComputeReduceSum(), and EluCompute(). Due to computational time constraints, the calculation of the numerator and denominator needs to be performed in parallel with the activation function update. Compute(), MatmulComputeP(), MatmulComputeZ(), and ComputeResult() correspond one-to-one with the four steps of this invention, EluCompute() is the function used for activation function update, and ComputeReduceSum corresponds to K. T The function performs row-by-row summation. CopyInP() is used to move intermediate values ​​P, and CopyOut() is used to move the result.

[0074] The remaining parts mainly include the definition of matrix multiplication classes, named mm1, mm2, and mm3 respectively; the setting of matrix multiplication execution information structures, named tiling1, tiling2, and tiling3; the definition of external input intermediate quantity outputs, that is, defining the inputs 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 and P matrices and the Q matrix updated by the activation function. The output is defined as resultGlobal. The declaration of the on-board TQue variables first sets Tposition to VECOUT, including zQueue, used to store the calculated result Z matrix of the denominator; and qOutQueue and kOutQueue, used to store the block Q and K after the activation function update. T The system includes a matrix B, consisting of a `bAllQueue` to store the complete B matrix and a `resultQueue` to store the calculated result matrix. Next, `Tposition` is set to `VECIN`, which includes `qQueue` and `kQueue` to store the loaded block inputs Q and K. T Includes pQueue, used to store intermediate values ​​P of the moved-in blocks; includes bQueue, used to store blocks K. TThe result is expanded by row-by-row summation. The TBuf variable on the board is declared; TBuf variables are mainly used in calculations and not for data transfer. expBuf and compareBuf are set for recording the bit positions of the exponentiation and comparison operations in the activation function, respectively; sumBuf and sumlBuf are set for temporary space and result recording in the row-by-row summation, respectively. Inputs N, d, and t are the number of blocks, and tileLength is the size of the block matrix.

[0075] The `Init()` function of this operator performs three main operations. The first operation: loads the partitioning information from the three matrix multiplications. The second operation: locates the external input / output memory positions using the `SetGlobalBuffer()` method. The third operation: initializes all allocated `Tque` and `Tbuf` objects, requiring explicit memory allocation.

[0076] The `Process()` function of this operator is mainly responsible for executing the computational logic function. First, it divides the inputs Q and K into blocks and performs multiple loops according to these blocks. During the loops, it uses the `DataCopy` interface of Ascend C to move the blocks of Q and K from global memory into the Vector computation unit. Then, it calls the `Compute()` function, which corresponds to updating the activation function and performing parallel computations on some parts of the numerator and denominator. After the loops finish, it calls the `MatmulComputeP()` function to calculate the numerator and the `MatmulComputeZ()` function to calculate the denominator. Finally, because the P matrix is ​​large, it divides the intermediate quantities P and Z into blocks and performs multiple loops according to these blocks. During the loops, it moves the blocks of the P matrix from global memory into the Vector computation unit. Then, it calls the `ComputeResult()` function to normalize the blocks of the computation result and moves the blocks of the result from the Vector computation unit back to global memory. The specific implementation is shown in pseudocode 1.

[0077] Pseudocode 1: Implementation of the Process() function

[0078]

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

[0080] (1) Implementation of the Compute() function

[0081] In the Compute() function, the K is first evaluated 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 asynchronously by setting the parameter in <> to false. While waiting for the result, the EluCompute() function is called to update the activation function of Q. The updated Q is then moved in blocks to an off-board intermediate quantity for storage by calling the DataCopy interface in Ascend C. Finally, the ComputeReduceSum() function is called to update K. T The summation is performed in blocks. The specific implementation of the Compute() function is shown in pseudocode 2.

[0082] Pseudocode 2: Implementation of the Compute() function

[0083]

[0084]

[0085] Within the function, `EluCompute()` is used to calculate the activation function. For the function's parameter `X` (LocalTensor), the `CompareScalar` interface in Ascend C compares each element of `X` with a scalar of 0. If the result is less than 0, the corresponding bit is set to 1; otherwise, it's 0. Compared to generating a `LocalTensor` of all zeros and then performing an element-wise comparison between two `LocalTensors`, this saves `d*N / t*sizeof(float)` of cache space. Next, the `Exp` interface in Ascend C obtains the intermediate result `expTensor` of `X`, which is the element-wise exponentiation of `X`. Then, the `Adds` interface in Ascend C obtains `x+1`, and the result is stored in `x`. Since the original value of `x` does not affect the calculation result, this again saves `d*N / t*sizeof(float)` of cache space. Finally, the `Select` interface in Ascend C selects elements based on the bit values ​​of the previously obtained comparison result `compareTensor`, achieving the effect of taking the exponent when `x` is less than 0 and taking `x+1` when `x` is greater than or equal to 0. The final result is stored in `y`. The specific implementation of the EluCompute() function is shown in pseudocode 3.

[0086] Pseudocode 3: Implementation of the EluCompute() function

[0087]

[0088] Within the function, `ComputeReduceSum()` performs a row-by-row summation. When the input N is less than 64, the `WholeReduceSum` interface in Ascend C can sum all data within each repeat. Since N is less than 64, the summation portion for each iteration can be limited to less than 8 data blocks. Therefore, the `WholeReduceSum` interface in Ascend C is used to perform d / t repeat calculations, with each repeat summing N data points. When the input N is greater than or equal to 64, the results are first accumulated into a (d / t, 64) empty space by calling the `Add` interface in Ascend C N / 64 times. Then, the `WholeReduceSum` interface in Ascend C is used to perform d / t repeat calculations, with each repeat summing 64 data points. The specific implementation of the `ComputeReduceSum()` function is shown in pseudocode 4.

[0089] Pseudocode 4: Implementation of the ComputeReduceSum() function

[0090]

[0091] (2) Implementation of the MatmulComputeP() function

[0092] The MatmulComputeP() function mainly calls Ascend C's IterateAll to perform matrix multiplication of the Q matrix and the A matrix to obtain the P matrix.

[0093] (3) Implementation of the MatmulComputeZ() function

[0094] The MatmulComputeZ() function mainly calls Ascend C's IterateAll to perform matrix multiplication of the Q matrix and the B matrix to obtain the Z matrix.

[0095] (4) Implementation of the ComputeResult() function

[0096] The `ComputeResult()` function primarily utilizes the hardware feature of the Ascend NPU's vector computing unit, which supports interval-based vector computation. It calls the Div interface in Ascend C to perform interval-based vector division. The address steps between different data blocks in a single iteration of the `result` vector, `p` vector, and `z` vector are configured as `d / 8`, `d / 8`, and `1`, respectively; simultaneously, the address steps between adjacent iterations of the same data block are configured as `1`, `1`, and `0`, respectively, to calculate the normalized result. The `result` and `p` vectors move the same coordinates (`d*8`) between each repeat. The `z` vector, due to its different shape, moves 64 coordinates each time. Depending on the computation block, the `p` vector only moves its block during processing, while the `z` vector remains complete and requires correction of its coordinates to block size `i*N / t*8`. A total of N / t / 8 repeat vector division operations are performed, with d / 8 vector division operations performed in each repeat. The specific implementation of the `ComputeResult()` function is shown in pseudocode 5.

[0097] Pseudocode 5: Implementation of the ComputeResult() function

[0098]

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

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

[0101] Pseudocode 6: Implementation of the operator kernel function

[0102]

[0103]

[0104] In the code above, line 1 requests the TPipe type, a memory management system specific to Ascend C, a common step for all Ascend C operators. Lines 2-7 call CopyTiling to load the matrix block information. Line 8 declares the LinearAttention class variable linearAttention. Line 9 executes the Init function. Line 10 performs the registration of matrix multiplication computation information. Line 11 executes the Process function. At this point, the kernel function of the Ascend_CoppersmithWinograd operator is complete. To use the advanced matrix multiplication API, host-side tiling code also needs to be written, which includes matrix multiplication operation registration information and computation scale settings.

[0105] In this operator, a total of three matrix multiplications are performed, namely K T Matrix partitioning and matrix multiplication of matrix V, matrix multiplication of matrix Q and matrix A, and matrix multiplication of matrix Q and matrix B. This example uses matrix multiplication of matrix Q and matrix A. 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 ND, the DataType is float, and the transpose enable parameter is false; the TPosition of the right matrix is ​​GM, the CubeFormat is ND, the DataType is float, and the transpose enable parameter is false; the TPosition of the result matrix is ​​GM, the CubeFormat is ND, and the DataType is float; since Bias is not needed, the isBias parameter is set to false.

[0106] 4. Operator correctness testing and performance testing

[0107] (1) Operator correctness test

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

[0109] Table 1 Common values ​​for single-head dimension in 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 of shape N×d and data type float as input data. The test then examines Q and K... T The matrices are processed using the `numpy.where` method. First, the activation function is updated; if x is less than 0, then it equals e. x If x is greater than or equal to 0, then it equals x+1. This yields the updated Q and K. T After the matrix is ​​generated, the updated K is calculated using numpy.matmul. T The matrix multiplication of matrix Q and matrix V results in matrix A. Then, using numpy.matmul, the matrix multiplication of the updated matrix Q and matrix A results in matrix P. Finally, numpy.sum is used to calculate the updated matrix K. TThe matrix is ​​summed along N dimensions to obtain matrix B. Then, the matrix multiplication result Z of the updated Q matrix and matrix B is calculated using numpy.matmul. Finally, matrix P and matrix Z are divided row-by-row using numpy.divide to obtain the golden matrix. Using golden as the truth value, the truth value calculated by Python is compared element-by-element with the result calculated by the developed operator. The comparison method is as follows:

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

[0113] 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 (4) is true, the element is calculated correctly; otherwise, it is incorrect.

[0114]

[0115] In formula (5), error num Calculate the error count for each element, all num `etol` represents the total number of elements in the result, and `etol` represents the error rate in element calculation. In the validation metrics, `etol ≤ 10`. -4 Then the operator is considered to be calculated accurately.

[0116] Table 2 Accuracy of Basic and Optimized Schemes

[0117]

[0118]

[0119] (2) Operator performance test

[0120] Under the premise of ensuring correctness, the running time of the linear approximate attention operator before and after the optimization of reducing cache usage and dividing into blocks was compared. N is 256, the x-axis is the size of d in the computation dimension, starting from the smaller 16 and doubling each time, from 16 to 128. Each group was calculated 10 times, and the y-axis is the average running time of the algorithm.

[0121] Depend on Figure 9 It can be seen that when d is relatively small, that is, when the input size is small, the optimization effect of block division is not obvious or even non-existent. This is because when the input size is small, even if the memory usage in the cache is not optimized, block division can be used directly for calculation. When d gradually increases, the optimization effect reaches up to 23.30%.

[0122] Next, under the same input size variation, the average running time of the basic scheme and the final optimized scheme were compared.

[0123] like Figure 10 As shown, the optimized solution's runtime is significantly shorter than the basic solution's runtime for any input size tested. Furthermore, the optimization efficiency gradually increases with the size of the computational input. The lowest optimization efficiency is 20.7%, while the highest reaches 46%.

[0124] After increasing N in the test, due to the limited cache size, the block partitioning of the basic and optimized solutions remained essentially the same, and the block partitioning optimization had no significant effect. Figure 11 The results demonstrate the effectiveness of the pipelining and vector computation optimization in the optimization scheme. In the test with N=512, especially with a large input size of d=128, the optimization efficiency dropped to 23.1% because the number of blocks could not be reduced. When d=32, the optimization scheme achieved the highest optimization efficiency of 33.9% in this set of experiments.

[0125] The efficiency of the CPU implementation and optimization schemes was also compared, such as... Figure 12 As shown, the computation time of the CPU solution is more than 15 times that of the optimized solution, demonstrating the significant acceleration of computation achieved by the hardware-friendly implementation of this operator presented in this paper.

Claims

1. A linear approximate attention operator acceleration method based on NPU, characterized in that the method includes the following steps: Step 1, The matrix is ​​divided into blocks of size d according to dimension d. The matrix is ​​moved from global memory into the Unified Buffer of the Vector computation unit, and the Q matrix is ​​divided into blocks of size N according to the N-dimensional dimensions. The matrix is ​​moved from global memory into the Unified Buffer of the vector computation unit, and then divided into blocks in the vector computation unit. =elu( +1 activation function update; Step 2: Perform pipelined parallel optimization: First, perform asynchronous matrix multiplication operations. During matrix multiplication, the vector computation unit simultaneously calculates the activation function update of the block Q, and also performs the operation on... The i-th block of the matrix The row-by-row summation process along the N-dimensional plane yields multiple matrices B, where, The row-by-row summation calculation is optimized for vector computation based on the hardware characteristics of the NPU: Step 2-1: Allocate blank space (m, 64); Step 2-2: Add the vector of the first block (m,64) and the vector of the second block, and record the resulting vector in the blank space of (m,64). The remaining blocks are also added to the vector in the blank space by vector addition. Steps 2-3: If If the column dimensions of the matrix are not divisible by 64, then the remainder is used as the tail block and accumulated again. The row-by-row summation of the matrix is ​​transformed into the row-by-row summation of the (m,64) space, which is convenient for the Vector computation unit; Step 3: Update the activation function. The matrix is ​​moved to global memory, and matrix multiplication QA is performed in the matrix multiplication unit to obtain matrix P; Step 4: Calculate matrix multiplication QB in the matrix multiplication calculation unit to obtain the normalized vector Z block; Step 5: Move P and Z from global memory to the local memory of the vector computation unit in the form of blocks. Expand the Z block in the vector computation unit and normalize the P block row by row according to the vector divided by the Z block to obtain the result matrix.

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

3. The NPU-based linear approximate attention operator acceleration method according to claim 1, characterized in that, in step 5, the cache memory usage for the expansion of the Z-block is optimized: Step 5-1: Utilize the vector computation unit to support spaced vector computation. 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 calculated directly. Step 5-2: By optimizing cache memory usage, the problem can be resolved by: Convert to 。

Images