Optimization method, device, storage medium and program product for attention mechanism calculation

Abstract

The embodiment of the present application provides an optimization method, device, storage medium and program product for attention mechanism calculation, which relates to the field of artificial intelligence technology. The method includes: executing a forward recalculation operator based on the query block and the key block to obtain a forward normalization result. Executing a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result, and simultaneously executing a dot product summation operator based on the forward output block and the output gradient block in parallel to obtain an intermediate tensor. Executing a first type of transformation operator on the value intermediate result to obtain a value gradient block; and simultaneously executing a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor and the forward normalization result in parallel to obtain a reverse normalization result. Finally, based on the reverse normalization result, the key block and the query block, the query gradient block and the key gradient block are obtained, thereby achieving effective pipeline processing of different operators, thereby improving the overall performance of the attention mechanism calculation.

Description

Technical Field

[0001] The embodiments of the present application relate to the field of artificial intelligence technology, and in particular to an optimization method, device, storage medium, and program product for attention mechanism calculation. Background Art

[0002] In the field of deep learning, attention mechanisms are widely used in various deep learning models. For example, the core computation of the Transformer model is attention mechanism calculation, and large language models generally use the Transformer model as the core infrastructure. Therefore, in the era of widespread application of large language, the execution performance of attention mechanism calculation has become particularly important.

[0003] In the related art, when executing the reverse attention mechanism calculation, multiple operators associated with the reverse attention mechanism calculation are executed sequentially in a serial manner. This serial approach leads to low operator execution efficiency, which in turn leads to poor performance of the reverse attention mechanism calculation. Summary of the Invention

[0004] The embodiments of the present application provide an optimization method, device, storage medium and program product for attention mechanism calculation, which are used to improve the performance of attention mechanism calculation.

[0005] On the one hand, the embodiments of the present application provide an optimization method for attention mechanism calculation, including:

[0006] Execute forward recalculation operators based on query blocks and key blocks to obtain forward normalized results;

[0007] Performing a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; performing a first type of transformation operator on the value intermediate result to obtain a value gradient block;

[0008] Executing a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product sum operator and the first matrix multiplication operator are executed in parallel;

[0009] Executing a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a backnormalization result, wherein the gradient backpropagation operator and the first type transformation operator are executed in parallel;

[0010] Based on the reverse normalization result and the key block, a query gradient block is obtained; and based on the reverse normalization result and the query block, a key gradient block is obtained.

[0011] On the one hand, an embodiment of the present application provides an optimization device for attention mechanism calculation, including:

[0012] The recalculation module is used to execute the forward recalculation operator based on the query block and the key block to obtain the forward normalized result;

[0013] A back propagation module is configured to perform a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; and perform a first type of transformation operator on the value intermediate result to obtain a value gradient block;

[0014] The back propagation module is further configured to execute a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product sum operator and the first matrix multiplication operator are executed in parallel;

[0015] The back propagation module is further configured to execute a gradient back propagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a back normalization result, wherein the gradient back propagation operator and the first type transformation operator are executed in parallel;

[0016] The back propagation module is further configured to obtain a query gradient block based on the back normalization result and the key block; and to obtain a key gradient block based on the back normalization result and the query block.

[0017] Optionally, the recalculation module is further configured to:

[0018] During execution of the forward recalculation operator, the output gradient blocks are loaded from the video memory into an on-chip cache in parallel.

[0019] Optionally, the forward recalculation operator includes: a second matrix multiplication operator and a normalization operator;

[0020] The recalculation module is specifically used for:

[0021] Dividing the query block into a plurality of first sub-blocks; and dividing the key block into a plurality of second sub-blocks, wherein each first sub-block corresponds to a second sub-block;

[0022] For each first sub-block and the corresponding second sub-block, respectively, performing: executing a second matrix multiplication operator based on the first sub-block and the corresponding second sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core and executing the normalization operator to obtain a sub-normalization result;

[0023] The forward normalization result is generated based on the obtained multiple sub-normalization results.

[0024] Optionally, the recalculation module is further configured to:

[0025] The normalization operator is executed by reading the sub-matrix multiplication result from the register through a vector core, and after obtaining the sub-normalization result, the sub-normalization result is stored in an on-chip cache.

[0026] Optionally, the recalculation module is further configured to:

[0027] During execution of the forward recalculation operator, register synchronization is adopted between the tensor core and the vector core.

[0028] Optionally, the gradient back-propagation operator includes: a third matrix multiplication operator and a dot product difference operator;

[0029] The back propagation module is specifically used for:

[0030] Dividing the output gradient block into a plurality of third sub-blocks; and dividing the value block into a plurality of fourth sub-blocks, wherein each third sub-block corresponds to one fourth sub-block;

[0031] For each third sub-block and the corresponding fourth sub-block, respectively, performing: executing a third matrix multiplication operator based on a third sub-block and the corresponding fourth sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core, and executing the dot product difference operator on the intermediate tensor and the forward normalization result to obtain a sub-dot product result;

[0032] The reverse normalization result is generated based on the obtained multiple sub-point product results.

[0033] Optionally, the back propagation module is further configured to:

[0034] The sub-matrix multiplication result is read from the register through the vector core, the dot product difference operator is executed with the intermediate tensor and the forward normalization result, and after the sub-dot product result is obtained, the sub-dot product result is stored in the on-chip cache.

[0035] Optionally, the back propagation module is specifically used to:

[0036] Executing a fourth matrix multiplication operator based on the inverse normalization result and the key block to obtain a query intermediate result; executing a second type transformation operator on the query intermediate result to obtain the query gradient block;

[0037] executing a fifth matrix multiplication operator based on the inverse normalization result and the query block to obtain a key intermediate result, wherein the fifth matrix multiplication operator and the second-type transformation operator are executed in parallel;

[0038] A third type of transformation operator is performed on the key intermediate result to obtain the key gradient block.

[0039] Optionally, the query blocks correspond to multiple rounds of inner loops, and the key blocks and the value blocks are input data of the current round of inner loops;

[0040] The third type transformation operator is executed in parallel with the forward recalculation operator in the next inner loop corresponding to the query block.

[0041] Optionally, a split module is also included;

[0042] The query block is a block in the target query tensor processed by a computing unit; the computing unit is any one of the multiple computing units included in the artificial intelligence chip;

[0043] The splitting module is specifically used for:

[0044] Get the original query tensor, where the shape of the original query tensor includes: the number of batches and attention heads combined, and the sequence length;

[0045] Obtaining a number of splits based on the sequence length and the hardware alignment granularity of the artificial intelligence chip;

[0046] When the number of merges is greater than the number of splits, and the ratio of the number of merges to the total number of computing units is an integer, splitting the original query tensor into target query tensors processed by each of the plurality of computing units according to the number of merges and the total number of computing units;

[0047] When the number of merges is not greater than the number of splits, and the ratio of the number of splits to the total number of computing units is an integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the sequence length and the total number of computing units.

[0048] Optionally, the splitting module is further configured to:

[0049] When the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is a non-integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the number of merges, the sequence length, and the total number of computing units.

[0050] In one aspect, an embodiment of the present application provides a computer device comprising a memory, an artificial intelligence chip, and a computer program stored in the memory and executable on the artificial intelligence chip. When the artificial intelligence chip executes the computer program, the method for optimizing the calculation of the attention mechanism described above is implemented, including:

[0051] The AI chip reads the query blocks and key blocks from the memory, and performs forward recalculation operators based on the query blocks and key blocks to obtain forward normalized results.

[0052] The artificial intelligence chip performs a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; performs a first type of transformation operator on the value intermediate result to obtain a value gradient block;

[0053] The artificial intelligence chip executes a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product sum operator and the first matrix multiplication operator are executed in parallel;

