An attention mechanism calculation method, device, medium and product

Abstract

The application discloses an attention mechanism calculation method, device, medium and product. The method comprises the following steps: controlling a second thread bundle group to load an i-th query matrix block; controlling the second thread bundle group to perform matrix multiplication and exponential operation on the i-th query matrix block and a transposed j-th key matrix block to obtain a j-th attention score matrix block; and controlling a first thread bundle group and the second thread bundle group to alternately perform attention mechanism operation on the j-th attention score matrix block by using different subblocks of a j-th value matrix block until a last subblock of a j-th attention result matrix block is obtained. The embodiment of the application can provide sufficient register resources for the calculation of the attention mechanism, thereby improving the calculation 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] The attention mechanism computation process mainly involves three types of matrices: query matrix, key matrix, and value matrix. However, as the matrix size increases significantly, the existing warpgroup configuration, which relies on a single warpgroup to handle all computational tasks, cannot provide sufficient register resources for the attention mechanism computation, thus limiting the computational efficiency of the attention mechanism.

[0003] For example, in an existing warpgroup configuration of "1 producer warpgroup + 2 consumer warpgroups", the producer warpgroup is only responsible for data transfer, while the computation work is entirely handled by the consumer warpgroups. However, each thread in the consumer warpgroup can only use a maximum of 240 registers. If the size of the query matrix is ​​too large, loading the query matrix alone will consume most of the register resources of the consumer warpgroup, leaving insufficient register resources to support the subsequent core attention mechanism computation requirements, ultimately resulting in inefficient execution of the attention mechanism computation. Summary of the Invention

[0004] This application provides an attention mechanism computation method, device, medium, and product to solve the problem that existing thread bundle configurations cannot provide sufficient register resources, thereby limiting the computational efficiency of the attention mechanism.

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

[0006] Control the second thread bundle to load the i-th query matrix block;

[0007] The second thread group is controlled 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;

[0008] The first thread group and the second thread group are controlled to alternately use different sub-blocks of the j-th value matrix block and the same j-th attention score matrix block to perform attention mechanism operations until the last sub-block of the j-th value matrix block is used to perform attention mechanism operations with the same j-th attention score matrix block, so as to obtain the last sub-block of the j-th attention result matrix block; wherein, the j-th key matrix block and the j-th value matrix block are pre-stored in shared memory.

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

[0010] The control of the first thread group and the second thread group to alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations, resulting in the last sub-block of the j-th attention result matrix block, includes:

[0011] The second thread bundle is controlled to perform 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 m-th sub-block of the j-th attention result matrix block;

[0012] After obtaining the m-th sub-block of the j-th attention result matrix block, the first thread bundle group is controlled to perform attention mechanism operation with 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 to obtain the m+1-th sub-block of the j-th attention result matrix block.

[0013] After obtaining the (m+1)th sub-block of the j-th attention result matrix block, let m=m+1, and return to the step of controlling the second thread group to perform 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 shared memory, until m is the preset total number of sub-blocks, and obtain the last sub-block of the j-th attention result matrix block.

[0014] As an improvement to the above scheme, the control of the second thread bundle group performs attention mechanism operations 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 m-th sub-block of the j-th attention result matrix block, including:

[0015] If j=1, then the second thread bundle group is controlled to perform 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 m-th sub-block of the first attention result matrix block and write it into the shared memory;

