A method, device, medium, and product for calculating attention mechanisms

Abstract

This application discloses an attention mechanism calculation method, device, medium, and product. The method includes: a first kernel controlling a second thread group to load the i-th query matrix block; the first kernel controlling the second thread group to calculate the j-th attention score matrix block; the first kernel controlling the first and second thread groups to alternately utilize different sub-blocks of the j-th value matrix block to perform a preset pre-processing stage of attention mechanism operation with the j-th attention score matrix block, until intermediate results of different sub-blocks of the last attention result matrix block of the i-th query matrix block are obtained; and the second kernel controlling a 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 perform a preset post-processing stage of attention mechanism operation. Using the embodiments of this application can provide sufficient register resources for attention mechanism calculation, thereby improving computational efficiency.

Description

Technical Field

[0001] This application relates to the field of artificial intelligence technology, and in particular to a computational method, device, medium and product for attention mechanism. Background Technology

[0002] In the process of attention mechanism computation, three types of matrices are mainly involved: query matrix, key matrix, and value matrix. However, as the matrix size increases significantly, existing technologies have revealed obvious defects: (1) At the warpgroup level, the entire attention mechanism computation is performed by relying on only a single warpgroup, which leads to excessive register pressure due to task concentration; (2) At the kernel scheduling level, the above warpgroups are controlled by only a single kernel, which coordinates the scheduling and execution of the entire computation task. The combination of these two problems makes it impossible for register resources to meet the needs of large-scale attention mechanism computation, ultimately limiting the computational efficiency of attention mechanism. Summary of the Invention

[0003] This application provides an attention mechanism computation method, device, medium, and product to solve the problems in the prior art where attention mechanism computation relies on a single thread bundle to execute the entire process, and where the thread bundle is controlled by a single kernel, resulting in insufficient register resources to meet large-scale attention mechanism computation and ultimately limiting the computational efficiency of the attention mechanism.

[0004] To achieve the above objectives, embodiments of this application provide an attention mechanism calculation method, including:

[0005] The first kernel controls the second thread bundle to load the i-th query matrix block;

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

[0007] 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;

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

[0009] As an improvement to the above scheme, the first kernel controls the second thread bundle to load the i-th query matrix block, including:

[0010] 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;

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

[0012] As an improvement to the above scheme, after obtaining the j-th attention score matrix block, the method further includes:

[0013] 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;

[0014] 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:

[0015] 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;

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

[0017] As an improvement to the above scheme, the j-th attention score matrix block is written into the first register group of the second thread bundle group and the first register group of the first thread bundle group;

[0018] 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:

[0019] 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;

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

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

[0022] As an improvement to the above scheme, 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:

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

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

[0025] As an improvement to the above scheme, 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:

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

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

[0028] As an improvement to the above solution, the method further includes:

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

[0030] As an improvement to the above solution, the method further includes:

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

[0032] As an improvement to 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.

[0033] As an improvement to the above solution, the method further includes:

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

[0035] Compared with the prior art, the attention mechanism calculation method, device, medium and product provided in this application embodiment, through the first kernel controlling the first thread bundle group and the second thread bundle group to alternately utilize different sub-blocks of the j-th value matrix block and the j-th attention score matrix block to perform a preset pre-sequence stage of attention mechanism calculation, and then the second kernel controlling the third thread bundle group to perform a preset post-sequence stage of attention mechanism calculation, not only avoids a single thread bundle group bearing the entire attention mechanism calculation process and effectively distributes the register pressure of a single thread bundle group, but also avoids a single kernel coordinating and scheduling all thread bundle groups, making the division of labor clearer, providing sufficient register resources for large-scale attention mechanism calculation, and ultimately effectively improving the calculation efficiency of attention mechanism. Attached Figure Description

[0036] Figure 1 This is a flowchart of an attention mechanism calculation method provided in an embodiment of this application;

[0037] Figure 2 This is a schematic diagram of the distribution of matrix elements and registers held by a thread, provided in an embodiment of this application;

[0038] Figure 3 This is a schematic diagram illustrating another thread-held matrix element and register distribution provided in an embodiment of this application;

[0039] Figure 4 This is a pipeline diagram of an attention mechanism calculation method provided in an embodiment of this application;

[0040] Figure 5This is a schematic diagram of an attention mechanism calculation provided in an embodiment of this application;

[0041] Figure 6 This is a pipeline diagram of a thread bundle group provided in an embodiment of this application;

[0042] Figure 7 This is another pipeline diagram of a thread bundle provided in an embodiment of this application;

[0043] Figure 8 This is yet another pipeline diagram of a thread bundle provided in the embodiments of this application;

[0044] Figure 9 This is a distribution diagram of shared memory usage in the early stage provided in an embodiment of this application;

[0045] Figure 10 This is a resource distribution diagram of shared memory usage in a later stage provided in an embodiment of this application;

[0046] Figure 11 This is a structural block diagram of an artificial intelligence chip provided in an embodiment of this application;

[0047] Figure 12 This is a structural block diagram of an attention mechanism computing device provided in an embodiment of this application. Detailed Implementation

[0048] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those of ordinary skill in the art without creative effort are within the scope of protection of this application.

[0049] In the description of this application, the sequence number of each process does not imply the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of this application.

[0050] In this application description, the terms "exemplary" or "for example" are used to indicate examples, illustrations, or explanations. Any embodiment or design described as "exemplary" or "for example" in the embodiments of this application should not be construed as being more preferred or advantageous than other embodiments or designs. Specifically, the use of terms such as "exemplary" or "for example" is intended to present the relevant concepts in a specific manner.

[0051] In this application description, the terms "first," "second," etc., are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of this application described herein can be implemented, for example, in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus. The term "based on" means "at least partially based on." The term "according to" means "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, such as A and / or B, indicating three cases: including only A, only B, and both A and B. Unless otherwise stated, the term "multiple" means two or more.

[0052] It is worth noting that this application applies to artificial intelligence chips, which can be AI chips such as graphics processing units (GPUs) and general-purpose graphics processing units (GPGPUs).

[0053] The Computing Unit (CU) is the core of an AI chip. The kernel within the CU is responsible for scheduling and executing warps, undertaking parallel processing tasks such as attention mechanism computations. Due to the limitations of AI chips, each thread within a warp group is limited to using a maximum of 256 Thread Level Registers (TLRs), and each thread within a thread block is limited to using a maximum of 512 registers.

[0054] For a single held matrix of size 16 A 16-thread bundle, where each thread holds a matrix element of (16) 16) / 32 = 8. The number of registers required for data held by one thread (i.e., the number of registers held by that thread) depends on the precision type of the data. For matrix elements of half-precision floating-point type (fp16 type), such as Figure 2As shown, a thread t0 holds 8 matrix elements, corresponding to 4 registers: r0, r1, r2, and r3. For single-precision floating-point (FP32) matrix elements, such as... Figure 3 As shown, a certain thread t0 holds 8 matrix elements, corresponding to 8 registers: r0, r1, r2, r3, r4, r5, r6, and r7.

[0055] Group Shared Memory (GSM) is a storage space shared by a group of threads. In GSM, the number of bytes occupied by matrix elements of different precision types varies. For a matrix element of type fp16, it occupies 2 bytes, i.e., sizeof(fp16) = 2. For a matrix element of type fp32, it occupies 4 bytes, i.e., sizeof(fp32) = 4.

[0056] Global memory (GLM) serves as a shared storage space across all AI chips, capable of storing various resources.

[0057] It is worth noting that, in the embodiments of this application, the j-th attention score matrix block is the j-th attention score matrix block of the currently loaded i-th query matrix block, and the j-th attention result matrix block is the j-th attention result matrix block of the currently loaded i-th query matrix block.