[0054] The artificial intelligence chip executes a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a backnormalization result, wherein the gradient backpropagation operator and the first type of transformation operator are executed in parallel;

[0055] The artificial intelligence chip obtains a query gradient block based on the reverse normalization result and the key block; and obtains a key gradient block based on the reverse normalization result and the query block.

[0056] On the one hand, an embodiment of the present application provides a computer-readable storage medium, which stores a computer program that can be executed by a computer device. When the computer program runs on the computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0057] On the one hand, an embodiment of the present application provides a computer program product, which includes a computer program stored on a computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer device, the computer device performs the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0058] In an embodiment of the present application, a forward recalculation operator is executed based on the query block and the key block to obtain a forward normalization result; then a first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain an intermediate result, wherein, when executing the first matrix multiplication operator, a dot product sum operator is executed in parallel based on the forward output block and the output gradient block to obtain an intermediate tensor, so that the calculation time of the first matrix multiplication operator and the dot product sum operator overlaps.

[0059] Then, the first type of transformation operator is executed on the value intermediate result to obtain the value gradient block; at the same time, the gradient backpropagation operator is executed in parallel based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain the reverse normalization result, so that the calculation time of the first type of transformation operator and the gradient backpropagation operator overlaps.

[0060] Finally, based on the reverse normalization results and key blocks, the query gradient blocks are obtained; and based on the reverse normalization results and query blocks, the key gradient blocks are obtained. This enables effective pipeline processing of different operators during the reverse attention mechanism calculation process. Compared with the serial method, the execution efficiency of the operators is improved, thereby improving the overall performance of the attention mechanism calculation. BRIEF DESCRIPTION OF THE DRAWINGS

[0061] 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 paying any creative work.

[0062] Figure 1 A schematic diagram of the structure of an artificial intelligence chip provided in an embodiment of the present application;

[0063] Figure 2 A schematic diagram of input and output data of the reverse attention mechanism operator provided in an embodiment of the present application;

[0064] Figure 3 Schematic diagram of a process for optimizing attention mechanism calculation provided in an embodiment of the present application Figure 1 ;

[0065] Figure 4 Schematic diagram of a process for optimizing attention mechanism calculation provided in an embodiment of the present application Figure 2 ;

[0066] Figure 5 A schematic diagram of the structure of an optimization device for attention mechanism calculation provided in an embodiment of the present application;

[0067] Figure 6 A schematic diagram of the structure of a computer device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0068] In order to make the purpose, technical solutions and beneficial effects of the present invention more clearly understood, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0069] The following is a brief introduction to the system architecture diagram applicable to the technical solution of the embodiment of the present application. It should be noted that the system architecture diagram introduced below is only used to illustrate the embodiment of the present application and is not limiting.

[0070] refer to Figure 1 , which is a structural diagram of an artificial intelligence chip applicable to an embodiment of the present application. The artificial intelligence chip 100 includes at least: a video memory 101, and multiple computing units 102, wherein the computing unit 102 can be a streaming processing cluster (SPC for short).

[0071] Each computing unit 102 includes: an on-chip cache 103 , a register 104 , a tensor core 105 and a vector core 106 , wherein the tensor core 105 and the vector core 106 are heterogeneous cores.

[0072] Video memory 101 can be high-bandwidth memory (HBM) or other types of memory. On-chip cache 103 is a temporary memory with a smaller capacity than video memory 101 but faster data transfer speed. On-chip cache 103 can be a general matrix main buffer (GMB) or other types of memory.

[0073] Register 104 is a shared register between tensor core 105 and vector core 106. Tensor core 105 and vector core 106 can exchange data through register 104, which can be a thread-local register (TLR). Compared to on-chip cache 103, register 104 has a smaller capacity but faster data exchange speed.

[0074] In an embodiment of the present application, multiple computing units 102 perform attention mechanism calculations in parallel. Taking one computing unit as an example, the tensor core 105 and the vector core 106 in the computing unit 102 perform a forward recalculation operator based on the query block and the key block to obtain a forward normalized result. In the process of executing the forward recalculation operator, the tensor core 105 and the vector core 106 transmit the results obtained by the intermediate calculation through the register 104, and the calculation results output by the forward recalculation operator and required by the subsequent operator are saved through the on-chip cache 103, without the need to write the results obtained by the intermediate calculation and the calculation results required by the subsequent operator back to the video memory. This makes full use of the register 104 and the on-chip cache 103, improves the storage resource utilization of the computing unit 102; at the same time, it avoids the computing unit 102 from frequently accessing the video memory, thereby improving the computing efficiency of the computing unit 102.

[0075] Then, the tensor core 105 performs a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain an intermediate result. While the tensor core 105 is performing the first matrix multiplication operator, the vector core 106 concurrently performs a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor. This allows the time for the tensor core 105 to perform the first matrix multiplication operator to overlap with the time for the vector core 106 to perform the dot product sum operator, avoiding the vector core 106 being in a waiting state while the tensor core 105 is performing the operator, thereby improving the utilization of computing resources in the computing unit 102 and improving the overall computing efficiency of the computing unit 102.

[0076] Next, the vector core 106 performs a first type of transformation operator on the intermediate value result to obtain a value gradient block. Simultaneously, the tensor core 105 and the vector core 106 perform a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a backpropagation result. During the execution of the backpropagation operator, the tensor core 105 and the vector core 106 transmit the intermediate calculation results through the register 104, and the calculation results output by the backpropagation operator and required by subsequent operators are stored through the on-chip cache 103, without having to write the intermediate calculation results and the calculation results required by subsequent operators back to the video memory. In this way, the register 104 and the on-chip cache 103 are fully utilized, thereby improving the storage resource utilization of the computing unit 102. At the same time, the computing unit 102 is prevented from frequently accessing the video memory, thereby improving the computational efficiency of the computing unit 102. Secondly, the computation time of the first type transformation operator and the gradient back-propagation operator overlaps, thereby improving the utilization of the computation resources in the computation unit 102 and improving the overall computation efficiency of the computation unit 102 .

[0077] Finally, based on the reverse normalization results and key blocks, the query gradient blocks are obtained; and based on the reverse normalization results and query blocks, the key gradient blocks are obtained, so that in the reverse attention mechanism calculation process, the computing unit 102 performs effective pipeline processing on different operators, which improves the execution efficiency of the operators and the overall performance of the artificial intelligence chip 100 compared to the serial method.

[0078] In addition to the above-mentioned structure, the artificial intelligence chip 100 in this application may also include other structures, which are not specifically limited in this application.

[0079] The artificial intelligence chip 100 can be: a graphics processing unit (GPU), a general-purpose computing on graphics processing units (GPGPU), a domain-specific architecture (DSA), etc.

[0080] This application is based on Figure 1 The architecture diagram of the artificial intelligence chip shown provides an optimization method for attention mechanism calculation.

[0081] Attention mechanism calculation can be applied to various scenarios, such as image processing, speech processing, and text processing.

[0082] In different application scenarios, the physical meaning of the input data calculated by the attention mechanism can be different. For example, in the speech processing scenario, the input data calculated by the attention mechanism can be speech data used in tasks such as speech enhancement, speech recognition, and speech synthesis.

[0083] For example, in an image processing scenario, the input data calculated by the attention mechanism can be image data used in tasks such as image preprocessing, image segmentation, and target detection.

[0084] For example, in a text processing scenario, the input data calculated by the attention mechanism can be text data used in tasks such as text generation and text recognition.

[0085] In practical applications, attention mechanism calculations include forward attention mechanism and reverse attention mechanism calculations. In the embodiments of this application, the focus is on optimizing reverse attention mechanism calculations. During iterative training of a model that includes attention mechanism calculations, each round of iteration requires backpropagation to calculate the model's gradient information, and then adjust the model parameters based on the gradient information. The backpropagation process includes reverse attention mechanism calculations.

