Attention mechanism calculation optimization method and device, storage medium and program product

Abstract

The embodiment of the invention provides an attention mechanism calculation optimization method and device, a storage medium and a program product, and relates to the technical field of artificial intelligence, and the method comprises the steps: executing a forward re-calculation operator based on a query block and a key block, and obtaining 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 executing a point multiplication summation operator in parallel based on the forward output block and the output gradient block to obtain an intermediate tensor; executing a first type of transformation operator on the value intermediate result to obtain value gradient blocks; and meanwhile, a gradient back propagation operator is executed in parallel based on the output gradient block, the value block, the intermediate tensor and the forward normalization result, and a reverse normalization result is obtained. And finally, on the basis of a reverse normalization result, the key blocks and the query blocks, obtaining query gradient blocks and key gradient blocks, so that effective pipeline processing is performed on different operators, and the overall performance of attention mechanism calculation is improved.

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 calculating an attention mechanism. Background Art

[0002] In the field of deep learning, the calculation of the attention mechanism is widely present in various models of deep learning. For example, the core calculation part of the Transformer model is the calculation of the attention mechanism, and large language models generally use the Transformer model as the core infrastructure. Therefore, in the era when large language models are widely used, the execution performance of the attention mechanism calculation becomes particularly important.

[0003] Under the related technology, in the process of performing the reverse attention mechanism calculation, a plurality of operators associated with the reverse attention mechanism calculation are sequentially executed in a serial manner. This serial manner will result in low execution efficiency of the operators, and thus 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 calculating an attention mechanism, which are used to improve the performance of the attention mechanism calculation.

[0005] On the one hand, the embodiments of the present application provide an optimization method for calculating an attention mechanism, including: Executing a forward recomputation operator based on a query block and a key block to obtain a forward normalization result; Executing a first matrix multiplication operator based on an output gradient block and the forward normalization result to obtain a value intermediate result; and executing a first type of transformation operator on the value intermediate result to obtain a value gradient block; Executing a dot product summation operator based on a forward output block and the output gradient block to obtain an intermediate tensor; the dot product summation operator and the first matrix multiplication operator are executed in parallel; Executing a gradient backpropagation operator based on the output gradient block, a value block, the intermediate tensor and the forward normalization result to obtain a reverse normalization result, and the gradient backpropagation operator and the first type of transformation operator are executed in parallel; Obtaining a query gradient block based on the reverse normalization result and the key block; and obtaining a key gradient block based on the reverse normalization result and the query block.

[0006] On the one hand, the embodiments of the present application provide an optimization device for calculating an attention mechanism, including: A recomputation module, configured to execute a forward recomputation operator based on a query block and a key block to obtain a forward normalization result; A backpropagation module, configured to perform a first matrix multiplication operator based on the output gradient chunks and the forward normalization result to obtain an intermediate value result; and perform a first type of transformation operator on the intermediate value result to obtain value gradient chunks. The backpropagation module is further configured to perform a dot product summation operator based on the forward output chunks and the output gradient chunks to obtain an intermediate tensor; the dot product summation operator and the first matrix multiplication operator are executed in parallel. The backpropagation module is further configured to perform a gradient backpropagation operator based on the output gradient chunks, value chunks, the intermediate tensor, and the forward normalization result to obtain a reverse normalization result; the gradient backpropagation operator and the first type of transformation operator are executed in parallel. The backpropagation module is further configured to obtain query gradient chunks based on the reverse normalization result and the key chunks; and obtain key gradient chunks based on the reverse normalization result and the query chunks.

[0007] Optionally, the recomputation module is further configured to: During the execution of the forward recomputation operator, load the output gradient chunks from the video memory to the on-chip cache in parallel.

[0008] Optionally, the forward recomputation operator includes: a second matrix multiplication operator and a normalization operator. The recomputation module is specifically configured to: Divide the query chunks into multiple first sub-chunks; and divide the key chunks into multiple second sub-chunks, where each first sub-chunk corresponds to a second sub-chunk. For each first sub-chunk and the corresponding second sub-chunk, respectively perform: perform a second matrix multiplication operator on one first sub-chunk and the corresponding second sub-chunk through a tensor core to obtain a sub-matrix multiplication result; save the sub-matrix multiplication result in a register; read the sub-matrix multiplication result from the register through a vector core to perform the normalization operator to obtain a sub-normalization result. Generate the forward normalization result based on the obtained multiple sub-normalization results.

[0009] Optionally, the recomputation module is further configured to: After obtaining the sub-normalization result by reading the sub-matrix multiplication result from the register through a vector core to perform the normalization operator, save the sub-normalization result in the on-chip cache.

[0010] Optionally, the recomputation module is further configured to: During the execution of the forward recomputation operator, use register synchronization between the tensor core and the vector core.

[0011] Optionally, the gradient backpropagation operator includes: a third matrix multiplication operator and a dot product difference operator; The backpropagation module is specifically configured to: Partition the output gradient into a plurality of third sub-blocks; and partition the value into a plurality of fourth sub-blocks, where each third sub-block corresponds to a fourth sub-block; For each third sub-block and the corresponding fourth sub-block, respectively perform: perform the third matrix multiplication operator on a third sub-block and the corresponding fourth sub-block based on a tensor core to obtain a sub-matrix multiplication result; save the sub-matrix multiplication result in a register; read the sub-matrix multiplication result from the register through a vector core, and perform the dot product difference operator on the sub-matrix multiplication result, the intermediate tensor, and the forward normalization result to obtain a sub-dot product result; Generate the back normalization result based on the obtained plurality of sub-dot product results.

[0012] Optionally, the backpropagation module is further configured to: After reading the sub-matrix multiplication result from the register through a vector core and performing the dot product difference operator on the sub-matrix multiplication result, the intermediate tensor, and the forward normalization result to obtain a sub-dot product result, save the sub-dot product result in the on-chip cache.

[0013] Optionally, the backpropagation module is specifically configured to: Perform a fourth matrix multiplication operator on the back normalization result and the key block to obtain a query intermediate result; perform a second type of transformation operator on the query intermediate result to obtain the query gradient block; Perform a fifth matrix multiplication operator on the back normalization result and the query block to obtain a key intermediate result, where the fifth matrix multiplication operator and the second type of transformation operator are executed in parallel; Perform a third type of transformation operator on the key intermediate result to obtain the key gradient block.

[0014] Optionally, the query block corresponds to multiple rounds of inner loops, and the key block and the value block are the input data for the current round of inner loops; The third type of transformation operator and the forward recomputation operator in the next round of inner loops corresponding to the query block are executed in parallel.