[0058] For example, the first kernel controls the second thread group to load the first query matrix block, and controls the second thread group to perform matrix multiplication and exponentiation operations to obtain the first attention score matrix block of the first query matrix block. The first and second thread groups are then controlled to alternately perform a preset pre-processing stage for the attention mechanism operation of the first attention score matrix block of the first query matrix block, respectively obtaining intermediate results of different sub-blocks of the first attention result matrix block of the first query matrix block. Next, the attention mechanism operation of the second attention score matrix block of the first 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 first query matrix block is obtained.

[0059] Then, 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 first query matrix block, and performs the preset post-stage of attention mechanism operation to obtain the last attention result matrix block of the first query matrix block.

[0060] Next, the first kernel controls the second thread group to load the next (i.e., the second) query matrix block, and controls the second thread group to perform matrix multiplication and exponentiation operations to obtain the j-th attention score matrix block of the second query matrix block. The first and second thread groups are controlled to alternately perform the preset pre-sequence stage of the attention mechanism operation of the j-th attention score matrix block of the second query matrix block until the intermediate results of different sub-blocks of the last attention result matrix block of the second query matrix block are obtained. 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 second query matrix block and perform the preset post-sequence stage of the attention mechanism operation. This process continues until the last query matrix block is loaded, and the last sub-block of the last attention result matrix block of the last query matrix block is obtained.

[0061] See Figure 1 The attention mechanism calculation method provided in this application includes:

[0062] S11, The first kernel controls the second thread bundle to load the i-th query matrix block;

[0063] S12. The first kernel controls the second thread group to perform matrix multiplication and exponential 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.

[0064] S13. 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, and obtain 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, and obtain 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;

[0065] S14. 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 attention mechanism operation to obtain the last attention result matrix block of the i-th query matrix block.

[0066] For example, the matrix multiplication operation between the i-th query matrix block and the transposed j-th key matrix block can be performed using the mma instruction. The mma instruction requires matrix A to be in a register and matrix B to be in shared memory so that matrices A and B can be multiplied. Therefore, in this embodiment, the i-th query matrix block can be preloaded into the register resources of the second thread bundle and the j-th key matrix block can be written into shared memory, so that the j-th key matrix block is stored in shared memory, thus enabling the matrix multiplication operation to be performed using the mma instruction.

[0067] This application does not impose specific restrictions on the thread group that writes the j-th key matrix block and the j-th value matrix block to shared memory; these restrictions can be set according to actual needs. For example, before the preset pre-stage of attention mechanism operation on the j-th attention score matrix block, the first thread group is controlled to write the j-th key matrix block and the j-th value matrix block from global memory to shared memory.

[0068] Generally, to improve computational efficiency and reduce data storage overhead, input data (including query matrix blocks, key matrix blocks, and value matrix blocks) is of low precision type. Simultaneously, the intermediate results of the attention mechanism computation (e.g., attention score matrix blocks) are also of low precision type, suitable for scenarios with low precision requirements. Furthermore, to improve computational precision, when the input data is of low precision type, the precision of various intermediate results in the attention mechanism computation can be converted, suitable for scenarios with high precision requirements. For example, the output precision type of the mma instruction can be specified to output a high-precision result. Assuming the input data is of first precision type and the output precision type of the mma instruction is of second precision type, then the j-th attention score matrix block output will be of second precision type; where the precision of the second precision type is greater than that of the first precision type.

[0069] This configuration of low-precision input and high-precision output can leverage the efficiency of low precision to improve computational speed and reduce memory usage, while retaining critical computational details through the high-fidelity characteristics of high precision. Of course, the output precision type of the mma instruction can also be the same as the input data precision type, for example, both being first-precision type. This eliminates the need for subsequent precision conversion, reducing time overhead and making it suitable for scenarios with low precision requirements. Therefore, the j-th attention score matrix block can be either first-precision type or second-precision type.

[0070] Before alternating attention mechanism operations, the embodiments of this application prepare the necessary data, such as the j-th attention score matrix block obtained by performing matrix multiplication and exponentiation operations on the i-th query matrix block and the transposed j-th key matrix block. This j-th attention score matrix block will be used in the preset pre-stage of each alternating attention mechanism operation of the j-th attention score matrix block of the i-th query matrix block.

[0071] During the pre-defined pre-processing phase of the alternating attention mechanism operation, the first thread group and the second thread group alternately utilize different sub-blocks of the j-th value matrix block. For example, for the first query matrix block, during the first alternating attention mechanism operation, the second thread group utilizes the first sub-block of the j-th value matrix block of the first query matrix block; during the second alternating attention mechanism operation, the first thread group utilizes the second sub-block of the j-th value matrix block of the first query matrix block; during the third alternating attention mechanism operation, the second thread group utilizes the third sub-block of the j-th value matrix block of the first query matrix block, and so on.

[0072] This application embodiment utilizes two kernels (a first kernel and a second kernel) to perform the attention mechanism operation. When the first kernel performs the preset pre-order stage for the j-th attention score matrix block, until it obtains the intermediate results of different sub-blocks of the last attention result matrix block of the i-th query matrix block, the second kernel is triggered to start executing the preset post-order stage of the attention mechanism operation. This preset pre-order stage and the preset post-order stage constitute the entire attention mechanism operation for the i-th query matrix block. For example, the two kernels can be synchronized through a preamble mechanism.

[0073] This application embodiment controls a first thread group and a second thread group to alternately utilize different sub-blocks of the j-th value matrix block and the j-th attention score matrix block for a preset pre-sequence stage of attention mechanism computation. Then, a second kernel controls a third thread group for a preset post-sequence stage of attention mechanism computation. This not only avoids a single thread group bearing the entire attention mechanism computation process, effectively distributing the load on a single thread group, but also avoids a single kernel coordinating and scheduling all thread groups, resulting in a clearer division of labor. This provides sufficient register resources for large-scale attention mechanism computation, ultimately improving the computational efficiency of the attention mechanism. Furthermore, adding a second kernel enables pipeline reordering, reducing waiting time for shared memory (reuse) release.

[0074] The embodiments of this application are particularly applicable to the multi-head latent attention (MLA) mechanism, because the attention head dimension of the MLA mechanism is significantly larger than that of the traditional attention mechanism, resulting in a significant increase in matrix size.

[0075] In one optional embodiment, the first kernel controls the second thread bundle to load the i-th query matrix block, including:

[0076] If i=1, 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 register group of the second thread group.

[0077] 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 into the register group of the second thread group; where I is the last query matrix block.

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

[0079] In related technologies, for each query matrix block, the first thread group first writes the i-th query matrix block from global memory to shared memory, and then the second thread group writes the i-th query matrix block from shared memory to the register group of the second thread group. However, this can lead to the query matrix block and the attention result matrix block reusing the same shared memory space. Therefore, before loading the next query matrix block, it is necessary to wait for all business logic of the previous inner loop to finish, that is, for the previous attention result matrix block to be completely written from shared memory to global memory before execution. This results in the tcore (tensor core) and vcore (vector core) hardware computing resources being idle during this period, with a large compute bubble, leading to reduced performance utilization and room for improvement. Figure 4WG0 is the first thread bundle, WG1 is the second thread bundle, O(J)_acc1 and O(J)_acc3 are the sub-blocks of the last attention result matrix block, O(1)_half is half of the attention result matrix block of the first query matrix block, Q(2) is the second query matrix block, KV(1) is the first key-value matrix block, K(1) is the first key matrix block, T is transpose, cp means writing data from global memory to shared memory, ld means loading data, specifically writing data from shared memory to registers, st means writing data from registers to shared memory, tma is used to implement data transfer between global memory and shared memory, here it means writing data from shared memory to global memory.

[0080] Based on this, the embodiments of this application use the relevant technology for the first query matrix block, while for the second, ..., I query matrix blocks, the second thread group is controlled to write the i-th query matrix block directly from global memory to the eleventh register group regQ of the second thread group. This avoids waiting for all business logic of the previous inner loop to finish before loading the next query matrix block, thereby avoiding the phenomenon of hardware computing resources being idle and improving performance utilization.

