Attention mechanism calculation method and device, medium and product

Abstract

The invention discloses an attention mechanism calculation method and device, a medium and a product. The method comprises the steps that a first kernel controls a second thread bundle group to load an ith query matrix block; the first kernel controls the second thread bundle group to calculate a jth attention score matrix block; and the first kernel controls the first thread bundle group and the second thread bundle group to alternately utilize different sub-blocks of the jth value matrix block and the jth attention score matrix block to perform a preset preorder stage of attention mechanism operation, an intermediate result of different sub-blocks of the last attention result matrix block of the ith query matrix block is obtained; and the second kernel controls the third thread bundle group to load intermediate results of all sub-blocks of the last attention result matrix block of the ith query matrix block, and carries out a preset subsequent stage of attention mechanism operation. By adopting the embodiment of the invention, sufficient register resources can be provided for the calculation of the attention mechanism, so that the calculation efficiency is improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence, and in particular to an attention mechanism calculation method, device, medium and product. BACKGROUND

[0002] In the attention mechanism calculation process, three types of matrices will be mainly involved: query matrix, key matrix and value matrix. However, as the matrix size increases significantly, the existing technology exposes obvious defects: (1) at the warpingroup level, only a single warpgroup is relied on to perform the full-process attention mechanism calculation, and the task concentration causes excessive register pressure; (2) at the kernel scheduling level, only a single kernel is used to control the above warpgroup, and the kernel is responsible for the scheduling and execution of the full-process calculation task. The superposition of the two problems makes the register resources unable to meet the large-scale attention mechanism calculation, and ultimately limits the calculation efficiency of the attention mechanism. SUMMARY

[0003] The present application provides an attention mechanism calculation method, device, medium and product to solve the problem that the existing technology only relies on a single warpgroup to perform the full-process attention mechanism calculation, and only a single kernel is used to control the above warpgroup, making the register resources unable to meet the large-scale attention mechanism calculation, and ultimately limiting the calculation efficiency of the attention mechanism.

[0004] To achieve the above-mentioned purpose, the present application provides an attention mechanism calculation method, comprising: The first kernel controls the second warpgroup to load the ith query matrix block; The first kernel controls the second warpgroup to perform matrix multiplication and exponential operation on the ith query matrix block and the transposed jth key matrix block, to obtain the jth attention score matrix block; The first kernel controls the first warpgroup and the second warpgroup to alternately use different sub-blocks of the jth value matrix block and the jth attention score matrix block in a preset pre-sequence stage of attention mechanism operation, to respectively obtain intermediate results of different sub-blocks of the jth attention result matrix block, let j=j+1, and return to the step of the first kernel controlling the second warpgroup to perform matrix multiplication and exponential operation on the ith query matrix block and the transposed jth key matrix block, until j=J, to obtain intermediate results of different sub-blocks of the last attention result matrix block of the ith query matrix block; wherein J is a preset total number of blocks, and the jth key matrix block and the jth value matrix block are pre-stored in the shared memory; The second kernel controls the third thread bundle group to load all sub-blocks of the last one attention result matrix block of the i-th query matrix block, and perform a preset post-processing stage of the attention mechanism operation, to obtain the last one attention result matrix block of the i-th query matrix block.

[0005] As an improvement of the above scheme, the first kernel controls the second thread bundle group to load the i-th query matrix block, including: If i = 1, the first kernel controls the first thread bundle group to write the first query matrix block from the global memory to the shared memory, and controls the second thread bundle group to write the first query matrix block from the shared memory to the eleventh register group of the second thread bundle group. If i = 2,...,I, the first kernel controls the second thread bundle group to write the i-th query matrix block from the global memory to the eleventh register group of the second thread bundle group; wherein I is the last one query matrix block.

[0006] As an improvement of the above scheme, after obtaining the j-th attention score matrix block, the method further includes: The first kernel controls the second thread bundle group to calculate the current row accumulation value of the last one attention score matrix block of the i-th query matrix block required for the attention mechanism operation, and write it to the global memory. The second kernel controls the third thread bundle group to load all sub-blocks of the last one attention result matrix block of the i-th query matrix block, and perform a preset post-processing stage of the attention mechanism operation, to obtain the last one attention result matrix block of the i-th query matrix block, including: The second kernel controls the third thread bundle group to write the current row accumulation value of the last one attention score matrix block of the i-th query matrix block from the global memory to the register of the third thread bundle group. The second kernel controls the third thread bundle group to load all sub-blocks of the last one attention result matrix block of the i-th query matrix block and perform splicing, to obtain the intermediate result of the last one attention result matrix block of the i-th query matrix block, divide it by the current row accumulation value of the last one attention score matrix block of the i-th query matrix block, to obtain the last one attention result matrix block of the i-th query matrix block.

[0007] As an improvement of the above scheme, the j-th attention score matrix block is written to the first register group of the second thread bundle group and the first register group of the first thread bundle group. The first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation on the mth sub-block of the jth value matrix block in the shared memory by using the jth attention score matrix block in the first register group of the second thread bundle group, to obtain an intermediate result of the mth sub-block of the jth attention result matrix block. The first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation on the mth sub-block of the jth value matrix block in the shared memory by using the jth attention score matrix block in the first register group of the second thread bundle group, to obtain an intermediate result of the mth sub-block of the jth attention result matrix block. After obtaining the intermediate result of the mth sub-block of the jth attention result matrix block, the first kernel controls the first thread bundle group to perform the preset preliminary stage of attention mechanism operation on the m+1th sub-block of the jth value matrix block in the shared memory by using the jth attention score matrix block in the first register group of the first thread bundle group, to obtain an intermediate result of the m+1th sub-block of the jth attention result matrix block. After obtaining the intermediate result of the m+1th sub-block of the jth attention result matrix block, let m=m+1, and return to the step of the first kernel controlling the second thread bundle group to perform the preset preliminary stage of attention mechanism operation on the mth sub-block of the jth value matrix block in the shared memory by using the jth attention score matrix block in the first register group of the second thread bundle group, until m is the preset total number of sub-blocks, to obtain an intermediate result of the last sub-block of the jth attention result matrix block.

[0008] As an improvement of the above scheme, the first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation on the mth sub-block of the jth value matrix block in the shared memory by using the jth attention score matrix block in the first register group of the second thread bundle group, to obtain an intermediate result of the mth sub-block of the jth attention result matrix block, including: If j=1, the first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation on the mth sub-block of the first value matrix block in the shared memory by using the first attention score matrix block in the first register group of the second thread bundle group, to obtain an intermediate result of the mth sub-block of the first attention result matrix block and write it into the shared memory. If j = 1, the first kernel controls the first thread bundle group to perform the preset pre-sequence stage of attention mechanism operation on the first attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the first attention result matrix block and write it into the shared memory.

[0009] As an improvement of the above scheme, the first kernel controls the first thread bundle group to perform the preset pre-sequence stage of attention mechanism operation on the first attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the first attention result matrix block, including: If j = 1, the first kernel controls the first thread bundle group to perform the preset pre-sequence stage of attention mechanism operation on the first attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the first attention result matrix block and write it into the shared memory. If j = 1, the first kernel controls the first thread bundle group to perform the preset pre-sequence stage of attention mechanism operation on the first attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the first attention result matrix block and write it into the shared memory.

[0010] As an improvement of the above scheme, the method further comprises: For the i-th query matrix block, after the preset pre-sequence stage of the last alternating attention mechanism operation of the second thread bundle group, the first kernel controls the second thread bundle group to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by the preset pre-sequence stage of attention mechanism operation of the second thread bundle group into the global memory.

[0011] As an improvement of the above scheme, the method further comprises: For the i-th query matrix block, after the preset pre-stage of the attention mechanism operation of the last alternation, the first kernel controls the first thread bundle group to write all the sub-blocks of the last attention result matrix block obtained by the first thread bundle group performing the preset pre-stage of the attention mechanism operation to the global memory.

[0012] As an improvement of the above scheme, in the shared memory, the j-th value matrix block and the j-th key matrix block share the same j-th key-value matrix block.

[0013] As an improvement of the above scheme, the method further comprises: For the i-th query matrix block, after the preset pre-stage of the last alternation of the attention mechanism operation, the first kernel controls the first thread bundle group to write the j+x-th key-value matrix block from the global memory to the position of the j-th key-value matrix block in the shared memory; wherein x is the number of key-value matrix blocks stored in the shared memory at the same time.

[0014] Compared with the prior art, the attention mechanism calculation method, device, medium and product provided by the embodiment of the application, by the first kernel controlling the first thread bundle group and the second thread bundle group to alternately use different sub-blocks of the j-th value matrix block to perform the preset pre-stage of the attention mechanism operation with the j-th attention score matrix block, and then by the second kernel controlling the third thread bundle group to perform the preset post-stage of the attention mechanism operation, not only avoids a single thread bundle group from bearing the whole process of the attention mechanism calculation, effectively disperses the register pressure of the single thread bundle group, but also avoids a single kernel from coordinating and scheduling all thread bundle groups, the division of labor is more clear, sufficient register resources can be provided for large-scale attention mechanism calculation, and finally the calculation efficiency of the attention mechanism is effectively improved. BRIEF DESCRIPTION OF DRAWINGS

[0015] Figure 1 is a flowchart of an attention mechanism calculation method provided by the embodiment of the application; Figure 2 is a schematic diagram of a thread holding a matrix element and a register distribution provided by the embodiment of the application; Figure 3 is another schematic diagram of a thread holding a matrix element and a register distribution provided by the embodiment of the application; Figure 4 is a flowchart of an attention mechanism calculation method of a related technology provided by the embodiment of the application; Figure 5 is a principle diagram of an attention mechanism calculation provided by the embodiment of the application; Figure 6 is a flowchart of a thread bundle group provided by the embodiment of the application; Figure 7 is another pipeline diagram of a thread bundle group provided by an embodiment of the present application; Figure 8 is still another pipeline diagram of a thread bundle group provided by an embodiment of the present application; Figure 9 is a resource distribution diagram of shared memory occupancy in an early stage provided by an embodiment of the present application; Figure 10 is a resource distribution diagram of shared memory occupancy in a late stage provided by an embodiment of the present application; Figure 11 is a structural block diagram of an artificial intelligence chip provided by an embodiment of the present application; Figure 12 is a structural block diagram of an attention mechanism computing device provided by an embodiment of the present application. DETAILED DESCRIPTION