[0015] Optionally, a splitting module is further included; The query block is a block in a target query tensor processed by a computing unit; the computing unit is any one of a plurality of computing units included in the artificial intelligence chip; The splitting module is specifically configured to: Obtain the original query tensor, where the shape of the original query tensor includes: the combined number of batches and attention heads, and the sequence length; Based on the sequence length and the hardware alignment granularity of the artificial intelligence chip, obtain the splitting number; When the combined number is greater than the splitting number, and the ratio of the combined number to the total number of computing units is an integer, split the original query tensor into target query tensors processed by each of the multiple computing units according to the combined number and the total number of computing units; When the combined number is not greater than the splitting number, and the ratio of the splitting number to the total number of computing units is an integer, split the original query tensor into target query tensors processed by each of the multiple computing units according to the sequence length and the total number of computing units.

[0016] Optionally, the splitting module is further configured to: When the ratio of the combined number to the total number of computing units is a non-integer, and the ratio of the splitting number to the total number of computing units is a non-integer, split the original query tensor into target query tensors processed by each of the multiple computing units according to the combined number, the sequence length, and the total number of computing units.

[0017] On the one hand, an embodiment of the present application provides a computer device, including 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 optimization method for calculating the attention mechanism is implemented, including: The artificial intelligence chip reads the query block and the key block from the memory, and executes a forward recomputation operator based on the query block and the key block to obtain a forward normalization result; The artificial intelligence chip executes a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; and executes a first type of transformation operator on the value intermediate result to obtain a value gradient block; The artificial intelligence chip executes a dot product summation operator based on the forward output block and the output gradient block to obtain an intermediate tensor; the dot product summation operator and the first matrix multiplication operator are executed in parallel; 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 backward normalization result, and the gradient backpropagation operator and the first type of transformation operator are executed in parallel; The artificial intelligence chip obtains a query gradient block based on the backward normalization result and the key block; and obtains a key gradient block based on the backward normalization result and the query block.

[0018] On the one hand, an embodiment of the present application provides a computer-readable storage medium storing a computer program executable by a computer device. When the computer program runs on the computer device, the computer device is caused to execute the steps of the optimization method for attention mechanism calculation described above.

[0019] On the one hand, an embodiment of the present application provides a computer program product. The computer program product includes a computer program stored on a computer-readable storage medium. The computer program includes program instructions. When the program instructions are executed by a computer device, the computer device is caused to execute the steps of the optimization method for attention mechanism calculation described above.

[0020] In an embodiment of the present application, a forward recomputation operator is executed based on query chunks and key chunks to obtain a forward normalization result; then a first matrix multiplication operator is executed based on the output gradient chunks and the forward normalization result to obtain a value intermediate result. When executing the first matrix multiplication operator, a dot product summation operator is executed in parallel based on the forward output chunks and the output gradient chunks to obtain an intermediate tensor, so that the calculation times of the first matrix multiplication operator and the dot product summation operator overlap.

[0021] Next, a first type of transformation operator is executed on the value intermediate result to obtain value gradient chunks; at the same time, a gradient backpropagation operator is executed in parallel based on the output gradient chunks, value chunks, intermediate tensor, and forward normalization result to obtain a backward normalization result, so that the calculation times of the first type of transformation operator and the gradient backpropagation operator overlap.

[0022] Finally, query gradient chunks are obtained based on the backward normalization result and key chunks; and key gradient chunks are obtained based on the backward normalization result and query chunks, realizing effective pipeline processing of different operators during the backward attention mechanism calculation. Compared with the serial method, the execution efficiency of the operators is improved, thereby enhancing the overall performance of the attention mechanism calculation. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] To more clearly illustrate the technical solutions in the embodiments of the present invention, the following briefly introduces the accompanying drawings required for the description of the embodiments. Obviously, the following described accompanying drawings are only some embodiments of the present invention. For those of ordinary skill in the art, other accompanying drawings can be obtained without creative efforts based on these drawings.

[0024] Figure 1 It is a schematic structural diagram of an artificial intelligence chip provided by an embodiment of the present application; Figure 2 It is a schematic diagram of the input and output data of the backward attention mechanism operator provided by an embodiment of the present application; Figure 3 Flow schematic of an optimization method for calculating an attention mechanism provided by an embodiment of the present application Figure 1 ; Figure 4 Flow schematic of an optimization method for calculating an attention mechanism provided by an embodiment of the present application Figure 2 ; Figure 5 Structural schematic diagram of an optimization device for calculating an attention mechanism provided by an embodiment of the present application; Figure 6 Structural schematic diagram of a computer device provided by an embodiment of the present application. Detailed implementation manners

[0025] In order to make the objectives, technical solutions and beneficial effects of the present invention clearer and more understandable, 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 used to limit the present invention.

[0026] Some simple introductions will be made below to the system architecture diagram applicable to the technical solutions of the embodiments of the present application. It should be noted that the system architecture diagrams introduced below are only used to illustrate the embodiments of the present application and are not limiting.

[0027] 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 at least includes: a video memory 101 and a plurality of computing units 102. Among them, the computing unit 102 can be a Streaming Processing Cluster (SPC for short).

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

[0029] The video memory 101 can be a High Bandwidth Memory (HBM for short), or other types of memories. The on-chip cache 103 is a temporary memory, its capacity is smaller than that of the video memory 101, but the data exchange speed is faster than that of the video memory 101. The on-chip cache 103 can be a Gemm Main Buffer (GMB for short), or other types of memories.

[0030] Register 104 is a register shared by tensor core 105 and vector core 106. The tensor core 105 and the vector core 106 can perform data interaction through the register 104. The register 104 can be a Thread-Local Register (TLR for short). Compared with the on-chip cache 103, the capacity of the register 104 is smaller than that of the on-chip cache 103, but the data exchange speed is faster than that of the on-chip cache 103.

[0031] In the embodiment of the present application, multiple computing units 102 execute the attention mechanism calculation in parallel. Taking one computing unit as an example, the tensor core 105 and the vector core 106 in the computing unit 102 execute the forward recomputation operator based on the query block and the key block to obtain the forward normalization result. During the execution of the forward recomputation operator, the tensor core 105 and the vector core 106 transfer the results obtained from the intermediate calculation through the register 104, and save the calculation results output by the forward recomputation operator and required by subsequent operators through the on-chip cache 103, without writing the results obtained from the intermediate calculation 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, improving the utilization rate of the storage resources 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.

[0032] Then, the tensor core 105 executes the first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain the value intermediate result. When the tensor core 105 executes the first matrix multiplication operator, the vector core 106 executes the dot product summation operator in parallel based on the forward output block and the output gradient block to obtain the intermediate tensor; this makes the time for the tensor core 105 to execute the first matrix multiplication operator overlap with the time for the vector core 106 to execute the dot product summation operator, avoiding the vector core 106 from being in a waiting state all the time when the tensor core 105 executes the operator, thereby improving the utilization rate of the computing resources in the computing unit 102 and at the same time improving the overall computing efficiency of the computing unit 102.