[0081] In an optional embodiment, after obtaining the j-th attention score matrix block, the method further includes:

[0082] 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;

[0083] 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:

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

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

[0086] It's worth noting that the preset preceding stage refers to the stage where the attention result matrix blocks are not updated (scaled) using the current row's accumulated value `accum_rsum_new`. Therefore, the preset following stage needs to use this value to update the intermediate results of the attention result matrix blocks to obtain the final attention result matrix blocks. Specifically, the current row's accumulated value `accum_rsum_new` is the accumulated value of each row in the current attention score matrix block, that is, the accumulated value of each row in the matrix formed by the 1st to jth attention score matrix blocks; each row corresponds to one accumulated value.

[0087] In this embodiment, after the first kernel obtains the intermediate results of different sub-blocks of the last attention result matrix block of the i-th query matrix block, it can trigger the second kernel to start executing a preset post-stage of the attention mechanism operation. Specifically, the first kernel can control the second thread group to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory, and control the first thread group to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory, and then trigger the second kernel to execute. 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 from global memory into its own register group, and concatenate them to obtain the intermediate result of the last attention result matrix block of the i-th query matrix block. This intermediate result is then divided by the current row accumulation value accum_rsum_new of the j-th attention score matrix block to obtain the last attention result matrix block of the i-th query matrix block.

[0088] In an optional embodiment, the j-th attention score matrix block is written into the first register group of the second thread bundle and the first register group of the first thread bundle.

[0089] 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:

[0090] 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;

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

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

[0093] It is worth noting that the preset pre-processing stage of the attention mechanism operation using the m-th sub-block of the j-th attention score matrix block and the j-th value matrix block involves matrix multiplication of the m-th sub-block of the j-th attention score matrix block and the j-th value matrix block. For example, matrix multiplication can be performed using the mma instruction. The mma instruction requires matrix A to be in a register and matrix B to be in shared memory so that matrices A and B can be multiplied. Therefore, in this embodiment, the j-th attention score matrix block is written to the first register group regP of the second thread bundle and the first register group regP of the first thread bundle, and the j-th value matrix block is written to shared memory. This allows the matrix multiplication of the j-th attention score matrix block and the j-th value matrix block to be performed using the mma instruction.

[0094] In this embodiment, different thread bundles are controlled to write the j-th attention score matrix in blocks into their respective first register groups regP. Specifically, the second thread bundle is controlled to write the j-th attention score matrix in blocks into its own first register group regP; the second thread bundle is controlled to write the j-th attention score matrix in blocks from its own first register group regP into shared memory; and the first thread bundle is controlled to write the j-th attention score matrix in blocks from shared memory into its own first register group regP.

[0095] In this embodiment, the second thread group is first controlled to perform a preset pre-stage of attention mechanism operation using the m-th sub-block of the j-th value matrix block in shared memory. Then, the first thread group is controlled to perform a preset pre-stage of attention mechanism operation using the (m+1)-th sub-block of the j-th value matrix block in shared memory. Then, m=m+1 is set, and this process is repeated until m equals the total number of preset sub-blocks. This indicates that the preset pre-stage of attention mechanism operation for the last sub-block of the j-th value matrix block has been reached. At this point, the preset pre-stage of attention mechanism operation for the m-th (m=total number of preset sub-blocks) sub-block of the j-th value matrix block, i.e., the last sub-block, is used to perform the preset pre-stage of attention mechanism operation, and the intermediate result of the last sub-block of the j-th attention result matrix block is obtained.

[0096] In an optional embodiment, the first kernel controls the second thread group to perform a preset pre-processing 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:

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

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

[0099] It is worth noting that J is the preset total number of blocks, i.e., J = total number of key matrix blocks = total number of value matrix blocks = total number of attention score matrix blocks; the embodiments of this application control the second thread bundle to use different steps for different attention score matrix blocks, specifically:

[0100] For the first attention score matrix block, in the preset pre-processing stage of controlling the second thread bundle to perform attention mechanism operations, after obtaining the intermediate result of the m-th sub-block of the first attention result matrix block, writing it into shared memory can avoid register occupation and provide sufficient register resources for subsequent calculations.

[0101] For the 2nd, ..., Jth attention score matrix blocks (i.e., excluding the 1st attention score matrix block), the second thread group needs to first write the intermediate result of the m-th sub-block of the previous (i.e., the (j-1th)th attention result matrix block) from shared memory into its own second register group regO to participate in subsequent calculations, thereby realizing the reuse of the second register group regO of the second thread group. This embodiment avoids the risk of performance degradation caused by register overflow by reusing register resources, ensuring the continuity and stability of the attention mechanism calculation process. Similarly, writing the intermediate result of the m-th sub-block of the j-th attention result matrix block into shared memory avoids register occupation and provides sufficient register resources for subsequent calculations. Moreover, writing the intermediate result of the m-th sub-block of the last attention result matrix block of the i-th query matrix block into shared memory supports the pre-set subsequent stage of writing data from shared memory.

[0102] In an optional embodiment, the first kernel controls the first thread group to perform a preset pre-processing 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:

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

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

[0105] It is worth noting that, in the embodiments of this application, the first thread bundle group uses different steps for different attention score matrix blocks, specifically:

[0106] For the first attention score matrix block, in the preset pre-processing stage of controlling the first thread bundle to perform attention mechanism operations, after obtaining the intermediate result of the (m+1)th sub-block of the first attention result matrix block, writing it into shared memory can avoid register occupation and provide sufficient register resources for subsequent calculations.

[0107] For the 2nd, ..., Jth attention score matrix blocks (i.e., excluding the 1st and last attention score matrix blocks), the second thread group needs to control the second thread group to first write the intermediate result of the (m+1)th sub-block of the previous (j-1)th attention result matrix block from shared memory into its own second register group regO to participate in subsequent calculations, thereby realizing the reuse of the second register group regO of the first thread group. This embodiment of the application avoids the risk of performance degradation caused by register overflow by reusing register resources, ensuring the continuity and stability of the attention mechanism calculation process. Similarly, writing the intermediate result of the (m+1)th sub-block of the jth attention result matrix block into shared memory avoids register occupation and provides sufficient register resources for subsequent calculations. Furthermore, writing the intermediate result of the (m+1)th sub-block of the last attention result matrix block into shared memory supports the pre-set subsequent stage of writing data from shared memory.

[0108] In an optional embodiment, the method further includes:

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

[0110] It is worth noting that, for the i-th query matrix block, the preset pre-order stage of the last alternating attention mechanism operation of the second thread bundle does not refer to the preset pre-order stage of the last attention mechanism operation of all alternating processes of the i-th query matrix block, but rather to the preset pre-order stage of the last attention mechanism operation of the second thread bundle during all alternating processes of the i-th query matrix block. This may be the last of all alternating processes, or it may not be the last of all alternating processes. In other words, after the preset pre-order stage of the last alternating attention mechanism operation of the second thread bundle ends, the second thread bundle will no longer perform the preset pre-order stage of the attention mechanism operation of the i-th query matrix block, but the first thread bundle may still perform the preset pre-order stage of the attention mechanism operation of the i-th query matrix block, that is, the first thread bundle may perform the preset pre-order stage of the last attention mechanism operation in step S13.

[0111] In this embodiment of the application, 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 second thread bundle will no longer participate in the attention mechanism operation. At this time, the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread bundle itself in the preset pre-stage of the attention mechanism operation are written into global memory for the second kernel to perform the preset post-stage of the attention mechanism operation.

[0112] It is understandable that the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread group itself in the preset pre-stage 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 group itself in the preset pre-stage of the alternating attention mechanism operation constitute the intermediate result of the last attention result matrix block of the i-th query matrix block.

[0113] In an optional embodiment, the method further includes:

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