[0016] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all the embodiments of the present application. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative work fall within the scope of protection of the present application.

[0017] In the description of the present application, the size of the serial number of the following processes does not mean the order of execution. The execution order of each process should be determined according to its function and inherent logic, and should not constitute any limitation on the implementation process of the embodiments of the present application.

[0018] In the description of the present application, the term "exemplary" or "for example" is used to mean by way of example, illustration or description. Any embodiment or design solution described as "exemplary" or "for example" in the embodiments of the present application should not be interpreted as more preferred or more advantageous than other embodiments or design solutions. Rather, the use of "exemplary" or "for example" is intended to present the relevant concept in a specific manner.

[0019] In the description of the present application, the terms "first", "second", etc. are used to distinguish similar objects, and do not necessarily have to be used to describe a specific order or sequence. It should be understood that the data thus used can be interchanged under appropriate circumstances, so that the embodiments of the present application described herein can be implemented, for example, in an order other than that illustrated or described herein. In addition, the terms "include" and "have" and any variations thereof are intended to cover non-exclusive inclusion, for example, a process, method, system, product or device that includes a series of steps or units does not have to be limited to those steps or units clearly listed, but can include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices. The term "based on" is "at least partially based on". The term "according to" is "at least partially according to". The term "one embodiment" means "at least one embodiment"; the term "another embodiment" means "at least one additional embodiment"; the term "some embodiments" means "at least some embodiments". The term "and / or" means at least one of the connected objects, for example, A and / or B means three cases including A alone, B alone, and A and B both. Unless otherwise stated, the term "a plurality of" means two or more.

[0020] It is worth noting that the present application is applicable to artificial intelligence chips, which can be graphics processing units (GPU), general-purpose graphics processing units (GPGPU) and other artificial intelligence chips.

[0021] The computing unit (CU) is the core of the artificial intelligence chip, and the kernel in the computing unit is responsible for scheduling and executing thread bundles (warp), and undertakes parallel processing tasks such as attention mechanism calculation tasks. Limited by the artificial intelligence chip, each thread in a thread bundle group is limited to use a maximum of 256 registers (Tread Level Register, TLR), and each thread in a thread block is limited to use a maximum of 512 registers.

[0022] For a 16 thread bundle, the number of matrix elements held by each thread is (16 16) / 32 = 8. The number of registers required by the data held by a thread (i.e. the number of registers held by the thread) is related to the precision type of the data. For a matrix element of half-precision floating point type (fp16 type), such as Figure 2As shown, a certain thread t0 holds 8 matrix elements, corresponding to holding 4 registers r0, r1, r2 and r3. For the matrix elements of the single-precision floating-point type (fp32 type), as shown in the following table, the number of bytes occupied by each matrix element is 4, that is, sizeof (fp32) = 4. Figure 3 As shown, a certain thread t0 holds 8 matrix elements, corresponding to holding 8 registers r0, r1, r2, r3, r4, r5, r6 and r7.

[0023] The group shared memory (GSM) is a storage space shared by the thread bundle group. In the group shared memory, the number of bytes occupied by the matrix elements of different precision types is different. For the matrix elements of the fp16 type, the number of bytes occupied is 2, that is, sizeof (fp16) = 2. For the matrix elements of the fp32 type, the number of bytes occupied is 4, that is, sizeof (fp32) = 4.

[0024] The global memory (GLM) is a storage space shared by the full artificial intelligence chip, and can store various resources.

[0025] It is worth noting that the jth attention score matrix block described in the embodiments of the present application is the jth attention score matrix block of the ith query matrix block currently loaded, and the jth attention result matrix block is the jth attention result matrix block of the ith query matrix block currently loaded.

[0026] For example, the first kernel controls the second thread bundle group to load the 1st query matrix block, controls the second thread bundle group to perform matrix multiplication and exponential operation, obtains the 1st attention score matrix block of the 1st query matrix block, controls the first thread bundle group and the second thread bundle group to alternately perform the preset pre-sequence stage of the attention mechanism operation of the 1st attention score matrix block of the 1st query matrix block, and respectively obtains the intermediate results of different sub-blocks of the 1st attention result matrix block of the 1st query matrix block. Then, the attention mechanism operation of the 2nd attention score matrix block of the 1st query matrix block can be performed, and so on, until the intermediate result of the last sub-block of the last attention result matrix block of the 1st query matrix block is obtained.

[0027] Then, the second kernel controls the third thread bundle group to load the intermediate results of all sub-blocks of the last attention result matrix block of the 1st query matrix block, and perform the preset post-sequence stage of the attention mechanism operation, to obtain the last attention result matrix block of the 1st query matrix block.

[0028] Then, the first kernel controls the second thread bundle group to load a next (i.e., a second) query matrix block, controls the second thread bundle group to perform matrix multiplication and exponential operation, obtains a jth attention score matrix block of the second query matrix block, controls the first thread bundle group and the second thread bundle group to alternately perform a preset pre-sequence stage of attention mechanism operation of the jth attention score matrix block of the second query matrix block, until intermediate results of different sub-blocks of a last attention result matrix block of the second query matrix block are obtained; the second kernel controls the third thread bundle group to load the intermediate results of all sub-blocks of the last attention result matrix block of the second query matrix block, and perform a preset post-sequence stage of attention mechanism operation, and so on, until a last query matrix block is loaded, and a last sub-block of a last attention result matrix block of the last query matrix block is obtained.

[0029] Referring to Figure 1 The attention mechanism calculation method provided by the embodiment of the application comprises: S11, a first kernel controls a second thread bundle group to load an ith query matrix block; S12, the first kernel controls the second thread bundle group to perform matrix multiplication and exponential operation on the ith query matrix block and a transposed jth key matrix block, to obtain a jth attention score matrix block; S13, the first kernel controls a first thread bundle group and the second thread bundle group to alternately perform a preset pre-sequence stage of attention mechanism operation on different sub-blocks of the jth value matrix block and the jth attention score matrix block, to respectively obtain intermediate results of different sub-blocks of a jth attention result matrix block, and let j=j+1, and return to the step of the first kernel controlling the second thread bundle group to perform matrix multiplication and exponential operation on the ith query matrix block and a transposed jth key matrix block, until j=J, and intermediate results of different sub-blocks of a last attention result matrix block of the ith query matrix block are obtained; wherein J is a preset total number of blocks, the jth key matrix block and the jth value matrix block are pre-stored in a shared memory; S14, a second kernel controls a third thread bundle group to load intermediate results of all sub-blocks of a last attention result matrix block of the ith query matrix block, and perform a preset post-sequence stage of attention mechanism operation, to obtain the last attention result matrix block of the ith query matrix block.

[0030] For example, the matrix multiplication of the ith query matrix block and the transposed jth key matrix block can be performed by an mma instruction. The mma instruction requires that the A matrix is in a register and the B matrix is in shared memory, so that the AB matrix can be multiplied. Therefore, the embodiment of the present application can load the ith query matrix block into the register resource of the second thread bundle group in advance, and write the jth key matrix block into the shared memory, so that the jth key matrix block is stored in the shared memory, so that the matrix multiplication can be performed by the mma instruction.

[0031] The embodiment of the present application does not specifically limit the thread bundle group that writes the jth key matrix block and the jth value matrix block into the shared memory, and can be set according to actual needs. For example, before the preset pre-sequence stage of the jth attention score matrix block performing the attention mechanism operation, the first thread bundle group is controlled to write the jth key matrix block and the jth value matrix block from the global memory to the shared memory.

[0032] Generally, in order to improve the calculation efficiency and reduce the data storage overhead, the input data (including the query matrix block, the key matrix block and the value matrix block) is of a low precision type, and the intermediate result (for example, the attention score matrix block) of the attention mechanism operation is also of a low precision type, which is suitable for a scene with low precision requirement. Further, in order to improve the calculation precision, in the case that the input data is of a low precision type, the various intermediate results in the attention mechanism calculation can be converted in precision, which is suitable for a scene with high precision requirement. For example, the output precision type of the mma instruction can be specified to output a high precision result. Assuming that the input data is of a first precision type and the output precision type of the mma instruction is of a second precision type, the output jth attention score matrix block is of the second precision type; wherein the precision of the second precision type is greater than that of the first precision type.

[0033] This configuration of low precision input and high precision output can not only improve the operation speed and reduce the memory occupation by virtue of the efficiency of low precision, but also retain key calculation details by virtue of the high fidelity characteristics of high precision. Of course, the output precision type of the mma instruction can also be the same as the precision type of the input data, for example, both are of the first precision type, so that subsequent precision conversion is not required, which reduces a certain time overhead and is suitable for a scene with low precision requirement. Therefore, the jth attention score matrix block can be of the first precision type or the second precision type.

[0034] The embodiment of the present application prepares the required data before alternately performing the attention mechanism operation, for example, the jth attention score matrix block obtained by performing the matrix multiplication operation and the exponential operation on the ith query matrix block and the transposed jth key matrix block will be used in the preset pre-stage of the jth attention score matrix block of the ith query matrix block in each alternating attention mechanism operation.

[0035] In the preset pre-stage of the alternating attention mechanism operation, the first thread bundle group and the second thread bundle group alternately use different sub-blocks of the jth value matrix block. For example, for the 1st query matrix block, in the first alternating attention mechanism operation, the second thread bundle group uses the first sub-block of the jth value matrix block of the 1st query matrix block; in the second alternating attention mechanism operation, the first thread bundle group uses the second sub-block of the jth value matrix block of the 1st query matrix block; in the third alternating attention mechanism operation, the second thread bundle group uses the third sub-block of the jth value matrix block of the 1st query matrix block, and so on.

[0036] The embodiment of the present application uses two kernels (the first kernel and the second kernel) to complete the attention mechanism operation. When the first kernel performs the preset pre-stage of the jth attention score matrix block, until the intermediate result of the different sub-blocks of the last attention result matrix block of the ith query matrix block is obtained, the second kernel is triggered to start the preset post-stage of the attention mechanism operation. The preset pre-stage and the preset post-stage constitute the entire attention mechanism operation of the ith query matrix block. For example, the synchronization of the two kernels can be performed through the preamble mechanism.