[0033] Next, the vector core 106 performs a first type of transformation operator on the value intermediate result to obtain a value gradient block; at the same time, 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 reverse normalization result. During the execution of the backpropagation operator, the tensor core 105 and the vector core 106 transfer the results obtained from intermediate calculations through the register 104, and save the calculation results output by the backpropagation operator and required by subsequent operators through the on-chip cache 103, without writing the results obtained from intermediate calculations and the calculation results required by subsequent operators back to the video memory. In this way, both the register 104 and the on-chip cache 103 are fully utilized, improving the utilization rate of the storage resources of the computing unit 102; at the same time, the computing unit 102 is prevented from frequently accessing the video memory, thereby improving the computing efficiency of the computing unit 102. Secondly, the calculation times of the first type of transformation operator and the gradient backpropagation operator overlap, thereby improving the utilization rate of the computing resources in the computing unit 102 and at the same time improving the overall computing efficiency of the computing unit 102.

[0034] Finally, 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, realizing that during the calculation process of the reverse attention mechanism, the computing unit 102 performs effective pipelining on different operators. Compared with the serial method, the execution efficiency of the operators is improved, and the overall performance of the artificial intelligence chip 100 is also improved.

[0035] In addition to the above structure, the artificial intelligence chip 100 in this application may further include other structures, and specific limitations are not made in this application.

[0036] The artificial intelligence chip 100 may be: a Graphics Processing Unit (GPU), a General-purpose computing on graphics processing units (GPGPU), a Domain Specific Architecture (DSA), etc.

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

[0038] Attention mechanism calculation can be applied to various scenarios. For example, image processing scenarios, speech processing scenarios, text processing scenarios, etc.

[0039] 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 the speech data used in tasks such as speech enhancement, speech recognition, and speech synthesis.

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

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

[0042] In practical applications, the attention mechanism calculation includes: forward attention mechanism and backward attention mechanism calculations. In the embodiments of the present application, the optimization is mainly aimed at the backward attention mechanism calculation. During the iterative training process of the model including the attention mechanism calculation, the gradient information of the model needs to be calculated through backpropagation in each round of iteration, and then the model parameters are adjusted based on the gradient information. Among them, the backward attention mechanism calculation is included in the backpropagation process.

[0043] In the embodiments of the present application, referring to Figure 2 , the backward attention mechanism (Attention bwd) operator includes: 6 input tensors, which are (Q, K, V, O, dO, M1) respectively, 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.

[0044] The backward attention mechanism operator also includes 3 intermediate tensors and 3 output tensors. The 3 intermediate tensors are (m, s, RD) respectively, where m and s are the original intermediate results generated during the forward attention mechanism calculation, and RD is the original random mask generated during the forward attention mechanism calculation. The 3 intermediate tensors are also used as the input of the backward attention mechanism operator for calculation.

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

[0046] The shape of any one 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.

[0047] In an embodiment of the present application, an artificial intelligence chip includes multiple computing units, and the multiple computing units execute attention mechanism calculations in parallel. However, in the related art, for an original query tensor of any shape, a fixed splitting method (i.e., according to the sequence length) is used to split the original query tensor onto multiple computing units for calculation. When the sequence length of the original query tensor is small, the target query tensor split onto 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.

[0048] In view of this, in an embodiment of the present application, the central processing unit obtains the target query tensor processed by each computing unit in the following manner. The specific process is as follows: Obtain the original query tensor. The shape of the original query tensor includes: the combined number of batches and attention heads, and the sequence length, where the combined number of batches and attention heads refers to: That is, the product of the batch size and the number of attention heads.

[0049] Based on the sequence length and the hardware alignment granularity of the artificial intelligence chip, obtain the splitting number. Here, the hardware alignment granularity (block_size) represents: the minimum sequence length processed by an instruction of the artificial intelligence chip. The specific value of the hardware alignment granularity is determined by the type of the artificial intelligence chip. In some embodiments, the ratio of the sequence length to the hardware alignment granularity of the artificial intelligence chip is used as the splitting number.

[0050] When the combined number is greater than the splitting number, and the ratio of the combined number to the total number of computing units is an integer, split the original query tensor into target query tensors processed by each of the multiple computing units according to the combined number and the total number of computing units.

[0051] At this time, the sequence length of the target query tensor is the same as that of the original query tensor, and the shape of the target query tensor is: ;

[0052] where N represents the total number of computing units.

[0053] For example, assume the shape of the original query tensor is , N = 16, block_size = 64. At this time, the combined number = 16, the splitting number = 512 / 64 = 8, that is, the combined number is greater than the splitting number, and the ratio of the combined number to the total number of computing units is an integer. Therefore, split the original query tensor into target query tensors processed by N computing units respectively in the combined number dimension. At this time, the shape of the target query tensor is [1, 512, head_size].

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

[0055] At this time, 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: ; where N represents the total number of computing units.

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

[0057] In some embodiments, when the ratio of the merging quantity to the total number of computing units is a non-integer, and the ratio of the splitting quantity to the total number of computing units is a non-integer, the original query tensor is split into target query tensors to be processed by each of the multiple computing units according to the merging quantity, the sequence length, and the total number of computing units.

[0058] At this time, the shape of the target query tensor is: ; where n1 represents the splitting ratio in the merging quantity dimension, n2 represents the splitting ratio in the sequence length dimension, , and N represents the total number of computing units.

[0059] For example, assume the shape of the original query tensor is , N = 16, and block_size = 64. At this time, the merging quantity = 8, the splitting quantity = 512 / 64 = 8, that is, the ratio of the merging quantity to the total number of computing units is a non-integer, and the ratio of the splitting quantity to the total number of computing units is a non-integer. Therefore, combining the merging quantity dimension and the sequence length dimension, the original query tensor is split into target query tensors to be processed by each of the N computing units. At this time, n1 = 8, n2 = 2, and the shape of the target query tensor is [1, 256, head_size].

[0060] It should be noted that when the merging quantity is greater than the splitting quantity, the ratio of the merging quantity to the total number of computing units is a non-integer, and the ratio of the splitting quantity 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 respectively in both the merging quantity dimension and the sequence length dimension.

[0061] When the merging quantity is not greater than the splitting quantity, the ratio of the splitting quantity to the total number of computing units is a non-integer, and the ratio of the merging quantity 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 respectively in both the merging quantity dimension and the sequence length dimension. The present application does not make specific limitations on this.

[0062] In the embodiments of the present application, according to the actual shape of the original query tensor, appropriate splitting dimensions are selected in the merging 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 dimensions. In this way, the larger dimensions can be preferentially split, so that each computing unit can be fully utilized, and the performance of the operator is improved.