[0016] If j=2,...,J-1, then the second thread group controls the m-th sub-block of the (j-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs 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 m-th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block;

[0017] If j=J, then the second thread group controls the m-th sub-block of the (J-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs attention mechanism operation with the last attention score matrix block in the first register group of the second thread group and the m-th sub-block of the last value matrix block in the shared memory to obtain the m-th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

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

[0019] For the i-th query matrix block, after the last alternating attention mechanism operation of the second thread group, the second thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

[0020] As an improvement to the above scheme, the control of the first thread bundle group to perform attention mechanism operation on 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 to obtain the (m+1)-th sub-block of the j-th attention result matrix block includes:

[0021] If j=1, then the first thread bundle group is controlled to perform attention mechanism operation by 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 (m+1)th sub-block of the first attention result matrix block and write it into the shared memory.

[0022] If j=2,...,J-1, control the first thread group to write 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 perform 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 (m+1)th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block;

[0023] If j=J, control the first thread group to write 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 perform attention mechanism operation with the last attention score matrix block in the first register group of the first thread group and the (m+1)th sub-block of the last value matrix block in the shared memory to obtain the (m+1)th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

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

[0025] For the i-th query matrix block, after the last alternating attention mechanism operation of the first thread group, the first thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

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

[0027] After the first thread group writes all sub-blocks of the last attention result matrix block obtained by itself through attention mechanism operation to global memory, the first thread group writes the (i+1)th query matrix to the position of the i-th query matrix in the shared memory.

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

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

[0030] For the i-th query matrix block, after the last alternating attention mechanism operation, the first thread group is controlled to write the j+x-th key-value matrix block from global memory to 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.

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

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

[0033] 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, S r2 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;

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

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

[0036] As an improvement to the above scheme, the control of the second thread group involves performing 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, including:

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

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

[0039] As an improvement to the above scheme, before the first thread group and the second thread group alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations and obtain the last sub-block of the j-th attention result matrix block, the method further includes:

[0040] The second thread group is controlled to calculate the row target value of the j-th attention score matrix block required for the attention mechanism operation, and write it into the first register block of the second thread group;

[0041] The second thread bundle is controlled to write the row target value of the j-th attention score matrix block in the first register block of the second thread bundle into the shared memory;

[0042] The first thread bundle is controlled to write the row target value of the j-th attention score matrix block in the shared memory into the first register block of the first thread bundle.

[0043] To achieve the above objectives, embodiments of this application also provide an attention mechanism computing device, including a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, wherein the processor implements the attention mechanism computing method as described above when executing the computer program.

[0044] To achieve the above objectives, embodiments of this application also provide a computer-readable storage medium, the computer-readable storage medium including 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 above.

[0045] To achieve the above objectives, embodiments of this application also provide a computer program product, including a computer program or instructions, which, when executed by a processor, implement the attention mechanism computation method as described above.

[0046] Compared with existing technologies, the present application provides an attention mechanism calculation method, device, medium, and product. This involves controlling a second thread group to load the i-th query matrix block; 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; and controlling the first and second thread groups to alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism calculations until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism calculations, obtaining the last sub-block of the j-th attention result matrix block. Therefore, by controlling the first and second thread groups to alternately perform attention mechanism calculations, the present application breaks the pure producer-consumer relationship between thread groups, avoiding the problem of insufficient register resources caused by relying on a single thread group to undertake all calculation tasks. This provides sufficient register resources for attention mechanism calculations, ultimately effectively improving the computational efficiency of the attention mechanism. Attached Figure Description

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

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

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

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

[0051] Figure 5 This is a pipeline diagram of the early stage of a thread bundle provided in an embodiment of this application;

[0052] Figure 6 This is a pipeline diagram of the intermediate stage of a thread bundle provided in an embodiment of this application;

[0053] Figure 7 This is a pipeline diagram of the late stage of a thread bundle provided in an embodiment of this application;

[0054] Figure 8 This is a pipeline diagram of the late stage and the next early stage of a thread bundle provided in an embodiment of this application;

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

[0056] Figure 10 This is a shared memory resource distribution diagram for a mid-term stage provided in an embodiment of this application;

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

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

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

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

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

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

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

[0064] 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).

[0065] The Computing Unit (CU), as the core of an AI chip, is responsible for scheduling and executing warps to handle 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.

[0066] 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 half-precision floating-point matrix elements (fp16 type), the distribution of matrix elements and registers held by one thread is as follows: 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, the distribution of matrix elements and registers held by one thread is as follows: 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.

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

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

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

[0070] For example, the first query matrix block is loaded, and the second thread group is controlled to perform matrix multiplication and exponentiation operations to obtain the j-th attention score matrix block of the first query matrix block. The first and second thread groups are then controlled to alternately perform attention mechanism operations on the first attention score matrix block of the first query matrix block until the last sub-block of the first attention result matrix block of the first query matrix block is obtained. Then, attention mechanism operations on the second attention score matrix block of the first query matrix block can be performed, and so on, until the last sub-block of the last attention result matrix block of the first query matrix block is obtained.

[0071] Then, the next (i.e. the second) query matrix block is loaded, and the second thread group is controlled to perform matrix multiplication and exponentiation operations to obtain the first attention score matrix block of the second query matrix block. The first thread group and the second thread group are controlled to alternately perform the attention mechanism operation of the first attention score matrix block of the second query matrix block until the last sub-block of the last attention result matrix block of the second query matrix block is obtained. This process is repeated 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.

[0072] See Figure 1 , Figure 1 This is a flowchart of an attention mechanism calculation method provided in an embodiment of this application. The attention mechanism calculation method includes:

[0073] S11, Control the second thread bundle to load the i-th query matrix block;

[0074] S12. Control 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.

[0075] S13. Control the first thread group and the second thread group to alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operation until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operation, and obtain the last sub-block of the j-th attention result matrix block; wherein, the j-th key matrix block and the j-th value matrix block are pre-stored in shared memory.

[0076] 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. Specifically, it means controlling the second thread bundle to write the i-th query matrix block from shared memory to the eleventh register group regQ of the second thread bundle.

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

[0078] This application does not impose specific restrictions on the thread group that writes the i-th query matrix block, 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, the first thread group can be controlled to write the i-th query matrix block, the j-th key matrix block, and the j-th value matrix block from global memory to shared memory.

[0079] 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 operation (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 intermediate results can be precision-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.

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

[0081] 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 each alternating attention mechanism operation of the j-th attention score matrix block of the i-th query matrix block.

[0082] During the alternating attention mechanism operations, the first and second thread bundles 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 bundle 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 bundle 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 bundle utilizes the third sub-block of the j-th value matrix block of the first query matrix block, and so on.

[0083] This application embodiment breaks the pure producer-consumer relationship between thread bundles by controlling the first thread bundle group and the second thread bundle group to alternately perform attention mechanism operations. Both thread bundle groups undertake a part of the computation task, avoiding the problem of insufficient register resources caused by relying on a single thread bundle group to undertake all computation tasks. This provides sufficient register resources for the computation of attention mechanism, and ultimately effectively improves the computational efficiency of attention mechanism.

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

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

[0086] The control of the first thread group and the second thread group to alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations, resulting in the last sub-block of the j-th attention result matrix block, includes:

[0087] The second thread bundle is controlled to perform 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 m-th sub-block of the j-th attention result matrix block;

[0088] After obtaining the m-th sub-block of the j-th attention result matrix block, the first thread bundle group is controlled to perform attention mechanism operation with 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 to obtain the m+1-th sub-block of the j-th attention result matrix block.

[0089] After obtaining the (m+1)th sub-block of the j-th attention result matrix block, let m=m+1, and return to the step of controlling the second thread group to perform 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 shared memory, until m is the preset total number of sub-blocks, and obtain the last sub-block of the j-th attention result matrix block.

[0090] It is worth noting that the attention mechanism operation performed 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.

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

[0092] In this embodiment, the second thread group is first controlled to perform 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 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 is the preset total number of sub-blocks. This indicates that the attention mechanism operation of the last sub-block of the j-th value matrix block has been performed. At this time, the m-th (m=preset total number of sub-blocks) sub-block of the j-th value matrix block, i.e., the last sub-block, is used to perform attention mechanism operation, and the last sub-block of the j-th attention result matrix block is obtained. At this time, the final result of the j-th attention result matrix block of the i-th query matrix block is obtained.

[0093] In an optional embodiment, the control of the second thread bundle group performs attention mechanism operations 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 m-th sub-block of the j-th attention result matrix block, including:

[0094] If j=1, then the second thread bundle group is controlled to perform 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 m-th sub-block of the first attention result matrix block and write it into the shared memory;

[0095] If j=2,...,J-1, then the second thread group controls the m-th sub-block of the (j-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs 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 m-th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block;

[0096] If j=J, then the second thread group controls the m-th sub-block of the (J-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs attention mechanism operation with the last attention score matrix block in the first register group of the second thread group and the m-th sub-block of the last value matrix block in the shared memory to obtain the m-th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

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

[0098] For the first attention score matrix block, after the second thread bundle performs the attention mechanism operation to obtain 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.

[0099] For the 2nd, ..., J-1th attention score matrix blocks (i.e., excluding the 1st and last attention score matrix blocks), the second thread group needs to first write the m-th sub-block of the previous (i.e., the j-1th) attention result matrix block from shared memory into its own second register group regO to participate in subsequent calculations, thereby achieving the reuse of the second register group regO of the second 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 m-th sub-block of the j-th attention result matrix block into shared memory avoids register occupancy, providing sufficient register resources for subsequent calculations.

[0100] For the last attention score matrix block, the second thread group also needs to first write the m-th sub-block of the previous (i.e., the (j-1)-th) attention result matrix block from shared memory into its own second register group regO to participate in subsequent calculations, thus 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. Since it is the last attention result matrix block, it does not need to participate in subsequent calculations; the m-th sub-block of the last attention result matrix block of the i-th query matrix block is written to shared memory in first precision type.

[0101] Specifically, the first precision type is the same as the precision type of the value matrix blocks in shared memory. Generally, the precision type of matrix data in shared memory is a low precision type, which can reduce data storage overhead. Therefore, if the attention mechanism operation uses the second precision type, then all sub-blocks of the resulting attention result matrix block will be of the second precision type. Thus, the m-th sub-block of the last attention result matrix block needs to be converted from the second precision type to the first precision type before being written to shared memory. If the attention mechanism operation uses the first precision type, then all sub-blocks of the resulting attention result matrix block will be of the first precision type. Therefore, no precision conversion is needed, and the m-th sub-block of the last attention result matrix block can be directly written to shared memory.

[0102] It is worth noting that for the attention mechanism operation of the 2nd, ..., Jth attention result matrix blocks, since it is not the 1st attention result matrix block, after performing matrix multiplication on the m-th sub-block of the j-th attention score matrix block and the j-th value matrix block, the result of this matrix multiplication needs to be used to update the m-th sub-block of the previous (i.e., the (j-1th)th attention result matrix block to obtain the m-th sub-block of the j-th attention result matrix block; the specific update method is superposition, that is, the j-th attention score matrix block... The m-th sub-block of the j-th value matrix block + the m-th sub-block of the (j-1)-th attention result matrix block. The mma instruction can perform matrix multiplication and superposition operations, so it can be used to achieve both to obtain the m-th sub-block of the j-th attention result matrix block.

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

[0104] For the i-th query matrix block, after the last alternating attention mechanism operation of the second thread group, the second thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

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

[0106] In this embodiment of the application, for the i-th query matrix block, after 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, all sub-blocks of the last attention result matrix block obtained by the second thread bundle itself through the attention mechanism operation are written into the global memory for subsequent global memory access to the attention result matrix block.

[0107] It is understandable that all sub-blocks of the last attention result matrix block obtained by the second thread bundle itself through attention mechanism operation and all sub-blocks of the last attention result matrix block obtained by the first thread bundle itself through alternating attention mechanism operation form the last attention result matrix block of the i-th query matrix block.

[0108] In an optional embodiment, controlling the first thread bundle group to perform attention mechanism operations 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 to obtain the (m+1)-th sub-block of the j-th attention result matrix block includes:

[0109] If j=1, then the first thread bundle group is controlled to perform attention mechanism operation by 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 (m+1)th sub-block of the first attention result matrix block and write it into the shared memory.

[0110] If j=2,...,J-1, control the first thread group to write 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 perform 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 (m+1)th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block;

[0111] If j=J, control the first thread group to write 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 perform attention mechanism operation with the last attention score matrix block in the first register group of the first thread group and the (m+1)th sub-block of the last value matrix block in the shared memory to obtain the (m+1)th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

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

[0113] For the first attention score matrix block, after the first thread bundle performs attention mechanism operations to obtain 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.

[0114] For the 2nd, ..., J-1th attention score matrix blocks (i.e., excluding the 1st and last attention score matrix blocks), the second thread group needs to first write the (m+1)th sub-block of the previous (j-1th) attention result matrix block from shared memory into its own second register group regO to participate in subsequent calculations, thus realizing the reuse of the second register group regO of the second 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 (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.

[0115] For the last attention score matrix block, the second thread group also needs to first write 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, thus 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. Since it is the last attention result matrix block, it does not need to participate in subsequent calculations; the (m+1)th sub-block of the obtained last attention result matrix block is written to shared memory in first precision type.

[0116] Specifically, the first precision type is the same as the precision type of the value matrix blocks in shared memory. Generally, the precision type of matrix data in shared memory is a low precision type, which can reduce data storage overhead. Therefore, if the attention mechanism operation uses the second precision type, then all sub-blocks of the resulting attention result matrix block will be of the second precision type. Thus, the (m+1)th sub-block of the last attention result matrix block of the i-th query matrix block needs to be converted from the second precision type to the first precision type before being written to shared memory. If the attention mechanism operation uses the first precision type, then all sub-blocks of the resulting attention result matrix block will be of the first precision type. Therefore, no precision conversion is needed, and the (m+1)th sub-block of the last attention result matrix block can be directly written to shared memory.

[0117] Understandably, for the attention mechanism operation of the 2nd, ..., Jth attention result matrix blocks, since it is not the 1st attention result matrix block, after performing matrix multiplication on the (m+1)th sub-block of the jth attention score matrix block and the jth value matrix block, the result of this matrix multiplication needs to be used to update the (m+1)th sub-block of the previous (i.e., the (j-1)th) attention result matrix block in order to obtain the (m+1)th sub-block of the jth attention result matrix block. The specific update method is superposition, that is, the jth attention score matrix block... The (m+1)th sub-block of the j-th value matrix block plus the (m+1)th sub-block of the (j-1)-th attention result matrix block. The mma instruction can perform both matrix multiplication and superposition operations; therefore, it can be used to implement both operations to obtain the m-th sub-block of the j-th attention result matrix block.

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

[0119] For the i-th query matrix block, after the last alternating attention mechanism operation of the first thread group, the first thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

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

[0121] In this embodiment of the application, for the i-th query matrix block, after the last alternating attention mechanism operation of the first thread bundle group, the first thread bundle group will no longer participate in the attention mechanism operation. At this time, all sub-blocks of the last attention result matrix block obtained by the first thread bundle group itself through the attention mechanism operation are written into the global memory for subsequent global memory access to the attention result matrix block.

[0122] 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, all the sub-blocks of the last attention result matrix block obtained by the second thread group itself through attention mechanism operation and all the sub-blocks of the last attention result matrix block obtained by the first thread group itself through alternating attention mechanism operation form the last attention result matrix block of the i-th query matrix block.

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

[0124] After the first thread group writes all sub-blocks of the last attention result matrix block obtained by itself through attention mechanism operation to global memory, the first thread group writes the (i+1)th query matrix to the position of the i-th query matrix in the shared memory.

[0125] It's worth noting that after the first thread group writes all sub-blocks of the last attention result matrix block obtained from its attention mechanism operation into global memory, it indicates that the attention mechanism operation of the i-th query matrix has ended or is about to end, and the attention mechanism operation of the next (i.e., the (i+1)-th) query matrix can begin. At this point, the first thread group can be controlled to write the (i+1)-th query matrix into the position of the i-th query matrix in shared memory, allowing the thread group to form a pipeline and continue the attention mechanism operation of the (i+1)-th query matrix block. In other words, after the (i+1)-th query matrix is ​​written into the position of the i-th query matrix in shared memory, let i = i+1, and continue to return to step S11 for execution.

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

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

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

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

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

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

[0132] For the i-th query matrix block, after the last alternating attention mechanism operation, the first thread group is controlled to write the j+x-th key-value matrix block from global memory to 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.

[0133] It is worth noting that, in this embodiment, after the last alternating attention mechanism operation of 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 attention mechanism operation on the i-th query matrix block.

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

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

[0136] 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, S r2 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;

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

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

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

[0140] 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. max1The 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.

[0141] (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.

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

[0143] 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... The relationship between them.

[0144] 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. u2 The relationship between them;

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

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

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

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

[0149] 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 restrictions 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.

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

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

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

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

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

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

[0156] In an optional embodiment, before the first thread group and the second thread group alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations, and the last sub-block of the j-th attention result matrix block is obtained, the method further includes:

[0157] The second thread group is controlled to calculate the row target value of the j-th attention score matrix block required for the attention mechanism operation, and write it into the first register block of the second thread group;

[0158] The second thread bundle is controlled to write the row target value of the j-th attention score matrix block in the first register block of the second thread bundle into the shared memory;

[0159] The first thread bundle is controlled to write the row target value of the j-th attention score matrix block in the shared memory into the first register block of the first thread bundle.

[0160] It is worth noting that both the first and second thread bundles are responsible for 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, including:

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

[0162] The current row's accumulated value, accum_rsum_new, is specifically 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; where each row corresponds to one accumulated value.

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

[0164] Specifically, the first register group of the second thread bundle includes the fourth register group, the fifth register group, and the sixth register group, which are used to store the historical row accumulated value `accum_rsum`, the current row accumulated value `accum_rsum_new`, and the modifier value, respectively. Similarly, the first register group of the first thread bundle includes the fourth register group, the fifth register group, and the sixth register group, which are used to store the historical row accumulated value `accum_rsum`, the current row accumulated value `accum_rsum_new`, and the modifier value, respectively.

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

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

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

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

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

[0170] Therefore, the remaining row values ​​required for the second thread bundle to calculate the aforementioned three row target values ​​are also controlled, including: the maximum row value `chunk_rmax`, the historical global maximum row value `accum_rmax`, the current global maximum row value `accum_rmax_new`, and the row accumulation value `chunk_rsum`. These values ​​are written to the second register group of the second thread bundle. These values ​​do not participate in the attention mechanism calculation of the first thread bundle, therefore, there is no need to write them to the first thread bundle via shared memory. Specifically, the second register group of the second thread bundle includes: the seventh register group, the eighth register group, the ninth register group, and the tenth register group, which are used to store: the maximum row value `chunk_rmax`, the historical global maximum row value `accum_rmax`, the current global maximum row value `accum_rmax_new`, and the row accumulation value `chunk_rsum`, respectively.

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

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

[0173] like Figure 4This is a schematic diagram illustrating the principle of attention mechanism calculation provided in an embodiment of this application. A matrix multiplication operation is performed between the query matrix block Q and the transposed key matrix block K 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. Alternating attention mechanism operations are then performed between the attention score matrix block P and a sub-block of the value matrix block V to obtain a sub-block of the attention result matrix block O. Figure 4 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 four sub-blocks of the attention result matrix block O, including the first sub-block O_0, the third sub-block O_1, the second sub-block O_2, and the fourth sub-block O_3. Then, the first thread bundle group and the second thread bundle group each use half of the value matrix block alternately, finally obtaining half of the attention result matrix block.

[0174] like Figures 5-8 This is a layout diagram of a thread bundle group provided in an embodiment of this application. First, it should be noted that cp indicates writing data from global memory to shared memory, ld indicates loading data, specifically writing data from shared memory to a register, and st indicates writing data from a register to shared memory.

[0175] The first thread group WG0 controls the writing of the first query matrix block Q1, the first key-value matrix block KV1 (including the first key matrix block K1 and the first value matrix block V1), and the second key-value matrix block KV1 (including the second key matrix block K1 and the second value matrix block V1) into shared memory.

[0176] The second thread group WG1 writes Q1 into its own register and uses Q1 to perform a matrix multiplication operation with K1 (i.e., Q1...). After K1), perform the softmax operation. Here, the softmax operation means that after Q1... After K1, the operations preceding the alternating attention mechanism operations, such as exponentiation, precision conversion, and calculation of row target values, are performed. Then, the first attention score matrix is ​​written into shared memory in block P1, and the alternating attention mechanism operations of WG1 can begin.

[0177] WG0 writes P1 into its own register, and can then begin WG0's alternating attention mechanism operation.

[0178] When j=1, WG1 uses P1 and the first sub-block V(1)_0 of the first value matrix block to perform attention mechanism operation, obtaining the first sub-block O(1)_acc0 of the first attention result matrix block and writing it into shared memory; WG0 uses P1 and the second sub-block V(1)_2 of the first value matrix block to perform attention mechanism operation, obtaining the second sub-block O(1)_acc2 of the first attention result matrix block; WG1 uses P1 and the third sub-block of the first value matrix block... V(1)_1 performs the attention mechanism operation, obtaining the third sub-block O(1)_acc1 of the first attention result matrix block and writing it into shared memory. At this point, for WG1, it is the last operation using Q1, K1, and V1, and the subsequent operations of Q1, K2, and V2 can begin. WG0 performs the attention mechanism operation using P1 and the fourth sub-block V(1)_3 of the first value matrix block, obtaining the fourth sub-block O(1)_acc3 of the first attention result matrix block. Since WG0 does not need to use V1 and P1 for operation after performing the attention mechanism operation using P1 and the fourth sub-block V(1)_3 of the first value matrix block, it can control the position of WG0 writing the third key-value matrix block KV3 into shared memory KV1.

[0179] In the process of alternating attention operations (j=2,...,J-1), before performing 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 sub-block O(j)_acc of the (j-1)-th attention result matrix block O(j) is first written to the required thread group. The remaining processes are similar to the above and will not be repeated here.

[0180] When j=J, after each alternating attention operation, the sub-block O(J)_acc of the Jth attention result matrix block O(J) is obtained. This sub-block is then precision-transformed to obtain the first sub-block O(J)_0, the second sub-block O(J)_2, the third sub-block O(J)_1, and the fourth sub-block O(J)_3 of the Jth attention result matrix block O(J). The remaining processes are similar to those described above and will not be repeated here.

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

[0182] Instruction logic for the first thread bundle:

[0183] Loop{ / Q outer loop /

[0184] tma cp Q / / Write the query matrix block Q from global memory to shared memory;

[0185] Loop{ / KV internal circulation /

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

[0187] ldmatrix regP / / Writes the attention score matrix p of fp16 from shared memory to the first register group regP of the first thread bundle;

[0188] ld.gsm acc_rsum, acc_rsum_new, modifier / / Writes the historical row accumulation value accum_rsum, the current row accumulation value accum_rsum_new, and the modifier from shared memory to the first register block of the first thread bundle;

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

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

[0191] regO = mma(regP, V, O)

[0192] regO = regO / acc_rsum_new / / Together with the previous instruction, this forms the m+1th sub-block of the first attention score matrix block in the first register group regP of the first thread bundle group, which is then used for attention mechanism operation with the first value matrix block in the shared memory.

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

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

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

[0196] regO = modifier accum_rsum regO / /

[0197] regO = mma(P, V, regO)

[0198] regO = regO / acc_rsum_new / / Together with the previous two instructions, this forms the m+1 sub-block 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+1 sub-block of the jth value matrix block in the shared memory, which are then used for attention mechanism operations.

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

[0200] / / Used for attention score matrix blocks in the middle (i.e., except for the first and last ones);

[0201] {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 first thread bundle;

[0202] regO = modifier accum_rsum regO

[0203] regO = mma(P, V, regO)

[0204] regO = regO / acc_rsum_new / / Together with the previous two instructions, this forms the m+1 sub-block 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+1 sub-block of the Jth value matrix block in the shared memory, which are then used for attention mechanism operations.

[0205] Typecast regO to regO_fp16 / / Precision conversion; Specifically, after performing precision conversion on the (m+1)th block of the last attention result matrix block in the second register group regO of the first thread beam group, it is written to the twelfth register group regO_fp16 of the first thread beam group, thus realizing the precision conversion from fp32 to fp16.

[0206] Stmatrix regO_fp16 to O(gsm) / / Writes the (m+1)th block of the last attention result matrix after precision transformation to shared memory;

[0207] / / Used for the attention score matrix block of the last one;

[0208] }

[0209] }

[0210] Tma half O / / Writes the sub-block of the last attention result matrix block obtained by performing attention mechanism operations on itself into global memory;

[0211] }

[0212] }

[0213] Instruction logic for the second thread bundle:

[0214] Loop{ / Q outer loop /

[0215] Ldmatrix Q ->regQ / / Writes the query matrix block Q from shared memory to the eleventh register group regQ of the second thread bundle;

[0216] Loop{ / KV internal circulation /

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

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

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

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

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

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

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

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

[0225] St.gsm acc_rsum, acc_rsum_new, modifier / / Writes the historical row accumulation value accum_rsum, the current row accumulation value accum_rsum_new, and the modifier value to shared memory;

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

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

[0228] regO = mma(regP, V, O)

[0229] regO = regO / acc_rsum_new / / Together with the previous instruction, this forms the m-th sub-block of the first attention score matrix block in the first register group regP of the second thread bundle group, which is then used for attention mechanism operation with the first value matrix block in the shared memory.

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

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

[0232] {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;

[0233] regO = modifier accum_rsum regO

[0234] regO = mma(P, V, regO)

[0235] regO = regO / acc_rsum_new / / Together with the previous two instructions, this forms the m-th sub-block of the (j-1)-th attention result matrix block of the second register group regO of the second thread bundle group, the m-th sub-block of the j-th attention score matrix block in the first register group regP of the second thread bundle group, and the j-th value matrix block in the shared memory, which are then subjected to attention mechanism operations.

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

[0237] / / Used for attention score matrix blocks in the middle (i.e., except for the first and last ones);

[0238] {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;

[0239] regO = modifier accum_rsum regO

[0240] regO = mma(P, V, regO)

[0241] regO = regO / acc_rsum_new / / Together with the previous two instructions, this forms the m-th sub-block of the (J-1)-th attention result matrix block of the second register group regO of the second thread bundle group, the J-th attention score matrix block in the first register group regP of the second thread bundle group, and the m-th sub-block of the J-th value matrix block in the shared memory, which are then used for attention mechanism operations.

[0242] Typecast regO to regO_fp16 / / Precision conversion; Specifically, after performing precision conversion on the m-th sub-block of the last attention result matrix block in the second register group regO of the second thread bundle, it is written to the twelfth register group regO_fp16 of the second thread bundle, thus realizing the precision conversion from fp32 to fp16.

[0243] Stmatrix regO_fp16 to O(gsm) / / Writes the m-th sub-block of the last attention result matrix after precision transformation to shared memory;

[0244] / / Used for the attention score matrix block of the last one;

[0245] }

[0246] }

[0247] Tma half O / / Writes the sub-block of the last attention result matrix block obtained by performing attention mechanism operations on itself into global memory;

[0248] }

[0249] }

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

[0251] Table 1

[0252]

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

[0254] For the fourth, fifth, and sixth register groups, since the matrix elements held by one thread are distributed across two rows, and each row needs to calculate and store the historical row accumulation value accum_rsum, the current row accumulation value accum_rsum_new, and the modifier, the fourth, fifth, and sixth register groups each require two registers.

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

[0256] 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).

[0257] The number of registers in register group regO_fp16 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).

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

[0259] Table 2

[0260]

[0261] 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).

[0262] 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).

[0263] 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).

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

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

