A method, device, equipment, medium and product for accelerating attention mechanism
By prefetching the key matrix and value matrix in the forward propagation and backward propagation stages of the attention mechanism and continuing to perform calculations at the calculation nodes, the traditional attention mechanism solves the problem of large calculations and high memory consumption when processing long sequences, and achieves more efficient computing performance and accuracy.
Patent Information
- Application Number
- CN202411994533.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-12-31
- Publication Date
- 2025-04-04
AI Technical Summary
When processing long sequences or large-scale data, the traditional attention mechanism has a large amount of computation and high memory consumption, resulting in inefficiency.
By prefetching the key matrix and value matrix in the forward propagation and backward propagation phases and continuing the calculations at the corresponding computing nodes, memory access latency and bandwidth requirements are reduced, and the computing process is optimized to improve computing efficiency.
The processing efficiency, computing performance and accuracy of the attention mechanism is significantly improved, especially when processing long sequences, which can maintain efficient computing speed and low latency.
Smart Images

Figure CN119761448A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of computer technology, and in particular to an attention mechanism acceleration method, device, equipment, medium and product. Background Art
[0002] In fields such as natural language processing and computer vision, attention mechanisms, especially self-attention mechanisms, have become key technologies. However, traditional attention mechanisms are computationally intensive and memory intensive, especially when processing long sequences or large-scale data, where computational and memory requirements increase rapidly. The main operations in the attention mechanism include matrix multiplication and normalization, and the computational complexity of these operations is closely related to the length of the input sequence and the dimension of the model. When processing long sequences, the computational complexity and memory usage of these operations increase significantly, resulting in low efficiency. Summary of the invention
[0003] The present invention provides an attention mechanism acceleration method, device, equipment, medium and product, which improves the processing efficiency, computing performance and accuracy of the attention mechanism in fields such as natural field processing and computer vision.
[0004] In a first aspect, an embodiment of the present disclosure provides an attention mechanism acceleration method, which is applied to any processing core of an electronic device, and the method includes a forward propagation stage and a backward propagation stage;
[0005] In the forward propagation stage, a query matrix is obtained, and forward calculation is performed according to the query matrix. In the forward calculation process, a key matrix and a value matrix are pre-fetched, and forward calculation is continued at the corresponding calculation node according to the key matrix and the value matrix to obtain the current forward output matrix;
[0006] In the backward propagation stage, the key matrix, value matrix and query matrix are obtained, and backward calculation is performed based on the key matrix, value matrix and query matrix. During the backward calculation process, the intermediate matrix is pre-fetched, and backward calculation is continued at the corresponding calculation node based on the intermediate matrix to obtain the current gradient matrix.
[0007] In a second aspect, the present disclosure provides an attention mechanism acceleration device, including:
[0008] A forward propagation module is used to obtain a query matrix in the forward propagation stage, perform forward calculation according to the query matrix, pre-fetch a key matrix and a value matrix during the forward calculation process, and continue to perform forward calculation at corresponding calculation nodes according to the key matrix and the value matrix to obtain a current forward output matrix;
[0009] The backward propagation module is used to obtain the key matrix, value matrix and query matrix in the backward propagation stage, perform backward calculation according to the key matrix, value matrix and query matrix, pre-fetch the intermediate matrix during the backward calculation process, and continue to perform backward calculation at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix.
[0010] In a third aspect, an embodiment of the present disclosure provides an electronic device, including:
[0011] at least one processor; and
[0012] a memory communicatively connected to at least one processor; wherein,
[0013] The memory stores a computer program that can be executed by at least one processor, and the computer program is executed by at least one processor so that the at least one processor can execute an attention mechanism acceleration method provided in the above-mentioned first aspect embodiment.
[0014] In a fourth aspect, an embodiment of the present disclosure provides a computer-readable storage medium, which stores computer instructions, and the computer instructions are used to enable a processor to implement an attention mechanism acceleration method provided in the embodiment of the first aspect above when executing.
[0015] In a fifth aspect, an embodiment of the present disclosure provides a computer program product, which includes a computer program, and when the computer program is executed by a processor, it implements an attention mechanism acceleration method provided in the embodiment of the first aspect above.
[0016] An attention mechanism acceleration method, device, equipment, medium and product of an embodiment of the present invention, in the forward propagation stage, obtains a query matrix, performs forward calculation according to the query matrix, pre-fetches a key matrix and a value matrix during the forward calculation process, and continues to perform forward calculation at the corresponding calculation node according to the key matrix and the value matrix to obtain the current forward output matrix; in the backward propagation stage, obtains a key matrix, a value matrix and a query matrix, performs backward calculation according to the key matrix, the value matrix and the query matrix, and pre-fetches an intermediate matrix during the backward calculation process, and continues to perform backward calculation at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix. The above technical scheme improves the processing efficiency, computing performance and accuracy of the attention mechanism in fields such as natural field processing and computer vision.
[0017] It should be understood that the contents described in this section are not intended to identify the key or important features of the embodiments of the present invention, nor are they intended to limit the scope of the present invention. Other features of the present invention will become easily understood through the following description. BRIEF DESCRIPTION OF THE DRAWINGS
[0018] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.
[0019] Figure 1 is a flow chart of an attention mechanism acceleration method provided in Embodiment 1 of the present invention;
[0020] Figure 2a This is an example diagram of load balancing of forward propagation calculation provided by Embodiment 1 of the present invention;
[0021] Figure 2b This is an example diagram of load balancing of a back-propagation calculation provided in the first embodiment of the present invention;
[0022] Figure 3 This is a flowchart of a forward calculation of an attention mechanism acceleration method provided in Example 1 of the present invention;
[0023] Figure 4 This is a flowchart of backward calculation of an attention mechanism acceleration method provided in Example 1 of the present invention;
[0024] Figure 5 is a schematic diagram of the structure of an attention mechanism acceleration device provided in Embodiment 2 of the present invention;
[0025] Figure 6 It is a structural schematic diagram of an electronic device provided in Embodiment 3 of the present invention. DETAILED DESCRIPTION
[0026] In order to enable those skilled in the art to better understand the scheme of the present invention, the technical scheme in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work should fall within the scope of protection of the present invention.
[0027] It should be noted that the terms "first", "second" and "target" in the specification and claims of the present invention and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the data used in this way can be interchanged where appropriate, so that the embodiments of the present invention described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions, for example, a process, method, system, product or device that includes a series of steps or units is not necessarily limited to those steps or units that are clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.
[0028] In fields such as natural language processing and computer vision, attention mechanisms (especially self-attention mechanisms) have become key application technologies. However, traditional attention mechanisms are computationally intensive and memory intensive, especially when processing long sequences or large-scale data, where computational and memory requirements increase rapidly. The main operations in the attention mechanism include matrix multiplication and normalization, and the computational complexity of these operations is closely related to the length of the input sequence and the dimension of the model. When processing long sequences, the computational complexity and memory usage of these operations increase significantly, resulting in low efficiency.
[0029] The attention acceleration algorithm (FlashAttention) reduces memory bandwidth consumption through efficient memory management. It uses block-wise processing and local storage optimization techniques to concentrate memory usage in local areas, thereby reducing memory access latency and bandwidth requirements. FlashAttention optimizes the computational process, especially in matrix multiplication and normalization operations. By improving the computational method and reducing redundant computations, it can significantly accelerate attention computations while maintaining high accuracy. The technology is optimized for modern hardware such as GPUs and TPUs, using their parallel processing capabilities to improve performance. The characteristics of the hardware architecture, such as memory access patterns and utilization of computing units, are taken into account during design to maximize computational efficiency, reduce computational complexity and memory requirements, and improve the computational method of the attention mechanism. It can process longer sequences and maintain efficient computational speed and low latency.
[0030] FlashAttention, a widely used acceleration method for the attention mechanism, is mainly based on data locality to divide the computing tasks into blocks to reduce the amount of data interaction from the main memory to the cache and increase the occupancy rate of the GPU computing power. Compared with the original attention calculation method, FlashAttention alleviates the communication pressure through methods such as block division and operator fusion, and also provides more optimization opportunities for higher computing parallelism. In addition, it can also save a lot of video memory space occupied by the intermediate result matrix. In general, FlashAttention will bring very high benefits to the scene of bandwidth bottleneck. For the scene of computing bottleneck, FlashAttention will also reduce the overall computing delay and the video memory occupied at runtime, which is of great significance when processing long sequences or large-scale models.
[0031] FlashAttention uses a very mature matrix multiplication library and nonlinear computing library, and takes advantage of its efficient block processing and efficient memory access mode to accelerate attention calculations, fully utilizing the parallel processing capabilities of hardware to improve model training and reasoning capabilities. FlashAttention not only supports the standard self-attention mechanism, but also supports various variants of attention mechanisms, such as cross attention, multi-head attention, etc.
[0032] Traditional self-attention mechanisms (such as the attention mechanism in Transform 2) involve matrix operations on the entire input sequence, with a computational complexity of (O(n^2*d)), where (n) is the sequence length and (d) is the model dimension. This computational approach becomes impractical for long sequences. To address the bottlenecks of traditional attention mechanisms, researchers have proposed various optimization techniques, such as sparse attention, low-rank approximation, and memory optimization techniques. These methods aim to reduce computational complexity or reduce memory usage, but usually face a trade-off between performance and implementation complexity.
[0033] FlashAttention combines the above technical backgrounds and provides an efficient solution to the high memory consumption and computing bottleneck of attention mechanism calculation. It significantly improves the computational efficiency of the attention mechanism through hardware optimization, efficient memory management, and improved computing process, especially when processing long sequences.
[0034] However, in the context of the continuous increase in the number of large model parameters and the continuous introduction of new hardware architectures, FlashAttention, which is optimized for a specific graphics processing unit (GPU), cannot support hardware of other architectures well. Since FlashAttention itself is completely implemented based on GPU, its hardware coupling is relatively high, and the complexity of supporting variant attention mechanisms is also relatively high, and there may not be the same performance improvement on different hardware platforms. In addition, compared with the original attention calculation method, FlashAttention does not improve the time complexity. Instead, it introduces new operators, such as logarithmic operations, which further increases the complexity of implementation. For hardware architectures with computing bottlenecks, it may lead to negative optimization; FlashAttention divides the computing tasks and requires the subtasks obtained by the division to be executed in the cache, which leads to greater register pressure and cache space pressure.
[0035] Therefore, an embodiment of the present invention provides an attention mechanism acceleration method to solve the above technical problems.
[0036] Embodiment 1
[0037] Figure 1 This is a flowchart of an attention mechanism acceleration method provided in Example 1 of the present invention. This embodiment is applicable to the situation of optimizing the attention mechanism involved in the field of natural language processing and computer vision. The method can be executed by an attention mechanism acceleration device, and the attention mechanism acceleration device can be implemented in the form of hardware and / or software.
[0038] The method is applied to any processing core of an electronic device. The electronic device may include multiple processors, each of which may be considered as a processing core, or the electronic device may include at least one processor, and the processor may include multiple processing cores.
[0039] Figure 2a This is an example diagram of load balancing of forward propagation calculation provided by Embodiment 1 of the present invention; Figure 2b This is an example diagram of load balancing of a back-propagation calculation provided by the first embodiment of the present invention. Specifically, the entire calculation load is divided into blocks, and the calculation load is divided according to the length direction of the sequence (sequence) dimension to avoid additional atomic additions (atomicadd), that is, Figure 2a As shown, each subtask in the forward computation contains a row for the computational load, such as Figure 2b As shown, each subtask of the backward calculation contains a column of computational load, and each processing core processes one subtask at a time. Each processing core can implement the computational execution of the task load in parallel.
[0040] When adding causal masks, the load of each subtask is not the same, and it conforms to the law of arithmetic progression. You can consider using task queues and other methods to achieve multi-core load balancing, but dynamic load balancing cannot pre-fetch data when switching tasks, and the registration and switching of tasks themselves also have additional overhead, and the computational load of causal masks has obvious rules, so here we consider using static task allocation to achieve load balancing. The overall idea is similar to distributed computing. For this load balancing problem that follows the law of arithmetic progression and is arranged in a triangular matrix, a greedy algorithm can be used to obtain the optimal solution. First, sort the task loads from high to low. For the forward propagation calculation process, this order is the reverse order of the row order, and for the backward process, it is the positive order of the column order; each time, assign the task to the core with the lowest total load. For example, there are four processing cores responsible for the calculation of a head (head of multi-head attention). The distribution is as follows, in order to achieve load balancing of the causal mask attention mechanism operation suitable for multi-core architecture during the forward propagation calculation process:
[0041]
[0042] The backward propagation calculation is similar, and can be allocated in positive order, which will not be described in detail in this embodiment.
[0043] It is understandable that since atomic addition (atomicadd) is required when calculating the gradient of the query matrix Q in the reverse direction, this performance will be affected by the number of threads that accumulate at the same time. The more threads that use atomic add, the more instructions that need to be waited. Without adding a causal mask, all tasks in the reverse direction will start from the first row, and the gradient of each row of Q calculated will need to use atomic add, resulting in a very long waiting time. It takes n-1 atomicadd dQ matrix times, and n is the number of cores allocated to each head. If the sequence is not too long, it will lead to a huge waste of computing power. Therefore, the calculation order without causal mask is adjusted to be consistent with the calculation order with causal mask, that is, thread 0 starts from the first row of the first column, thread 1 starts from the second row of the second column, and so on. This will avoid the occupation of atomic add resources by threads to the greatest extent.
[0044] The present invention is adapted to a multi-core hardware architecture. For the multi-core architecture, the computing load of each processing core may be different, and the final performance is determined by the core with the longest running time. The present invention proposes a load balancing method based on a greedy algorithm to reduce cavitation in multi-core operation and achieve load balancing of causal mask attention mechanism operations suitable for multi-core architecture.
[0045] like Figure 1 As shown, the method includes:
[0046] S101. In the forward propagation stage, a query matrix is obtained, and forward calculation is performed according to the query matrix. In the forward calculation process, a key matrix and a value matrix are pre-fetched, and forward calculation is continued at the corresponding computing node according to the key matrix and the value matrix to obtain the current forward output matrix.
[0047] In this embodiment, the query matrix is the Q matrix, the key matrix is the K matrix, the value matrix is the V matrix, and the current forward output matrix can be understood as the output matrix (Output) formed in the current forward calculation iteration, that is, the O matrix.
[0048] Specifically, in the forward propagation stage, the forward propagation calculation of the current iteration first obtains the query matrix, and initializes the forward output matrix O, the first intermediate matrix L and the second intermediate matrix M. The first intermediate matrix and the second intermediate matrix are both intermediate parameters in the forward calculation process and are an operator. The product of the transpose of the Q matrix and the K matrix is calculated to obtain the third intermediate matrix S. The third intermediate matrix is also an intermediate parameter in the forward calculation process and is an operator. The maximum value of S and M is taken as the N matrix, and the exponent of the N matrix in two adjacent forward calculations is calculated, that is, exp_sub. After the determination of the N matrix is completed, it also includes pre-taking the next K matrix, and returning to recalculate the S matrix based on the new K matrix, and iterative execution. After determining exp_sub, the fifth intermediate matrix P is calculated based on the S matrix and the M matrix. The fifth intermediate matrix is also an intermediate parameter in the forward calculation process and is an operator. After the calculation of the P matrix is completed, the row accumulation of the P matrix is immediately calculated, and then the sum of the row accumulation of the P matrix and exp_sub is calculated to obtain the first intermediate matrix L. After the first intermediate matrix is determined, the P matrix is converted from the fp32 data format to the fp16 data format. Based on the product of the P matrix in fp16 format and the V matrix and the output matrix of the previous iteration, the current forward output matrix of the current iteration round is calculated.
[0049] In one embodiment, Figure 3 is a flowchart of a forward calculation of an attention mechanism acceleration method provided in Example 1 of the present invention, such as Figure 3 As shown, the query matrix is obtained, and forward calculation is performed according to the query matrix. During the forward calculation process, the key matrix and the value matrix are pre-fetched, and forward calculation is continued at the corresponding calculation node according to the key matrix and the value matrix to obtain the current forward output matrix, including:
[0050] S1011, obtain the query matrix Q, initialize the forward output matrix O, the first intermediate matrix L and the second intermediate matrix M.
[0051] First, execute "load Q" to obtain the Q matrix, and then execute "init O i (0) , L i (0) 、M i (0) ", initialize the forward output matrix O, the first intermediate matrix L and the second intermediate matrix M.
[0052] S1012, obtaining the current key matrix K, determining the third intermediate matrix S according to the Q matrix and the current K matrix, determining the current K matrix as the historical K matrix, and re-obtaining the current K matrix.
[0053] In this embodiment, at the beginning of the iteration, the K matrix at the current moment is taken from the outer loop memory, and "S = QK T ", calculate the product of the transpose of the Q matrix and the K matrix to obtain the S matrix. After the S matrix is determined, the current K matrix is determined as the historical K matrix to facilitate the re-acquisition of a new current K matrix. After the S matrix is calculated, "load K" is executed to obtain a new current K matrix. Alternatively, after the S matrix is calculated and the maximum value between the S matrix and the M matrix is determined based on the S matrix calculation, "load K" is executed.
[0054] S1013. Determine a fourth intermediate matrix N according to the S matrix and the M matrix.
[0055] In this embodiment, the fourth intermediate matrix is also an intermediate parameter in the forward calculation process and is an operator.
[0056] After the S matrix is determined in each iteration, "N=max(M, rowmax(S))" is executed to compare the cumulative sum of the determined S matrix with the M matrix, and the maximum value of S and M is taken as the N matrix. After the N matrix is determined, "loadK" can be executed to pre-fetch the new current K matrix.
[0057] Since the calculation process of FlashAttention is relatively complicated, we consider using the asynchronous transmission component Direct Memory Access (DMA) to pre-fetch data in advance when calculating the current data. Since the first operation of the calculation is to calculate S = QK T , so the first K can be read into the cache in the outer loop, and S = QK in the current loop T After the calculation is completed, the next K matrix is read in advance to asynchronously implement calculation communication masking.
[0058] After the N matrix is determined, execute "exp_sub = exp(N i j-1 -Ni j )”, based on the N matrix N in the previous iteration i j-1 and the N matrix N of the current round i j The difference between them is subjected to exponential function operation to obtain the exponential vector parameter exp_sub.
[0059] S1014. Obtain the current fifth intermediate matrix P according to the N matrix calculation, and sum each P matrix.
[0060] Execute "P=exp(SN)", perform exponential function operation on the difference between the S matrix and the N matrix to obtain the P matrix. After determining the P matrix, immediately execute the row accumulation "rowsum(P)" of the P matrix, that is, "P=exp(SN)&rowsum(P)".
[0061] In the forward and backward calculations of FlashAttention, the exponential operation is a relatively complex operator. However, in the activation function softmax, the input range of the exponential function is (-∞,0], so we can consider fitting the negative semi-axis of the exponential function with Taylor expansion. Since the dynamic range of the exponential function in the negative semi-axis is very small, only the second-order Taylor expansion is required to achieve the data accuracy of FP32, and the coefficients of the Taylor expansion are stored in the on-chip cache, which greatly increases the data processing speed. Furthermore, the exponential function calculation process is expanded in the loop using instruction-level parallel technology, and ROWSUM (summing the matrix by row dimension) is loop-fused, that is, the data output by the exponential is accumulated in the row direction in the loop, which further improves the instruction throughput rate and computing unit occupancy rate, and completes the exponential function implementation with faster performance and more accurate numerical values.
[0062] S1015. Pre-fetch value matrix V.
[0063] Before calculating PV, "load V" is executed to pre-fetch the V matrix from the outer loop memory so that it can be directly called when the P*V calculation is applied subsequently, realizing the calculation communication masking asynchronously.
[0064] S1016. Determine the current forward output matrix according to the P matrix and the V matrix.
[0065] The first intermediate matrix L is calculated based on the sum of the row accumulation rowsun(P) of the P matrix and the exponential vector parameter exp_sub. The first intermediate matrix L is also an intermediate parameter in the forward calculation process and is an operator. Execute "L = exp_sub + rowsum(P)". Convert the data format of the P matrix, and calculate the current forward output matrix based on the V matrix, the P matrix after format conversion, exp_sub, and the historical forward output matrix in the previous iteration (if the current iteration is the first round, the historical forward output matrix is the initialization value.
[0066] In one embodiment, determining the current forward output matrix according to the sum of each P matrix and the N matrix includes:
[0067] The data type of the P matrix is converted, and the current forward output matrix is determined based on the product of the converted P matrix and the V matrix and the reference forward output matrix.
[0068] The reference forward output matrix is the product of the exponential vector parameter of the N matrix and the historical forward output matrix, that is, "exp_sub*O i j-1 ”.
[0069] The data type of the initial output of the P matrix is fp32. The fp32 data type is required when updating the L matrix. However, when calculating PV, the P matrix is required to be of fp16 data type. Therefore, the cast conversion operation of the P matrix from fp32 is fused with the calculation of the matrix multiplication by vector. Matrix-vector multiplication is an operator with low computational complexity, and its bottleneck lies in the communication bandwidth. The component occupancy rate of the computing core is low, and cast is also a memory access bottleneck operator. Therefore, the operator parallel method of mixed precision attention calculation is adopted to fuse the two operators to improve the occupancy rate and instruction throughput.
[0070] Specifically, after the L matrix is determined, execute "Cast_fp32_fp16(P)" to convert the P matrix from fp32 type to fp16 type, and then execute "O i j =PV+exp_sub*O i j-1 ", calculate the product of the format-converted P matrix and the V matrix, calculate the product of exp_sub and the historical forward output matrix of the previous iteration, calculate the sum of the two products, and get the current forward output matrix. Execute "output O" to output the current forward output matrix O.
[0071] O i j The calculation steps refer to the forward output matrix "exp_sub*O ij-1 The determination of " and the conversion of the P matrix data type "Cast_fp32_fp16(P)" can be executed in parallel. After obtaining the outputs of the two, the O is calculated based on the output results of the two. i j The final result is the operator parallelism of mixed-precision attention calculation.
[0072] In the above technical content, the subscript i of each matrix represents the task load identifier, the superscript j represents the current iteration, and j-1 represents the previous iteration.
[0073] S102. In the backward propagation stage, the key matrix, the value matrix and the query matrix are obtained, and backward calculation is performed according to the key matrix, the value matrix and the query matrix. In the backward calculation process, the intermediate matrix is pre-fetched, and the backward calculation is continued at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix.
[0074] In one embodiment, Figure 4 is a flowchart of backward calculation of an attention mechanism acceleration method provided in Example 1 of the present invention, such as Figure 4 As shown, the key matrix, value matrix and query matrix are obtained, and backward calculation is performed according to the key matrix, value matrix and query matrix. During the backward calculation process, the intermediate matrix is pre-fetched, and the backward calculation is continued at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix, including:
[0075] S1021. Obtain the current key matrix K, the current value matrix V and the current query matrix Q, and pre-fetch the first intermediate matrix L, the sixth intermediate matrix D and the seventh intermediate matrix dO.
[0076] In this embodiment, the sixth intermediate matrix and the seventh intermediate matrix can be understood as intermediate parameters in backward calculation, which is a kind of operator.
[0077] In the reverse calculation, the data D, L, and dO are all read into the cache in advance before the calculation. The calculation of the matrix P and the reading of the Q matrix are slightly more complicated. The first block of the Q matrix is read into the cache outside the loop, and then the Q matrix block for the next loop is pre-fetched in the inner loop. The calculation and communication in the inner loop are completely masked by asynchronous implementation.
[0078] Specifically, execute "load K, V, Q0" and "init K, V, Q0" to obtain the initial current K matrix, current Q matrix and current V matrix and perform initialization. Since the Q matrix needs to be pre-fetched and updated continuously in the future, the initially obtained Q matrix is represented as Q0. In addition to obtaining the K matrix, Q matrix and V matrix, it is also necessary to execute "load D, L" and "load dO" to pre-acquire the L matrix, D matrix and dO matrix.
[0079] S1022. Determine a third intermediate matrix S according to the current Q matrix and the current K matrix, determine the current Q matrix as the historical Q matrix, and reacquire the current Q matrix.
[0080] At the beginning of the iteration, execute "S = QK T ", calculate the product of the transpose of the Q matrix and the K matrix to obtain the S matrix. After the S matrix is determined, the current Q matrix is determined as the historical Q matrix to facilitate the re-obtaining of the new current Q matrix. After the calculation of the S matrix is completed, execute "load Q" to obtain the new current Q matrix.
[0081] S1023. Perform calculations based on the S matrix and the L matrix to obtain a fifth intermediate matrix P, and determine a first current gradient matrix dV corresponding to the current V matrix and a second current gradient matrix dK corresponding to the current K matrix based on the P matrix.
[0082] The P matrix is obtained according to the exponential function of the difference between the S matrix and the L matrix, and P matrices of different data types are obtained. The first current gradient matrix dV corresponding to the current V matrix is calculated based on the fp16 type P matrix; the second current gradient matrix dK corresponding to the current K matrix is calculated based on the fp32 type P matrix.
[0083] In one embodiment, calculation is performed according to the S matrix and the L matrix to obtain a fifth intermediate matrix P, and a first current gradient matrix dV corresponding to the current K matrix and a second current gradient matrix dK corresponding to the current K matrix are determined according to the P matrix, including:
[0084] a. Perform calculations based on the S matrix and the L matrix to obtain a fifth intermediate matrix P of the first data type and a fifth intermediate matrix P of the second data type, respectively.
[0085] In this embodiment, the first data type may be understood as fp16 type, and the second data type may be understood as fp32 type.
[0086] Specifically, "P=exp(SL)" is executed to calculate a P matrix, and an fp16-type P matrix and an fp32-type P matrix are output.
[0087] In the above technical solution, the P matrix requires two data precisions, FP32 and FP16. Therefore, two data formats are output when calculating the P matrix, saving the time of one cast and realizing the operator parallelism of mixed-precision attention calculation.
[0088] b. Determine a first current gradient matrix dV corresponding to the current V matrix based on the dO matrix, the P matrix of the first data type, and the gradient matrix corresponding to the historical V matrix.
[0089] In this embodiment, the first current gradient matrix may be understood as a gradient matrix corresponding to the current V matrix in the current iteration.
[0090] Specifically, execute "dV i j =dV i j-1 +P T dO”, where dV i j-1 is the gradient matrix corresponding to the history V matrix, according to the dO matrix, the transpose of the fp16 type P matrix, and dV i j-1 , calculate the first current gradient matrix dV corresponding to the current V matrix i j .
[0091] c. Determine the eighth intermediate matrix dP based on the dO matrix and the current V matrix.
[0092] In this embodiment, the eighth intermediate matrix can be understood as an intermediate parameter in the backward calculation, which is an operator.
[0093] Execute "dP = dOV T ", calculate the product of dO and the transpose of the current V matrix to get the dP matrix.
[0094] d. Determine a ninth intermediate matrix dS according to the dP matrix, the D matrix, and the P matrix of the second data type.
[0095] In this embodiment, the ninth intermediate matrix can be understood as an intermediate parameter in the backward calculation, which is an operator.
[0096] Execute "dS = P*(dP-D)" to determine the difference between the dP matrix and the D matrix, and take the product of the difference and the fp16 type P matrix to obtain the dS matrix.
[0097] e. Determine a second current gradient matrix dK corresponding to the current K matrix based on the dS matrix, the current Q matrix, and the gradient matrix corresponding to the historical K matrix.
[0098] In this embodiment, the second current gradient matrix can be understood as a gradient matrix corresponding to the current K matrix in the current iteration.
[0099] Execute "dK i j =dK i j-1 +dS T Q", where dK i j-1is the gradient matrix corresponding to the historical K matrix, and the Q matrix is obtained. If this is the first round of iteration, the Q matrix is still Q 0 If it is not the first iteration, the Q matrix is the pre-fetched Q matrix. Determine the product of the transpose of the dS matrix and the Q matrix, and take the product and dK i j-1 The sum of the values is calculated to obtain the second current gradient matrix dK corresponding to the current K matrix i j .
[0100] S1024: Determine a current third gradient matrix dQ corresponding to the current Q matrix according to the S matrix and the current K matrix, as well as the gradient matrix corresponding to the historical Q matrix.
[0101] In this embodiment, the third current gradient matrix can be understood as a gradient matrix corresponding to the current Q matrix in the current iteration.
[0102] Execute "dQ i j =dQ i j-1 +dSK", where dQ i j-1 is the gradient matrix corresponding to the historical Q matrix, determine the product of the dS matrix and the K matrix, and take the product and dQ i j-1 The sum of the values is calculated to obtain the third current gradient matrix dQ corresponding to the current Q matrix i j .
[0103] Execute “output dQ, KV, dV” to output the current backward output matrix of the backward calculation, including the current first gradient matrix dV, the current second gradient matrix dK, and the current third gradient matrix dQ.
[0104] In the above technical solution, the K matrix and the dS matrix are both transposed and non-transposed in the reverse calculation process. If a matrix stored in row-major order in the memory is accessed in column-major order, a bank conflict will occur. Similarly, a matrix stored in column-major order is read in row-major order, which will also cause a bank conflict, thereby increasing the memory access time. Therefore, this embodiment proposes a data arrangement method without bank conflicts. For example, for an 8*8 matrix, data is placed in the same column to represent the same cache, as shown in the following table:
[0105] 0 8 16 24 32 40 48 56 57 1 9 17 25 33 41 49 50 58 2 10 18 26 34 42 43 51 59 3 11 19 27 35 36 44 52 60 4 12 20 28 29 37 45 53 61 5 13 21 22 30 38 46 54 62 6 14 15 23 31 39 47 55 63 7
[0106] Conflicts are avoided by arranging data in a diagonal pattern. In this way, data with the same layout can be read in either row-major order or column-major order. This does not incur additional storage time and does not require data rearrangement to transpose the matrix, thus achieving matrix transposition optimization without memory access conflicts.
[0107] An attention mechanism acceleration method provided by an embodiment of the present invention can be applied to the fields of natural language processing and computer vision, such as the construction, training and application of natural language processing models in the field of natural language processing, especially the Transformer model and its derivative architecture, and the construction, training and application of image processing models in the field of computer vision, so as to improve the model's attention to key features of the image and reduce the interference of irrelevant information, thereby improving the performance and efficiency of the task. The method is applied to any processing core of an electronic device, and the method includes a forward propagation stage and a backward propagation stage; in the forward propagation stage, a query matrix is obtained, and forward calculation is performed according to the query matrix. In the forward calculation process, a key matrix and a value matrix are pre-fetched, and forward calculation is continued at the corresponding calculation node according to the key matrix and the value matrix to obtain the current forward output matrix; in the backward propagation stage, a key matrix, a value matrix and a query matrix are obtained, and backward calculation is performed according to the key matrix, the value matrix and the query matrix. In the backward calculation process, an intermediate matrix is pre-fetched, and backward calculation is continued at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix. The above technical solution is adapted to a multi-core hardware architecture. For a multi-core architecture, the computing load of each core may be different, and the final performance is determined by the core with the longest running time, so as to achieve load balancing of the causal mask attention mechanism operation of the multi-core architecture; the operators in the calculation process are optimized in a targeted manner to improve the calculation performance and accuracy; by efficiently reusing the cache space, the data originally to be stored in the register is placed in the cache, which relieves the register pressure and cache pressure, and has better versatility and scalability; compared with FlashAttention, the degree of operator fusion is further improved, the data precision conversion and multiplication are parallelized, and the row accumulation and exponential operation are parallelized, which increases the utilization rate of the hardware and the instruction throughput. In addition, an asynchronous transmission unit is used to realize the mutual masking of calculation and communication, and the bubbles in the calculation process are further eliminated by designing a software pipeline.
[0108] Embodiment 2
[0109] Figure 5 It is a structural schematic diagram of an attention mechanism acceleration device provided in Example 2 of the present invention.
[0110] like Figure 5 As shown, the device comprises:
[0111] A forward propagation module 21 is used to obtain a query matrix in the forward propagation stage, perform forward calculation according to the query matrix, pre-fetch a key matrix and a value matrix during the forward calculation process, and continue to perform forward calculation at corresponding calculation nodes according to the key matrix and the value matrix to obtain a current forward output matrix;
[0112] The backward propagation module 22 is used to obtain the key matrix, the value matrix and the query matrix in the backward propagation stage, perform backward calculation according to the key matrix, the value matrix and the query matrix, pre-fetch the intermediate matrix during the backward calculation process, and continue to perform backward calculation at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix.
[0113] The attention mechanism acceleration device adopted in this technical solution improves the processing efficiency, computing performance and accuracy of the attention mechanism in fields such as natural field processing and computer vision.
[0114] Optionally, the forward propagation module 21 includes:
[0115] An initialization unit, used to obtain a query matrix Q, initialize a forward output matrix O, a first intermediate matrix L, and a second intermediate matrix M;
[0116] A first S matrix determination unit is used to obtain a current key matrix K, determine a third intermediate matrix S according to the Q matrix and the current K matrix, determine the current K matrix as a historical K matrix, and re-acquire the current K matrix;
[0117] An N matrix determining unit, used for determining a fourth intermediate matrix N according to the S matrix and the M matrix;
[0118] a first P matrix determining unit, configured to obtain a current fifth intermediate matrix P according to the N matrix calculation, and to sum each P matrix;
[0119] A V matrix pre-fetch unit, used for pre-fetching a value matrix V;
[0120] The output matrix determination unit is used to determine the current forward output matrix according to the P matrix and the V matrix.
[0121] Optionally, an output matrix determination unit is specifically used to:
[0122] The data type of the P matrix is converted, and the current forward output matrix is determined based on the product of the converted P matrix and the V matrix and with reference to the forward output matrix.
[0123] Optionally, the reference forward output matrix is the product of the exponential vector parameter of the N matrix and the historical forward output matrix.
[0124] Optionally, the back propagation module 22 includes:
[0125] An initial matrix acquisition unit, used to acquire a current key matrix K, a current value matrix V and a current query matrix Q, and pre-fetch a first intermediate matrix L, a sixth intermediate matrix D and a seventh intermediate matrix dO;
[0126] a second S matrix determining unit, configured to determine a third intermediate matrix S according to a current Q matrix and a current K matrix, determine the current Q matrix as a historical Q matrix, and reacquire the current Q matrix;
[0127] a second P matrix determining unit, configured to perform calculations according to the S matrix and the L matrix to obtain a fifth intermediate matrix P, and determine, according to the P matrix, a first current gradient matrix dV corresponding to the current V matrix and a second current gradient matrix dK corresponding to the current K matrix;
[0128] The gradient matrix determining unit is used to determine a current third gradient matrix dQ corresponding to the current Q matrix according to the S matrix, the current K matrix, and the gradient matrix corresponding to the historical Q matrix.
[0129] Optional. The second P matrix determination unit is specifically used for:
[0130] Perform calculations based on the S matrix and the L matrix to obtain a fifth intermediate matrix P of the first data type and a fifth intermediate matrix P of the second data type, respectively;
[0131] Determine a first current gradient matrix dV corresponding to the current V matrix according to the dO matrix, the P matrix of the first data type, and the gradient matrix corresponding to the historical V matrix;
[0132] Determine an eighth intermediate matrix dP according to the dO matrix and the current V matrix;
[0133] Determine a ninth intermediate matrix dS according to the dP matrix, the D matrix, and the P matrix of the second data type;
[0134] According to the dS matrix and the current Q matrix, and the gradient matrix corresponding to the historical K matrix, a second current gradient matrix dK corresponding to the current K matrix is determined.
[0135] The attention mechanism acceleration device provided in the embodiment of the present invention can execute the attention mechanism acceleration method provided in any embodiment of the present invention, and has the corresponding functional modules and beneficial effects of the execution method.
[0136] Embodiment 3
[0137] Figure 6: is a schematic diagram of the structure of an electronic device provided in Embodiment 3 of the present invention. The electronic device is intended to represent various forms of digital computers, such as laptop computers, desktop computers, workbenches, personal digital assistants, servers, blade servers, mainframe computers, and other suitable computers. The electronic device may also represent various forms of mobile devices, such as personal digital processing, cellular phones, smart phones, wearable devices (such as helmets, glasses, watches, etc.) and other similar computing devices. The components shown herein, their connections and relationships, and their functions are merely examples and are not intended to limit the implementation of the present invention described and / or required herein.
[0138] like Figure 6 As shown, the electronic device 30 includes at least one processor 31, and a memory connected to the at least one processor 31, such as a read-only memory (ROM) 32, a random access memory (RAM) 33, etc., wherein the memory stores a computer program that can be executed by at least one processor, and the processor 31 can perform various appropriate actions and processes according to the computer program stored in the read-only memory (ROM) 32 or the computer program loaded from the storage unit 38 to the random access memory (RAM) 33. In the RAM 33, various programs and data required for the operation of the electronic device 30 can also be stored. The processor 31, the ROM 32, and the RAM 33 are connected to each other through a bus 34. An input / output (I / O) interface 35 is also connected to the bus 34.
[0139] A number of components in the electronic device 30 are connected to the I / O interface 35, including: an input unit 36, such as a keyboard, a mouse, etc.; an output unit 37, such as various types of displays, speakers, etc.; a storage unit 38, such as a disk, an optical disk, etc.; and a communication unit 39, such as a network card, a modem, a wireless communication transceiver, etc. The communication unit 39 allows the electronic device 30 to exchange information / data with other devices through a computer network such as the Internet and / or various telecommunication networks.
[0140] The processor 31 may be a variety of general and / or special processing components with processing and computing capabilities. Some examples of the processor 31 include, but are not limited to, a central processing unit (CPU), a graphics processing unit (GPU), various dedicated artificial intelligence (AI) computing chips, various processors running machine learning model algorithms, a digital signal processor (DSP), and any appropriate processor, controller, microcontroller, etc. The processor 31 executes the various methods and processes described above, such as the attention mechanism acceleration method.
[0141] In some embodiments, the attention mechanism acceleration method may be implemented as a computer program, which is tangibly contained in a computer-readable storage medium, such as a storage unit 38. In some embodiments, part or all of the computer program may be loaded and / or installed on the electronic device 30 via the ROM 32 and / or the communication unit 39. When the computer program is loaded into the RAM 33 and executed by the processor 31, one or more steps of the attention mechanism acceleration method described above may be performed. Alternatively, in other embodiments, the processor 31 may be configured to execute the attention mechanism acceleration method in any other appropriate manner (e.g., by means of firmware).
[0142] Various implementations of the systems and techniques described above herein can be implemented in digital electronic circuit systems, integrated circuit systems, field programmable gate arrays (FPGAs), application specific integrated circuits (ASICs), application specific standard products (ASSPs), systems on chips (SOCs), load programmable logic devices (CPLDs), computer hardware, firmware, software, and / or combinations thereof. These various implementations can include: being implemented in one or more computer programs that can be executed and / or interpreted on a programmable system including at least one programmable processor, which can be a special purpose or general purpose programmable processor that can receive data and instructions from a storage system, at least one input device, and at least one output device, and transmit data and instructions to the storage system, the at least one input device, and the at least one output device.
[0143] Computer programs for implementing the methods of the present invention may be written in any combination of one or more programming languages. These computer programs may be provided to a processor of a general-purpose computer, a special-purpose computer, or other programmable data processing device, so that when the computer program is executed by the processor, the functions / operations specified in the flow chart and / or block diagram are implemented. The computer program may be executed entirely on the machine, partially on the machine, partially on the machine and partially on a remote machine as a stand-alone software package, or entirely on a remote machine or server.
[0144] In the context of the present invention, a computer-readable storage medium may be a tangible medium that may contain or store a computer program for use by or in combination with an instruction execution system, device or equipment. A computer-readable storage medium may include, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, devices or equipment, or any suitable combination of the foregoing. Alternatively, a computer-readable storage medium may be a machine-readable signal medium. A more specific example of a machine-readable storage medium may include an electrical connection based on one or more lines, a portable computer disk, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), an optical fiber, a portable compact disk read-only memory (CD-ROM), an optical storage device, a magnetic storage device, or any suitable combination of the foregoing.
[0145] To provide interaction with a user, the systems and techniques described herein may be implemented on an electronic device having: a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user; and a keyboard and a pointing device (e.g., a mouse or trackball) through which the user can provide input to the electronic device. Other types of devices may also be used to provide interaction with the user; for example, the feedback provided to the user may be any form of sensory feedback (e.g., visual feedback, auditory feedback, or tactile feedback); and input from the user may be received in any form (including acoustic input, voice input, or tactile input).
[0146] The systems and techniques described herein may be implemented in a computing system that includes backend components (e.g., as a data server), or a computing system that includes middleware components (e.g., an application server), or a computing system that includes frontend components (e.g., a user computer with a graphical user interface or a web browser through which a user can interact with implementations of the systems and techniques described herein), or a computing system that includes any combination of such backend components, middleware components, or frontend components. The components of the system may be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include: a local area network (LAN), a wide area network (WAN), a blockchain network, and the Internet.
[0147] A computing system may include a client and a server. The client and the server are generally remote from each other and usually interact through a communication network. The client and server relationship is generated by computer programs running on the corresponding computers and having a client-server relationship with each other. The server may be a cloud server, also known as a cloud computing server or cloud host, which is a host product in the cloud computing service system to solve the defects of difficult management and weak business scalability in traditional physical hosts and VPS services.
[0148] It should be understood that the various forms of processes shown above can be used to reorder, add or delete steps. For example, the steps described in the present invention can be executed in parallel, sequentially or in different orders, as long as the desired results of the technical solution of the present invention can be achieved, and this document does not limit this.
[0149] The above specific implementations do not constitute a limitation on the protection scope of the present invention. It should be understood by those skilled in the art that various modifications, combinations, sub-combinations and substitutions can be made according to design requirements and other factors. Any modification, equivalent substitution and improvement made within the spirit and principle of the present invention should be included in the protection scope of the present invention.
Claims
1. A method for accelerating an attention mechanism, characterized in that: Applied to any processing core of an electronic device, the method comprises a forward propagation stage and a backward propagation stage; In the forward propagation stage, a query matrix is obtained, and forward calculation is performed according to the query matrix. In the forward calculation process, a key matrix and a value matrix are pre-fetched, and forward calculation is continued at the corresponding calculation node according to the key matrix and the value matrix to obtain the current forward output matrix; In the backward propagation stage, the key matrix, value matrix and query matrix are obtained, and backward calculation is performed based on the key matrix, value matrix and query matrix. During the backward calculation process, the intermediate matrix is pre-fetched, and backward calculation is continued at the corresponding calculation node based on the intermediate matrix to obtain the current gradient matrix.
2. The method according to claim 1, characterized in that: The obtaining of the query matrix, performing forward calculation according to the query matrix, pre-fetching a key matrix and a value matrix during the forward calculation process, and continuing forward calculation at corresponding computing nodes according to the key matrix and the value matrix to obtain a current forward output matrix, includes: Obtain the query matrix Q, initialize the forward output matrix O, the first intermediate matrix L and the second intermediate matrix M; Obtain the current key matrix K, determine the third intermediate matrix S according to the Q matrix and the current K matrix, determine the current K matrix as the historical K matrix, and re-obtain the current K matrix; Determine a fourth intermediate matrix N according to the S matrix and the M matrix; The current fifth intermediate matrix P is calculated according to the N matrix, and each P matrix is summed; Pre-fetch matrix V; According to the P matrix and the V matrix, the current forward output matrix is determined.
3. The method according to claim 2, characterized in that Determining the current forward output matrix according to the P matrix and the V matrix includes: The data type of the P matrix is converted, and the current forward output matrix is determined based on the product of the converted P matrix and the V matrix and with reference to the forward output matrix.
4. The method according to claim 3, characterized in that The reference forward output matrix is the product of the exponential vector parameter of the N matrix and the historical forward output matrix.
5. The method according to claim 1, characterized in that The step of obtaining a key matrix, a value matrix, and a query matrix, performing backward calculation according to the key matrix, the value matrix, and the query matrix, pre-fetching an intermediate matrix during the backward calculation, and continuing to perform backward calculation at a corresponding calculation node according to the intermediate matrix to obtain a current gradient matrix includes: Obtain the current key matrix K, the current value matrix V and the current query matrix Q, and pre-fetch the first intermediate matrix L, the sixth intermediate matrix D and the seventh intermediate matrix dO; Determine a third intermediate matrix S according to the current Q matrix and the current K matrix, determine the current Q matrix as the historical Q matrix, and re-acquire the current Q matrix; Calculate according to the S matrix and the L matrix to obtain a fifth intermediate matrix P, and determine according to the P matrix a first current gradient matrix dV corresponding to the current V matrix and a second current gradient matrix dK corresponding to the current K matrix; A current third gradient matrix dQ corresponding to the current Q matrix is determined according to the S matrix, the current K matrix, and the gradient matrix corresponding to the historical Q matrix.
6. The method according to claim 5, characterized in that The step of calculating according to the S matrix and the L matrix to obtain a fifth intermediate matrix P, and determining according to the P matrix a first current gradient matrix dV corresponding to the current K matrix and a second current gradient matrix dK corresponding to the current K matrix, comprises: Perform calculations based on the S matrix and the L matrix to obtain a fifth intermediate matrix P of the first data type and a fifth intermediate matrix P of the second data type, respectively; Determine a first current gradient matrix dV corresponding to the current V matrix according to the dO matrix, the P matrix of the first data type, and the gradient matrix corresponding to the historical V matrix; Determine an eighth intermediate matrix dP according to the dO matrix and the current V matrix; Determine a ninth intermediate matrix dS according to the dP matrix, the D matrix, and the P matrix of the second data type; According to the dS matrix and the current Q matrix, and the gradient matrix corresponding to the historical K matrix, a second current gradient matrix dK corresponding to the current K matrix is determined.
7. An attention mechanism acceleration device, characterized in that: include: A forward propagation module is used to obtain a query matrix in the forward propagation stage, perform forward calculation according to the query matrix, pre-fetch a key matrix and a value matrix during the forward calculation process, and continue to perform forward calculation at corresponding calculation nodes according to the key matrix and the value matrix to obtain a current forward output matrix; The backward propagation module is used to obtain the key matrix, value matrix and query matrix in the backward propagation stage, perform backward calculation according to the key matrix, value matrix and query matrix, pre-fetch the intermediate matrix during the backward calculation process, and continue to perform backward calculation at the corresponding calculation node according to the intermediate matrix to obtain the current gradient matrix.
8. An electronic device, characterized in that: include: at least one processor; as well as a memory communicatively connected to the at least one processor; wherein, The memory stores a computer program executable by the at least one processor, and the computer program is executed by the at least one processor so that the at least one processor can execute an attention mechanism acceleration method as described in any one of claims 1-6.
9. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores computer instructions, and the computer instructions are used to enable a processor to implement an attention mechanism acceleration method described in any one of claims 1-6 when executed.
10. A computer program product, characterized in that The computer program product comprises a computer program which, when executed by a processor, implements an attention mechanism acceleration method according to any one of claims 1-6.