[0063] Based on the sequence length and the hardware alignment granularity of the artificial intelligence chip, the splitting quantity is obtained. When the ratio of the splitting quantity to the total number of computing units is an integer, the original query tensor is split into target query tensors processed by N computing units respectively according to the sequence length dimension, so that each computing unit can process a target query tensor not less than the hardware alignment granularity. In this way, the computing unit does not need to perform granularity filling for overly fine (i.e., less than the hardware alignment granularity) target query tensors during the calculation process, thereby improving the overall computing performance of the computing unit.

[0064] 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 is split into N target key tensors, the original value tensor V is split into N target value tensors, the original forward output O is split into N target forward outputs, the original output gradient dO is split into N target output gradients, the original input mask M1 is split into N target input masks, the original intermediate result is split into N target intermediate results, the original random mask is split into N target random masks, etc., which will not be elaborated here.

[0065] Multiple computing units perform attention mechanism calculations in parallel based on the data obtained by each of them. Hereinafter, one computing unit will be taken as an example for detailed introduction: In the embodiments of the present application, the reverse attention mechanism operator can specifically be a reverse fast attention mechanism (FlashAttention bwd) operator. During the execution of the reverse fast attention mechanism operator, each operator inside the reverse fast attention mechanism operator performs calculations in a block manner.

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

[0067] The execution process of the reverse fast attention mechanism operator includes: an outer loop, a middle loop, and an 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.

[0068] One round of the upper loop contains multiple rounds of the middle loop, and the Q_chunk input in each round of the middle loop is different; each round of the middle loop contains multiple rounds of the inner loop, and the K_chunk and V_chunk input in each round of the inner loop are different, but the Q_chunk is the same. The intermediate results generated during the above three-level loop process can be reused.

[0069] The embodiments of this application take one round of the inner loop as an example to introduce the flow of the optimization method for attention mechanism calculation, which is executed by an artificial intelligence chip. See Figure 3 , and specifically includes the following steps: Step 301, execute a forward recomputation operator based on the query chunk and the key chunk to obtain a forward normalization result.

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

[0071] In some embodiments, the forward recomputation operator includes: a second matrix multiplication operator and a normalization operator. The execution process of the forward recomputation operator includes the following steps: Divide the query chunk into multiple first sub-chunks; and, divide the key chunk into multiple second sub-chunks, where each first sub-chunk corresponds to a second sub-chunk.

[0072] For each first sub-chunk and the corresponding second sub-chunk, respectively execute: Execute a second matrix multiplication operator on a first sub-block and a corresponding second sub-block based on a tensor core to obtain a sub-matrix multiplication result; then save the sub-matrix multiplication result in a register.

[0073] In some cases, in order to fully utilize the register resources and avoid the sub-matrix multiplication result exceeding the storage space of the register, in the embodiments of the present application, based on the pre-allocated storage space of the register, 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 result obtained by multiplying the first sub-block and the corresponding second sub-block can fully occupy the register resources, while reducing the number of split sub-blocks, thereby reducing the number of loop calculations.

[0074] Next, read the sub-matrix multiplication result from the register through a vector core to execute a normalization operator to obtain a sub-normalization result.

[0075] In specific implementation, obtain an intermediate result block m_chunk and an 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.

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

[0077] After calculating the sub-normalization result, since the subsequent calculation needs to use the sub-normalization result, therefore, the sub-normalization result is saved in the on-chip cache instead of writing it back to the video memory. In this way, during the subsequent calculation process, the sub-normalization result can be directly read from the on-chip cache without accessing the video memory, thereby improving the calculation efficiency and making full use of the on-chip cache.

[0078] In some embodiments, during the execution of the forward recomputation operator, register synchronization is adopted between the tensor core and the vector core.

[0079] 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 the registers. After the tensor core has produced 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. At this time, the producer is the tensor core and the consumer is the vector core.

[0080] The vector core consumes the data of the preset number of registers, that is, reads the sub-matrix multiplication result from the registers, and executes a normalization operator on the sub-matrix multiplication result to obtain a sub-normalization result. After the vector core has consumed the data of the preset number of registers, it sends a buffer ready instruction to the tensor core.

[0081] In addition, during the process of the vector core reading the sub-matrix multiplication result from the registers and executing the normalization operator, the tensor core concurrently executes the second matrix multiplication operator based on the next first sub-block and the corresponding second sub-block to obtain the sub-matrix multiplication result, causing the calculation times of the tensor core and the vector core to overlap.

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

[0083] In the embodiments of the present application, the query block and the key block are further divided into multiple sub-blocks; then the tensor core executes the second matrix multiplication operator on the sub-blocks to obtain the sub-matrix multiplication result and stores the sub-matrix multiplication result in the registers; the vector core reads the sub-matrix multiplication result from the registers and executes the normalization operator to obtain the sub-normalization result. At the same time, the tensor core concurrently executes the second matrix multiplication operator based on the next sub-block, causing the calculation times of the tensor core and the vector core to overlap, realizing effective pipelining for different operators, thereby improving the execution efficiency; Secondly, data interaction between the tensor core and the vector core is carried out through registers, avoiding frequent entry and exit of the interaction data from the on-chip cache, thereby improving the overall performance of the operator.

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

[0085] Specifically, the tensor core transposes the forward normalization result and then inputs it together with the output gradient block into the first matrix multiplication operator for calculation to obtain the value intermediate result.

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

[0087] In some embodiments, during the execution of the forward recomputation operator, the output gradient block is loaded into the on-chip cache in parallel from the video memory, where the output gradient block is a block in the target output gradient processed by a computing unit. In this way, after obtaining the forward normalization result, the forward normalization result and the output gradient block can be directly read from the on-chip cache to execute the first matrix multiplication operator, thereby improving the operator execution efficiency.

[0088] Step 303: Execute a dot product summation operator based on the forward output block and the output gradient block to obtain an intermediate tensor.

[0089] Specifically, perform a dot product operation on the forward output block and the output gradient block through a vector kernel, and sum the obtained dot product results by row to obtain an intermediate tensor. After calculating the intermediate tensor, since the subsequent calculation requires the use of the intermediate tensor, the intermediate tensor is stored in the on-chip cache.

[0090] In the embodiments of the present application, the dot product summation operator and the first matrix multiplication operator are executed in parallel, that is, during the execution of the first matrix multiplication operator based on the output gradient block and the forward normalization result, the dot product summation operator based on the forward output block and the output gradient block is executed in parallel.

[0091] In some embodiments, after both the first matrix multiplication operator and the dot product summation operator are executed, other operators are continued to be executed backward to avoid execution errors.

[0092] Step 304: Execute a gradient backpropagation operator based on the output gradient block, the value block, the intermediate tensor, and the forward normalization result to obtain a backward normalization result.