[0086] In the examples of this application, see Figure 2 The reverse attention mechanism (Attention bwd) operator includes: 6 input tensors, namely (Q, K, V, O, dO, M1), where Q represents the original query tensor, K represents the original key tensor, V represents the original value tensor, O represents the original forward output, dO represents the original output gradient, and M1 represents the original input mask.

[0087] The reverse attention operator also includes three intermediate tensors and three output tensors. The three intermediate tensors are (m, s, RD), where m and s are the original intermediate results produced by the forward attention mechanism calculation process, and RD is the original random mask produced by the forward attention mechanism calculation process. The three intermediate tensors are also used as inputs for the reverse attention mechanism calculation.

[0088] The three output tensors are (dQ, dK, dV), where dQ represents the query gradient, dK represents the key gradient, and dV represents the value gradient.

[0089] The shape of any of the above input tensors, intermediate tensors, and output tensors includes the following dimensions: , where batch_size represents the batch size, head_num represents the number of attention heads, sequence_length represents the sequence length, and head_size represents the size of the attention head.

[0090] In an embodiment of the present application, the artificial intelligence chip includes multiple computing units, and the multiple computing units perform attention mechanism calculations in parallel. However, in related technologies, for raw query tensors of any shape, a fixed splitting method (i.e., according to sequence length) is used to split the raw query tensor into multiple computing units for calculation. When the sequence length of the raw query tensor is short, the target query tensor split into each computing unit will be small, and may even be smaller than the hardware alignment granularity of the artificial intelligence chip, resulting in performance degradation.

[0091] In view of this, in an embodiment of the present application, the target query tensor processed by each computing unit is obtained by the central processing unit in the following manner. The specific process is as follows:

[0092] Get the original query tensor. The shape of the original query tensor includes: the number of batches and attention heads combined, and the sequence length. The number of batches and attention heads combined refers to: , which is the product of the batch size and the number of attention heads.

[0093] The number of splits is determined based on the sequence length and the AI chip's hardware alignment granularity. The hardware alignment granularity (block_size) represents the minimum sequence length that can be processed by a single instruction on the AI chip. The specific value of the hardware alignment granularity depends on the type of AI chip. In some embodiments, the ratio of the sequence length to the AI chip's hardware alignment granularity is used as the number of splits.

[0094] When the number of merges is greater than the number of splits, and the ratio of the number of merges to the total number of computational units is an integer, the original query tensor is split into target query tensors processed by multiple computational units according to the number of merges and the total number of computational units.

[0095] At this point, the sequence length of the target query tensor is the same as the sequence length of the original query tensor, and the shape of the target query tensor is:

[0096] ;

[0097] Where N represents the total number of computing units.

[0098] For example, setting the shape of the original query tensor to , N=16, block_size=64. In this case, the number of merges = 16, and the number of splits = 512 / 64 = 8. That is, the number of merges is greater than the number of splits, and the ratio of the number of merges to the total number of computing units is an integer. Therefore, the original query tensor is split into target query tensors processed by N computing units in the merge number dimension. In this case, the shape of the target query tensor is [1, 512, head_size].

[0099] When the number of merges is not greater than the number of splits, and the ratio of the number of splits to the total number of computational units is an integer, the original query tensor is split into target query tensors processed by multiple computational units according to the sequence length and the total number of computational units.

[0100] At this point, the sequence length of the target query tensor is the ratio of the sequence length of the original query tensor to the total number of computing units, that is, the shape of the target query tensor is:

[0101] ;

[0102] Where N represents the total number of computing units.

[0103] For example, setting the shape of the original query tensor to , N=16, block_size=64. At this time, the number of merges = 16, the number of splits = 2048 / 64 = 32, that is, the number of merges is not greater than the number of splits, and the ratio of the number of splits to the total number of computing units is an integer. Therefore, in the sequence length dimension, the original query tensor is split into target query tensors processed by N computing units respectively. At this time, the shape of the target query tensor is .

[0104] In some embodiments, when the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is a non-integer, the original query tensor is split into target query tensors processed by multiple computing units according to the number of merges, sequence length and total number of computing units.

[0105] At this point, the shape of the target query tensor is:

[0106] ;

[0107] Among them, n1 represents the split ratio in the dimension of merge quantity, and n2 represents the split ratio in the dimension of sequence length. , N represents the total number of computing units.

[0108] For example, setting the shape of the original query tensor to , N=16, block_size=64. In this case, the number of merges = 8, and the number of splits = 512 / 64 = 8. That is, the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is also a non-integer. Therefore, combining the merge number dimension and the sequence length dimension, the original query tensor is split into target query tensors processed by N computing units, each of which is processed by the target query tensor. In this case, n1=8, n2=2, and the shape of the target query tensor is [1, 256, head_size].

[0109] It should be noted that when the number of merges is greater than the number of splits, the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is an integer, the original query tensor can also be split into target query tensors processed by N computing units separately in both the merge number dimension and the sequence length dimension.

[0110] When the number of merges is not greater than the number of splits, the ratio of the number of splits to the total number of computing units is a non-integer, and the ratio of the number of merges to the total number of computing units is an integer, the original query tensor can also be split into target query tensors processed by N computing units, each in both the merge number dimension and the sequence length dimension. This application does not impose specific limitations on this.

[0111] In an embodiment of the present application, according to the actual shape of the original query tensor, a suitable splitting dimension is selected in the merge quantity dimension and the sequence length dimension, and the original query tensor is split into target query tensors processed by N computing units respectively according to the selected splitting dimension. In this way, larger dimensions can be split preferentially, so that each computing unit can be fully utilized, thereby improving operator performance.

[0112] The number of splits is determined based on the sequence length and the hardware alignment granularity of the AI chip. Only when the ratio of the number of splits to the total number of compute units is an integer is the original query tensor split according to the sequence length dimension into target query tensors processed by each of the N compute units. This allows each compute unit to process target query tensors that are no smaller than the hardware alignment granularity. This eliminates the need to padded target query tensors that are too fine (i.e., smaller than the hardware alignment granularity) during computation, thereby improving the overall computing performance of the compute unit.

[0113] It should be noted that other input tensors and intermediate tensors can be split into N computing units for calculation in the same way. For example, the original key tensor K can be split into N target key tensors, the original value tensor V can be split into N target value tensors, the original forward output O can be split into N target forward outputs, the original output gradient dO can be split into N target output gradients, the original input mask M1 can be split into N target input masks, the original intermediate result can be split into N target intermediate results, and the original random mask can be split into N target random masks, etc., which will not be repeated here.

[0114] Multiple computing units perform attention mechanism calculations in parallel based on the data they obtain. The following is an example of a computing unit:

[0115] In an embodiment of the present application, the reverse attention mechanism operator may be a reverse fast attention mechanism (FlashAttention bwd) operator. During the execution of the reverse fast attention mechanism operator, each operator in the reverse fast attention mechanism operator is calculated in a block manner.

[0116] Therefore, for each computing unit, after obtaining the target query tensor, the target query tensor is split into multiple query blocks Q_chunk; similarly, the obtained target key tensor is split into multiple key blocks K_chunk; the obtained target value tensor is split into multiple value blocks V_chunk; the obtained target forward output is split into multiple forward output blocks O_chunk; the obtained target output gradient is split into multiple output gradient blocks dO_chunk; the obtained target input mask is split into multiple input mask blocks M1_chunk; the obtained target intermediate result is split into multiple intermediate result blocks m_chunk and multiple intermediate result blocks s_chunk; the obtained target random mask is split into multiple random mask blocks RD_chunk, etc., which will not be repeated here.