[0115] It is worth noting that, for the i-th query matrix block, the preset pre-order stage of the last alternating attention mechanism operation of the first thread bundle does not refer to the preset pre-order stage of the last attention mechanism operation of all alternating processes of the i-th query matrix block, but rather to the preset pre-order stage of the last attention mechanism operation of the first thread bundle during all alternating processes of the i-th query matrix block. This may be the last of all alternating processes, or it may not be the last of all alternating processes. In other words, after the preset pre-order stage of the last alternating attention mechanism operation of the first thread bundle ends, the first thread bundle will no longer perform the preset pre-order stage of the attention mechanism operation of the i-th query matrix block, but the second thread bundle may still perform the preset pre-order stage of the attention mechanism operation of the i-th query matrix block, that is, the second thread bundle may perform the preset pre-order stage of the last attention mechanism operation in step S13.

[0116] In this embodiment of the application, 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 thread bundle will no longer participate in the preset pre-stage of the attention mechanism operation. At this time, the intermediate results of all sub-blocks of the last attention result matrix block obtained by the first thread bundle itself in the preset pre-stage of the attention mechanism operation are written into global memory for the second kernel to perform the preset post-stage of the attention mechanism operation.

[0117] It is understandable that the last attention result matrix block refers to the last attention result matrix block of the i-th query matrix block. Then, the intermediate results of all sub-blocks of the last attention result matrix block obtained by the second thread group itself in the preset pre-stage of attention mechanism operation and the intermediate results of all sub-blocks of the last attention result matrix block obtained by the first thread group itself in the preset pre-stage of alternating attention mechanism operation constitute the intermediate result of the last attention result matrix block of the i-th query matrix block.

[0118] In one alternative embodiment, 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.

[0119] In this embodiment, the j-th value matrix block and the j-th key matrix block share the same key-value matrix block. Specifically, they share the same j-th key-value matrix block. This means that a common portion of the j-th key-value matrix block is shared by both the j-th value matrix block and the j-th key matrix block, while the remaining portion is shared by each of the j-th value matrix block and the j-th key matrix block. In other words, the j-th key-value matrix block includes both the j-th value matrix block and the j-th key matrix block. This clearly shows that the data storage overhead is reduced from the original "overhead of the value matrix block alone + overhead of the key matrix block alone" to "overhead of the remaining portion of the value matrix block + overhead of the common portion + overhead of the remaining portion of the key matrix block," effectively reducing data storage overhead and further alleviating the problem of hardware resource constraints.

[0120] This application does not impose specific restrictions on the shared portion of the j-th value matrix block and the j-th key matrix block; the shared portion and the exclusive portion can be set according to actual needs.

[0121] Furthermore, to reduce data storage overhead, the entire key-value matrix is ​​divided into key matrix blocks or value matrix blocks, with a preset portion within each block being a shared portion. For example, the j-th key-value matrix block is designated as the j-th key matrix block; a preset portion within the j-th key-value matrix block is designated as the j-th value matrix block. Specifically, the preset portion is a preset column area. Assuming the j-th key-value matrix block has 576 columns, and the preset column area is the first 512 columns, then the j-th key matrix block and the j-th value matrix block share the same 576-column key-value matrix block. The complete 576-column area is taken as the j-th key matrix block, and the first 512 columns within it are taken as the j-th value matrix block.

[0122] Specifically, before controlling 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, the first thread group is controlled to write the j-th key-value matrix block into the shared memory to provide data support for the subsequent matrix multiplication and exponentiation operations performed by the second thread group.

[0123] In an optional embodiment, the method further includes:

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

[0125] It is worth noting that, in this embodiment, after the preset pre-stage of the last alternating attention mechanism operation for the i-th query matrix block, the first thread group is controlled to write the j+x-th key-value matrix block from global memory to shared memory. This provides data support for the subsequent control of the second thread group and the first thread group to alternately use different sub-blocks of the j+x-th value matrix block for attention mechanism operation, so that the thread group forms a pipeline process and continues to perform the attention mechanism operation for the i-th query matrix block.

[0126] In an optional embodiment, after 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-processing stage, and after obtaining intermediate results of different sub-blocks of the j-th attention result matrix block, the method further includes:

[0127] 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+1-th key matrix block to obtain the j+1-th attention score matrix block.

[0128] In this embodiment of the application, after obtaining the intermediate results of different sub-blocks of the j-th attention result matrix block, the first kernel can control the second thread bundle to perform matrix multiplication and exponentiation operations using the i-th query matrix block and the transposed j+1-th key matrix block, providing data support for the preset pre-stage of attention mechanism operation of the j+1-th attention score matrix block.

[0129] 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 through the following steps:

[0130] The first objective function is min(S) max1 -S r1 -S u1 ); where S max1 To limit the maximum number of registers used by each thread within the first thread bundle, S r1 The number of registers reserved for the first preset number, S u1 This represents the actual number of registers that each thread within the first thread bundle needs to use.

[0131] Construct the second objective function as min(S) max2 -S r2 -S u2 ); where S max2 To limit the maximum number of registers used by each thread within the second thread bundle, Sr2 The number of spaces reserved for the second preset register, S u2 This represents the actual number of registers that each thread within the second thread bundle needs to use;

[0132] Construct the third objective function as min(M) max1 -M max2 ); where M max1 M represents the maximum capacity of the shared memory. max2 This represents the maximum total amount of shared memory that the first thread bundle and the second thread bundle need to occupy at the same time.

[0133] By jointly solving the first objective function, the second objective function, and the third objective function, 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 obtained.

[0134] It is worth noting that if the control thread group directly processes the query matrix and key-value matrix for attention mechanism calculation, the large size of the matrix (especially in long sequence scenarios) will exhaust hardware resources, consume a large amount of bandwidth, and easily cause blocking. The storage of intermediate results may exceed the limit, leading to calculation interruption. Therefore, this embodiment uses query matrix block and key-value matrix block for attention mechanism calculation, allowing each thread to focus on the local operation of an independent block, achieving efficient and stable execution of the attention mechanism under hardware resource constraints.

[0135] It's worth noting that the third thread bundle only requires register resources for the intermediate results of all sub-blocks of the last attention result matrix block of the i-th query matrix block and the current row's accumulated value of the last attention score matrix block of the i-th query matrix block. This requires far fewer than 512 registers and does not occupy shared memory. Therefore, the register and shared memory resources used by the third thread bundle are not discussed here. The following discussion focuses only on the hardware resources required by the first and second thread bundles.

[0136] Due to limitations in AI chip hardware resources: (1) When each thread bundle executes a task, the number of registers used by each thread within each thread bundle must not exceed the maximum number of registers that each thread within that thread bundle is limited to use; specifically, when the first thread bundle executes the attention mechanism, the first actual number of registers S that each thread within it needs to use. u1 It should be less than or equal to the maximum number of registers S that each thread in the first thread bundle can use. max1 The number of actual registers S that each thread within the second thread bundle needs to use when implementing the attention mechanism. u2 It should be less than or equal to the maximum number of registers S that each thread in the second thread bundle can use.max2 Furthermore, some registers can be reserved for each thread bundle for auxiliary operations by the compiler, ensuring that the thread bundles complete tasks efficiently and stably. For example, S can be reserved for the first thread bundle. r1 One register is reserved for the second thread bundle. r2 There are 1 register; among them, S r1 The number of registers reserved for the first preset number, S r2 The number of spaces reserved for the second preset register, S r1 and S r2 They can be equal or unequal.

[0137] (2) The first thread group and the second thread group perform attention mechanism calculations in parallel, so they share one shared memory when processing in parallel. The first thread group and the second thread group perform attention mechanism calculations alternately, so the total amount of actual shared memory required at different times may be different. Therefore, the maximum value M of the total amount of actual shared memory required at the same time when the first thread group and the second thread group perform attention mechanism calculations in parallel is required. max2 The capacity must not exceed the maximum capacity M of the shared memory. max1 Restrictions.