[0093] In specific implementation, the gradient backpropagation operator and the first type of transformation operator are executed in parallel, that is, during the execution of 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, the value block, the intermediate tensor, and the forward normalization result.

[0094] 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: Divide the output gradient block into multiple third sub-blocks; and divide the value block into multiple fourth sub-blocks, where each third sub-block corresponds to a fourth sub-block.

[0095] For each third sub - block and the corresponding fourth sub - block, the following operations are respectively performed: Execute a third matrix multiplication operator based on a third sub - block and the corresponding fourth sub - block through a tensor core to obtain a sub - matrix multiplication result, and store the sub - matrix multiplication result in a register.

[0096] In some cases, based on the storage space of pre - allocated registers, the output gradient block is divided into multiple third sub - blocks; and the value block is divided into multiple fourth sub - blocks. In this way, the sub - matrix multiplication results obtained by multiplying the third sub - block and the corresponding fourth sub - block can fully utilize the register resources, while reducing the number of split sub - blocks, thereby reducing the number of loop calculations.

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

[0098] In 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'. Select a sub - tensor d' and a sub - normalization result p' associated with the third sub - block and the fourth sub - block of the current round of calculation.

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

[0100] After calculating the sub - dot - product result ds', since the subsequent calculation needs to use the sub - dot - product result ds', the sub - dot - product result ds' is stored in the on - chip cache instead of writing it back to the video memory. In this way, during the subsequent calculation process, the sub - dot - product result ds' can be directly read from the on - chip cache, thereby improving the calculation efficiency.

[0101] In some embodiments, during the execution of the gradient backpropagation operator, register synchronization is adopted between the tensor core and the vector core.

[0102] Specifically, the tensor core executes the third matrix multiplication operator, obtains the sub - matrix multiplication result and stores the sub - matrix multiplication result in the register; then, notifies the vector core to read the sub - matrix multiplication result from the register, and perform a dot - product difference operator with the intermediate tensor and the forward normalization result to obtain a sub - dot - product result.

[0103] In addition, when reading the sub-matrix multiplication result from the register through the vector core and performing the dot product difference operator with the intermediate tensor and the forward normalization result, the tensor core executes the third matrix multiplication operator in parallel based on the next third sub-block and the corresponding fourth sub-block to obtain the sub-matrix multiplication result, so that the calculation times of the tensor core and the vector core overlap.

[0104] Finally, based on the obtained multiple sub-dot product results, a reverse normalization result is generated; specifically, the multiple sub-dot product results can be concatenated to obtain the reverse normalization result.

[0105] In the embodiment of the present application, the output gradient block and the value block are further divided into multiple sub-blocks; then the tensor core executes the third matrix multiplication operator on the sub-blocks to obtain the sub-matrix multiplication result and stores the sub-matrix multiplication result in the register; the vector core reads the sub-matrix multiplication result from the register to execute the dot product difference operator to obtain the sub-dot product result. At the same time, the tensor core executes the second matrix multiplication operator in parallel based on the next sub-block, so that the calculation times of the tensor core and the vector core overlap, realizing effective pipelining processing for different operators, thereby improving the execution efficiency; Secondly, data interaction between the tensor core and the vector core is carried out through the register, avoiding the frequent entry and exit of the interaction data from the on-chip cache, thereby improving the overall performance of the operator.

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

[0107] Specifically, the key block is a block in the target key tensor processed by a computing unit. Read the reverse normalization result and the key block from the on-chip cache to perform matrix multiplication calculation to obtain the query gradient block; read the reverse normalization result and the query block from the on-chip cache to perform matrix multiplication calculation to obtain the key gradient block.

[0108] The value gradient block obtained through step 302, as well as the query gradient block and the key gradient block obtained through step 305, are the results obtained in one round of the inner loop.

[0109] In the embodiment of the present application, the forward recalculation operator is executed based on the query block and the key block to obtain the forward normalization result; then the first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain the value intermediate result. When executing the first matrix multiplication operator, the dot product summation operator is executed in parallel based on the forward output block and the output gradient block to obtain the intermediate tensor, so that the calculation times of the first matrix multiplication operator and the dot product summation operator overlap.

[0110] Next, a first type of transformation operator is executed on the value intermediate result to obtain value gradient blocks; at the same time, a gradient backpropagation operator is executed in parallel based on the output gradient blocks, value blocks, intermediate tensors, and forward normalization results to obtain a reverse normalization result, so that the calculation times of the first type of transformation operator and the gradient backpropagation operator overlap.

[0111] Finally, based on the reverse normalization result and key blocks, query gradient blocks are obtained; and based on the reverse normalization result and query blocks, key gradient blocks are obtained, realizing effective pipelining of different operators during the calculation process of the reverse attention mechanism. Compared with the serial method, the execution efficiency of the operators is improved, thereby enhancing the overall performance of the attention mechanism calculation.

[0112] In some embodiments, the embodiments of the present application also provide another implementation manner for obtaining query gradient blocks and key gradient blocks: Execute a fourth matrix multiplication operator based on the reverse normalization result and key blocks to obtain a query intermediate result; execute a second type of transformation operator on the query intermediate result to obtain query gradient blocks.

[0113] Next, execute a fifth matrix multiplication operator based on the reverse normalization result and query blocks to obtain a key intermediate result, where the fifth matrix multiplication operator and the second type of transformation operator are executed in parallel; execute a third type of transformation operator on the key intermediate result to obtain key gradient blocks.

[0114] Specifically, the reverse normalization result is read from the on-chip cache through a tensor core, and then the reverse normalization result and key blocks are input into the fourth matrix multiplication operator for calculation to obtain a query intermediate result. After the execution of the fourth matrix multiplication operator ends, the second type of transformation operator is started to be executed on the query intermediate result through a vector core to obtain query gradient blocks.

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

[0116] After the execution of the second type of transformation operator and the fifth matrix multiplication operator ends, the third type of transformation operator is started to be executed on the key intermediate result through a vector core to obtain key gradient blocks.

[0117] In the embodiments of the present application, during the process of executing the second type of transformation operator through a vector core, the reverse normalization result and query blocks are input into the fifth matrix multiplication operator in parallel through a tensor core for calculation to obtain a key intermediate result, so that the calculation times of the second type of transformation operator and the fifth matrix multiplication operator overlap, realizing effective pipelining of different operators, thereby improving the execution efficiency of the operators.

[0118] In some embodiments, the query chunk corresponds to multiple rounds of inner loops, and the key chunk and value chunk are the input data for the current round of inner loops. The third type of transformation operator and the forward recomputation operator in the next round of inner loops corresponding to the query chunk are executed in parallel.