[0117] The execution process of the reverse fast attention mechanism operator includes: outer loop, middle loop, and inner loop. The outer loop is a batch dimension loop with a step size of 1; the middle loop is a chunk Q dimension loop with a step size of Q_chunk; the inner loop is a chunk KV dimension loop with a step size of K_chunk or V_chunk.

[0118] A top-level loop consists of multiple middle-level loops, each with a different Q_chunk input. Each middle-level loop consists of multiple inner-level loops, each with a different K_chunk and V_chunk input, but the same Q_chunk. The intermediate results generated by these three levels of loops can be reused.

[0119] This embodiment of the application takes one inner loop as an example to introduce the process of the optimization method of attention mechanism calculation, which is executed by an artificial intelligence chip. Figure 3 , specifically including the following steps:

[0120] Step 301: Execute a forward recalculation operator based on the query block and the key block to obtain a forward normalization result.

[0121] Specifically, a query chunk is a chunk in the target query tensor processed by a computing unit, and a key chunk is a chunk in the target key tensor processed by a computing unit.

[0122] In some embodiments, the forward recalculation operator includes: a second matrix multiplication operator and a normalization operator. The execution process of the forward recalculation operator includes the following steps:

[0123] The query block is divided into a plurality of first sub-blocks; and the key block is divided into a plurality of second sub-blocks, wherein each first sub-block corresponds to one second sub-block.

[0124] For each first sub-block and the corresponding second sub-block, execute:

[0125] A second matrix multiplication operator is executed based on a first sub-block and a corresponding second sub-block by the tensor core to obtain a sub-matrix multiplication result; and then the sub-matrix multiplication result is stored in a register.

[0126] In some cases, to fully utilize register resources while preventing the sub-matrix multiplication results from exceeding the register storage space, in embodiments of the present application, based on the pre-allocated register storage space, the query block is divided into multiple first sub-blocks; and the key block is divided into multiple second sub-blocks. In this way, the sub-matrix multiplication results obtained by performing matrix multiplication on the first sub-block and the corresponding second sub-block can fully utilize register resources, while reducing the number of split sub-blocks and, therefore, the number of loop calculations.

[0127] Next, the sub-matrix multiplication result is read from the register by the vector core and the normalization operator is executed to obtain the sub-normalized result.

[0128] In a specific implementation, obtain the intermediate result block m_chunk and the intermediate result block s_chunk, then further split the intermediate result block m_chunk into multiple intermediate sub-blocks m'_chunk, and further split the intermediate result block s_chunk into multiple intermediate sub-blocks s'_chunk. Select an intermediate sub-block m'_chunk and an intermediate sub-block s'_chunk associated with the first sub-block and the second sub-block of the current round of calculation.

[0129] Perform a mask operation on the sub-matrix multiplication result to obtain the sub-mask result f; then, based on the sub-mask result f, the intermediate sub-block m'_chunk and the intermediate sub-block s'_chunk, obtain the sub-normalization result p', as shown in the following formula (1):

[0130]

[0131] After the subnormalized results are obtained, they are stored in the on-chip cache instead of being written back to the video memory because they are needed for subsequent calculations. This allows subsequent calculations to read the subnormalized results directly from the on-chip cache without accessing the video memory, improving computational efficiency and making full use of the on-chip cache.

[0132] In some embodiments, register synchronization is employed between the tensor cores and the vector cores during execution of the forward recalculation operator.

[0133] Specifically, the synchronization granularity is set to a preset number of registers. The Tensor Core executes the second matrix multiplication operator to obtain the sub-matrix multiplication result and stores the sub-matrix multiplication result in a register. After the Tensor Core produces data for the preset number of registers (i.e., the sub-matrix multiplication result), it sends a data ready synchronization instruction to the Vector Core. In this case, the producer is the Tensor Core and the consumer is the Vector Core.

[0134] The vector core consumes data from a preset number of registers, reading the sub-matrix multiplication results from the registers and performing a normalization operator on the sub-matrix multiplication results to obtain the sub-normalized results. After consuming data from the preset number of registers, the vector core sends a buffer ready instruction to the tensor core.

[0135] In addition, in the process of reading the sub-matrix multiplication result from the register through the vector core to execute the normalization operator, the second matrix multiplication operator is executed in parallel based on the next first sub-block and the corresponding second sub-block through the tensor core to obtain the sub-matrix multiplication result, so that the calculation time of the tensor core and the vector core overlaps.

[0136] Finally, based on the multiple sub-normalization results obtained, a forward normalization result is generated; specifically, the multiple sub-normalization results can be spliced together to obtain a forward normalization result.

[0137] In an embodiment of the present application, the query block and the key block are further divided into multiple sub-blocks; the tensor core then performs a second matrix multiplication operator on the sub-block to obtain a sub-matrix multiplication result, and stores the sub-matrix multiplication result in a register; the vector core reads the sub-matrix multiplication result from the register and performs a normalization operator to obtain a sub-normalization result. At the same time, the tensor core executes the second matrix multiplication operator based on the next sub-block in parallel. This allows the computation time of the tensor core and the vector core to overlap, achieving effective pipeline processing of different operators, thereby improving execution efficiency.

[0138] Secondly, data exchange between tensor cores and vector cores is carried out through registers, avoiding frequent entry and exit of interactive data in the on-chip cache, thereby improving the overall performance of the operator.

[0139] Step 302: Execute a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; execute a first type of transformation operator on the value intermediate result to obtain a value gradient block.

[0140] Specifically, the forward normalization result is transposed by the tensor core and then multiplied by the first matrix of the output gradient block input to obtain the intermediate result.

[0141] The data type of the intermediate value result is converted by using the first type transformation operator through the vector kernel to obtain the value gradient block, where the data types include: FP32 (32-bit floating point number), BF16 (16-bit floating point number), etc.

[0142] In some embodiments, during the execution of the forward recalculation operator, an output gradient block is loaded from video memory into an on-chip cache in parallel, where the output gradient block is a block of target output gradients processed by a computation unit. This allows the forward normalization result and the output gradient block to be directly read from the on-chip cache to execute the first matrix multiplication operator after obtaining the forward normalization result, thereby improving operator execution efficiency.

[0143] Step 303: Perform a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor.

[0144] Specifically, a vector kernel is used to perform a dot product operation on the forward output block and the output gradient block, and the dot product results are summed row by row to obtain an intermediate tensor. After the intermediate tensor is calculated, it is stored in the on-chip cache because it is needed for subsequent calculations.

[0145] In an embodiment of the present application, the dot product sum operator and the first matrix multiplication operator are executed in parallel, that is, in the process of executing the first matrix multiplication operator based on the output gradient block and the forward normalization result, the dot product sum operator based on the forward output block and the output gradient block is executed in parallel.

[0146] In some embodiments, execution of other operators is continued only after the first matrix multiplication operator and the dot product summation operator are completed, so as to avoid execution errors.

[0147] Step 304 , performing a gradient backpropagation operator based on the output gradient blocks, the value blocks, the intermediate tensors, and the forward normalization results to obtain a reverse normalization result.

[0148] In a specific implementation, the gradient backpropagation operator and the first type of transformation operator are executed in parallel, that is, in the process of executing the first type of transformation operator on the value intermediate result, the gradient backpropagation operator is executed in parallel based on the output gradient block, value block, intermediate tensor and forward normalization result.

[0149] In some embodiments, the gradient backpropagation operator includes: a third matrix multiplication operator and a dot product difference operator, and the execution process of the gradient backpropagation operator includes the following steps:

[0150] The output gradient block is divided into a plurality of third sub-blocks; and the value block is divided into a plurality of fourth sub-blocks, wherein each third sub-block corresponds to one fourth sub-block.

[0151] For each third sub-block and the corresponding fourth sub-block, execute:

[0152] A third matrix multiplication operator is executed based on a third sub-block and a corresponding fourth sub-block through the tensor core to obtain a sub-matrix multiplication result, and the sub-matrix multiplication result is stored in a register.

[0153] In some cases, based on pre-allocated register storage space, the output gradient block is divided into multiple third sub-blocks, and the value block is divided into multiple fourth sub-blocks. This allows the matrix multiplication of the third sub-block with the corresponding fourth sub-block to fully utilize register resources, while reducing the number of sub-blocks and, consequently, the number of loop calculations.

[0154] Next, the sub-matrix multiplication result is read from the register through the vector core, and the dot multiplication difference operator is performed with the intermediate tensor and the forward normalization result to obtain the sub-dot multiplication result.

[0155] In the specific implementation, the intermediate tensor is further split into multiple sub-tensors d', and the forward normalization result is further split into multiple sub-normalization results p'. A sub-tensor d' and a sub-normalization result p' associated with the third and fourth sub-blocks of this round of calculation are selected.

[0156] Subtract the sub-matrix multiplication result dp' from the corresponding sub-tensor d', and perform a dot multiplication operation on the sub-normalized result p' and the obtained subtraction result to obtain the sub-dot multiplication result ds', as shown in the following formula (2):

[0157]

[0158] After calculating the sub-dot product result ds', since it is needed in subsequent calculations, it is saved in the on-chip cache instead of being written back to the video memory. This allows subsequent calculations to directly read the sub-dot product result ds' from the on-chip cache, improving computational efficiency.

[0159] In some embodiments, during the execution of the gradient backpropagation operator, register synchronization is employed between the tensor cores and the vector cores.

[0160] Specifically, the tensor core executes the third matrix multiplication operator to obtain the sub-matrix multiplication result and saves the sub-matrix multiplication result in the register; then, it notifies the vector core to read the sub-matrix multiplication result from the register, and performs the point multiplication difference operator with the intermediate tensor and the forward normalization result to obtain the sub-point multiplication result.

[0161] In addition, the sub-matrix multiplication result is read from the register through the vector core, and in the process of performing the dot multiplication and difference operator with the intermediate tensor and the forward normalization result, the third matrix multiplication operator is executed in parallel based on the next third sub-block and the corresponding fourth sub-block through the tensor core to obtain the sub-matrix multiplication result, so that the calculation time of the tensor core and the vector core overlaps.

[0162] Finally, based on the multiple sub-point multiplication results obtained, a reverse normalization result is generated; specifically, the multiple sub-point multiplication results can be spliced to obtain a reverse normalization result.

[0163] In the embodiment of the present application, the output gradient block and the value block are further divided into multiple sub-blocks; the tensor core then performs a third matrix multiplication operator on the sub-block to obtain a sub-matrix multiplication result, and stores the sub-matrix multiplication result in a register; the vector core reads the sub-matrix multiplication result from the register and performs a dot multiplication and difference operator to obtain a sub-dot multiplication result. At the same time, the tensor core executes a second matrix multiplication operator based on the next sub-block in parallel. This allows the computation time of the tensor core and the vector core to overlap, achieving effective pipeline processing of different operators, thereby improving execution efficiency.

[0164] Secondly, data exchange between tensor cores and vector cores is carried out through registers, avoiding frequent entry and exit of interactive data in the on-chip cache, thereby improving the overall performance of the operator.

[0165] Step 305: Obtain query gradient blocks based on the reverse normalization result and the key blocks; and obtain key gradient blocks based on the reverse normalization result and the query blocks.

[0166] Specifically, a key block is a block in the target key tensor processed by a computation unit. The denormalized result is read from the on-chip cache and the key block is subjected to matrix multiplication to obtain the query gradient block. The denormalized result is read from the on-chip cache and the query block is subjected to matrix multiplication to obtain the key gradient block.

[0167] The value gradient block obtained in step 302 and the query gradient block and key gradient block obtained in step 305 are the results obtained in one inner loop.

[0168] In an embodiment of the present application, a forward recalculation operator is executed based on the query block and the key block to obtain a forward normalization result; then a first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain an intermediate result, wherein, when executing the first matrix multiplication operator, a dot product sum operator is executed in parallel based on the forward output block and the output gradient block to obtain an intermediate tensor, so that the calculation time of the first matrix multiplication operator and the dot product sum operator overlaps.

[0169] Then, the first type of transformation operator is executed on the value intermediate result to obtain the value gradient block; at the same time, the gradient backpropagation operator is executed in parallel based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain the reverse normalization result, so that the calculation time of the first type of transformation operator and the gradient backpropagation operator overlaps.

[0170] Finally, based on the reverse normalization results and key blocks, the query gradient blocks are obtained; and based on the reverse normalization results and query blocks, the key gradient blocks are obtained. This enables effective pipeline processing of different operators in the reverse attention mechanism calculation process. Compared with the serial method, the execution efficiency of the operators is improved, thereby improving the overall performance of the attention mechanism calculation.

[0171] In some embodiments, the present application also provides another implementation method for obtaining query gradient blocks and key gradient blocks:

[0172] A fourth matrix multiplication operator is executed based on the inverse normalization result and the key block to obtain a query intermediate result; and a second type of transformation operator is executed on the query intermediate result to obtain a query gradient block.

[0173] Next, the fifth matrix multiplication operator is executed based on the inverse normalization result and the query block to obtain the key intermediate result, wherein the fifth matrix multiplication operator and the second type of transformation operator are executed in parallel; the third type of transformation operator is executed on the key intermediate result to obtain the key gradient block.

[0174] Specifically, the Tensor Core reads the denormalized result from the on-chip cache. The denormalized result and the key block are then fed into the fourth matrix multiplication operator for computation, yielding the query intermediate result. After the fourth matrix multiplication operator completes, the Vector Core performs the second type of transformation on the query intermediate result, yielding the query gradient block.

[0175] In the process of executing the second type of transformation operator through the vector core, the inverse normalization result is read from the on-chip cache in parallel through the tensor core, and then the inverse normalization result and the query block are input into the fifth matrix multiplication operator for calculation to obtain the key intermediate result.

[0176] After the second type transformation operator and the fifth matrix multiplication operator are executed, the third type transformation operator is executed by vector checking the key intermediate result to obtain the key gradient block.

[0177] In an embodiment of the present application, in the process of executing the second type of transformation operator through the vector core, the inverse normalization result and the query block are input into the fifth matrix multiplication operator in parallel through the tensor core for calculation to obtain the key intermediate result, so that the calculation time of the second type of transformation operator and the fifth matrix multiplication operator overlaps, realizing effective pipeline processing of different operators, thereby improving the execution efficiency of the operators.

[0178] In some embodiments, the query block corresponds to multiple inner loops, and the key block and the value block are input data for the current inner loop. The third type of transformation operator is executed in parallel with the forward recalculation operator in the next inner loop corresponding to the query block.

[0179] Specifically, each query block corresponds to a middle-level loop, and a middle-level loop includes multiple inner-level loops. The key blocks and value blocks input into the multiple inner-level loops are different.

[0180] When the current inner loop executes the third type of transformation operator, the next inner loop starts in parallel. That is, the forward recalculation operator is executed based on the query block (the same query block as the query block used in the current inner loop) and the next key block (different from the key block used in the current inner loop) to obtain the forward normalization result.

[0181] It should be noted that if this round of inner loop is the last inner loop in multiple rounds of inner loops corresponding to the query block, then when this round of inner loop executes the third type of transformation operator, the first inner loop in the next round of middle loop will be executed in parallel, that is, based on the next query block (different from the query block used in this round of inner loop) and the key block input by the first inner loop (different from the key block used in this round of inner loop), the forward recalculation operator is executed to obtain the forward normalization result.