[0037] The embodiment of the present application controls the first thread bundle group and the second thread bundle group to alternately use different sub-blocks of the jth value matrix block through the first kernel to perform the preset pre-stage of the attention mechanism operation with the jth attention score matrix block, and then controls the third thread bundle group to perform the preset post-stage of the attention mechanism operation through the second kernel. Not only does this avoid a single thread bundle group from bearing the full-process attention mechanism calculation, effectively dispersing the load pressure of the single thread bundle group, but also avoids a single kernel from overall scheduling all thread bundle groups, making the division of labor more clear, providing sufficient register resources for large-scale attention mechanism calculation, and finally effectively improving the calculation efficiency of the attention mechanism. In addition, the increase of the second kernel can rearrange the pipeline and reduce the waiting for the release of the shared memory (reuse).

[0038] The embodiments of the present application are particularly suitable for a multi-head latent attention (MLA, Multi-Head Latent Attention) mechanism, because the attention head dimension of the MLA mechanism is obviously larger than that of the traditional attention mechanism, so that the matrix size is obviously increased.

[0039] In an alternative embodiment, the first kernel controls the second thread bundle group to load the i-th query matrix block, comprising: If i = 1, the first kernel controls the first thread bundle group to write the first query matrix block from the global memory to the shared memory, and controls the second thread bundle group to write the first query matrix block from the shared memory to the register group of the second thread bundle group. If i = 2,..., I, the first kernel controls the second thread bundle group to write the i-th query matrix block from the global memory to the register group of the second thread bundle group; wherein I is the last query matrix block.

[0040] It is worth noting that controlling the second thread bundle group to load the i-th query matrix block means controlling the second thread bundle group to load the i-th query matrix block into its own register resource.

[0041] In the related art, for each query matrix block, the first thread bundle group is first controlled to write the i-th query matrix block from the global memory to the shared memory, and then the second thread bundle group is controlled to write the i-th query matrix block from the shared memory to the register group of the second thread bundle group. However, this may cause the query matrix block and the attention result matrix block to share the same shared memory space, so that before the next query matrix block is loaded, all business logic in the inner loop needs to be waited to end, that is, the last attention result matrix block is completely written from the shared memory to the global memory before execution, which makes the tcore (tensor core) and vcore (vector core) hardware computing resources idle (idle state) during this period, resulting in a large compute bubble, which reduces the performance utilization and has room for improvement, such as Figure 4, WG0 is a first thread bundle group, WG1 is a second thread bundle group, O(J)_acc1, O(J)_acc3 are each sub-block of the last 1 attention result matrix block, O(1)_half is a half of the attention result matrix block of the 1st query matrix block, Q(2) is the 2nd query matrix block, KV(1) is the 1st key-value matrix block, K(1) is the 1st key matrix block, T is a transpose, cp represents writing data from global memory to shared memory, ld represents loading data, specifically, writing data from shared memory to a register, st represents writing data from a register to shared memory, tma is used to realize data transfer between global memory and shared memory, here, it refers to writing data from shared memory to global memory.

[0042] Based on this, the embodiments of the present application continue to use related technologies for the 1st query matrix block, and for the 2nd,...,Ith query matrix block, the second thread bundle group is controlled to write the ith query matrix block from global memory directly to the eleventh register group regQ of the second thread bundle group, avoiding the need to wait for all business logic of the last inner loop to end before loading the next query matrix block, thereby avoiding the phenomenon of idle hardware computing resources and improving performance utilization.

[0043] In an optional embodiment, after obtaining the jth attention score matrix block, the method further comprises: The first kernel controls the second thread bundle group to calculate the current row accumulation value of the last 1 attention score matrix block of the ith query matrix block required for attention mechanism operation and write it to global memory; The second kernel controls the third thread bundle group to load the intermediate results of all sub-blocks of the last 1 attention result matrix block of the ith query matrix block and perform a preset post-processing stage of the attention mechanism operation to obtain the last 1 attention result matrix block of the ith query matrix block, including: The second kernel controls the third thread bundle group to write the current row accumulation value of the last 1 attention score matrix block of the ith query matrix block from the global memory to the registers of the third thread bundle group; The second kernel controls the third thread bundle group to load the intermediate results of all sub-blocks of the last 1 attention result matrix block of the ith query matrix block and perform splicing to obtain the intermediate results of the last 1 attention result matrix block of the ith query matrix block, divide it by the current row accumulation value of the last 1 attention score matrix block of the ith query matrix block, and obtain the last 1 attention result matrix block of the ith query matrix block.

[0044] It is worth mentioning that the preset pre-stage refers to a stage in which the current row accumulated value accum_rsum_new is not used to update the (scaled) attention result matrix block, and therefore, the preset post-stage needs to use the value to update the intermediate result of the attention result matrix block to obtain the final attention result matrix block. The current row accumulated value accum_rsum_new is specifically an accumulated value of each row of the current attention score matrix block, that is, an accumulated value of each row of the matrix formed by the first to jth attention score matrix blocks. Each row corresponds to an accumulated value.

[0045] The embodiment of the present application can trigger the second kernel to start performing the preset post-stage of the attention mechanism operation after the first kernel obtains the intermediate results of different sub-blocks of the last attention result matrix block of the ith query matrix block. Specifically, the second kernel can be triggered to execute after the first kernel controls the second thread group to write all sub-blocks of the last attention result matrix block obtained by the second thread group through attention mechanism operation into the global memory, and controls the first thread group to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by the first thread group through attention mechanism operation into the global memory. The second kernel controls the third thread group to load the intermediate results of all sub-blocks of the last attention result matrix block of the ith query matrix block from the global memory to the register group of the third thread group, and performs splicing to obtain the intermediate result of the last attention result matrix block of the ith query matrix block, and divides the intermediate result by the current row accumulated value accum_rsum_new of the jth attention score matrix block to obtain the last attention result matrix block of the ith query matrix block.

[0046] In an optional embodiment, the jth attention score matrix block is written into the first register group of the second thread group and the first register group of the first thread group; The first kernel controls the first thread group and the second thread group to alternately perform the preset pre-stage of the attention mechanism operation on different sub-blocks of the jth value matrix block and the jth attention score matrix block, respectively, to obtain the intermediate results of different sub-blocks of the jth attention result matrix block, including: The first kernel controls the second thread group to perform the preset pre-stage of the attention mechanism operation on the mth sub-block of the jth value matrix block in the shared memory and the jth attention score matrix block in the first register group of the second thread group to obtain the intermediate result of the mth sub-block of the jth attention result matrix block; After obtaining the intermediate result of the mth sub-block of the jth attention result matrix block, the first kernel controls the first thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the jth attention score matrix block in the first register group of the first thread bundle group and the m+1th sub-block of the jth value matrix block in the shared memory, to obtain an intermediate result of the m+1th sub-block of the jth attention result matrix block. After obtaining the intermediate result of the m+1th sub-block of the jth attention result matrix block, let m = m + 1, and return to the step of controlling the second thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the jth attention score matrix block in the first register group of the second thread bundle group and the mth sub-block of the jth value matrix block in the shared memory, until m is the preset total number of sub-blocks, to obtain an intermediate result of the last sub-block of the jth attention result matrix block.

[0047] It is worth noting that the preset pre-sequence stage of attention mechanism operation on the jth attention score matrix block and the mth sub-block of the jth value matrix block involves matrix multiplication of the jth attention score matrix block and the mth sub-block of the jth value matrix block. For example, the matrix multiplication can be performed by an mma instruction. The mma instruction requires that the A matrix is in the register and the B matrix is in the shared memory, so that the AB matrix can be multiplied. Therefore, in the embodiment of the application, the jth attention score matrix block is written into the first register group regP of the second thread bundle group and the first register group regP of the first thread bundle group, and the jth value matrix block is written into the shared memory, so that the matrix multiplication of the jth attention score matrix block and the jth value matrix block can be performed by the mma instruction.

[0048] The embodiment of the application controls different thread bundle groups to write the jth attention score matrix block into the respective first register group regP. Specifically, the second thread bundle group is controlled to write the jth attention score matrix block into the first register group regP of the second thread bundle group; the second thread bundle group is controlled to write the jth attention score matrix block from the first register group regP of the second thread bundle group into the shared memory; and the first thread bundle group is controlled to write the jth attention score matrix block from the shared memory into the first register group regP of the first thread bundle group.

[0049] The embodiment of the application first controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation by using the mth sub-block of the jth value matrix block in the shared memory, then controls the first thread bundle group to perform the preset preliminary stage of attention mechanism operation by using the (m+1)th sub-block of the jth value matrix block in the shared memory, and then sets m=m+1. The process is repeated until m is the total number of preset sub-blocks. It is explained that the preset preliminary stage of attention mechanism operation of the last sub-block of the jth value matrix block is performed at this time. The preset preliminary stage of attention mechanism operation is performed by using the mth (m=total number of preset sub-blocks) sub-block of the jth value matrix block, that is, the last sub-block, to obtain the intermediate result of the last sub-block of the jth attention result matrix block.

[0050] In an optional embodiment, the first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation by using the jth attention score matrix block in the first register group of the second thread bundle group and the mth sub-block of the jth value matrix block in the shared memory to obtain the intermediate result of the mth sub-block of the jth attention result matrix block, including: If j=1, the first kernel controls the second thread bundle group to perform the preset preliminary stage of attention mechanism operation by using the first attention score matrix block in the first register group of the second thread bundle group and the mth sub-block of the first value matrix block in the shared memory to obtain the intermediate result of the mth sub-block of the first attention result matrix block and write it into the shared memory. If j=2,...,J, the first kernel controls the second thread bundle group to write the intermediate result of the mth sub-block of the (j-1)th attention result matrix block from the shared memory into the second register group of the second thread bundle group, and perform the preset preliminary stage of attention mechanism operation by using the jth attention score matrix block in the first register group of the second thread bundle group and the mth sub-block of the jth value matrix block in the shared memory to obtain the intermediate result of the mth sub-block of the jth attention result matrix block and write it into the shared memory.