[0119] Specifically, each query chunk corresponds to one round of middle loops. One round of middle loops includes multiple rounds of inner loops, and the key chunks and value chunks input in multiple rounds of inner loops are different.

[0120] When the current round of inner loops executes to the third type of transformation operator, the next round of inner loops starts to be executed in parallel, that is, based on the query chunk (the same as the query chunk used in the current round of inner loops) and the next key chunk (different from the key chunk used in the current round of inner loops), the forward recomputation operator is executed to obtain the forward normalization result.

[0121] It should be noted that if the current round of inner loops is the last round of inner loops among the multiple rounds of inner loops corresponding to the query chunk, when the current round of inner loops executes to the third type of transformation operator, the first inner loop in the next round of middle loops starts to be executed in parallel, that is, based on the next query chunk (different from the query chunk used in the current round of inner loops) and the key chunk input in the first inner loop (different from the key chunk used in the current round of inner loops), the forward recomputation operator is executed to obtain the forward normalization result.

[0122] In the embodiments of the present application, effective pipelining processing is implemented for different operators between multiple loops of the reverse attention mechanism calculation. Compared with the serial method, the execution efficiency of the operators is improved, thereby enhancing the overall performance of the attention mechanism calculation.

[0123] To better explain the embodiments of the present application, the following introduces the process of an optimization method for attention mechanism calculation provided by the embodiments 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 shown as: 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.

[0124] The first operator combination includes: the forward recomputation operator 401 and the third type of transformation operator 402 that are executed in parallel. The forward recomputation operator 401 includes the following operations: MMA (matrix multiplication operation), masking operation, sub (subtraction), exp (exponential function), div (division).

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

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

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

[0128] The third operator combination includes: a first type transformation operator 405 and a gradient backpropagation operator 406 that are executed in parallel.

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

[0130] The fourth operator combination includes: a fourth matrix multiplication operator 407; The fourth matrix multiplication operator 407 includes the following operations: MMA.

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

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

[0133] Next, taking one round of the inner loop as an example, the execution process of each operator combination in the reverse fast attention mechanism operator is introduced: Load the output gradient block dO_chunk 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: Input the query block Q_chunk, key block K_chunk, intermediate result block m_chunk, and intermediate result block s_chunk into the forward recomputation operator 401 for calculation to obtain the forward normalization result P.

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

[0135] Perform T rounds of calculations on the data obtained by splitting. The input data for each round of calculation includes: a first sub - chunk q_chunk and an associated second sub - chunk k_chunk, an intermediate sub - chunk m'_chunk, and an intermediate sub - chunk s'_chunk. The specific process for each round of calculation is as follows: Execute MMA on the first sub - chunk q_chunk and the second sub - chunk k_chunk through a tensor core to obtain the sub - matrix multiplication result dp1, and save the sub - matrix multiplication result dp1 in a register.

[0136] Then, notify the vector core as a consumer to read the sub - matrix multiplication result dp1 from the register, and then sequentially perform a masking operation, sub, exp, and div based on the sub - matrix multiplication result dp1, the intermediate sub - chunk m'_chunk, and the intermediate sub - chunk s'_chunk according to the above formula (1) to obtain the sub - normalization result p', and save the sub - normalization result p' in the on - chip cache.

[0137] Concatenate the T sub - normalization results p' obtained to obtain the forward normalization result P.

[0138] During the execution of the forward recomputation operator 401, concurrently perform the third - type transformation operator 402 on the key intermediate result obtained in the previous inner - loop to obtain the key gradient chunk output in the previous inner - loop.

[0139] Only after both the forward recomputation operator 401 and the third - type transformation operator 402 have completed their calculations, start to execute the subsequent second operator combination.

[0140] The execution process of the second operator combination includes: Input the forward normalization result P and the output gradient chunk dO_chunk into the first matrix multiplication operator 403 through the tensor core for calculation to obtain the value intermediate result v1.

[0141] Input the forward output chunk O_chunk and the output gradient chunk dO_chunk into the dot - product summation operator 404 through the vector core for calculation to obtain the intermediate tensor D, and save the intermediate tensor D in the on - chip cache.

[0142] Only after both the first matrix multiplication operator 403 and the dot - product summation operator 404 have completed their calculations, start to execute the subsequent third operator combination.

[0143] The execution process of the third operator combination: Input the value intermediate result v1 into the first - type transformation operator 405 through the vector core for calculation to obtain the value gradient chunk dV_chunk of this round's inner - loop.

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

[0145] After both the first type of transformation operator 405 and the gradient backpropagation operator 406 have completed their calculations, the subsequent fourth operator combination is executed.

[0146] Execution process of the fourth operator combination: The reverse 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.

[0147] Execution process of the fifth operator combination: The second type of transformation operator 408 is executed on the query intermediate result q1 through vector checking to obtain the query gradient block dQ_chunk for the current round of inner loop.

[0148] The reverse 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.

[0149] After the fifth operator combination has been executed, the third type of transformation operator 402 is executed on the key intermediate result k1 through vector checking to obtain the key gradient block dK_chunk for the current round of inner loop; at the same time, the forward recomputation operator 401 in the next round of inner loop is executed in parallel, and the execution process of the next round of inner loop is the same as that of the current round of inner loop, which will not be elaborated here.

[0150] In the embodiments of the present application, the forward recomputation operator is executed based on the query block and the key block to obtain the forward normalization result; then, the first matrix multiplication operator is executed based on the output gradient block and the forward normalization result to obtain the value intermediate result. Among them, when executing the first matrix multiplication operator, the dot product summation operator is executed in parallel based on the forward output block and the output gradient block to obtain the intermediate tensor, so that the calculation times of the first matrix multiplication operator and the dot product summation operator overlap.

[0151] Next, 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, value block, intermediate tensor, and forward normalization result to obtain the reverse normalization result, so that the calculation times of the first type of transformation operator and the gradient backpropagation operator overlap.

[0152] Finally, based on the reverse normalization result and the key chunks, query gradient chunks are obtained; and based on the reverse normalization result and the query chunks, key gradient chunks are obtained, realizing effective pipelining of different operators during the calculation process of the reverse attention mechanism. Compared with the serial method, the execution efficiency of the operators is improved, thereby enhancing the overall performance of the attention mechanism calculation.