[0138] S u1 S u2 and M max2 The size of the query matrix block, key-value matrix block, key matrix block, and value matrix block are related to this. Based on this, the embodiments of this 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 This allows us to obtain appropriate sizes for query matrix blocks, key-value matrix blocks, and key and value matrix blocks, ensuring that when using matrix blocks for attention mechanism operations, we can fully utilize hardware resources while avoiding register overflows and shared memory contention.

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

[0140] Based on the second preset relation, a second objective function is constructed; the second preset relation includes at least the size of the query matrix block, the size of the key matrix block, the size of the value matrix block, and S. u2The relationship between them;

[0141] Based on the third predefined relation, a third objective function is constructed; the third predefined relation includes at least 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 The relationship between them;

[0142] Furthermore, certain dimensions of the matrix blocks can be pre-defined, and then the first, second, and third objective functions can be solved jointly. For example, the attention head dimensions of the key matrix blocks and the value matrix blocks can be determined based on the task complexity.

[0143] In an optional embodiment, the third preset relationship further includes: target quantity and M. max2 The relationship between them;

[0144] The target number includes the number of key-value matrix blocks stored simultaneously in the shared memory and the number of query matrix blocks stored simultaneously in the shared memory.

[0145] It is understandable that the more key-value matrix blocks and query matrix blocks stored simultaneously in shared memory, the more shared memory is consumed. Therefore, M max2 It is also related to the number of key-value matrix blocks and query matrix blocks stored simultaneously in the shared memory. This application does not impose specific limitations on the number of key-value matrix blocks and query matrix blocks stored simultaneously in the shared memory; these can be set according to actual needs, and only need to be positive integers.

[0146] In an optional embodiment, controlling the second thread group to perform matrix multiplication and exponentiation operations 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:

[0147] The second thread group is controlled to perform matrix multiplication with 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, which is written to the third register group of the second thread group. The j-th attention score matrix block in the third register group of the second thread group is then subjected to exponential operation to obtain the j-th attention score matrix block, which is overwritten and written to the third register group of the second thread group. The j-th attention score matrix block in the third register group of the second thread group is then converted to higher precision and written to the first register group of the second thread group. Finally, the j-th attention score matrix block in the first register group of the second thread group is written to shared memory.

[0148] The first thread bundle group is controlled to write the j-th attention score matrix in the shared memory into the first register group of the first thread bundle group in blocks.

[0149] It is worth noting that in this embodiment, the j-th key matrix block in the shared memory is first transposed to obtain the transposed j-th key matrix block. Then, a matrix multiplication operation is performed between 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. Next, an exponential operation is performed on the j-th attention score matrix block before scaling to achieve scaling of the attention score matrix block, thus obtaining the j-th attention score matrix block. Specifically, according to exp((S-accum_rmax)... sf), which performs an exponential operation on the attention score matrix blocks before scaling; where exp() represents the natural exponential function, S is the j-th attention score matrix block before scaling, accum_rmax is the historical global row maximum value, and sf is the scaling factor.

[0150] This application embodiment uses the third register group regS to perform matrix multiplication and exponentiation operations, avoiding the risk of performance degradation caused by register overflow and ensuring the continuity and stability of the attention mechanism calculation process.

[0151] Furthermore, to improve computational accuracy, the j-th attention score matrix block obtained after matrix multiplication and exponentiation is of second precision type. Therefore, to reduce data storage overhead in shared memory, the j-th attention score matrix block in the third register group regS of the second thread bundle is precision-converted and written to the first register group regP of the second thread bundle. The j-th attention score matrix block in the first register group regP of the second thread bundle is then written to shared memory, ensuring that the j-th attention score matrix block written to shared memory is of first precision type.

[0152] In an optional embodiment, after obtaining the j-th attention score matrix block, the method further includes:

[0153] The first kernel controls the second thread group to calculate the row target value of the j-th attention score matrix block required for the attention mechanism operation, and writes it into the first register group of the second thread group; wherein, the row target value includes: historical row accumulated value, correction value for correcting historical row accumulated value, current row accumulated value, row maximum value, historical global row maximum value, current global row maximum value, and row accumulated value;

[0154] The first kernel controls the second thread bundle to write the correction value of the j-th attention score matrix block in the first register block of the second thread bundle into the shared memory;

[0155] The first kernel controls the first thread bundle to write the correction value of the j-th attention score matrix block in the shared memory into the first register block of the first thread bundle.

[0156] It is worth noting that both the first and second thread bundles undertake the attention mechanism computation. The data required for the attention mechanism computation, in addition to the attention score matrix blocks, also includes the row target values ​​of the attention score matrix blocks. The row target values, besides the aforementioned current row accumulation value `accum_rsum_new`, also include:

[0157] The historical row accumulation value accum_rsum is specifically the accumulation value of each row of the historical attention score matrix block, that is, the accumulation value of each row of the matrix formed by the 1st to j-1st attention score matrix blocks; where each row corresponds to one accumulation value.

[0158] The modifier is used to correct the historical row's accumulated value. Specifically, the modifier is used to correct the historical row's accumulated value accum_rsum to obtain the current row's accumulated value accum_rsum_new.

[0159] Furthermore, the calculation of the current row's accumulated value accum_rsum_new and the modifier also requires the following row target values:

[0160] The maximum value of a row, chunk_rmax, is specifically the maximum value of each row in the attention score matrix block, that is, the maximum value of each row in the j-th attention score matrix block; where each row corresponds to a maximum value.

[0161] The historical global row maximum value, accum_rmax, is specifically the maximum value of each row in the historical global attention score matrix block, that is, the maximum value of each row in the matrix formed by the 1st to j-1st attention score matrix blocks; where each row corresponds to a maximum value.

[0162] The current global row maximum value, accum_rmax_new, is specifically the maximum value of each row in the current global attention score matrix block, that is, the maximum value of each row in the matrix formed by the 1st to jth attention score matrix blocks; where each row corresponds to a maximum value.

[0163] The cumulative value of a row, chunk_rsum, is specifically the cumulative value of each row in the attention score matrix block, that is, the cumulative value of the j-th attention score matrix block in each row; where each row corresponds to one cumulative value.

[0164] Therefore, the second thread bundle calculates the aforementioned seven row target values, including: historical row accumulation value `accum_rsum`, current row accumulation value `accum_rsum_new`, modifier, maximum row value `chunk_rmax`, historical global maximum row value `accum_rmax`, current global maximum row value `accum_rmax_new`, and row accumulation value `chunk_rsum`, and stores them sequentially in the fourth, fifth, sixth, seventh, eighth, ninth, and tenth register groups of the second thread bundle. The first thread bundle only needs the modifier, which can be stored in the sixth register group of the first thread bundle. The fourth, fifth, and sixth register groups belong to the first register group; the seventh, eighth, ninth, and tenth register groups belong to the second register group.

[0165] In one specific embodiment, the hardware resources of a certain artificial intelligence chip are as follows: (1) Each thread in each thread bundle is limited to using a maximum of 256 registers; (2) The shared memory has two storage banks, including a first storage bank and a second storage bank, which allows simultaneous access to data in both banks and improves memory access efficiency. Furthermore, the shared memory can be divided into a first storage bank and a second storage bank according to address. For example, the first storage bank is a high-address storage bank (highbank), i.e., the high-address area in the shared memory, allocated to 96KB; the second storage bank is a low-address storage bank (lowbank), i.e., the low-address area in the shared memory, allocated to 128KB. Of course, other methods can also be used for division, which are not specifically limited here.

[0166] For this AI chip, each matrix block is set according to the attention head dimension (headdim). The sequence length (seqlen) is calculated using the attention mechanism provided in this application embodiment: the size of the query matrix block Q is 576. 64; The size of the key-value matrix block KV is 576. 32; Taking the complete key-value matrix block as the key matrix block K, the size of the key matrix block is: 576 32; If the first 512 columns are taken as the value matrix block V, then the size of the value matrix block is: 512 32.