[0051] It is worth noting that J is the total number of preset blocks, that is, J=total number of key matrix blocks=total number of value matrix blocks=total number of attention score matrix blocks. The embodiment of the application controls the second thread bundle group to use different steps for different attention score matrix blocks. Specifically: For the first attention score matrix block, after the preset preliminary stage of attention mechanism operation is performed by the second thread bundle group to obtain the intermediate result of the mth sub-block of the first attention result matrix block, writing it into the shared memory can avoid register occupation and provide sufficient register resources for subsequent calculation.

[0052] For the 2nd,..., Jth attention score matrix block (i.e. except for the 1st attention score matrix block), the second thread bundle group needs to control to write the intermediate result of the mth sub-block of the previous (i.e. the (j-1)th) attention result matrix block from the shared memory to the second register group regO of the second thread bundle group to participate in the subsequent calculation, so as to realize the reuse of the second register group regO of the second thread bundle group. The embodiment of the present application avoids the risk of performance degradation caused by register overflow by reusing the register resource, and ensures the continuity and stability of the attention mechanism calculation process. Similarly, writing the intermediate result of the mth sub-block of the jth attention result matrix block to the shared memory can avoid register occupation and provide sufficient register resources for subsequent calculation. Moreover, writing the intermediate result of the mth sub-block of the last 1st attention result matrix block of the i th query matrix block to the shared memory provides support for writing data from the shared memory in the preset post-processing stage.

[0053] In an optional embodiment, the first kernel controls the preset pre-processing stage of the first thread bundle group using the (j)th attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the (j)th value matrix block in the shared memory to perform attention mechanism operation, to obtain the intermediate result of the (m+1)th sub-block of the (j)th attention result matrix block, including: If j=1, the first kernel controls the preset pre-processing stage of the first thread bundle group using the 1st attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the 1st value matrix block in the shared memory to perform attention mechanism operation, to obtain the intermediate result of the (m+1)th sub-block of the 1st attention result matrix block and write it to the shared memory; If j=2,...,J, the first kernel controls the first thread bundle group to write the intermediate result of the (m+1)th sub-block of the (j-1)th attention result matrix block from the shared memory to the second register group of the first thread bundle group, and controls the first thread bundle group to perform the preset pre-processing stage of attention mechanism operation using the (j)th attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the (j)th value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the (j)th attention result matrix block and write it to the shared memory.

[0054] It is worth noting that the embodiment of the present application controls the first thread bundle group to use different steps for different attention score matrix blocks, specifically: For the first attention score matrix block, in the preset pre-phase of controlling the first thread bundle group to perform the attention mechanism operation, writing the intermediate result of the m+1th sub-block of the first attention result matrix block to the shared memory can avoid register occupation and provide sufficient register resources for subsequent calculation.

[0055] For the second,..., Jth attention score matrix block (i.e., except for the first and last attention score matrix blocks), the second thread bundle group needs to write the intermediate result of the m+1th sub-block of the (j-1)th attention result matrix block from the shared memory to its second register group regO to participate in subsequent calculation, thereby realizing the reuse of the second register group regO of the first thread bundle group. The embodiments of the present application reuse the register resources to avoid the risk of performance degradation caused by register overflow and ensure the continuity and stability of the attention mechanism calculation process. Similarly, writing the intermediate result of the m+1th sub-block of the jth attention result matrix block to the shared memory can avoid register occupation and provide sufficient register resources for subsequent calculation. Moreover, writing the intermediate result of the m+1th sub-block of the last attention result matrix block to the shared memory provides support for writing data from the shared memory in the preset post-phase.

[0056] In an optional embodiment, the method further comprises: For the ith query matrix block, after the preset pre-phase of the last alternating attention mechanism operation of the second thread bundle group, the first kernel controls the second thread bundle group to write the intermediate results of all sub-blocks of the last attention result matrix block obtained in the preset pre-phase of the attention mechanism operation of the second thread bundle group to the global memory.

[0057] It is worth noting that for the ith query matrix block, the preset pre-phase of the last alternating attention mechanism operation of the second thread bundle group does not refer to the preset pre-phase of the last attention mechanism operation of all alternating processes of the ith query matrix block, but refers to the preset pre-phase of the last attention mechanism operation of the second thread bundle group in all alternating processes of the ith query matrix block, which may be the last one of all alternating processes or may not be the last one of all alternating processes. In other words, after the preset pre-phase of the end of the last alternating attention mechanism operation of the second thread bundle group, the second thread bundle group will not perform the preset pre-phase of the attention mechanism operation of the ith query matrix block again, but the first thread bundle group may still perform the preset pre-phase of the attention mechanism operation of the ith query matrix block, i.e., the first thread bundle group may perform the preset pre-phase of the last attention mechanism operation in step S13.

[0058] In the preset pre-phase of the attention mechanism operation, the second thread bundle group writes the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread bundle group itself into the global memory, for the preset post-phase of the attention mechanism operation of the second kernel.

[0059] It can be understood that the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread bundle group itself in the preset pre-phase of the attention mechanism operation and the intermediate results of all sub-blocks of the last attention result matrix block obtained by the first thread bundle group itself in the preset pre-phase of the alternating attention mechanism operation constitute the intermediate results of the last attention result matrix block of the ith query matrix block.

[0060] In an optional embodiment, the method further comprises: In the preset pre-phase of the attention mechanism operation, the second thread bundle group writes the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread bundle group itself into the global memory, for the preset post-phase of the attention mechanism operation of the second kernel.

[0061] It is worth noting that for the ith query matrix block, the preset pre-phase of the last alternating attention mechanism operation of the first thread bundle group does not refer to the preset pre-phase of the last attention mechanism operation of all alternating processes of the ith query matrix block, but refers to the preset pre-phase of the last attention mechanism operation of the first thread bundle group in all alternating processes of the ith query matrix block, which may be the last one of all alternating processes or may not be the last one of all alternating processes. In other words, after the preset pre-phase of the last alternating attention mechanism operation of the first thread bundle group ends, the first thread bundle group will not perform the preset pre-phase of the attention mechanism operation of the ith query matrix block any more, but the second thread bundle group may still perform the preset pre-phase of the attention mechanism operation of the ith query matrix block, that is, the second thread bundle group may perform the last preset pre-phase of the attention mechanism operation in step S13.

[0062] The first thread bundle group will not participate in the preset pre-phase of the attention mechanism operation any more after the preset pre-phase of the last one of the attention mechanism operation of the first thread bundle group, and at this time, all sub-blocks of the last one of the attention result matrix blocks obtained by the first thread bundle group itself performing the preset pre-phase of the attention mechanism operation are written into the global memory for the preset post-phase of the attention mechanism operation of the second kernel.

[0063] It can be understood that the last one of the attention result matrix blocks refers to the last one of the attention result matrix blocks of the ith query matrix block, and then the intermediate results of all sub-blocks of the last one of the attention result matrix blocks obtained by the second thread bundle group itself performing the preset pre-phase of the attention mechanism operation and the intermediate results of all sub-blocks of the last one of the attention result matrix blocks obtained by the first thread bundle group itself performing the preset pre-phase of the alternating attention mechanism operation are combined to form the intermediate results of the last one of the attention result matrix blocks of the ith query matrix block.

[0064] In an optional embodiment, in the shared memory, the jth value matrix block and the jth key matrix block share the same jth key-value matrix block.

[0065] In the embodiment of the present application, the jth value matrix block and the jth key matrix block share the same key-value matrix block, specifically, the same jth key-value matrix block, which means that the jth key-value matrix block has a common part that is shared by the jth value matrix block and the jth key matrix block, and the remaining part is respectively owned by the jth value matrix block and the jth key matrix block, that is, the jth key-value matrix block includes the jth value matrix block and the jth key matrix block. It can be obviously seen that the data storage overhead is reduced from the original "value matrix block exclusive overhead + key matrix block exclusive overhead" to "value matrix block remaining part overhead + common part overhead + key matrix block remaining part overhead", which effectively reduces the data storage overhead and further alleviates the problem of hardware resource shortage.

[0066] The embodiment of the present application does not specifically limit the part shared by the jth value matrix block and the jth key matrix block, and the shared part and the exclusive part can be set according to actual needs.

[0067] Further, in order to reduce the overhead of data storage, the whole key-value matrix block is divided into a key matrix block or a value matrix block, and a preset part of the key-value matrix block is taken as a shared part. For example, the jth key-value matrix block is taken as the jth key matrix block, and a preset part of the jth key-value matrix block is taken as the jth value matrix block. Specifically, the preset part is a preset column region. Assuming that the number of columns of the jth key-value matrix block is 576, and the preset column region is the first 512 columns, the jth key matrix block and the jth value matrix block share the same 576-column key-value matrix block, the whole 576-column region is taken as the jth key matrix block, and the first 512 columns of the 576-column region are taken as the jth value matrix block.

[0068] Specifically, before the control of the second thread bundle group to perform the matrix multiplication and the exponential operation on the ith query matrix block and the transposed jth key matrix block, the first kernel controls the first thread bundle group to write the jth key-value matrix block into the shared memory, so as to provide data support for the subsequent matrix multiplication and exponential operation of the second thread bundle group.

[0069] In an optional embodiment, the method further comprises: For the ith query matrix block, after a preset previous stage of the last alternating attention mechanism operation, the first kernel controls the first thread bundle group to write the j+xth key-value matrix block from the global memory into the position of the jth key-value matrix block in the shared memory; wherein x is the number of key-value matrix blocks simultaneously stored in the shared memory.

[0070] It is worth noting that, after a preset previous stage of the last alternating attention mechanism operation of the ith query matrix block, the first kernel controls the first thread bundle group to write the j+xth key-value matrix block from the global memory into the shared memory, so as to provide data support for the subsequent control of the second thread bundle group and the first thread bundle group to alternately perform the attention mechanism operation on different subblocks of the j+xth value matrix block, so that the thread bundle group forms a pipeline processing and continues to perform the attention mechanism operation of the ith query matrix block.

[0071] In an optional embodiment, after the first kernel controls the first thread bundle group and the second thread bundle group to alternately perform the attention mechanism operation on different subblocks of the jth value matrix block and the jth attention score matrix block in a preset previous stage, and obtains the intermediate results of different subblocks of the jth attention result matrix block, the method further comprises: The first kernel controls the second thread bundle group to perform the matrix multiplication and the exponential operation on the ith query matrix block and the transposed j+1th key matrix block, to obtain the j+1th attention score matrix block.