[0182] In an embodiment of the present application, effective pipeline processing of different operators is implemented between multiple layers of loops in the reverse attention mechanism calculation, which improves the execution efficiency of the operators compared to the serial method, thereby improving the overall performance of the attention mechanism calculation.

[0183] In order to better explain the embodiment of the present application, the following describes the process of an optimization method for attention mechanism calculation provided by the embodiment of the present application in combination with the specific implementation scenario of the reverse fast attention mechanism operator. The process of this method is executed by an artificial intelligence chip, such as Figure 4 As shown:

[0184] The reverse fast attention mechanism operator is divided into multiple operator combinations, namely the first operator combination, the second operator combination, the third operator combination, the fourth operator combination, and the fifth operator combination.

[0185] The first operator combination includes: a forward recalculation operator 401 and a third type transformation operator 402 executed in parallel, and the forward recalculation operator 401 includes the following operations: MMA (matrix multiplication operation), mask operation, sub (subtraction), exp (exponential function), div (division).

[0186] The third type transformation operator 402 includes the following operations: type conversion, specifically converting FP32 to BF16.

[0187] The second operator combination includes: a first matrix multiplication operator 403 and a dot product summation operator 404 executed in parallel.

[0188] The first matrix multiplication operator 403 includes the following operations: MMA. The dot product summation operator 404 includes the following operations: mul (dot product) and row summation operation.

[0189] The third operator combination includes: a first type transformation operator 405 and a gradient back-propagation operator 406 executed in parallel.

[0190] The first type transformation operator 405 includes the following operations: type conversion, specifically converting FP32 to BF16. The gradient back propagation operator 406 includes the following operations: MMA, sub, and mul.

[0191] The fourth operator combination includes: a fourth matrix multiplication operator 407;

[0192] The fourth matrix multiplication operator 407 includes the following operations: MMA.

[0193] The fifth operator combination includes: a second type transformation operator 408 and a fifth matrix multiplication operator 409 executed in parallel.

[0194] The second type transformation operator 408 includes the following operations: type conversion, specifically converting FP32 to BF16. The fifth matrix multiplication operator 409 includes the following operations: MMA.

[0195] Next, we will take one inner loop as an example to introduce the execution process of each operator combination in the reverse fast attention mechanism operator:

[0196] The output gradient block dO_chunk is loaded from the video memory into the on-chip cache. During the loading process, the first operator combination is executed in parallel. The execution process of the first operator combination includes:

[0197] The query chunk Q_chunk, the key chunk K_chunk, the intermediate result chunk m_chunk and the intermediate result chunk s_chunk are input into the forward recalculation operator 401 for calculation to obtain a forward normalization result P.

[0198] Specifically, the query block Q_chunk is divided into T first sub-blocks q_chunk; the key block K_chunk is divided into T second sub-blocks k_chunk; the intermediate result block m_chunk is further split into T intermediate sub-blocks m'_chunk, and the intermediate result block s_chunk is further split into T intermediate sub-blocks s'_chunk, where T is greater than 1.

[0199] Perform T rounds of calculation based on the split data. The input data for each round of calculation includes: a first sub-block q_chunk and an associated second sub-block k_chunk, an intermediate sub-block m'_chunk, and an intermediate sub-block s'_chunk. The specific process of each round of calculation is:

[0200] MMA is performed on the first sub-block q_chunk and the second sub-block k_chunk through the tensor core to obtain a sub-matrix multiplication result dp1, and the sub-matrix multiplication result dp1 is saved in a register.

[0201] Next, the vector core is notified to read the sub-matrix multiplication result dp1 from the register as a consumer, and then performs mask operations, sub, exp, div in sequence based on the sub-matrix multiplication result dp1, the intermediate sub-block m'_chunk and the intermediate sub-block s'_chunk according to the above formula (1) to obtain the sub-normalized result p', and save the sub-normalized result p' in the on-chip cache.

[0202] The T sub-normalization results p' are concatenated to obtain the forward normalization result P.

[0203] During the execution of the forward recalculation operator 401 , the third type transformation operator 402 is executed in parallel on the key intermediate result obtained in the previous inner loop to obtain the key gradient block output by the previous inner loop.

[0204] After the forward recalculation operator 401 and the third type transformation operator 402 are both completed, the subsequent second operator combination is executed.

[0205] The execution process of the second operator combination includes:

[0206] The forward normalization result P and the output gradient block dO_chunk are input into the first matrix multiplication operator 403 through the tensor core for calculation to obtain the intermediate result v1.

[0207] The forward output block O_chunk and the output gradient block dO_chunk are input to the dot product summation operator 404 through the vector core to obtain an intermediate tensor D, and the intermediate tensor D is stored in the on-chip cache.

[0208] After the first matrix multiplication operator 403 and the dot product summation operator 404 are both calculated, the subsequent third operator combination is executed.

[0209] The execution process of the third operator combination:

[0210] The intermediate value result v1 is input into the first type transformation operator 405 through the vector kernel for calculation to obtain the value gradient block dV_chunk of the current inner loop.

[0211] The output gradient block dO_chunk, the value block V_chunk, the intermediate tensor D and the forward normalization result P are input into the gradient back-propagation operator 406 for calculation to obtain the reverse normalization result dS.

[0212] After the calculation of the first type transformation operator 405 and the gradient back propagation operator 406 is completed, the subsequent fourth operator combination is executed.

[0213] The execution process of the fourth operator combination:

[0214] The inverse normalization result dS and the key block K_chunk are input into the fourth matrix multiplication operator 407 for calculation to obtain the query intermediate result q1.

[0215] The execution process of the fifth operator combination:

[0216] The second type transformation operator 408 is executed on the vector check query intermediate result q1 to obtain the query gradient block dQ_chunk of the current inner loop.

[0217] The inverse normalization result dS and the query block Q_chunk are input into the fifth matrix multiplication operator 409 for calculation to obtain the key intermediate result k1.

[0218] After the execution of the fifth operator combination is completed, the third type transformation operator 402 is executed through the vector verification key intermediate result k1 to obtain the key gradient block dK_chunk of this inner loop; at the same time, the forward recalculation operator 401 in the next inner loop process is executed in parallel. The execution process of the next inner loop is the same as that of the current inner loop, and will not be repeated here.

[0219] In an embodiment of the present application, a forward recalculation operator is executed based on the query block and the key block to obtain a forward normalization result; then a first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain an intermediate result, wherein, when executing the first matrix multiplication operator, a dot product sum operator is executed in parallel based on the forward output block and the output gradient block to obtain an intermediate tensor, so that the calculation time of the first matrix multiplication operator and the dot product sum operator overlaps.

[0220] Then, the first type of transformation operator is executed on the value intermediate result to obtain the value gradient block; at the same time, the gradient backpropagation operator is executed in parallel based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain the reverse normalization result, so that the calculation time of the first type of transformation operator and the gradient backpropagation operator overlaps.

[0221] Finally, based on the reverse normalization results and key blocks, the query gradient blocks are obtained; and based on the reverse normalization results and query blocks, the key gradient blocks are obtained. This enables effective pipeline processing of different operators in the reverse attention mechanism calculation process. Compared with the serial method, the execution efficiency of the operators is improved, thereby improving the overall performance of the attention mechanism calculation.

[0222] Based on the same technical concept, the embodiment of the present application provides a structural diagram of an optimization device for attention mechanism calculation, such as Figure 5 As shown, the optimization device 500 for attention mechanism calculation includes:

[0223] A recalculation module 501 is configured to execute a forward recalculation operator based on the query block and the key block to obtain a forward normalized result;

[0224] A back propagation module 502 is configured to perform a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; and perform a first type of transformation operator on the value intermediate result to obtain a value gradient block.