[0167] like Figure 5The 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 then subjected to exponential operation and precision conversion to obtain the attention score matrix block P. The attention score matrix block P and the sub-blocks of the value matrix block V are used in a preset pre-stage of alternating attention mechanism operation to obtain the intermediate result of the sub-block of the attention result matrix block O. Figure 5 In this process, the value matrix is ​​evenly divided into four sub-blocks along the direction of the attention head dimension: the first sub-block V_0, the third sub-block V_1, the second sub-block V_2, and the fourth sub-block V_3. Each sub-block has a size of 128. 32, thus obtaining the intermediate results of four sub-blocks of the attention result matrix block O, including the intermediate results O_acc0 of the first sub-block, O_acc1 of the third sub-block, O_acc2 of the second sub-block, and 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 alternately to finally obtain the intermediate results of half of the attention result matrix block.

[0168] like Figures 6-8 First, let me explain: cp means writing data from global memory to shared memory; ld means loading data, specifically writing data from shared memory to registers; st means writing data from registers to shared memory; tma is used to move data between global memory and shared memory, here referring to writing data from shared memory to global memory; and ldbulk means loading data from global memory into registers in batches.

[0169] First kernel (specifically, ComputeKernel, the computing kernel):

[0170] The first thread 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 shared memory.

[0171] The second thread group WG1 writes Q(1) into its own register and performs a matrix multiplication operation (i.e., Q(1)) with the transposed K(1). K(1) T After T denotes transpose, a softmax operation is performed. Here, the softmax operation is represented in Q(1). K(1) TThen, operations preceded by alternating attention mechanism operations, such as exponentiation, precision conversion, and calculation of row target values, are performed. The first attention score matrix is ​​then written into shared memory in block P(1), which allows the pre-set pre-processing stage of alternating attention mechanism operations of WG1 to begin.

[0172] WG0 writes P(1) into its own register, which can start the preset pre-processing stage of WG0's alternating attention mechanism operation.

[0173] When 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-order stage of the attention mechanism operation, and obtains the intermediate result of the first sub-block O(1)_acc0 of the first attention result matrix block and writes it to 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-order stage of the attention mechanism operation, and obtains the intermediate result O(1)_acc2 of the second sub-block of the first attention result matrix block; WG1 uses P(1) and the third sub-block V(1)_0 of the first value matrix block to perform the preset pre-order stage of the attention mechanism operation, and obtains the intermediate result O(1)_acc2 of the second sub-block of the first attention result matrix block; In the preset pre-stage of the attention mechanism operation in block V(1)_1, the intermediate result O(1)_acc1 of the third sub-block of the first attention result matrix block is obtained and written to shared memory. At this time, for WG1, it is the last operation using Q(1), K(1) and V(1), and the subsequent operation of Q(1), K(2) and V(2) can begin. 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 performs the attention mechanism operation using P(1) and the fourth sub-block V(1)_3 of the first value matrix block, it is not necessary to use V(1) and P(1) for operation, at this time, WG0 can be controlled to write the third key-value matrix block KV(3) to the position of KV1 in shared memory.

[0174] In the process of alternating attention operations (j=2,...,J-1), before the pre-preceding stage of attention mechanism operations using the sub-block V(j)_acc of the j-th attention score matrix block P(j) and the j-th 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 to the required thread group. The remaining processes are similar to the above and will not be repeated here.

[0175] After the preset pre-order stage of each alternating attention operation is completed, 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, that is, each has O(i)_half), and is written to global memory through shared memory, and a preamble trigger signal is sent to the second kernel (specifically the PostprocessKernel).

[0176] After the second kernel detects and receives the preamble trigger signal, it starts controlling the third thread group to load the intermediate results of all sub-blocks of the Jth attention result matrix block and concatenate them to obtain the intermediate result O(i) of the Jth attention result matrix block of the i-th query block. Then, it starts the preset post-stage of the attention mechanism operation, including: rescale O(i), which specifically means dividing the intermediate result of the last attention result matrix block of the i-th query matrix by the current row accumulation value of the last attention score matrix block of the i-th query matrix; typecast O(i), which means performing data type conversion on the result obtained by rescale O(i); and tmaO(i), which means writing the result of typecast O(i) to global memory. At this time, a preamble complete signal is issued to indicate that the attention mechanism calculation of the i-th query matrix block is complete.

[0177] The above process is executed according to the following instructions:

[0178] Instruction logic for the first thread bundle:

[0179] Loop{ / Q outer loop /

[0180] {tma cp Q / / Write the first query matrix block from global memory to shared memory (only for the first query matrix block);

[0181] }

[0182] Loop{ / KV internal circulation /

[0183] tma cp KV / / Writes the key-value matrix into blocks (KV) from global memory to shared memory;

[0184] ld.gsm modifier / / Writes the modifier value from shared memory to the first register block of the first thread bundle;

[0185] Ldmatrix P / / Writes the attention score matrix block p from shared memory to the first register group regP of the first thread bundle;

[0186] Loop { / half seqlen of V (half of V) /

[0187] {init regO / / Initialize the second register group regO of the first thread bundle;

[0188] regO = mma(P, V, O) / / Pre-set pre-stage for performing attention mechanism operations on the first attention score matrix block in the first register group regP of the first thread bundle and the (m+1)th sub-block of the first value matrix block in the shared memory;

[0189] Stmatrix regO to O_acc(gsm) / / Writes the intermediate results of the (m+1)th sub-block of the first attention result matrix into shared memory;

[0190] / / Used for the first attention score matrix block;

[0191] {ldmatrix regO / / Write the intermediate result of the (m+1)th sub-block of the (j-1)th attention result matrix block from shared memory to the second register group regO of the first thread bundle;

[0192] regO = modifier regO

[0193] regO = mma(P, V, regO) / / Together with the previous instruction, this forms a preset pre-stage for performing attention mechanism operations on the intermediate result of the (j-1)th attention result matrix block of the second register group regO of the first thread bundle, the jth attention score matrix block in the first register group regP of the first thread bundle, and the m+1th sub-block of the jth value matrix block in the shared memory.

[0194] Stmatrix regO to O_acc(gsm) / / Writes the intermediate result of the (m+1)th block of the j-th attention result matrix into shared memory;

[0195] / / Used for other (i.e., attention score matrix blocks besides the first one);

[0196] }

[0197] Tma half O / / Writes the intermediate results of the sub-blocks of the last attention result matrix block obtained by the attention mechanism operation to global memory (Tma half O of the first thread bundle will complete faster than that of the second thread bundle).

[0198] supaTriggerProgrammaticLaunchCompletion() / / Preamble trigger, specifically used to trigger the second kernel to execute the code following the preamble;

[0199] }

[0200] }

[0201] Instruction logic for the second thread bundle:

[0202] Loop{ / Q outer loop /

[0203] Ldmatrix Q ->regQ / / Writes the query matrix block Q from shared memory to the eleventh register group regQ of the second thread bundle (only for the first query matrix block);

[0204] ldbulk Q to reg + shuffle reg layout / / Write the other (except the first) query matrices from global memory into blocks and write them to the eleventh register group regQ of the second thread bundle. Process the data in the registers through a data shuffling operation.

[0205] Loop{ / KV internal circulation /

[0206] regS = mma(regQ, smemK, regS) / / Perform matrix multiplication on the query matrix block in the eleventh register group regQ of the second thread bundle and the transposed key matrix block to obtain the attention score matrix block S of fp32 and write it into the third register group regS of the second thread bundle.

[0207] calc chunk_rmax, accum_rmax, accum_rmax_new / / Calculate the maximum value of a row, chunk_rmax, the historical global maximum value of a row, accum_rmax, and the current global maximum value of a row, accum_rmax_new;

[0208] exp((S-accum_rmax) sf ), reuse regS / / Perform exponential operation on the attention score matrix blocks in the third register group regS of the second thread bundle, and write the resulting attention score matrix blocks overwrite the third register group regS of the second thread bundle;

[0209] calc chunk_rsum, accum_rsum, accum_rsum_new / / Calculate the cumulative value of the current row, chunk_rsum, the cumulative value of the previous row, accum_rsum, and the cumulative value of the current row, accum_rsum_new;

[0210] Modifier=exp((accum_rmax – accum_rmax_new) sf ) / / Calculate the Modifier value;

[0211] accum_rsum_new = accum_rsum Modifier + chunk_rsum / / Update the current row's accumulated value accum_rsum_new;

[0212] Typecast regS to regP / / This converts the attention score matrix in the third register group regS of the second thread bundle into blocks, performs precision conversion, and writes it to the first register group regP of the second thread bundle, thus achieving precision conversion from fp32 to fp16.

[0213] Stmatrix regP to smemP / / Writes the attention score matrix in the first register group regP of the second thread bundle into shared memory in blocks;

[0214] St.gsm modifier / / Writes the modifier value to shared memory;

[0215] Loop { / half seqlen of V (half of V) /

[0216] {init regO / / Initialize the second register group regO of the second thread bundle;

[0217] regO = mma(P, V, O) / / Pre-set pre-stage for performing attention mechanism operations on the first attention score matrix block in the first register group regP of the second thread bundle and the m-th sub-block of the first value matrix block in the shared memory;

[0218] Stmatrix regO to O_acc(gsm) / / Writes the intermediate results of the m-th sub-block of the first attention result matrix into shared memory;

[0219] / / Used for the first attention score matrix block;

[0220] {ldmatrix regO / / Write 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;

[0221] regO = modifier regO

[0222] regO = mma(P, V, regO) / / Together with the previous instruction, this forms a preset pre-stage for performing attention mechanism operations on 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, the j-th attention score matrix block in the first register group regP of the second thread bundle, and the m+1-th sub-block of the j-th value matrix block in the shared memory.

[0223] Stmatrix regO to O_acc(gsm) / / Writes the intermediate results of the m-th sub-block of the j-th attention result matrix into shared memory;

[0224] / / Used for other (i.e., attention score matrix blocks besides the first one);

[0225] }

[0226] st.glm acc_rsum_new / / Writes the current row's accumulated value acc_rsum_new to global memory;

[0227] Tma half O / / Writes the intermediate results of the sub-blocks of the last attention result matrix block obtained by the attention mechanism operation to global memory;

[0228] }

[0229] }

[0230] Instruction logic for the third thread bundle:

[0231] Loop{ / Q outer loop /

[0232] supaGridDependencySynchronize() / / preamble synchronization, specifically used to receive triggers from the first kernel;

[0233] ldbulk O / / Loads the intermediate results of the sub-blocks of the attention result matrix directly from global memory into the register group of the third thread bundle, and concatenates them into the intermediate results of the attention result matrix blocks;

[0234] ld.glm acc_rsum_new / / Writes the current row's accumulated value acc_rsum_new from global memory to the register group of the third thread bundle;

[0235] regO = regO / acc_rsum_new / / Divide the intermediate result of the attention result matrix block by the current row's accumulated value acc_rsum_new to obtain the attention result matrix block;

[0236] Typecase regO to regO_f16 / / This converts the data type of the attention result matrix into blocks, for example, from fp32 to fp16;

[0237] Stbulk regO_f16 / / Writes the attention result matrix after data type conversion in blocks from the register group of the third thread bundle to global memory in batches;

[0238] }

[0239] The register resources used by the first thread bundle in this scenario are shown in Table 1.

[0240] Table 1

[0241]

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

[0243] For the sixth register group, since the matrix elements held by one thread are distributed across two rows, and each row needs to calculate and store the modifier value, the sixth register group requires two registers.

[0244] The matrix data stored in the first register group regP is written to shared memory after the attention score matrix is ​​divided into blocks and its precision is converted by the second thread beam group, and then read from the shared memory. Therefore, the number of registers in the first register group regP = the number of registers in the second thread beam group regS / 2 = 8.

[0245] The number of registers in the second register group regO is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. The size of the sub-blocks in the attention result matrix / (number of thread bundles in the first thread bundle group) (The size of the matrix held by one thread bundle).

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

[0247] Table 2

[0248]

[0249] In Table 2, the actual number of registers required by each thread in the second thread bundle includes the number of registers in register group 11 (regQ), register group 3 (regS), register group 7, register group 8, register group 9, register group 10, register group 4, register group 5, register group 6, register group 1 (regP), register group 2 (regO), and register group 12 (regO_fp16).

[0250] The number of registers in the eleventh register group (regQ) is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. Query the size of the matrix block / (number of thread bundles in the second thread bundle group) (Matrix size held by one thread bundle).

[0251] The number of registers in the third register group regS is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. Attention score matrix blocks / (number of thread bundles in the second thread bundle group) (Matrix size held by one thread bundle).

[0252] Similar to the first thread bundle, the seventh, eighth, ninth, tenth, fourth, fifth, and sixth register bundles each contain two registers.

[0253] The matrix data stored in the first register group regP is obtained by the second thread beam group after dividing the attention score matrix into blocks and performing precision conversion. Therefore, the number of registers in the first register group regP = the number of registers in the second thread beam group regS / 2 = 8.

[0254] The number of registers in the second register group regO is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. The size of the sub-blocks in the attention result matrix / (number of thread bundles in the second thread bundle group) (The size of the matrix held by one thread bundle).

[0255] The register resources required by the third thread bundle are far less than 512 registers, so the register resources used by the third thread bundle will not be discussed here.

[0256] The shared memory resources used in this scenario are shown in Table 3, and the specific distribution of shared memory resources at each stage is as follows: Figures 9-10 Shared memory (GSM) is divided into a high-address memory area (highbank) and a low-address memory area (lowbank). The highbank occupies 96KB and the lowbank occupies 128KB. The white part inside the black solid box in the figure represents the unused shared memory space. Among them, the space occupied by the row target value and the other row values ​​is too small, only 8B, and is not included here. Figure 9 For the phase from start to end j=1, only the query matrix block Q (72KB), the key-value matrix block KV (each KV occupies 36KB), and the attention score matrix block P (2KB) occupy shared memory; Figure 10 During the period from the start of writing the i-th query matrix into shared memory to the end of j=1 to the end of j=J, the query matrix block Q will be released (indicated by dashed lines) and will not occupy shared memory. At this time, the sub-blocks O(j)_acc0, O(j)_acc1, O(j)_acc2, and O(j)_acc3 (each occupying 32KB), the key-value matrix block KV, and the attention score matrix block P will occupy shared memory.

[0257] Table 3

[0258]

[0259] As can be seen, the attention mechanism computation provided in this application embodiment can accurately adapt to the characteristics of the register and shared memory resources of the artificial intelligence chip, the hardware resources used do not exceed the usage limits, and mobilize as many hardware resources as possible to participate in the computation, fully release the computing potential of the hardware, ensure the continuity and stability of the attention mechanism computation process, and achieve high computing performance on the basis of efficient resource utilization.

[0260] See Figure 11This application provides an artificial intelligence chip including a global memory (GLM) and one or more computing units. Each computing unit includes: a first core, a second core, shared memory (GSM), one or more first thread bundles, one or more second thread bundles, and one or more third thread bundles. In the register space, regQ 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.

[0261] Furthermore, this application also provides a computer-readable storage medium, which includes a stored computer program; wherein, when the computer program is executed, it controls the device where the computer-readable storage medium is located to perform the attention mechanism calculation method as described in any of the above embodiments.

[0262] Furthermore, this application also provides a computer program product, including a computer program or instructions, which, when executed by a processor, implement the attention mechanism computation method as described in any of the above embodiments.

[0263] See Figure 12 This application provides an attention mechanism computing device 20, which includes a processor 21, a memory 22, and a computer program stored in the memory 22 and executable on the processor 21. When the processor 21 executes the computer program, it implements the steps described in the attention mechanism computing method embodiments above. Alternatively, when the processor 21 executes the computer program, it implements the functions of each module / unit in the various device embodiments described above.

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

[0265] The attention mechanism computing device 20 may include, but is not limited to, a processor 21 and a memory 22. Those skilled in the art will understand that the schematic diagram is merely an example of the attention mechanism computing device 20 and does not constitute a limitation on the attention mechanism computing device 20. It may include more or fewer components than illustrated, or combine certain components, or different components. For example, the attention mechanism computing device 20 may also include input / output devices, network access devices, buses, etc.

[0266] The processor 21 can be a Central Processing Unit (CPU), or other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor can be a microprocessor or any conventional processor. The processor 21 is the control center of the attention mechanism computing device 20, connecting all parts of the attention mechanism computing device 20 via various interfaces and lines.

[0267] The memory 22 can be used to store the computer programs and / or modules. The processor 21 implements 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 the data stored in the memory 22. The memory 22 may mainly include a program storage area and a data storage area. The program storage area may store the operating system, at least one application program required for a function (such as sound playback function, image playback function, etc.), etc.; the data storage area may store data created according to the use of the mobile phone (such as audio data, phonebook, etc.). In addition, the memory 22 may include high-speed random access memory, and may also include non-volatile memory, such as hard disk, memory, plug-in hard disk, smart media card (SMC), secure digital (SD) card, flash card, at least one disk storage device, flash memory device, or other volatile solid-state storage device.

[0268] If the modules / units integrated in the attention mechanism computing device 20 are implemented as software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, all or part of the processes in the methods of the above embodiments can also be implemented by a computer program instructing related hardware. The computer program can be stored in a computer-readable storage medium, and when executed by the processor 21, it can implement the steps of the various method embodiments described above. The computer program includes computer program code, which can be in the form of source code, object code, executable files, or certain intermediate forms. The computer-readable medium can include: any entity or device capable of carrying the computer program code, recording media, USB flash drives, portable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signals, telecommunication signals, and software distribution media, etc.

[0269] It should be noted that the device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Furthermore, in the accompanying drawings of the device embodiments provided in this application, the connection relationships between modules indicate that they have communication connections, which can be specifically implemented as one or more communication buses or signal lines. Those skilled in the art can understand and implement this without any creative effort.

[0270] The above description is the preferred embodiment of this application. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the principle of this application, and these improvements and modifications are also considered to be within the scope of protection of this application.

Claims

1. An attention mechanism calculation method, characterized in that, The method comprises the following steps: The first kernel controls the second thread bundle group to load the i-th query matrix block; The first kernel controls the second thread bundle group to perform matrix multiplication and 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; The first kernel controls 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 attention mechanism operation on the j-th attention score matrix block, to obtain the intermediate results of different sub-blocks of the j-th attention result matrix block, and let j = j + 1, and return to the step of controlling the second thread bundle group to perform matrix multiplication and exponential operation on the i-th query matrix block and the transposed j-th key matrix block, until j = J, to obtain the intermediate results of different sub-blocks of the last attention result matrix block of the i-th query matrix block; wherein J is a preset total number of blocks, and the j-th key matrix block and the j-th value matrix block are pre-stored in the shared memory; 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 i-th query matrix block, and perform the preset post-stage of attention mechanism operation, to obtain the last attention result matrix block of the i-th query matrix block.

2. The attention mechanism computation method of claim 1, wherein, 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 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 query matrix block. 3.The attention mechanism calculation method of claim 1, wherein, After obtaining the j-th 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 attention score matrix block of the i-th query matrix block required for attention mechanism operation, and writes the current row accumulation value into the global memory; 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 i-th query matrix block, and perform the preset post-stage of attention mechanism operation, to obtain the last attention result matrix block of the i-th query matrix block, comprising: The second kernel controls the third thread bundle group 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 group; The second kernel controls the third thread bundle group to load all sub-blocks of the last 1 attention result matrix block of the i-th query matrix block and splice intermediate results of all sub-blocks to obtain the intermediate result of the last 1 attention result matrix block of the i-th query matrix block, and divide the intermediate result by the current row accumulated value of the last 1 attention score matrix block of the i-th query matrix block to obtain the last 1 attention result matrix block of the i-th query matrix block. 4.The attention mechanism calculation method of claim 1, wherein, The j-th attention score matrix block is written into 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 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 a preset pre-sequence stage of attention mechanism operation on the j-th attention score matrix block to respectively obtain intermediate results of different sub-blocks of the j-th attention result matrix block, including: The first kernel controls the second thread bundle group to use the j-th attention score matrix block in the first register group of the second thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the m-th sub-block of the j-th value matrix block in the shared memory to obtain an 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 use the j-th attention score matrix block in the first register group of the first thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the m+1-th sub-block of the j-th value matrix block in the shared memory to obtain an 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, m is set to m+1, and the step of the first kernel controlling the second thread bundle group to use the j-th attention score matrix block in the first register group of the second thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the m-th sub-block of the j-th value matrix block in the shared memory is returned to until m is a preset total number of sub-blocks, and an intermediate result of the last 1 sub-block of the j-th attention result matrix block is obtained.

5. The attention mechanism computation method of claim 4, wherein, The first kernel controls the second thread bundle group to use the j-th attention score matrix block in the first register group of the second thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the m-th sub-block of the j-th value matrix block in the shared memory to obtain an 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 group to use the first attention score matrix block in the first register group of the second thread bundle group to perform a preset pre-sequence stage of attention mechanism operation on the m-th sub-block of the first value matrix block in the shared memory to obtain an intermediate result of the m-th sub-block of the first attention result matrix block and write the intermediate result 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 to the second register group of the second thread bundle group, and to perform the preset preliminary stage of the attention mechanism operation with 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.

6. The attention mechanism computation method of claim 4, wherein, The first kernel controls the first thread bundle group to perform the preset preliminary stage of the attention mechanism operation with the jth attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the jth value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the jth attention result matrix block, including: If j = 1, the first kernel controls the first thread bundle group to perform the preset preliminary stage of the attention mechanism operation with 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 = 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 to perform the preset preliminary stage of the attention mechanism operation with the jth attention score matrix block in the first register group of the first thread bundle group and the (m+1)th sub-block of the jth value matrix block in the shared memory, to obtain the intermediate result of the (m+1)th sub-block of the jth attention result matrix block and write it into the shared memory.

7. The attention mechanism computation method of claim 1, wherein, The method further comprises: For the ith query matrix block, after the preset preliminary 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 preliminary stage of the attention mechanism operation of the second thread bundle group into the global memory.

8. The attention mechanism computation method of claim 1, wherein, The method further comprises: For the ith query matrix block, after the preset preliminary stage of the last alternating attention mechanism operation of the first thread bundle group, the first kernel controls the first thread bundle group to write the intermediate results of all sub-blocks of the last attention result matrix block obtained by the preset preliminary stage of the attention mechanism operation of the first thread bundle group into the global memory.

9. The attention mechanism computation method of claim 1, wherein, In the shared memory, the jth value matrix block and the jth key matrix block share the same jth key-value matrix block.

10. The attention mechanism computation method of claim 9, wherein, The method further comprises: For the i-th query matrix block, after a preset presequence stage of the last alternating 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 simultaneously stored in the shared memory.

11. An attention mechanism computing device, comprising: A processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, the processor implements the attention mechanism calculation method according to any one of claims 1-10 when executing the computer program.

12. A computer-readable storage medium, characterized in that, The computer readable storage medium includes 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 claims 1-10 when running.

13. A computer program product, characterised in that, A computer program or instructions, which, when executed by a processor, implement the attention mechanism calculation method according to any one of claims 1-10.

Images