[0153] Based on the same technical concept, an embodiment of the present application provides a structural schematic diagram of an optimized device for attention mechanism calculation, as Figure 5 shown. The optimized device 500 for attention mechanism calculation includes: A recomputation module 501, configured to execute a forward recomputation operator based on the query chunks and the key chunks to obtain a forward normalization result; A backpropagation module 502, configured to execute a first matrix multiplication operator based on the output gradient chunks and the forward normalization result to obtain a value intermediate result; and execute a first type of transformation operator on the value intermediate result to obtain value gradient chunks; The backpropagation module 502 is further configured to execute a dot product summation operator based on the forward output chunks and the output gradient chunks to obtain an intermediate tensor; the dot product summation operator and the first matrix multiplication operator are executed in parallel; The backpropagation module 502 is further configured to execute a gradient backpropagation operator based on the output gradient chunks, the value chunks, the intermediate tensor, and the forward normalization result to obtain a reverse normalization result, and the gradient backpropagation operator and the first type of transformation operator are executed in parallel; The backpropagation module 502 is further configured to obtain query gradient chunks based on the reverse normalization result and the key chunks; and obtain key gradient chunks based on the reverse normalization result and the query chunks.

[0154] Optionally, the recomputation module 501 is further configured to: During the execution of the forward recomputation operator, load the output gradient chunks from the video memory to the on-chip cache in parallel.

[0155] Optionally, the forward recomputation operator includes: a second matrix multiplication operator and a normalization operator; The recomputation module 501 is specifically configured to: Divide the query chunks into multiple first sub-chunks; and divide the key chunks into multiple second sub-chunks, where each first sub-chunk corresponds to a second sub-chunk; For each first sub-block and the corresponding second sub-block, the following operations are respectively performed: perform a second matrix multiplication operator on a first sub-block and the corresponding second sub-block based on a tensor core to obtain a sub-matrix multiplication result; save the sub-matrix multiplication result in a register; read the sub-matrix multiplication result from the register by a vector core to perform the normalization operator to obtain a sub-normalization result; Generate the forward normalization result based on the obtained multiple sub-normalization results.

[0156] Optionally, the recomputation module 501 is further configured to: After reading the sub-matrix multiplication result from the register by a vector core to perform the normalization operator to obtain a sub-normalization result, save the sub-normalization result in the on-chip cache.

[0157] Optionally, the recomputation module 501 is further configured to: During the execution of the forward recomputation operator, register synchronization is adopted between the tensor core and the vector core.

[0158] Optionally, the gradient backpropagation operator includes: a third matrix multiplication operator and a dot product difference operator; The backpropagation module 502 is specifically configured to: Divide the output gradient block into multiple third sub-blocks; and divide the value block into multiple fourth sub-blocks, where each third sub-block corresponds to a fourth sub-block; For each third sub-block and the corresponding fourth sub-block, the following operations are respectively performed: perform a third matrix multiplication operator on a third sub-block and the corresponding fourth sub-block based on a tensor core to obtain a sub-matrix multiplication result; save the sub-matrix multiplication result in a register; read the sub-matrix multiplication result from the register by a vector core, and perform the dot product difference operator with the intermediate tensor and the forward normalization result to obtain a sub-dot product result; Generate the backward normalization result based on the obtained multiple sub-dot product results.

[0159] Optionally, the backpropagation module 502 is further configured to: After reading the sub-matrix multiplication result from the register by a vector core, and performing the dot product difference operator with the intermediate tensor and the forward normalization result to obtain a sub-dot product result, save the sub-dot product result in the on-chip cache.

[0160] Optionally, the backpropagation module 502 is specifically configured to: Perform a fourth matrix multiplication operator based on the backward normalization result and the key block to obtain a query intermediate result; perform a second type of transformation operator on the query intermediate result to obtain the query gradient block; Execute a fifth matrix multiplication operator based on the reverse normalization result and the query chunk to obtain a key intermediate result, where the fifth matrix multiplication operator is executed in parallel with the second type of transformation operator; Execute a third type of transformation operator on the key intermediate result to obtain the key gradient chunk.

[0161] Optionally, the query chunk corresponds to multiple rounds of inner loops, and the key chunk and the value chunk are the input data for the current round of inner loops; The third type of transformation operator is executed in parallel with the forward recomputation operator in the next round of inner loop corresponding to the query chunk.

[0162] Optionally, it further includes a splitting module 503; The query chunk is a chunk in the target query tensor processed by one computing unit; the one computing unit is any one of the multiple computing units included in the artificial intelligence chip; The splitting module 503 is specifically used for: Obtain an original query tensor, and the shape of the original query tensor includes: the combined number of batches and attention heads, and the sequence length; Based on the sequence length and the hardware alignment granularity of the artificial intelligence chip, obtain the splitting number; When the combined number is greater than the splitting number, and the ratio of the combined number to the total number of computing units is an integer, split the original query tensor into target query tensors processed by the multiple computing units respectively according to the combined number and the total number of computing units; When the combined number is not greater than the splitting number, and the ratio of the splitting number to the total number of computing units is an integer, split the original query tensor into target query tensors processed by the multiple computing units respectively according to the sequence length and the total number of computing units.

[0163] Optionally, the splitting module 503 is further used for: When the ratio of the combined number to the total number of computing units is a non-integer, and the ratio of the splitting number to the total number of computing units is a non-integer, split the original query tensor into target query tensors processed by the multiple computing units respectively according to the combined number, the sequence length, and the total number of computing units.

[0164] In the embodiment of the present application, a forward recomputation operator is executed based on query chunks and key chunks to obtain a forward normalization result; then, a first matrix multiplication operator is executed based on the output gradient chunks and the forward normalization result to obtain a value intermediate result. When executing the first matrix multiplication operator, a dot product summation operator is executed in parallel based on the forward output chunks and the output gradient chunks to obtain an intermediate tensor, so that the calculation times of the first matrix multiplication operator and the dot product summation operator overlap.

[0165] Next, a first type of transformation operator is executed on the value intermediate result to obtain value gradient chunks; at the same time, a gradient backpropagation operator is executed in parallel based on the output gradient chunks, value chunks, intermediate tensor, and forward normalization result to obtain a backward normalization result, so that the calculation times of the first type of transformation operator and the gradient backpropagation operator overlap.

[0166] Finally, query gradient chunks are obtained based on the backward normalization result and key chunks; and key gradient chunks are obtained based on the backward normalization result and query chunks, realizing effective pipelining of different operators during the calculation process of the reverse attention mechanism. Compared with the serial method, the execution efficiency of the operators is improved, thereby enhancing the overall performance of the attention mechanism calculation.

[0167] Based on the same technical concept, the embodiment of the present application provides a computer device, as Figure 6 shown, including at least one artificial intelligence chip 100 and a memory 601 connected to the at least one artificial intelligence chip 100. In the embodiment of the present application, the specific connection medium between the artificial intelligence chip 100 and the memory 601 is not limited. Figure 6 Taking the connection between the artificial intelligence chip 100 and the memory 601 through a bus as an example. The bus can be divided into an address bus, a data bus, a control bus, etc.