[0266] 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).

[0267] The number of registers in register group regO_fp16 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) (size of the matrix held by one thread bundle) = 4 (64 128) / (4 16 16) = 32. Since 32 is less than 144, the attention result matrix blocks after precision conversion can be stored in the eleventh register group regQ. Therefore, the twelfth register group regO_fp16 is not needed, and the number of registers in the twelfth register group regO_fp16 is 0.

[0268] 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-11 , respectively corresponding Figures 5-7 The shared memory resource distribution at each stage; 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 unoccupied 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 at 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 10During 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-1, 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. Figure 11 The period from the end of j=J-1 to the end of j=J is the time when the sub-blocks O(J)_0, O(J)_1, O(J)_2, and O(J)_3 (each occupying 16KB), the key-value matrix block KV, and the attention score matrix block P of the precision-converted attention score matrix block O(J) will occupy shared memory.

[0269] Table 3

[0270]

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

[0272] See Figure 12 This is a structural block diagram of an artificial intelligence chip provided in an embodiment of this application. The artificial intelligence chip includes global memory (GLM) and one or more computing units; each computing unit includes: shared memory (GSM), one or more first thread bundles, and one or more second 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.

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

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

[0275] See Figure 13 , Figure 13 This is a structural block diagram of an attention mechanism computing device 20 provided in an embodiment of this application. The attention mechanism computing device 20 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 in the above-described attention mechanism computing method embodiments. Alternatively, when the processor 21 executes the computer program, it implements the functions of each module / unit in the above-described device embodiments.

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

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

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

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

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

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