[0225] The back propagation module 502 is further configured to execute a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product sum operator and the first matrix multiplication operator are executed in parallel;

[0226] The back propagation module 502 is further configured to execute a gradient back propagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a back normalization result, wherein the gradient back propagation operator and the first type transformation operator are executed in parallel;

[0227] The back propagation module 502 is further configured to obtain a query gradient block based on the back normalization result and the key block; and obtain a key gradient block based on the back normalization result and the query block.

[0228] Optionally, the recalculation module 501 is further configured to:

[0229] During execution of the forward recalculation operator, the output gradient blocks are loaded from the video memory into an on-chip cache in parallel.

[0230] Optionally, the forward recalculation operator includes: a second matrix multiplication operator and a normalization operator;

[0231] The recalculation module 501 is specifically used for:

[0232] Dividing the query block into a plurality of first sub-blocks; and dividing the key block into a plurality of second sub-blocks, wherein each first sub-block corresponds to a second sub-block;

[0233] For each first sub-block and the corresponding second sub-block, respectively, performing: executing a second matrix multiplication operator based on the first sub-block and the corresponding second sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core and executing the normalization operator to obtain a sub-normalization result;

[0234] The forward normalization result is generated based on the obtained multiple sub-normalization results.

[0235] Optionally, the recalculation module 501 is further configured to:

[0236] The normalization operator is executed by reading the sub-matrix multiplication result from the register through a vector core, and after obtaining the sub-normalization result, the sub-normalization result is stored in an on-chip cache.

[0237] Optionally, the recalculation module 501 is further configured to:

[0238] During execution of the forward recalculation operator, register synchronization is adopted between the tensor core and the vector core.

[0239] Optionally, the gradient back-propagation operator includes: a third matrix multiplication operator and a dot product difference operator;

[0240] The back propagation module 502 is specifically used for:

[0241] Dividing the output gradient block into a plurality of third sub-blocks; and dividing the value block into a plurality of fourth sub-blocks, wherein each third sub-block corresponds to one fourth sub-block;

[0242] For each third sub-block and the corresponding fourth sub-block, respectively, performing: executing a third matrix multiplication operator based on a third sub-block and the corresponding fourth sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core, and executing the dot product difference operator on the intermediate tensor and the forward normalization result to obtain a sub-dot product result;

[0243] The reverse normalization result is generated based on the obtained multiple sub-point product results.

[0244] Optionally, the back propagation module 502 is further configured to:

[0245] The sub-matrix multiplication result is read from the register through the vector core, the dot product difference operator is executed with the intermediate tensor and the forward normalization result, and after the sub-dot product result is obtained, the sub-dot product result is stored in the on-chip cache.

[0246] Optionally, the back propagation module 502 is specifically configured to:

[0247] Executing a fourth matrix multiplication operator based on the inverse normalization result and the key block to obtain a query intermediate result; executing a second type transformation operator on the query intermediate result to obtain the query gradient block;

[0248] executing a fifth matrix multiplication operator based on the inverse normalization result and the query block to obtain a key intermediate result, wherein the fifth matrix multiplication operator and the second-type transformation operator are executed in parallel;

[0249] A third type of transformation operator is performed on the key intermediate result to obtain the key gradient block.

[0250] Optionally, the query blocks correspond to multiple rounds of inner loops, and the key blocks and the value blocks are input data of the current round of inner loops;

[0251] The third type transformation operator is executed in parallel with the forward recalculation operator in the next inner loop corresponding to the query block.

[0252] Optionally, a splitting module 503 is further included;

[0253] The query block is a block in the target query tensor processed by a computing unit; the computing unit is any one of the multiple computing units included in the artificial intelligence chip;

[0254] The splitting module 503 is specifically used for:

[0255] Get the original query tensor, where the shape of the original query tensor includes: the number of batches and attention heads combined, and the sequence length;

[0256] Obtaining a number of splits based on the sequence length and the hardware alignment granularity of the artificial intelligence chip;

[0257] When the number of merges is greater than the number of splits, and the ratio of the number of merges to the total number of computing units is an integer, splitting the original query tensor into target query tensors processed by each of the plurality of computing units according to the number of merges and the total number of computing units;

[0258] When the number of merges is not greater than the number of splits, and the ratio of the number of splits to the total number of computing units is an integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the sequence length and the total number of computing units.

[0259] Optionally, the splitting module 503 is further configured to:

[0260] When the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is a non-integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the number of merges, the sequence length, and the total number of computing units.

[0261] In an embodiment of the present application, a forward recalculation operator is executed based on the query block and the key block to obtain a forward normalization result; then a first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain an intermediate result, wherein, when executing the first matrix multiplication operator, a dot product sum operator is executed in parallel based on the forward output block and the output gradient block to obtain an intermediate tensor, so that the calculation time of the first matrix multiplication operator and the dot product sum operator overlaps.

[0262] Then, the first type of transformation operator is executed on the value intermediate result to obtain the value gradient block; at the same time, the gradient backpropagation operator is executed in parallel based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain the reverse normalization result, so that the calculation time of the first type of transformation operator and the gradient backpropagation operator overlaps.

[0263] Finally, based on the reverse normalization results and key blocks, the query gradient blocks are obtained; and based on the reverse normalization results and query blocks, the key gradient blocks are obtained. This enables effective pipeline processing of different operators in the reverse attention mechanism calculation process. Compared with the serial method, the execution efficiency of the operators is improved, thereby improving the overall performance of the attention mechanism calculation.

[0264] Based on the same technical concept, the embodiment of the present application provides a computer device, such as Figure 6 As shown, it includes at least one artificial intelligence chip 100 and a memory 601 connected to the at least one artificial intelligence chip 100. The specific connection medium between the artificial intelligence chip 100 and the memory 601 is not limited in the embodiment of the present application. Figure 6For example, the artificial intelligence chip 100 and the memory 601 are connected via a bus. The bus can be divided into an address bus, a data bus, a control bus, etc.

[0265] In an embodiment of the present application, the memory 601 stores instructions that can be executed by at least one artificial intelligence chip 100. At least one artificial intelligence chip 100 can execute the steps of the above-mentioned attention mechanism calculation optimization method by executing the instructions stored in the memory 601.

[0266] Among them, the artificial intelligence chip 100 is the control center of the computer device. It can use various interfaces and lines to connect various parts of the computer device. By running or executing instructions stored in the memory 601 and calling data stored in the memory 601, it can optimize the calculation of the attention mechanism. Optionally, the artificial intelligence chip 100 may include one or more processing units. The artificial intelligence chip 100 may integrate an application processor and a modem processor, wherein the application processor mainly processes the operating system, user interface, and application programs, etc., and the modem processor mainly processes wireless communications. It is understandable that the above-mentioned modem processor may not be integrated into the artificial intelligence chip 100. In some embodiments, the artificial intelligence chip 100 and the memory 601 can be implemented on the same chip. In some embodiments, they can also be implemented separately on independent chips.

[0267] The artificial intelligence chip 100 can be a general-purpose processor, such as a central processing unit (CPU), a digital signal processor, an application-specific integrated circuit (ASIC), a field-programmable gate array or other programmable logic device, a discrete gate or transistor logic device, or a discrete hardware component, and can implement or execute the various methods, steps, and logic block diagrams disclosed in the embodiments of this application. A general-purpose processor can be a microprocessor or any conventional processor. The steps of the methods disclosed in the embodiments of this application can be directly implemented and executed by a hardware processor, or by a combination of hardware and software modules in the processor.