[0168] In the embodiment of the present application, the memory 601 stores instructions executable by the at least one artificial intelligence chip 100. By executing the instructions stored in the memory 601, the at least one artificial intelligence chip 100 can execute the steps of the above-mentioned optimization method for attention mechanism calculation.

[0169] Among them, the artificial intelligence chip 100 is the control center of the computer device. It can connect various parts of the computer device through various interfaces and circuits, and optimize the calculation of the attention mechanism by running or executing the instructions stored in the memory 601 and calling the data stored in the memory 601. 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. Among them, the application processor mainly processes the operating system, user interface, application programs, etc., and the modem processor mainly processes wireless communication. It can be understood that the above modem processor may not be integrated into the artificial intelligence chip 100. In some embodiments, the artificial intelligence chip 100 and the memory 601 may be implemented on the same chip. In some embodiments, they may also be separately implemented on independent chips.

[0170] The artificial intelligence chip 100 may 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 devices, discrete gate or transistor logic devices, discrete hardware components, and can implement or execute the various methods, steps and logic block diagrams disclosed in the embodiments of the present application. The general-purpose processor may be a microprocessor or any conventional processor, etc. The steps of the method disclosed in combination with the embodiments of the present application may be directly embodied as being executed by a hardware processor, or executed by a combination of hardware and software modules in the processor.

[0171] The memory 601, as a non-volatile computer-readable storage medium, can be used to store non-volatile software programs, non-volatile computer-executable programs, and modules. The memory 601 can include at least one type of storage medium. For example, it can include flash memory, hard disks, multimedia cards, card-type memories, random access memory (RAM), static random access memory (SRAM), programmable read-only memory (PROM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), magnetic memories, magnetic disks, optical disks, and so on. The memory 601 is any other medium that can be used to carry or store the 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 can also be a circuit or any other device capable of implementing a storage function, for storing program instructions and / or data.

[0172] Based on the same inventive concept, embodiments of the present application provide a computer-readable storage medium that stores a computer program executable by a computer device. When the computer program runs on the computer device, it causes the computer device to execute the steps of the above-mentioned optimization method for attention mechanism calculation.

[0173] Based on the same inventive concept, embodiments of the present application provide a computer program product. The computer program product includes a computer program stored on a computer-readable storage medium. The computer program includes program instructions. When the program instructions are executed by a computer device, they cause the computer device to execute the steps of the above-mentioned optimization method for attention mechanism calculation.

[0174] Those skilled in the art should understand that the embodiments of the present invention can be provided as a method, or a computer program product. Therefore, the present invention can take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware aspects. Moreover, the present invention can take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk memories, CD-ROMs, optical memories, etc.) that contain computer-usable program code.

[0175] The present invention is described with reference to the flowcharts and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the invention. It should be understood that each flow and / or block in the flowchart and / or block diagram, and combinations of flows and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to the processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing device to produce a machine, such that the instructions executed by the processor of the computer device or other programmable data processing device produce means for implementing the functions specified in one flow Figure 1 one flow or multiple flows and / or blocks Figure 1 or multiple blocks.

[0176] These computer program instructions can 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, such that the instructions stored in the computer-readable memory produce a manufactured article including instruction means that implement the functions specified in one flow Figure 1 one flow or multiple flows and / or blocks Figure 1 or multiple blocks.

[0177] These computer program instructions can also be loaded onto a computer device or other programmable data processing device, such that a series of operational steps are executed on the computer device or other programmable device to produce a process implemented by the computer device, so that the instructions executed on the computer device or other programmable device provide steps for implementing the functions specified in one flow Figure 1 one flow or multiple flows and / or blocks Figure 1 or multiple blocks.

[0178] Although the preferred embodiments of the present invention have been described, additional changes and modifications can be made by those skilled in the art once they learn the basic creative concepts. Therefore, the appended claims are intended to be construed to include the preferred embodiments and all changes and modifications that fall within the scope of the present invention.

[0179] Obviously, those skilled in the art can make various changes and variations to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalent technologies, the present invention is also intended to include these modifications and variations.

Claims

1. A method for optimizing attention mechanism calculation, characterized in that: include: Execute the forward recalculation operator based on the query block and the key block to obtain the forward normalized result; Execute a first matrix multiplication operator based on the output gradient block and the forward normalization result to obtain a value intermediate result; Executing 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 back-propagation operator based on the output gradient block, the value block, the intermediate tensor and the forward normalization result to obtain a reverse normalization result, wherein the gradient back-propagation operator and the first type transformation operator are executed in parallel; 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.

2. The method according to claim 1, characterized in that 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, characterized in that 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: executing the second matrix multiplication operator based on the first sub-block and the corresponding second sub-block by the tensor core to obtain a sub-matrix multiplication result; saving the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register by the vector core to execute the normalization operator to obtain a sub-normalization result; The forward normalization result is generated based on the obtained multiple sub-normalization results.

4. The method according to claim 3, characterized in that 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.

5. The method according to claim 3, characterized in that Also includes: During the execution of the forward recalculation operator, register synchronization is adopted between the tensor core and the vector core.

6. The method according to claim 1, characterized in that The gradient back propagation operator includes: a third matrix multiplication operator and a point multiplication difference operator; The performing of 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 reverse normalization 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 perform: executing a third matrix multiplication operator based on a third sub-block and the corresponding fourth sub-block through a tensor core to obtain a sub-matrix multiplication result; saving the sub-matrix multiplication result in a register; reading the sub-matrix multiplication result from the register through a vector core, executing the point multiplication difference operator with the intermediate tensor and the forward normalization result to obtain a sub-point multiplication result; The inverse normalization result is generated based on the obtained multiple sub-point multiplication results.

7. The method according to claim 6, characterized in that The method further comprises: reading the sub-matrix multiplication result from the register through the vector core, performing the point multiplication difference operator with the intermediate tensor and the forward normalization result, and obtaining the sub-point multiplication result; The sub-point product results are stored in an on-chip cache.

8. The method according to claim 1, characterized in that The step of obtaining a query gradient block based on the reverse normalization result and the key block; 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.

9. The method according to claim 8, characterized in that 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.

10. The method according to any one of claims 1 to 9, characterized in that: The query block is a block in a target query tensor processed by a computing unit; the computing unit is any one of a plurality of 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 the number of splits based on the sequence length and the hardware alignment granularity of the artificial intelligence chip; When the merging number is greater than the splitting number, and the ratio of the merging number to the total number of computing units is an integer, splitting the original query tensor into target query tensors processed by each of the multiple computing units according to the merging number 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.

11. The method according to claim 10, characterized in that 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.

12. 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 as described in any one of claims 1 to 11 are implemented.

13. 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 as claimed in any one of claims 1 to 11.

14. 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, and when the program instructions are executed by a computer device, the computer device is caused to execute the steps of the method according to any one of claims 1 to 11.

Images