[0282] 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. A method for calculating an attention mechanism, characterized in that, include: Control the second thread bundle to load the i-th query matrix block; The second thread group is controlled 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; The first thread group and the second thread group are controlled to alternately use different sub-blocks of the j-th value matrix block and the same j-th attention score matrix block to perform attention mechanism operations until the last sub-block of the j-th value matrix block is used to perform attention mechanism operations with the same j-th attention score matrix block, so as to obtain the last sub-block of the j-th attention result matrix block; wherein, the j-th key matrix block and the j-th value matrix block are pre-stored in shared memory.

2. The attention mechanism calculation method as described in claim 1, characterized in that, The j-th attention score matrix block is written to the first register group of the second thread bundle and the first register group of the first thread bundle; The control of the first thread group and the second thread group to alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations, resulting in the last sub-block of the j-th attention result matrix block, includes: The second thread bundle is controlled to perform 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 m-th sub-block of the j-th attention result matrix block; After obtaining the m-th sub-block of the j-th attention result matrix block, the first thread bundle group is controlled to perform attention mechanism operation with 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 to obtain the m+1-th sub-block of the j-th attention result matrix block. After obtaining the (m+1)th sub-block of the j-th attention result matrix block, let m=m+1, and return to the step of controlling the second thread group to perform 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, until m is the preset total number of sub-blocks, and obtain the last sub-block of the j-th attention result matrix block.