[0072] After obtaining the intermediate results of different sub-blocks of the jth attention result matrix block, the first kernel can control the second thread bundle group to perform matrix multiplication and exponential operation on the ith query matrix block and the transposed (j+1)th key matrix block, thereby providing data support for the preset presequence stage of the subsequent attention mechanism operation of the (j+1)th attention score matrix block.

[0073] In an optional embodiment, the size of the query matrix block, the size of the key-value matrix block, the size of the key matrix block and the size of the value matrix block are determined by the following steps: The first objective function is constructed as min(S max1 -S r1 -S u1 ); wherein S max1 is the maximum number of registers used by each thread in the first thread bundle group, S r1 is the first preset number of reserved registers, and S u1 is the actual number of registers used by each thread in the first thread bundle group. The second objective function is constructed as min(S max2 -S r2 -S u2 ); wherein S max2 is the maximum number of registers used by each thread in the second thread bundle group, S r2 is the second preset number of reserved registers, and S u2 is the actual number of registers used by each thread in the second thread bundle group. The third objective function is constructed as min(M max1 -M max2 ); wherein M max1 is the maximum capacity of the shared memory, and M max2 is the maximum value of the actual total capacity of the shared memory occupied by the first thread bundle group and the second thread bundle group at the same time. The first objective function, the second objective function and the third objective function are jointly solved to obtain the size of the query matrix block, the size of the key-value matrix block, the size of the key matrix block and the size of the value matrix block.

[0074] It is worth noting that if the control thread bundle group directly processes the query matrix and the key-value matrix to perform attention mechanism calculation, the hardware resources will be exhausted due to the large size of the matrix (especially in the case of long sequences), and a large amount of bandwidth will be occupied, which will easily block, and the intermediate result storage may exceed the limit, causing the calculation to be interrupted. Therefore, the query matrix block and the key-value matrix block are used for attention mechanism calculation in the embodiments of the present application, so that each thread focuses on local operation of an independent block, and efficient and stable execution of the attention mechanism is realized under hardware resource constraints.

[0075] It is worth noting that the third thread bundle group only needs the register resources of the intermediate results of all sub-blocks of the last 1 attention result matrix block of the i-th query matrix block and the current row accumulation value of the last 1 attention score matrix block of the i-th query matrix block, which requires register resources far less than 512 registers, and does not occupy shared memory. Therefore, the register resources and shared memory resources used by the third thread bundle group are not discussed here. Only the hardware resources occupied by the first thread bundle group and the second thread bundle group are discussed below.

[0076] Due to the hardware resource limitation of the artificial intelligence chip: (1) When each thread bundle group executes a task, the number of registers used by each thread in each thread bundle group should not exceed the maximum number of registers limited for each thread in the thread bundle group; specifically, the first actual number of registers S u1 used by each thread in the first thread bundle group when the first thread bundle group executes the attention mechanism should be less than or equal to the maximum number of registers S max1 limited for each thread in the first thread bundle group; the first actual number of registers S u2 used by each thread in the second thread bundle group when the second thread bundle group executes the attention mechanism should be less than or equal to the maximum number of registers S max2 limited for each thread in the second thread bundle group. In addition, some registers can be reserved for each thread bundle group for auxiliary operations by the compiler to ensure that the thread bundle group efficiently and stably completes the task. For example, S r1 registers are reserved for the first thread bundle group, and S r2 registers are reserved for the second thread bundle group; wherein S r1 is the first preset number of reserved registers, S r2 is the second preset number of reserved registers, S r1 and S r2 may be equal or not equal.

[0077] (2) The first thread bundle group and the second thread bundle group perform attention mechanism calculation in parallel, so they share 1 shared memory when processing in parallel. The first thread bundle group and the second thread bundle group are alternate attention mechanisms, so the actual total capacity of the shared memory that needs to be occupied at different times may be different. Therefore, when the first thread bundle group and the second thread bundle group perform attention mechanism calculation in parallel, the maximum value M of the actual total capacity of the shared memory that needs to be occupied at the same time max2 , cannot exceed the maximum capacity M of the shared memory max1 .

[0078] S u1 , S u2 and M max2 are related to the sizes of the query matrix block, the key-value matrix block, the key matrix block and the value matrix block. Based on this, the embodiments of the application construct and solve the first objective function min(S max1 -S r1 -S u1 ), the second objective function min(S max2 -S r2 -S u2 ) and the third objective function min(M max1 -M max2 ) to obtain appropriate sizes of the query matrix block, the key-value matrix block, the key matrix block and the value matrix block, so as to ensure that when the attention mechanism operation is performed using the matrix block, the hardware resources are fully utilized while avoiding register overflow and shared memory contention.

[0079] Specifically, based on a first preset relationship, a first objective function is constructed; the first preset relationship at least includes a relationship between the size of the query matrix block, the size of the key matrix block, the size of the value matrix block and S u1 .

[0080] Based on a second preset relationship, a second objective function is constructed; the second preset relationship at least includes a relationship between the size of the query matrix block, the size of the key matrix block, the size of the value matrix block and S u2 . Based on a third preset relationship, a third objective function is constructed; the third preset relationship at least includes a relationship between the size of the query matrix block, the size of the key-value matrix block, the size of the key matrix block, the size of the value matrix block and M max2 . Further, some sizes of the matrix block can be preset, and then the first objective function, the second objective function and the third objective function are jointly solved. For example, the attention head dimension of the key matrix block and the attention head dimension of the value matrix block can be determined according to the task complexity.

[0081] In an optional embodiment, the third preset relationship further includes a relationship between the target number and M max2 . The target number includes a number of key-value matrix blocks simultaneously stored in the shared memory and a number of query matrix blocks simultaneously stored in the shared memory.

[0082] It can be understood that the more the number of key-value matrix blocks and the number of query matrix blocks simultaneously stored in the shared memory, the more the shared memory occupied, and therefore, M max2 is also related to the number of key-value matrix blocks and the number of query matrix blocks simultaneously stored in the shared memory. Embodiments of the present application do not specifically limit the number of key-value matrix blocks and the number of query matrix blocks simultaneously stored in the shared memory, which can be set according to actual needs and only needs to be a positive integer.

[0083] In an optional embodiment, the control of the second thread bundle group to perform the matrix multiplication operation and the exponential operation on the i th query matrix block and the transposed j th key matrix block to obtain the j th attention score matrix block includes: controlling the second thread bundle group to perform the matrix multiplication operation on the i th query matrix block and the transposed j th key matrix block to obtain the j th attention score matrix block before scaling and writing the j th attention score matrix block before scaling into the third register group of the second thread bundle group, performing the exponential operation on the j th attention score matrix block in the third register group of the second thread bundle group to obtain the j th attention score matrix block and cover-write the j th attention score matrix block into the third register group of the second thread bundle group, writing the j th attention score matrix block in the third register group of the second thread bundle group into the first register group of the second thread bundle group after precision conversion, and writing the j th attention score matrix block in the first register group of the second thread bundle group into the shared memory; controlling the first thread bundle group to write the j th attention score matrix block in the shared memory into the first register group of the first thread bundle group.