[0268] The memory 601 is a non-volatile computer-readable storage medium that can be used to store non-volatile software programs, non-volatile computer executable programs, and modules. The memory 601 may include at least one type of storage medium, such as a flash memory, a hard disk, a multimedia card, a card-type memory, a random access memory (RAM), a static random access memory (SRAM), a programmable read-only memory (PROM), a read-only memory (ROM), an electrically erasable programmable read-only memory (EEPROM), a magnetic storage device, a magnetic disk, an optical disk, and the like. The memory 601 is any other medium that can be used to carry or store desired program code in the form of instructions or data structures and can be accessed by a computer device, but is not limited thereto. The memory 601 in the embodiments of the present application may also be a circuit or any other device capable of performing a storage function, used to store program instructions and / or data.

[0269] Based on the same inventive concept, an embodiment of the present application provides a computer-readable storage medium, which stores a computer program that can be executed by a computer device. When the computer program runs on the computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0270] Based on the same inventive concept, an embodiment of the present application provides a computer program product, which includes a computer program stored on a computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer device, the computer device performs the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0271] Those skilled in the art will appreciate that embodiments of the present invention may be provided as methods or computer program products. Thus, the present invention may take the form of an entirely hardware embodiment, an entirely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, the present invention may take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to magnetic disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0272] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowcharts and / or block diagrams, as well as combinations of processes and / or blocks in the flowcharts and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer device or other programmable data processing device generate instructions for implementing the processes in the flowcharts and / or block diagrams. Figure 1 a process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0273] These computer program instructions may also be stored in a computer readable memory that can direct a computer device or other programmable data processing device to work in a specific manner, so that the instructions stored in the computer readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 a process or multiple processes and / or boxes Figure 1 The function specified in one or more boxes.

[0274] These computer program instructions can also be loaded onto a computer device or other programmable data processing device so that a series of operating steps are executed on the computer device or other programmable device to produce a process implemented by the computer device, thereby providing instructions for implementing the process. Figure 1 a process or multiple processes and / or boxes Figure 1 A step that specifies a function in one or more boxes.

[0275] Although the preferred embodiments of the present invention have been described, those skilled in the art may make additional changes and modifications to these embodiments once they have learned the basic creative concept. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications that fall within the scope of the present invention.

[0276] Obviously, those skilled in the art may make various changes and modifications to the present invention without departing from the spirit and scope of the present invention. Thus, if such changes and modifications fall within the scope of the claims and their equivalents, the present invention is intended to include such changes and modifications.

Claims

1. A method for optimizing attention mechanism calculation, characterized in that: include: Execute forward recalculation operators based on query blocks and key blocks to obtain forward normalized results; Performing a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; Performing a first type of transformation operator on the value intermediate result to obtain value gradient blocks; Executing a dot product sum operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product sum operator and the first matrix multiplication operator are executed in parallel; Executing a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a backnormalization result, wherein the gradient backpropagation operator and the first type transformation operator are executed in parallel; Obtaining query gradient blocks based on the reverse normalization result and the key blocks; and, obtaining a key gradient block based on the reverse normalization result and the query block; The forward recalculation operator includes: a second matrix multiplication operator and a normalization operator; the forward recalculation operator is executed based on the query block and the key block to obtain a forward normalization result, including: Dividing the query block into a plurality of first sub-blocks; and dividing the key block into a plurality of second sub-blocks, wherein each first sub-block corresponds to a second sub-block; For each first sub-block and the corresponding second sub-block, respectively, performing: executing a second matrix multiplication operator based on the first sub-block and the corresponding second sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core to execute the normalization operator, and in the process of obtaining the sub-normalization result, executing the second matrix multiplication operator based on the next sub-block in parallel by the tensor core; The forward normalization result is generated based on the obtained multiple sub-normalization results.

2. The method according to claim 1, wherein Also includes: During execution of the forward recalculation operator, the output gradient blocks are loaded from the video memory into an on-chip cache in parallel.

3. The method according to claim 1, wherein After the sub-matrix multiplication result is read from the register by the vector core and the normalization operator is executed to obtain the sub-normalization result, the method further includes: The sub-normalized results are stored in an on-chip cache.

4. The method according to claim 1, wherein Also includes: During execution of the forward recalculation operator, register synchronization is adopted between the tensor core and the vector core.

5. The method according to claim 1, wherein The gradient back propagation operator includes: a third matrix multiplication operator and a dot product difference operator; The performing a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain a backnormalization result includes: Dividing the output gradient block into a plurality of third sub-blocks; and dividing the value block into a plurality of fourth sub-blocks, wherein each third sub-block corresponds to one fourth sub-block; For each third sub-block and the corresponding fourth sub-block, respectively, performing: executing a third matrix multiplication operator based on a third sub-block and the corresponding fourth sub-block by a tensor core to obtain a sub-matrix multiplication result; storing the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by a vector core, and executing the dot product difference operator on the intermediate tensor and the forward normalization result to obtain a sub-dot product result; The reverse normalization result is generated based on the obtained multiple sub-point product results.

6. The method according to claim 5, wherein The method further includes: reading the sub-matrix multiplication result from the register through the vector core, performing the dot product difference operator on the intermediate tensor and the forward normalization result, and obtaining the sub-dot product result; The sub-point product result is stored in an on-chip cache.

7. The method according to claim 1, wherein Obtaining query gradient blocks based on the reverse normalization result and the key blocks; And, based on the reverse normalization result and the query block, obtaining a key gradient block, including: Executing a fourth matrix multiplication operator based on the inverse normalization result and the key block to obtain a query intermediate result; executing a second type transformation operator on the query intermediate result to obtain the query gradient block; executing a fifth matrix multiplication operator based on the inverse normalization result and the query block to obtain a key intermediate result, wherein the fifth matrix multiplication operator and the second-type transformation operator are executed in parallel; A third type of transformation operator is performed on the key intermediate result to obtain the key gradient block.

8. The method according to claim 7, wherein The query blocks correspond to multiple rounds of inner loops, and the key blocks and the value blocks are input data of this round of inner loops; The third type transformation operator is executed in parallel with the forward recalculation operator in the next inner loop corresponding to the query block.

9. The method according to any one of claims 1 to 8, characterized in that: The query block is a block in the target query tensor processed by a computing unit; the computing unit is any one of the multiple computing units included in the artificial intelligence chip, and the target query tensor processed by each computing unit is obtained in the following manner: Get the original query tensor, where the shape of the original query tensor includes: the number of batches and attention heads combined, and the sequence length; Obtaining a number of splits based on the sequence length and the hardware alignment granularity of the artificial intelligence chip; When the number of merges is greater than the number of splits, and the ratio of the number of merges to the total number of computing units is an integer, splitting the original query tensor into target query tensors processed by each of the plurality of computing units according to the number of merges and the total number of computing units; When the number of merges is not greater than the number of splits, and the ratio of the number of splits to the total number of computing units is an integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the sequence length and the total number of computing units.

10. The method according to claim 9, wherein Also includes: When the ratio of the number of merges to the total number of computing units is a non-integer, and the ratio of the number of splits to the total number of computing units is a non-integer, the original query tensor is split into target query tensors processed by each of the multiple computing units according to the number of merges, the sequence length, and the total number of computing units.

11. A computer device comprising a memory, an artificial intelligence chip, and a computer program stored in the memory and running on the artificial intelligence chip, characterized in that: When the artificial intelligence chip executes the computer program, the steps of the method according to any one of claims 1 to 10 are implemented.

12. A computer-readable storage medium, characterized in that It stores a computer program executed by a computer device, and when the computer program is run on the computer device, the computer device executes the steps of the method according to any one of claims 1 to 10.

13. A computer program product, characterized in that The computer program product comprises a computer program stored on a computer-readable storage medium, wherein the computer program comprises program instructions. When the program instructions are executed by a computer device, the computer device is caused to perform the steps of the method according to any one of claims 1 to 10.

Images