3. The attention mechanism calculation method as described in claim 2, characterized in that, The control of the second thread bundle group involves performing an 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 m-th sub-block of the j-th attention result matrix block, including: If j=1, then the second thread bundle group is controlled to perform 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 m-th sub-block of the first attention result matrix block and write it into the shared memory; If j=2,...,J-1, then the second thread group controls the m-th sub-block of the (j-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs 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 m-th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block; If j=J, then the second thread group controls the m-th sub-block of the (J-1)-th attention result matrix block to be written from the shared memory to the second register group of the second thread group, and performs attention mechanism operation with the last attention score matrix block in the first register group of the second thread group and the m-th sub-block of the last value matrix block in the shared memory to obtain the m-th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

4. The attention mechanism calculation method as described in any one of claims 1 to 3, characterized in that, The method further includes: For the i-th query matrix block, after the last alternating attention mechanism operation of the second thread group, the second thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

5. The attention mechanism calculation method as described in claim 2, characterized in that, The control of the first thread bundle group to perform 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 to obtain the (m+1)-th sub-block of the j-th attention result matrix block includes: If j=1, then the first thread bundle group is controlled to perform attention mechanism operation by 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 (m+1)th sub-block of the first attention result matrix block and write it into the shared memory. If j=2,...,J-1, control the first thread group to write 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 perform 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 (m+1)th sub-block of the j-th attention result matrix block and write it to the shared memory; J is the last attention score matrix block of the i-th query matrix block; If j=J, control the first thread group to write 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 perform attention mechanism operation with the last attention score matrix block in the first register group of the first thread group and the (m+1)th sub-block of the last value matrix block in the shared memory to obtain the (m+1)th sub-block of the last attention result matrix block and write it to the shared memory with the first precision type.

6. The attention mechanism calculation method as described in any one of claims 1, 2, and 5, characterized in that, The method further includes: For the i-th query matrix block, after the last alternating attention mechanism operation of the first thread group, the first thread group is controlled to write all sub-blocks of the last attention result matrix block obtained by itself through the attention mechanism operation into global memory.

7. The attention mechanism calculation method as described in claim 6, characterized in that, The method further includes: After the first thread group writes all sub-blocks of the last attention result matrix block obtained by itself through attention mechanism operation to global memory, the first thread group writes the (i+1)th query matrix to the position of the i-th query matrix in the shared memory.

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

9. The attention mechanism calculation method as described in claim 8, characterized in that, The method further includes: For the i-th query matrix block, after the last alternating attention mechanism operation, the first thread group is controlled to write the j+x-th key-value matrix block from global memory to 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.

10. The attention mechanism calculation method as described in claim 8, characterized in that, The dimensions of the query matrix block, the key-value matrix block, the key matrix block, and the value matrix block are determined through the following steps: 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. 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, S r2 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; 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. 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.

11. The attention mechanism calculation method as described in claim 1, characterized in that, The control of the second thread group involves performing 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, including: 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. 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.

12. The attention mechanism calculation method as described in claim 1, characterized in that, Before the first thread group and the second thread group alternately use different sub-blocks of the j-th value matrix block with the same j-th attention score matrix block to perform attention mechanism operations, until the last sub-block of the j-th value matrix block is used with the same j-th attention score matrix block to perform attention mechanism operations, and the last sub-block of the j-th attention result matrix block is obtained, the method further includes: The second thread group is controlled to calculate the row target value of the j-th attention score matrix block required for the attention mechanism operation, and write it into the first register block of the second thread group; The second thread bundle is controlled to write the row target value of the j-th attention score matrix block in the first register block of the second thread bundle into the shared memory; The first thread bundle is controlled to write the row target value of the j-th attention score matrix block in the shared memory into the first register block of the first thread bundle.

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

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

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

Images