[0084] It is worth noting that, in the embodiments of the present application, the j th key matrix block in the shared memory is first transposed to obtain the transposed j th key matrix block, then the i th query matrix block and the transposed j th key matrix block are used to perform the matrix multiplication operation to obtain the j th attention score matrix block before scaling, and then the exponential operation is performed on the j th attention score matrix block before scaling to obtain the j th attention score matrix block. Specifically, according to exp((S-accum_rmax) exp (Sj / sf), wherein exp() represents a natural exponential function, Sj represents the jth attention score matrix block before scaling, accum_rmax represents a historical global row maximum value, and sf represents a scaling factor.

[0085] The embodiments of the present application perform matrix multiplication and exponential operation by multiplexing the third register group regS, avoid the risk of performance degradation caused by register overflow, and ensure the continuity and stability of the attention mechanism calculation process.

[0086] Further, in order to improve the calculation accuracy, the jth attention score matrix block obtained after the matrix multiplication and the exponential operation is of a second precision type. Therefore, in order to reduce the data storage overhead in the shared memory, the jth attention score matrix block in the third register group regS of the second thread bundle group is converted in precision and written into the first register group regP of the second thread bundle group, and the jth attention score matrix block in the first register group regP of the second thread bundle group is written into the shared memory, so that the jth attention score matrix block written into the shared memory is of a first precision type.

[0087] In an alternative embodiment, after obtaining the jth attention score matrix block, the method further comprises: The first kernel controls the second thread bundle group to calculate the row target value of the jth attention score matrix block required for the attention mechanism operation and write it into the first register group regP of the second thread bundle group; wherein the row target value includes: a historical row accumulation value, a correction value for correcting the historical row accumulation value, a current row accumulation value, a row maximum value, a historical global row maximum value, a current global row maximum value, and a row accumulation value; The first kernel controls the second thread bundle group to write the correction value of the jth attention score matrix block in the first register group regP of the second thread bundle group into the shared memory; The first kernel controls the first thread bundle group to write the correction value of the jth attention score matrix block in the shared memory into the first register group regP of the first thread bundle group.

[0088] It is worth noting that the first thread bundle group and the second thread bundle group both undertake the attention mechanism operation, and the data required for the attention mechanism operation, in addition to the attention score matrix block, also includes the row target value of the attention score matrix block, which, in addition to the above-mentioned current row accumulation value accum_rsum_new, also includes: accum_rsum, which is an accumulated value of each row of the historical attention score matrix blocks, i.e. an accumulated value of each row of a matrix formed by the first to the j-1th attention score matrix blocks; wherein each row corresponds to an accumulated value.

[0089] modifier, which is used to modify the historical row accumulated value accum_rsum to obtain a current row accumulated value accum_rsum_new.

[0090] Further, the calculation of the current row accumulated value accum_rsum_new and the modifier also needs to use the following row target value: chunk_rmax, which is a maximum value of each row of the attention score matrix blocks, i.e. a maximum value of each row of the jth attention score matrix block; wherein each row corresponds to a maximum value.

[0091] accum_rmax, which is a maximum value of each row of the historical global attention score matrix blocks, i.e. a maximum value of each row of a matrix formed by the first to the j-1th attention score matrix blocks; wherein each row corresponds to a maximum value.

[0092] accum_rmax_new, which is a maximum value of each row of the current global attention score matrix blocks, i.e. a maximum value of each row of a matrix formed by the first to the jth attention score matrix blocks; wherein each row corresponds to a maximum value.

[0093] chunk_rsum, which is an accumulated value of each row of the attention score matrix blocks, i.e. an accumulated value of each row of the jth attention score matrix block; wherein each row corresponds to an accumulated value.

[0094] Therefore, the second thread bundle group is controlled to calculate the above-mentioned 7 row target values, including the historical row accumulation value accum_rsum, the current row accumulation value accum_rsum_new, the correction value modifier, the row maximum value chunk_rmax, the historical global row maximum value accum_rmax, the current global row maximum value accum_rmax_new, and the row accumulation value chunk_rsum, which are sequentially stored in the fourth register group, the fifth register group, the sixth register group, the seventh register group, the eighth register group, the ninth register group, and the tenth register group of the second thread bundle group. The first thread bundle group only needs the correction value modifier, which can be stored in the sixth register group of the first thread bundle group. Among them, the fourth register group, the fifth register group, and the sixth register group belong to the first register group; the seventh register group, the eighth register group, the ninth register group, and the tenth register group belong to the second register group.

[0095] In a specific embodiment, the hardware resource condition of a certain artificial intelligence chip is: (1) each thread in each thread bundle group is limited to use at most 256 registers; (2) the shared memory is provided with two memory bank regions, including a first memory bank region and a second memory bank region, so that data in the two regions can be accessed at the same time, improving the memory access efficiency. Further, the shared memory can divide the first memory bank region and the second memory bank region according to the address, for example, the first memory bank region is a high bank, i.e., a high address region in the shared memory, which is divided by 96 KB; the second memory bank region is a low bank, i.e., a low address region in the shared memory, which is divided by 128 KB. Of course, it can also be divided according to other methods, which is not limited here.

[0096] For the artificial intelligence chip, each matrix block is set to perform the attention mechanism provided in the application according to the attention head dimension (head dim) The size of the sequence length (seqlen); the size of the query matrix block Q is: 576 64; the size of the key-value matrix block KV is 576 32; take the complete key-value matrix block as the key matrix block K, and the size of the key matrix block is: 576 32; take the front 512 column region as the value matrix block V, and the size of the value matrix block is: 512 32.

[0097] For example, Figure 5, the query matrix block Q and the transposed key matrix block K are multiplied to obtain the attention score matrix block S before scaling, the attention score matrix block S before scaling is subjected to exponential operation and precision conversion to obtain the attention score matrix block P, and the preset presequence stage of the alternate attention mechanism operation of the sub-blocks of the attention score matrix block P and the sub-blocks of the value matrix block V is used to obtain the intermediate results of the sub-blocks of the attention result matrix block O. Figure 5 In the embodiment, the value matrix block is evenly divided into four sub-blocks according to the direction of the attention head dimension, including a first sub-block V_0, a third sub-block V_1, a second sub-block V_2 and a fourth sub-block V_3, and the size of each sub-block is 128 32, thereby obtaining the intermediate results of the four sub-blocks of the attention result matrix block O, including the intermediate result O_acc0 of the first sub-block, the intermediate result O_acc1 of the third sub-block, the intermediate result O_acc2 of the second sub-block and the intermediate result O_acc3 of the fourth sub-block. Then, the first thread bundle group and the second thread bundle group each use half of the value matrix block to alternate, and finally obtain the intermediate results of half of the attention result matrix block.

[0098] As Figure 6-8 , first of all, cp means writing data from global memory to shared memory, ld means loading data, specifically writing data from shared memory to register, st means writing data from register to shared memory, tma is used to realize data transfer between global memory and shared memory, here it means writing data from shared memory to global memory, and ldbulk means batch loading data from global memory to register.

[0099] The first kernel (specifically, ComputeKernel, a computing kernel): The first thread bundle group WG0 controls the writing of the first query matrix block Q(1), the first key-value matrix block KV(1) (including the first key matrix block K(1) and the first value matrix block V(1)), and the second key-value matrix block KV(2) (including the second key matrix block K(2) and the second value matrix block V(2)) into the shared memory.

[0100] The second thread bundle group WG1 writes Q(1) into its own register, and performs matrix multiplication (i.e., Q(1) K(1) T , T represents transposition) of Q(1) K(1) TAfter that, the operation before the alternating attention mechanism operation, such as the exponential operation, the precision conversion, the calculation of the row target value, etc. Then the first attention score matrix block P(1) is written into the shared memory, and the preset pre-phase of the alternating attention mechanism operation of WG1 can be started.

[0101] WG0 writes P(1) into its own register, and the preset pre-phase of the alternating attention mechanism operation of WG0 can be started.

[0102] j=1, WG1 uses P(1) and the first sub-block V(1)_0 of the first value matrix block to perform the preset pre-phase of the attention mechanism operation, obtains the intermediate result of the first sub-block O(1)_acc0 of the first attention result matrix block and writes it into the shared memory; WG0 uses P(1) and the second sub-block V(1)_2 of the first value matrix block to perform the preset pre-phase of the attention mechanism operation, obtains the intermediate result of the second sub-block O(1)_acc2 of the first attention result matrix block; WG1 uses P(1) and the third sub-block V(1)_1 of the first value matrix block to perform the preset pre-phase of the attention mechanism operation, obtains the intermediate result of the third sub-block O(1)_acc1 of the first attention result matrix block and writes it into the shared memory, at this time, for WG1, it is the last time of using Q(1), K(1) and V(1) to operate, and the subsequent operation of Q(1), K(2) and V(2) can be started; WG0 uses P(1) and the fourth sub-block V(1)_3 of the first value matrix block to perform the attention mechanism operation, and obtains the fourth sub-block O(1)_acc3 of the first attention result matrix block. Since after WG0 uses P(1) and the fourth sub-block V(1)_3 of the first value matrix block to perform the attention mechanism operation, it is not necessary to use V(1) and P(1) to perform the operation, at this time, the WG0 can be controlled to write the third key-value matrix block KV(3) into the position of KV1 of the shared memory.

[0103] j=2,...,J-1, in the process of alternating attention operation, before the preset pre-phase of the attention mechanism operation using the jth attention score matrix block P(j) and the sub-block V(j)_acc of the jth value matrix block V(j), the intermediate result O(j-1)_acc of the sub-block of the (j-1)th attention result matrix block O(j-1) is written into the required thread group. The remaining process is similar to the above and will not be described here.

[0104] j=J, after the preset preamble phase of attention operation is completed each time, the sub-block O(J)_acc of the Jth attention result matrix block O(J) is obtained (the first thread bundle group and the second thread bundle group each have half of the intermediate result of the Jth attention result matrix block, i.e., each has O(i)_half), and the preamble trigger signal is sent to the second kernel (specifically, the PostprocessKernel, the post-processing kernel) by writing to the global memory through the shared memory; After the second kernel detects that the preamble trigger signal is received, the third thread bundle group is controlled to load the intermediate results of all sub-blocks of the Jth attention result matrix block to splice the intermediate result O(i) of the Jth attention result matrix block of the ith query block, and a preset post-processing phase of the attention mechanism operation is started, including: rescale O(i) specifically refers to dividing the intermediate result of the last 1 attention result matrix block of the ith query matrix by the current row cumulative value of the last 1 attention score matrix block of the ith query matrix, typecaste O(i) refers to data type conversion of the result obtained by rescale O(i), and tmaO(i) refers to writing the result of typecaste O(i) to the global memory, at this time, the preamble complete signal is sent to indicate that the attention mechanism calculation of the ith query matrix block is completed.

[0105] The above process is specifically performed according to the following instructions: Instruction logic of the first thread bundle group: Loop{ / Q outer loop / {tma cp Q / / write the 1st query matrix block from global memory to shared memory (only for the 1st query matrix block); } Loop{ / KV inner loop / tma cp KV / / write the key-value matrix block KV from global memory to shared memory; ld.gsm modifier / / write the modifier from shared memory to the first register group of the first thread bundle group; Ldmatrix P / / write the attention score matrix block p from shared memory to the first register group regP of the first thread bundle group; Loop { / half seqlen of V (half of V) / {init regO / / initialize the second register group regO of the first thread bundle group; regO = mma(P, V, O) / / preset preamble phase of attention mechanism operation on the first attention score matrix block in the first register group regP of the first thread bundle group and the m+1 sub-block of the first value matrix block in the shared memory; Stmatrix regO to O_acc(gsm) / / write the intermediate result of the m+1 sub-block of the first attention result matrix block to the shared memory; } / / for the first attention score matrix block; {ldmatrix regO / / write the intermediate result of the m+1 sub-block of the j-1 attention result matrix block from the shared memory to the second register group regO of the first thread bundle group; regO = modifier regO regO = mma(P, V, regO) / / and the previous instruction, form a preset preamble phase of attention mechanism operation on the intermediate result of the m+1 sub-block of the j-1 attention result matrix block in the second register group regO of the first thread bundle group, the j attention score matrix block in the first register group regP of the first thread bundle group, and the m+1 sub-block of the j value matrix block in the shared memory; Stmatrix regO to O_acc(gsm) / / write the intermediate result of the m+1 block of the j attention result matrix block to the shared memory; } / / for other (i.e. except the first) attention score matrix block; } Tma half O / / write the intermediate result of the last attention result matrix block sub-block obtained by attention mechanism operation on itself to the global memory (the Tma half O of the first thread bundle group will be completed faster than that of the second thread bundle group); supaTriggerProgrammaticLaunchCompletion() / / preamble trigger, specifically used to trigger the preamble subsequent code of the second kernel; } } Instruction logic of the second thread bundle group: Loop{ / Q outer loop / Ldmatrix Q ->regQ / / write query matrix chunk Q from shared memory to the 11th register group regQ of the second thread bundle group (only for the 1st query matrix chunk); ldbulk Q to reg + shuffle reg layout / / write other (except the 1st) query matrix chunks from global memory to the 11th register group regQ of the second thread bundle group, and process the data of the registers through a data shuffle operation; Loop{ / KV inner loop / regS = mma(regQ, smemK, regS) / / perform matrix multiplication between the query matrix chunk in the 11th register group regQ of the second thread bundle group and the transposed key matrix chunk to obtain an attention score matrix chunk S in fp32, and write it to the 3rd register group regS of the second thread bundle group; calc chunk_rmax, accum_rmax, accum_rmax_new / / calculate the row maximum value chunk_rmax, the historical global row maximum value accum_rmax, and the current global row maximum value accum_rmax_new; exp((S-accum_rmax) sf ), reuse regS / / perform exponential operation on the attention score matrix chunk in the 3rd register group regS of the second thread bundle group to obtain an attention score matrix chunk, and overwrite it to the 3rd register group regS of the second thread bundle group; calc chunk_rsum, accum_rsum, accum_rsum_new / / calculate the row accumulation value chunk_rsum, the historical row accumulation value accum_rsum, and the current row accumulation value accum_rsum_new; Modifier=exp((accum_rmax – accum_rmax_new) sf ) / / calculate the modifier value Modifier; accum_rsum_new = accum_rsum Modifier + chunk_rsum / / update the current row accumulation value accum_rsum_new; Typecast regS to regP / / precision conversion of attention score matrix in the third register group regS of the second thread bundle group to the first register group regP of the second thread bundle group, implementing precision conversion from fp32 to fp16; Stmatrix regP to smemP / / write of attention score matrix block in the first register group regP of the second thread bundle group to shared memory; St.gsm modifier / / write of modifier to shared memory; Loop { / half seqlen of V (half of V) / {init regO / / initialization of the second register group regO of the second thread bundle group; regO = mma(P, V, O) / / preset pre-phase of attention mechanism operation of the first attention score matrix block in the first register group regP of the second thread bundle group and the m-th sub-block of the first value matrix block in the shared memory; Stmatrix regO to O_acc(gsm) / / intermediate result of the m-th sub-block of the first attention result matrix block and write to shared memory; } / / for the first attention score matrix block; {ldmatrix regO / / write of the m-th sub-block of the j-1-th attention result matrix block from shared memory to the second register group regO of the second thread bundle group; regO = modifier regO regO = mma(P, V, regO) / / and the previous instruction, form the preset pre-phase of attention mechanism operation of the intermediate result of the m-th sub-block of the j-1-th attention result matrix block of the second register group regO of the second thread bundle group, the j-th attention score matrix block in the first register group regP of the second thread bundle group and the m+1-th sub-block of the j-th value matrix block in the shared memory; Stmatrix regO to O_acc(gsm) / / write of the intermediate result of the m-th sub-block of the j-th attention result matrix block to shared memory; } / / for other (i.e. except the first) attention score matrix block; } st.glm acc_rsum_new / / write the current row accumulation value acc_rsum_new to global memory; Tma half O / / write the intermediate result of the sub-block of the last attention result matrix block obtained by attention mechanism operation on itself to global memory; } } Instruction logic of the third thread bundle group: Loop{ / Q outer loop / supaGridDependencySynchronize() / / preamble synchronization, specifically used to accept the trigger of the first kernel; ldbulk O / / directly bulk load the intermediate result of the sub-block of the attention result matrix block from global memory to the register group of the third thread bundle group, and splice it into the intermediate result of the attention result matrix block; ld.glm acc_rsum_new / / write the current row accumulation value acc_rsum_new from global memory to the register group of the third thread bundle group; regO = regO / acc_rsum_new / / divide the intermediate result of the attention result matrix block by the current row accumulation value acc_rsum_new to obtain the attention result matrix block; Typecase regO to regO_f16 / / perform data type conversion on the attention result matrix block, for example, from fp32 to fp16; Stbulk regO_f16 / / bulk write the attention result matrix block after data type conversion from the register group of the third thread bundle group to global memory; } The register resources used by the first thread bundle group in this scenario are as shown in Table 1.

[0106] Table 1 In Table 1, the actual number of registers required by each thread of the first thread bundle group includes the number of registers of the sixth register group, the first register group regP, and the second register group regO.

[0107] For the sixth register group, since the matrix elements held by 1 thread are distributed in two rows, and each row needs to be calculated and store the modifier, the sixth register group needs to include 2 registers.

[0108] The matrix data stored in the first register group regP is written into the shared memory after precision conversion by the second thread bundle group, and is read from the shared memory, so the number of registers of the first register group regP = the number of registers of regS of the second thread bundle group / 2 = 8.

[0109] The number of registers of the second register group regO = the number of registers required by 1 thread to hold the matrix elements stored in the register group The size of the sub-block of the attention result matrix block / (the number of thread bundles of the first thread bundle group The matrix size held by 1 thread bundle).

[0110] The register resources used by the second thread bundle group in this scenario are shown in Table 2; wherein stmatrix_num is the minimum number of registers used when executing the instruction stmatrix, which is generally 4.

[0111] Table 2 In Table 2, the actual number of registers required by each thread of the second thread bundle group includes the number of registers of the eleventh register group regQ, the third register group regS, the seventh register group, the eighth register group, the ninth register group, the tenth register group, the fourth register group, the fifth register group, the sixth register group, the first register group regP, the second register group regO, and the twelfth register group regO_fp16.

[0112] The number of registers of the eleventh register group regQ = the number of registers required by 1 thread to hold the matrix elements stored in the register group The size of the query matrix block / (the number of thread bundles of the second thread bundle group The matrix size held by 1 thread bundle); The number of registers of the third register group regS = the number of registers required by 1 thread to hold the matrix elements stored in the register group The size of the attention score matrix block / (the number of thread bundles of the second thread bundle group The matrix size held by 1 thread bundle); Like the first thread bundle group, the seventh register group, the eighth register group, the ninth register group, the tenth register group, the fourth register group, the fifth register group, and the sixth register group each include 2 registers.

[0113] The matrix data stored in the first register group regP is obtained after the second thread bundle group performs precision conversion on the attention score matrix block, and thus the number of registers of the first register group regP = the number of registers of regS of the second thread bundle group / 2 = 8.

[0114] The number of registers of the second register group regO = the number of registers required by one thread to hold the matrix element stored in the register group The size of the sub-block of the attention result matrix block / the number of thread bundles of the second thread bundle group The size of the matrix held by one thread bundle.

[0115] The third thread bundle group needs to occupy far less than 512 registers, and thus the register resources used by the third thread bundle group are not discussed here.

[0116] The shared memory resources used in this scenario are as shown in Table 3, and the shared memory resource distribution of each stage is as shown in Figure 9-10 The shared memory (GSM) is divided into a high bank region and a low bank region, the high bank region occupies 96 KB, and the low bank region occupies 128 KB. The white part in the black solid line box in the figure represents the unoccupied shared memory space. The space occupied by the row target value and the remaining row values is too small, only 8B, and is not counted here. Figure 9 For the stage from the beginning to j = 1, only the query matrix block Q (occupying 72 KB), the key-value matrix block KV (each KV occupying 36 KB), and the attention score matrix block P (occupying 2 KB) occupy the shared memory. Figure 10 For the stage from the beginning to j = 1 to the beginning to j = J before the end, the query matrix block Q is released (indicated by a dashed line) and does not occupy the shared memory. At this time, the sub-blocks O( j )_acc0, O( j )_acc1, O( j )_acc2, O( j )_acc3 of the attention result matrix block O( j )_acc (each occupying 32 KB), the key-value matrix block KV, and the attention score matrix block P occupy the shared memory.

[0117] Table 3 It can be seen that the attention mechanism calculation provided in the embodiments of the present application can accurately adapt to the register and shared memory resource characteristics of the artificial intelligence chip, the hardware resources used do not exceed the use limit, and as much as possible hardware resources are mobilized to participate in calculation, the algorithm potential of the hardware is fully released, the continuity and stability of the attention mechanism calculation process are ensured, and high calculation performance is realized on the basis of efficient resource utilization.

[0118] Referring to Figure 11 The artificial intelligence chip provided by the embodiment of the present application comprises a global memory (GLM) and one or more computing units; each computing unit comprises a first core, a second core, a shared memory (GSM), one or more first thread bundle groups, one or more second thread bundle groups, and one or more third thread bundle groups. regQ in the register space is the eleventh register group, regS is the third register group, regP is the first register group, regO is the second register group, and regO_fp16 is the twelfth register group.

[0119] In addition, the embodiment of the present application further provides a computer readable storage medium comprising a stored computer program; wherein the computer program controls the device where the computer readable storage medium is located to execute the attention mechanism calculation method according to any one of the above embodiments when running.

[0120] In addition, the embodiment of the present application further provides a computer program product comprising a computer program or instructions, which, when executed by a processor, implement the attention mechanism calculation method according to any one of the above embodiments.

[0121] Referring to Figure 12 The attention mechanism calculation device 20 provided by the embodiment of the present application comprises a processor 21, a memory 22, and a computer program stored in the memory 22 and executable on the processor 21. The processor 21 implements the steps in the above attention mechanism calculation method embodiments when executing the computer program. Alternatively, the processor 21 implements the functions of the modules / units in the above device embodiments when executing the computer program.

[0122] For example, the computer program can be divided into one or more modules / units, which are stored in the memory 22 and executed by the processor 21 to complete the present application. The one or more modules / units can be a series of computer program instruction segments capable of completing a specific function, which are used to describe the execution process of the computer program in the attention mechanism calculation device 20.

[0123] The attention mechanism calculation device 20 can include, but is not limited to, the processor 21 and the memory 22. Those skilled in the art can understand that the schematic diagram is only an example of the attention mechanism calculation device 20 and does not constitute a limitation on the attention mechanism calculation device 20, which can include more or fewer components than the diagram, or combine certain components, or different components, for example, the attention mechanism calculation device 20 can also include an input / output device, a network access device, a bus, etc.

[0124] The processor 21 can be a central processing unit (CPU), and can also be other general-purpose processors, digital signal processors (DSP), application specific integrated circuits (ASIC), field-programmable gate arrays (FPGA) or other programmable logic devices, discrete gates or transistor logic components, discrete hardware components, or the like. The general-purpose processor can be a microprocessor or the processor can also be any conventional processor. The processor 21 is a control center of the attention mechanism computing device 20, and is connected with various parts of the attention mechanism computing device 20 through various interfaces and lines.

[0125] The memory 22 can be used to store computer programs and / or modules. The processor 21 realizes various functions of the attention mechanism computing device 20 by running or executing the computer programs and / or modules stored in the memory 22, and calling data stored in the memory 22. The memory 22 can mainly include a program storage area and a data storage area. The program storage area can store an operating system, at least one application program required for a function (such as a sound playing function, an image playing function, etc.), and the like. The data storage area can store data created according to use of the mobile phone (such as audio data, a phone book, etc.), and the like. In addition, the memory 22 can include a high-speed random access memory, and can also include a nonvolatile memory, such as a hard disk, a memory, a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, at least one disk storage device, a flash memory device, or other volatile solid-state memory device.

[0126] The modules / units integrated in the attention mechanism calculation device 20 can be stored in a computer readable storage medium if they are implemented in the form of software function units and sold or used as independent products. Based on this understanding, all or part of the processes in the above-mentioned embodiment methods can also be completed by a computer program instructing related hardware. The computer program can be stored in a computer readable storage medium. When the computer program is executed by the processor 21, the steps of the above-mentioned various method embodiments can be implemented. The computer program includes computer program code, which can be in the form of source code, object code, executable files, or some intermediate forms, etc. The computer readable medium can include any entity or device capable of carrying the computer program code, recording medium, U disk, mobile hard disk, magnetic disk, optical disk, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signal, telecommunication signal, and software distribution medium, etc.

[0127] It should be noted that the apparatus embodiments described above are merely illustrative, and the units described as separate components can or can not be physically separated, and the components shown as units can or can not be physical units, i.e., they can be located in one place or distributed on multiple network units. Part or all of the modules can be selected to achieve the purpose of the embodiment scheme according to actual needs. In addition, the connection relationship between the modules in the apparatus embodiment provided by the present application indicates that there is a communication connection between them, which can be implemented as one or more communication buses or signal lines. Those of ordinary skill in the art can understand and implement it without creative labor.

[0128] The above is the preferred embodiment of the present application. It should be noted that for those of ordinary skill in the art, without departing from the principles of the present application, a number of improvements and refinements can be made, which are also considered within the scope of protection of the present application.

Claims

1. A method for calculating an attention mechanism, characterized in that, include: The first kernel controls the second thread bundle to load the i-th query matrix block; The first kernel controls the second thread group to perform matrix multiplication and exponentiation operations using the i-th query matrix block and the transposed j-th key matrix block to obtain the j-th attention score matrix block; The first kernel controls the first thread group and the second thread group to alternately use different sub-blocks of the j-th value matrix block to perform attention mechanism operations with the j-th attention score matrix block in a preset pre-stage, obtaining intermediate results of different sub-blocks of the j-th attention result matrix block respectively. Let j = j + 1, and return to the step where the first kernel controls the second thread group to perform matrix multiplication and exponentiation operations using the i-th query matrix block and the transposed j-th key matrix block, until j = J, obtaining intermediate results of different sub-blocks of the last attention result matrix block of the i-th query matrix block; where J is the preset total number of blocks, and the j-th key matrix block and the j-th value matrix block are pre-stored in shared memory; The second kernel controls the third thread bundle to load the intermediate results of all sub-blocks of the last attention result matrix block of the i-th query matrix block, and performs the preset post-stage of the attention mechanism operation to obtain the last attention result matrix block of the i-th query matrix block.

2. The attention mechanism calculation method as described in claim 1, characterized in that, The first kernel controls the second thread group to load the i-th query matrix block, including: If i=1, then the first kernel controls the first thread group to write the first query matrix block from global memory to the shared memory, and controls the second thread group to write the first query matrix block from the shared memory to the eleventh register group of the second thread group; If i = 2, ..., I, then the first kernel controls the second thread group to write the i-th query matrix block from the global memory to the eleventh register group of the second thread group; where I is the last query matrix block.

3. The attention mechanism calculation method as described in claim 1, characterized in that, After obtaining the j-th attention score matrix block, the method further includes: The first kernel controls the second thread group to calculate the current row accumulation value of the last attention score matrix block of the i-th query matrix block required for attention mechanism operation, and writes it to global memory; The second kernel controls the third thread group to load the intermediate results of all sub-blocks of the last attention result matrix block of the i-th query matrix block, and performs a preset post-processing stage of the attention mechanism operation to obtain the last attention result matrix block of the i-th query matrix block, including: The second kernel controls the third thread bundle to write the current row accumulation value of the last attention score matrix block of the i-th query matrix block from the global memory to the register of the third thread bundle; The second kernel controls the third thread bundle to load and concatenate the intermediate results of all sub-blocks of the last attention result matrix block of the i-th query matrix block, and then divides the intermediate result of the last attention result matrix block of the i-th query matrix block by the current row accumulation value of the last attention score matrix block of the i-th query matrix block to obtain the last attention result matrix block of the i-th query matrix block.

4. The attention mechanism calculation method as described in claim 1, characterized in that, The j-th attention score matrix block is written to the first register group of the second thread bundle and the first register group of the first thread bundle; The first kernel controls the first thread group and the second thread group to alternately utilize different sub-blocks of the j-th value matrix block to perform attention mechanism operations with the j-th attention score matrix block in a preset pre-stage, thereby obtaining intermediate results of different sub-blocks of the j-th attention result matrix block, including: The first kernel controls the second thread bundle to perform a preset pre-stage of attention mechanism operation using the j-th attention score matrix block in its first register group and the m-th sub-block of the j-th value matrix block in the shared memory, so as to obtain the intermediate result of the m-th sub-block of the j-th attention result matrix block; After obtaining the intermediate result of the m-th sub-block of the j-th attention result matrix block, the first kernel controls the first thread bundle group to perform a preset pre-stage of attention mechanism operation using the j-th attention score matrix block in its own first register group and the m+1-th sub-block of the j-th value matrix block in the shared memory, so as to obtain the intermediate result of the m+1-th sub-block of the j-th attention result matrix block. After obtaining the intermediate result of the (m+1)th sub-block of the j-th attention result matrix block, let m=m+1, and return to the preset pre-stage step in which the first kernel controls the second thread group to perform attention mechanism operation using the j-th attention score matrix block in its own first register group and the m-th sub-block of the j-th value matrix block in the shared memory, until m is the preset total number of sub-blocks, and obtain the intermediate result of the last sub-block of the j-th attention result matrix block.

5. The attention mechanism calculation method as described in claim 4, characterized in that, The first kernel controls the second thread group to perform a preset pre-stage of attention mechanism operation using the j-th attention score matrix block in its first register group and the m-th sub-block of the j-th value matrix block in the shared memory, to obtain the intermediate result of the m-th sub-block of the j-th attention result matrix block, including: If j=1, the first kernel controls the second thread bundle to perform a preset pre-stage of attention mechanism operation using the first attention score matrix block in its first register group and the m-th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the m-th sub-block of the first attention result matrix block and write it into the shared memory. If j=2,...,J, then the first kernel controls the second thread group to write the intermediate result of the m-th sub-block of the (j-1)-th attention result matrix block from the shared memory to the second register group of the second thread group, and performs a preset pre-stage of attention mechanism operation with the j-th attention score matrix block in the first register group of the second thread group and the m-th sub-block of the j-th value matrix block in the shared memory, to obtain the intermediate result of the m-th sub-block of the j-th attention result matrix block and write it to the shared memory.

6. The attention mechanism calculation method as described in claim 4, characterized in that, The first kernel controls the first thread group to perform a preset pre-stage of attention mechanism operation using the j-th attention score matrix block in its first register group and the (m+1)-th sub-block of the j-th value matrix block in the shared memory, to obtain the intermediate result of the (m+1)-th sub-block of the j-th attention result matrix block, including: If j=1, the first kernel controls the first thread bundle to perform a preset pre-stage of attention mechanism operation using the first attention score matrix block in its first register group and the (m+1)th sub-block of the first value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the first attention result matrix block and write it into the shared memory. If j=2,...,J, then the first kernel controls the first thread group to write the intermediate result of the (m+1)th sub-block of the (j-1)th attention result matrix block from the shared memory to the second register group of the first thread group, and performs a preset pre-stage of attention mechanism operation with the j-th attention score matrix block in the first register group of the first thread group and the (m+1)th sub-block of the j-th value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the j-th attention result matrix block and write it to the shared memory.

7. The attention mechanism calculation method as described in claim 1, characterized in that, The method further includes: For the i-th query matrix block, after the preset pre-stage of the last alternating attention mechanism operation of the second thread bundle, the first kernel controls the second thread bundle to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by itself in the preset pre-stage of the attention mechanism operation into global memory.

8. The attention mechanism calculation method as described in claim 1, characterized in that, The method further includes: For the i-th query matrix block, after the preset pre-stage of the last alternating attention mechanism operation of the first thread bundle, the first kernel controls the first thread bundle to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by itself in the preset pre-stage of the attention mechanism operation into global memory.

9. The attention mechanism calculation method as described in claim 1, characterized in that, In the shared memory, the j-th value matrix block and the j-th key matrix block share the same j-th key-value matrix block.

10. The attention mechanism calculation method as described in claim 9, characterized in that, The method further includes: For the i-th query matrix block, after the preset pre-processing stage of the last alternating attention mechanism operation, the first kernel controls the first thread group to write the j+x-th key-value matrix block from global memory to the position of the j-th key-value matrix block in the shared memory; where x is the number of key-value matrix blocks stored simultaneously in the shared memory.

11. A computing device with an attention mechanism, characterized in that, The system includes a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, wherein the processor, when executing the computer program, implements the attention mechanism computation method as described in any one of claims 1 to 10.

12. A computer-readable storage medium, characterized in that, The computer-readable storage medium includes a stored computer program; wherein, when the computer program is executed, it controls the device on which the computer-readable storage medium is located to perform the attention mechanism calculation method as described in any one of claims 1 to 10.

13. A computer program product, characterized in that, It includes a computer program or instructions that, when executed by a processor, implement the attention mechanism computation method as described in any one of claims 1 to 10.

Images