An attention mechanism calculation method, device, storage medium and product

Abstract

The application discloses an attention mechanism calculation method and device, a storage medium and a product. The method comprises the following steps: performing matrix multiplication operation on a query matrix block and a key matrix block to obtain a first product matrix block of a first precision type, performing exponential operation on the first product matrix block written in a second register group to obtain a numerator matrix block; writing the quantized numerator matrix block into a third register group according to a second precision type; writing the numerator matrix block into a fourth register group according to a third precision type and performing layout conversion through a shared memory; performing softmax operation on the layout-converted numerator matrix block and a value matrix block to obtain an attention result matrix block of the first precision type and write the attention result matrix block into a fifth register group; and writing the attention result matrix block into the shared memory according to the third precision type. The embodiment of the application can reduce the cumulative error of an intermediate operation result.

Description

Technical Field

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

[0002] Currently, attention mechanisms often employ a uniform low-precision type throughout the entire process: from input data, matrix multiplication, softmax normalization to output data, core steps all rely on a single low-precision paradigm. However, this approach has the following problems: during the iterative accumulation of intermediate computation results (such as the results of matrix multiplication and softmax normalization), the error is continuously amplified due to the low-precision characteristics, resulting in a large-scale cumulative error, which ultimately leads to a decrease in the accuracy of the output result. Summary of the Invention

[0003] This application provides an attention mechanism calculation method, device, storage medium, and product to solve the problem in the prior art that the error of intermediate calculation results is continuously amplified due to the low precision characteristics during the iterative accumulation process, resulting in a large-scale cumulative error and ultimately reducing the accuracy of the output results.

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

[0005] Control at least one first thread bundle to perform the following steps:

[0006] Write the query matrix in blocks from shared memory into the first register group;

[0007] The query matrix block of the first register group is multiplied with the key matrix block that has been copied into the shared memory in advance to obtain a first product matrix block of the first precision type and write it into the second register group.

[0008] The first product matrix block of the second register group is used to perform exponential operation to obtain the numerator matrix block of the softmax operation and overwrite it into the second register group.

[0009] The molecular matrix blocks of the second register group are quantized, and the quantized molecular matrix blocks are written into the third register group according to the second precision type.

[0010] The molecular matrix blocks of the third register group are written into the fourth register group according to the third precision type, and the layout transformation of the molecular matrix blocks of the fourth register group is performed through the shared memory.

[0011] The softmax operation is performed using the molecular matrix blocks after layout transformation of the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain the attention result matrix blocks of the first precision type and write them into the fifth register group;

[0012] The attention result matrix of the fifth register group is quantized in blocks, and the quantized attention result matrix blocks are written into the third register group according to the second precision type.

[0013] The attention result matrix of the third register group is divided into blocks according to the third precision type and written to the shared memory through the fourth register group; wherein the precision of the first precision type is greater than the precision of the second precision type, and the precision of the second precision type is greater than the precision of the third precision type.

[0014] As an improvement to the above scheme, the step of writing the molecular matrix blocks of the third register group into the fourth register group according to the third precision type, and performing layout transformation on the molecular matrix blocks of the fourth register group through the shared memory, includes:

[0015] The molecular matrix of the third register group is divided into blocks and written into the fourth register group according to the third precision type;

[0016] The molecular matrix of the fourth register group is written into the shared memory in blocks;

[0017] The molecular matrix of the shared memory is written into the fourth register group in blocks.

[0018] As an improvement to the above scheme, the step of quantizing the molecular matrix blocks of the second register group includes:

[0019] The attention result matrix of the second register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type.

[0020] The step of quantizing the attention result matrix of the fifth register group in blocks includes:

[0021] The attention result matrix of the fifth register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type.

[0022] As an improvement to the above scheme, the step of performing exponential operations on the first product matrix block of the second register group to obtain the numerator matrix block of the softmax operation and overwriting it into the second register group includes:

[0023] Calculate the maximum row value of the first product matrix block in the second register group and write it to the sixth register group;

[0024] Write the historical global row maximum value that is in the same row as the first product matrix block of the second register group into the seventh register group;

[0025] Compare the maximum row value of the sixth register group with the historical global row maximum value of the seventh register group, and write the larger value as the current global row maximum value into the eighth register group;

[0026] The first product matrix block of the second register group and the current global row maximum value of the eighth register group are subjected to exponential operation to obtain the numerator matrix block, which is then overwritten into the second register group.

[0027] As an improvement to the above scheme, the step of performing a softmax operation on the molecular matrix blocks after layout transformation of the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain the attention result matrix blocks of the first precision type and writing them into the fifth register group includes:

[0028] Calculate the row accumulation value of the first product matrix block in the second register group and write it to the ninth register group;

[0029] Write the accumulated value of the historical row that is in the same row as the first product matrix block of the second register group into the tenth register group;

[0030] The current historical global row maximum value of the eighth register group and the historical global row maximum value of the seventh register group are used to perform an exponential operation to obtain a correction value, which is then written into the eleventh register group. The correction value is used to correct the attention result matrix block obtained by the softmax operation performed on the query matrix block in the past.

[0031] The current row accumulation value is obtained by performing scaling operations using the correction value of the eleventh register group, the historical row accumulation value of the tenth register group, and the row accumulation value of the ninth register group, and then writing it into the twelfth register group.

[0032] The correction value of the eleventh register group is used to correct the current multiply-accumulate matrix block of the fifth register group, and then overwrites the fifth register group.

[0033] Matrix multiplication and addition operations are performed using the molecular matrix block after layout transformation of the fourth register group, the value matrix block pre-copied into the shared memory, and the multiply-add matrix block after correction of the fifth register group to obtain the multiply-add matrix block and overwrite it into the fifth register group.

[0034] The attention result matrix block is obtained by performing a division operation between the multiply-accumulate matrix block of the fifth register group and the current row accumulated value of the twelfth register group, and then overwrites and writes it into the fifth register group.

[0035] As an improvement to the above scheme, the size of the query matrix block is the attention head dimension. The sequence length of the query matrix block; the size of the key matrix block is the attention head dimension. The sequence length of the key matrix block; the size of the value matrix block is the attention head dimension. The sequence length of the value matrix block; the sequence length of the key matrix block is equal to the sequence length of the value matrix block;

[0036] The dimensions of the query matrix block, the key matrix block, and the value matrix block are determined through the following steps:

[0037] Based on the first preset relationship, the first objective function is established as min(S) max -S1-S2); where S max The maximum number of registers used by each thread in the first thread bundle is limited, S1 is the preset number of reserved registers, and S2 is the actual number of registers that each thread in the first thread bundle needs to use when the attention mechanism is executed; the first preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block and the actual number of registers;

[0038] Based on the second presupposed relationship, the second objective function is established as min(M) max -M); where M max M represents the maximum capacity of the shared memory, where M is the actual shared memory capacity required for all parallel execution of the attention mechanism by the first thread bundle in one computing unit; the second preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block, and the actual shared memory capacity.

[0039] Given a fixed attention head dimension, the first objective function and the second objective function are jointly solved to obtain the sequence length of the query matrix block, the sequence length of the key matrix block, and the sequence length of the value matrix block.

[0040] As an improvement to the above scheme, the second preset relationship also includes: the relationship between the first quantity, the second quantity, the third quantity and the actual shared memory capacity;

[0041] Wherein, the first quantity is the number of key matrix blocks that can be copied into the shared memory simultaneously, the second quantity is the number of value matrix blocks that can be copied into the shared memory simultaneously, and the third quantity is the number of parallel groups of the first thread bundle in one computing unit, and the first quantity and the second quantity are equal.

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

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

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

[0045] Compared with the prior art, the present application provides an attention mechanism calculation method, device, storage medium, and product, which executes the following steps by controlling at least one first thread bundle group: writing a query matrix block from shared memory to a first register group; performing matrix multiplication on the query matrix block of the first register group and a key matrix block pre-copied into the shared memory to obtain a first product matrix block of a first precision type and writing it to a second register group; performing exponential operation on the first product matrix block of the second register group to obtain a numerator matrix block of a softmax operation and overwriting it in the second register group; quantizing the numerator matrix block of the second register group and writing the quantized numerator matrix block to a third register group according to a second precision type; and writing the numerator matrix block of the third register group according to a second precision type. The third precision type is written to the fourth register group, and the layout of the molecular matrix blocks of the fourth register group is transformed through the shared memory. A softmax operation is performed on the layout-transformed molecular matrix blocks of the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain the attention result matrix blocks of the first precision type, which are then written to the fifth register group. The attention result matrix blocks of the fifth register group are quantized, and the quantized attention result matrix blocks are written to the third register group according to the second precision type. The attention result matrix blocks of the third register group are written to the shared memory according to the third precision type through the fourth register group. The precision of the first precision type is greater than that of the second precision type, and the precision of the second precision type is greater than that of the third precision type. Therefore, this embodiment controls intermediate operations to be performed in the first precision type, avoiding the use of a single low precision, thereby reducing the cumulative error of intermediate operation results and ultimately improving the accuracy of the attention output results. Furthermore, by reusing register resources, this embodiment not only avoids the risk of performance degradation caused by register overflow but also mobilizes as many hardware resources as possible for computation, fully releasing the potential of hardware computing power and achieving high computing performance based on efficient resource utilization. Attached Figure Description

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

[0047] Figure 2 This is a schematic diagram of the first register layout provided in the embodiments of this application;

[0048] Figure 3 This is a schematic diagram of the second register layout provided in the embodiments of this application;

[0049] Figure 4 This is a schematic diagram of the third register layout provided in the embodiments of this application;

[0050] Figure 5 This is a schematic diagram of the fourth register layout provided in the embodiments of this application;

[0051] Figure 6 This is a schematic diagram of the fifth register layout provided in the embodiments of this application;

[0052] Figure 7 This is a schematic diagram illustrating the distribution of shared memory resources according to an embodiment of this application;

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

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

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

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

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

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

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

[0060] 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 computation tasks. Generally, each CU can only schedule and execute one thread block at a time, enabling multiple warp groups to process tasks in parallel. A thread block includes 3 or 4 warp groups, a warp group includes 4 warps, and a warp includes 32 threads.

[0061] Due to limitations in AI chips, each thread within a single thread warp group is limited to a maximum of 256 Thread Level Registers (TLRs), and each thread within a single thread block is limited to a maximum of 512 registers. Thread warp groups are typically divided into producer warp groups and consumer warp groups. Each thread within a producer warp group is limited to a maximum of 32 registers, and each thread within a consumer warp group is limited to a maximum of 240 registers. Thus, a single thread block can include one producer warp group and two consumer warp groups, with each thread within this block limited to a total of 512 registers.

[0062] For a single held matrix of size 16 A 16-thread bundle, where each thread holds a matrix element of (16) 16) / 32 = 8. Currently, register layouts (i.e., the positional mapping between threads, matrix elements, and registers) include: normal layout and interleave layout. The normal layout is zigzag-shaped, and the interleave layout is N-shaped. Furthermore, considering the precision type, there are five specific register layouts:

[0063] (1) fp8 normal layout (8-bit floating-point precision normal layout): The precision type of each element in the matrix is ​​fp8, 4 consecutive elements are held by the same thread, and the register order is Z-shaped; such as Figure 2 A certain thread t0 holds 8 matrix elements, corresponding to holding 2 registers r0 and r2.

[0064] (2) fp16 interleave layout: The precision type of each element in the matrix is ​​fp16, and four consecutive elements are held by the same thread, with the registers arranged in an N-shape; for example Figure 3 A certain thread t0 holds 8 matrix elements, corresponding to 4 registers r0, r1, r2 and r3.

[0065] (3) fp16 normal layout (16-bit floating-point precision normal layout): The precision type of each element in the matrix is ​​fp16, two consecutive elements are held by the same thread, and the register order is Z-shaped. For example Figure 4A certain thread t0 holds 8 matrix elements, corresponding to 4 registers r0, r1, r2 and r3.

[0066] (4) fp32 interleave layout: The precision type of each element in the matrix is ​​fp32, two consecutive elements are held by the same thread, and the register order is N-shaped; such as Figure 5 A certain thread t0 holds 8 matrix elements, corresponding to 8 registers: r0, r1, r2, r3, r4, r5, r6, and r7.

[0067] (5) fp32 normal layout (32-bit floating-point precision normal layout): Each element in the matrix has a precision type of fp32, one element is held by one thread, and the registers are arranged in a Z-shape; for example Figure 6 A certain thread t0 holds 8 matrix elements, corresponding to 8 registers: r0, r1, r2, r3, r4, r5, r6, and r7.

[0068] Register placement translation has limitations. Placement translation between registers depends on position mapping, such as... Figure 2 and Figure 3 Since the matrix elements held by the same thread are in the same position, it supports direct conversion from fp16 interleave layout to fp8 normal layout. Figure 4 and Figure 5 Since the matrix elements held by the same thread are in the same position, direct conversion from fp32 interleave layout to fp16 normal layout is supported. Additionally, by using shared memory, conversion from fp32 normal layout to fp16 normal layout, and from fp16 normal layout to fp8 normal layout, can be supported without interleave layout.

[0069] 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 an fp8 matrix element, sizeof(fp8) occupies 1 byte. For an fp16 matrix element, sizeof(fp16) occupies 2 bytes. For an fp32 matrix element, sizeof(fp32) occupies 4 bytes.

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

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

[0072] Control at least one first thread bundle to perform the following steps:

[0073] S11. Write the query matrix into blocks from shared memory to the first register group;

[0074] S12. Perform matrix multiplication using the query matrix block of the first register group and the key matrix block pre-copied into the shared memory to obtain a first product matrix block of the first precision type and write it into the second register group.

[0075] S13. Perform exponential operation using the first product matrix block of the second register group to obtain the numerator matrix block of the softmax operation and overwrite it into the second register group.

[0076] S14. Quantize the molecular matrix blocks of the second register group, and write the quantized molecular matrix blocks into the third register group according to the second precision type.

[0077] S15. Write the molecular matrix blocks of the third register group into the fourth register group according to the third precision type, and perform layout conversion on the molecular matrix blocks of the fourth register group through the shared memory.

[0078] S16. Perform softmax operation using the molecular matrix block after layout conversion of the fourth register group and the value matrix block pre-copied into the shared memory to obtain the attention result matrix block of the first precision type and write it into the fifth register group.

[0079] S17. The attention result matrix of the fifth register group is quantized into blocks, and the quantized attention result matrix blocks are written into the third register group according to the second precision type.

[0080] S18. Divide the attention result matrix of the third register group into blocks according to the third precision type and write it into the shared memory through the fourth register group; wherein, the precision of the first precision type is greater than the precision of the second precision type, and the precision of the second precision type is greater than the precision of the third precision type.

[0081] It is worth noting that the embodiments of this application are applied to the computing unit, which controls at least one first thread bundle to execute steps S11 to S18 to complete the calculation of the attention mechanism.

[0082] Specifically, when performing matrix multiplication of the query matrix block and the key matrix block, a matrix multiplication-addition instruction, such as the mma instruction, is used. 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 query matrix block is written from shared memory to the first register group regQ, and then multiplied with the key matrix block that has been pre-copied into shared memory to obtain the first product matrix block, which is then written to the second register group regS. Here, during the matrix multiplication operation, the query matrix block is multiplied by the transpose of the key matrix block.

[0083] Furthermore, to prevent intermediate errors from being amplified by low precision and thus accumulating into large-scale errors, this embodiment specifies the output precision type of the matrix multiplication-addition instruction, enabling it to output high-precision results. Specifically, the output precision type of the matrix multiplication-addition instruction is set to the first precision type. Assuming the input query matrix block, key matrix block, and value matrix block use the third precision type, the precision type of the output first product matrix block is also the first precision type. In other words, each element in the first product matrix block is of the first precision type. This configuration of low-precision input and high-precision output not only leverages the efficiency of low precision to improve computational speed and reduce memory usage, but also preserves key computational details through the high-fidelity characteristics of high precision, ultimately improving the accuracy of the output results. For example, if the first precision type is fp32, then further, in order to convert the fp32 interleave layout to the fp16 normal layout by relying solely on registers, the first product matrix block of the second register group regS is stored with the first precision type and the first preset layout. Specifically, the first precision type is fp32 and the first preset layout is interleavelayout, that is, the layout of the first product matrix block of the second register group regS is fp32 interleave layout.

[0084] Next, the first product matrix block of the second register group regS is used to perform exponential operation, and the numerator matrix block of the softmax operation is obtained and overwritten into the second register group regS, thus realizing the reuse of the second register group regS.

[0085] When the numerator matrix block in the second register group regS undergoes softmax operation, it needs to be multiplied and added with the input value matrix block. Therefore, the numerator matrix block needs to maintain the same precision type as the value matrix block. Thus, the precision type of the numerator matrix block needs to be converted from the first precision type to the third precision type. However, since the register does not support direct conversion from the first precision type to the third precision type, it needs to be converted from the first precision type to the second precision type, and then to the third precision type. Specifically:

[0086] First, the molecular matrix of the second register group regS is divided into blocks and quantized. Here, in order to ensure that subsequent low-precision data can be written to the register, the molecular matrix is ​​first divided into blocks and quantized so that its numerical representation is adapted to the subsequent low-precision data representation, thus ensuring the feasibility of low-precision writing.

[0087] Next, the quantized molecular matrix blocks are written to the third register group regTrans according to the second precision type. For example, if the second precision type is fp16, then further, the quantized molecular matrix blocks are written to the third register group regTrans according to the second precision type and the second preset layout, where the second preset layout is normal layout. Here, it is assumed that the layout of the molecular matrix blocks in the second register group regS is fp32 interleave layout. In this way, the conversion from fp32 interleave layout to fp16 normal layout can be achieved using only registers, enabling the molecular matrix blocks to be converted from fp32 interleave layout to fp16 normal layout and written to the third register group regTrans.

[0088] Then, the numerator matrix blocks from the third register group regTrans are written to the fourth register group regP according to the third precision type, and the layout of the numerator matrix blocks in the fourth register group regP is transformed through shared memory. Here, since direct layout transformation between registers is not supported, such as the conversion from fp16 normal layout to fp8 normal layout, the register layout transformation is achieved by using shared memory, so that the numerator matrix blocks are transformed from fp16 normal layout to fp8 normal layout and written to the fourth register group regP.

[0089] After layout transformation, a softmax operation is performed on the transformed molecular matrix blocks in register group regP of the fourth register group to obtain attention result matrix blocks of the first precision type, which are then written to register group regO of the fifth register group. Here, to avoid intermediate errors being amplified due to low precision characteristics, resulting in large-scale cumulative errors, the attention result matrix blocks obtained by the softmax operation are of high precision type, which improves the accuracy of the attention output results. For example, for the i-th query matrix block, steps S11~S16 are executed with it, along with the key matrix block and the value matrix block. After each execution of step S16, the latest attention result matrix block for that query matrix block is obtained, and the latest attention result matrix block is written to register group regO of the fifth register group, ensuring that register group regO stores only the latest attention result matrix block for that query matrix block. Furthermore, to enable subsequent conversion from fp32 interleave layout to fp16 normal layout using only registers, the layout of the attention result matrix blocks in register group regO of the fifth register group is an fp32 interleave layout.

[0090] Finally, to output the molecular matrix blocks in third-precision type, the attention result matrix blocks in the fifth register group regO are first quantized to ensure the feasibility of low-precision writing. Then, the quantized attention result matrix blocks are written to the third register group regTrans in second-precision type. This allows the attention result matrix blocks to be converted from fp32 interleave layout to fp16 normal layout and written to the third register group regTrans using only registers. Then, the attention result matrix blocks in the third register group regTrans are written to shared memory in third-precision type via the fourth register group regP. Here, to effectively reduce memory access latency, for the i-th query matrix block, its final attention result matrix block can be obtained, and then this final attention result matrix block can be converted to higher precision before being written to shared memory.

[0091] In this embodiment of the application, register resources are reused multiple times during precision conversion, which can mobilize as many hardware resources as possible to participate in the calculation, fully release the potential of hardware computing power, thereby avoiding the risk of performance drop caused by register overflow, ensuring the continuity and stability of the attention mechanism calculation process, and achieving high computing performance on the basis of efficient resource utilization.

[0092] In an optional embodiment, the step of writing the molecular matrix blocks of the third register group into the fourth register group according to the third precision type, and performing layout transformation on the molecular matrix blocks of the fourth register group through the shared memory, includes:

[0093] The molecular matrix of the third register group is divided into blocks and written into the fourth register group according to the third precision type;

[0094] The molecular matrix of the fourth register group is written into the shared memory in blocks;

[0095] The molecular matrix of the shared memory is written into the fourth register group in blocks.

[0096] It is worth noting that direct layout conversion between registers is not supported here. Therefore, this embodiment of the application uses shared memory to achieve register layout conversion. Specifically, firstly, the molecular matrix blocks of the third register group regTrans are written to the fourth register group regP according to the third precision type. Here, only the low precision type is written. Next, the molecular matrix blocks of the fourth register group regP are written to the shared memory. At the same time, layout conversion is performed. That is, the molecular matrix blocks in the shared memory are layout converted while writing the molecular matrix blocks to the shared memory. Finally, the molecular matrix blocks in the shared memory are written back to the fourth register group regP for subsequent softmax operations.

[0097] In an optional embodiment, the quantization processing of the molecular matrix blocks of the second register group includes:

[0098] The attention result matrix of the second register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type.

[0099] The step of quantizing the attention result matrix of the fifth register group in blocks includes:

[0100] The attention result matrix of the fifth register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type.

[0101] In order to ensure the feasibility of subsequent low-precision writing, the embodiments of this application map the corresponding matrix blocks from the numerical representation of the first precision type to the numerical representation of the third precision type to realize the quantization processing of the matrix blocks. Compared with first mapping the numerical representation of the first precision type to the numerical representation of the second precision type, and then mapping the numerical representation of the second precision type to the numerical representation of the third precision type, one quantization process is reduced, which is more convenient.

[0102] In one optional embodiment, the size of the query matrix block is the attention head dimension. The sequence length of the query matrix block; the size of the key matrix block is the attention head dimension. The sequence length of the key matrix block; the size of the value matrix block is the attention head dimension. The sequence length of the value matrix block; the sequence length of the key matrix block is equal to the sequence length of the value matrix block;

[0103] The dimensions of the query matrix block, the key matrix block, and the value matrix block are determined through the following steps:

[0104] Based on the first preset relationship, the first objective function is established as min(S) max -S1-S2); where S max The maximum number of registers used by each thread within the first thread bundle is limited, S1 is the preset number of reserved registers, and S2 is the actual number of registers that each thread within the first thread bundle needs to use when the attention mechanism is executed; the first preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block, and the actual number of registers;

[0105] Based on the second presupposed relationship, the second objective function is established as min(M) max -M); where M max M represents the maximum capacity of the shared memory, where M is the actual shared memory capacity required for all parallel execution of the attention mechanism by the first thread bundle in one computing unit; the second preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block, and the actual shared memory capacity.

[0106] Given a fixed attention head dimension, the first objective function and the second objective function are jointly solved to obtain the sequence length of the query matrix block, the sequence length of the key matrix block, and the sequence length of the value matrix block.

[0107] It is worth noting that if the first thread group directly processes the query matrix, key matrix, and 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. Furthermore, the storage of intermediate results may exceed limits, leading to computational interruption. Therefore, this embodiment utilizes block-based calculation of the query matrix, key matrix, and value matrix for attention mechanism computation, allowing each thread to focus on the local computation of its independent block, achieving efficient and stable execution of the attention mechanism under hardware resource constraints.

[0108] The size of the matrix block above is determined by the attention head dimension and the sequence length. Specifically, to query the size of the matrix block: attention head dimension d Sequence length m; Key matrix block size: Attention head dimension d Sequence length n; Size of value matrix blocks: Attention head dimension d Sequence length n.

[0109] Due to limitations of 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 the thread bundle is limited to use; for example, if the thread bundle is specifically the first thread bundle, then the actual number of registers S2 that each thread within the first thread bundle needs to use when executing the attention mechanism should be less than or equal to the maximum number of registers S that each thread within the first thread bundle is limited to use. max In addition, S1 registers need to be reserved for each thread bundle for auxiliary operations by the compiler to ensure that the thread bundles complete the task efficiently and stably. Optionally, S1=10.

[0110] (2) Multiple parallel thread bundles controlled by the same computing unit share one shared memory. This requires that the capacity used by the shared memory (e.g., the actual shared memory capacity M required for all parallel first thread bundles in a computing unit to execute the attention mechanism) must not exceed the maximum capacity M of the shared memory. max Restrictions.

[0111] Based on this, the embodiments of this application construct and solve the first objective function min(S) max -S1-S2) and the second objective function min(M max -M) is used to obtain appropriate sizes for the query matrix, key matrix, and value matrix blocks, ensuring that when using matrix blocks for attention mechanism operations, hardware resources are fully utilized while avoiding register overflows and shared memory contention. Furthermore, during joint solving, the values ​​of the first objective function and the second objective function are constrained to be greater than or equal to 0, thus guaranteeing the feasibility of the solution.

[0112] In this embodiment, the attention head dimension headdim can be preset and set according to actual needs. For example, the attention head dimension can be determined according to the task complexity. For complex tasks (such as long text understanding and multimodal reasoning), a higher attention head dimension (e.g., 96 or 128) is used to capture finer-grained features; for simple tasks (such as short text classification), the attention head dimension can be reduced (e.g., 32) to reduce redundancy.

[0113] In an optional embodiment, the second preset relationship further includes: the relationship between the first quantity, the second quantity, the third quantity, and the actual shared memory capacity;

[0114] Wherein, the first quantity is the number of key matrix blocks that can be copied into the shared memory simultaneously, the second quantity is the number of value matrix blocks that can be copied into the shared memory simultaneously, and the third quantity is the number of parallel groups of the first thread bundle in one computing unit, and the first quantity and the second quantity are equal.

[0115] It is worth noting that the actual shared memory capacity required for all parallel first thread bundles executing the attention mechanism in a computing unit is also related to the number of value matrix blocks that can be copied into the shared memory simultaneously (i.e., the first number), the number of value matrix blocks that can be copied into the shared memory simultaneously (i.e., the second number), and the number of parallel groups of the first thread bundles in a computing unit (i.e., the third number N). This application embodiment does not impose specific limitations on the first, second, and third numbers; they can be set according to actual needs, and only need to be positive integers.

[0116] The more first and second threads there are, the more shared memory they occupy, and the more data they can provide for the first thread bundle, thus preventing the first thread bundle from becoming idle while waiting for data.

[0117] The third quantity N is closely related to the number of query matrix blocks in shared memory. A first thread bundle can execute the attention mechanism calculation for one query matrix block. After the attention mechanism calculation for that query matrix block is completed, the next query matrix block's attention mechanism calculation begins. Therefore, N parallel first thread bundles in one computation unit can simultaneously execute the attention mechanism calculation for N query matrix blocks. To avoid idle first thread bundles, the number of query matrix blocks that can be stored simultaneously in shared memory must be the same as the third quantity N. Therefore, the third quantity N also affects the usage of shared memory.

[0118] Optionally, if one computational unit has two first thread bundles, then the third quantity N=2. In this case, to reduce the waiting time of the first thread bundles, the quantity can be incremented sequentially starting from the first quantity = second quantity = 2, solving the first objective function and the second objective function to obtain the size of the query matrix block, the key matrix block, and the value matrix block.

[0119] Furthermore, considering that if the first objective function is not satisfied, for example, if the maximum number of registers that can be used per thread is exceeded, the compiler will help buffer the overflow to the thread's local memory. In this case, the calculation result is generally correct, only leading to a performance decrease. However, if the second objective function is not satisfied, data overwriting or allocation out-of-bounds situations may occur, leading to incorrect calculation results. Therefore, the priority of the second objective function is set higher than that of the first objective function, so that the joint solution follows the principle of satisfying the second objective function first, and then satisfying the first objective function.

[0120] Moreover, with the second objective function min(M) max When -M) is the primary optimization objective, and multiple segmented scenarios satisfy both the first and second objective functions, the scenario with minimum (M) will be prioritized. max -M) is smaller, which is also the case where M is larger, so that the blocks are divided into blocks to improve the utilization of shared memory.

[0121] For example, by jointly solving the first objective function and the second objective function, we obtain two first and second objective functions that meet the requirements for both partitioned scenarios: Partitioned Scenario 1: Size of the query matrix partition: d1 m1; Size of the key matrix block: d1 n1; Size of the value matrix block: d1 n1. Block partitioning scenario 2: Size of the query matrix block: d2 m2; Size of the key matrix block: d2 n2; Size of the value matrix block: d2 n2. If the M of block scenario 1 is greater than the M of block scenario 2, then block scenario 1 is selected for matrix partitioning, and the attention mechanism is executed.

[0122] Specifically, the second preset relationship is: M = the total shared memory capacity occupied by each matrix block; these matrix blocks include at least: query matrix block, key matrix block, value matrix block, and numerator matrix block;

[0123] The shared memory capacity occupied by the query matrix block is N. m d Query the number of bytes occupied by the precision type of the matrix block;

[0124] Shared memory capacity occupied by key matrix blocks = first quantity n d The number of bytes occupied by the precision type of the key matrix block;

[0125] Shared memory capacity occupied by value matrix blocks = second quantity n d The number of bytes occupied by the precision type of the value matrix block;

[0126] Shared memory capacity occupied by molecular matrix blocks = N m d The number of bytes occupied by the precision type of the value matrix block;

[0127] The shared memory capacity occupied by the attention result matrix blocks = N m d The number of bytes occupied by the precision type of the attention result matrix block.

[0128] Specifically, the first preset relationship is: S2 = the total number of registers in each register group; these register groups include at least: the first register group regQ, the second register group regS, the third register group regP, and the fourth register group regO;

[0129] The number of registers in the first register group regQ is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. m d / (Number of thread bundles in the first thread bundle group) (Matrix size held by one thread bundle).

[0130] The number of registers in the second register group regS is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. m n / (Number of thread bundles in the first thread bundle group) (Matrix size held by one thread bundle).

[0131] The number of registers in the third register group regTrans = max(the number of registers in the second register group regS / 2, the number of registers in the fifth register group regO / 2).

[0132] The number of registers in the fourth register group regP = max(the number of registers in the second register group regS / 4, the number of registers in the fourth register group regO / 4).

[0133] The number of registers in register group regO of the fifth register group is equal to the number of registers required for one thread to hold the matrix elements stored in this register group. m d / (Number of thread bundles in the first thread bundle group) (The size of the matrix held by one thread bundle).

[0134] In one optional embodiment, the shared memory includes a first storage area and a second storage area; wherein the first storage area is used to store at least one of the query matrix block, the value matrix block, the key matrix block, and the attention result matrix block; and the second storage area is used to store at least one of the query matrix block, the value matrix block, the key matrix block, and the attention result matrix block.

[0135] It is worth noting that the shared memory provided in this embodiment includes a first storage bank region and a second storage bank region. These two regions are used to store various matrix blocks, allowing simultaneous access to data in both regions and improving memory access efficiency. Specifically, the first and second storage bank regions are divided according to their addresses. For example, the first storage bank region is a high-address storage bank region (high bank), i.e., the high-address region in the shared memory; and the second storage bank region is a low-address storage bank region (low bank), i.e., the low-address region in the shared memory. Of course, other methods can also be used for division, and no specific limitations are made here.

[0136] In an optional embodiment, the step of performing exponential operations on the first product matrix block of the second register group to obtain the numerator matrix block of the softmax operation and overwriting it into the second register group includes:

[0137] Calculate the maximum row value of the first product matrix block in the second register group and write it to the sixth register group;

[0138] Write the historical global row maximum value that is in the same row as the first product matrix block of the second register group into the seventh register group;

[0139] Compare the maximum row value of the sixth register group with the historical global row maximum value of the seventh register group, and write the larger value as the current global row maximum value into the eighth register group;

[0140] The first product matrix block of the second register group and the current global row maximum value of the eighth register group are subjected to exponential operation to obtain the numerator matrix block, which is then overwritten into the second register group.

[0141] It is worth noting that 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 the j-th attention score matrix block in each row; where each row corresponds to a maximum value.

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

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

[0144] In this embodiment, the maximum row value of the first product matrix block, chunk_rmax, the historical global row value accum_rmax that is in the same row as the first product matrix block, and the larger of the maximum row value chunk_rmax and the historical global row value accum_rmax are used as the current global row value accum_rmax_new. These values ​​are used to perform exponential operations with the first product matrix block to obtain the numerator matrix block of the softmax operation and overwrite it into the second register group regS.

[0145] Of course, these values ​​also need to be stored in registers, so the maximum row value of the first product matrix block chunk_rmax, the historical global maximum row value accum_rmax, and the current global maximum row value accum_rmax_new are stored in different register groups: the sixth register group, the seventh register group, and the eighth register group.

[0146] Therefore, in the first preset relationship S2 = the total number of registers in each register group, these register groups also include: the sixth register group, the seventh register group, and the eighth register group. For example... Figures 2-6 Regardless of whether the matrix is ​​of type fp8, fp16, or fp32, the matrix elements held by one thread are distributed across two rows. Each row needs to calculate and store the maximum value of the first product matrix block, chunk_rmax, the historical global maximum value of the row, accum_rmax, and the current global maximum value of the row, accum_rmax_new. Therefore, the required sixth, seventh, and eighth register groups each include two registers.

[0147] In an optional embodiment, the step of performing a softmax operation on the molecular matrix blocks after layout transformation using the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain attention result matrix blocks of the first precision type and writing them into the fifth register group includes:

[0148] Calculate the row accumulation value of the first product matrix block in the second register group and write it to the ninth register group;

[0149] Write the accumulated value of the historical row that is in the same row as the first product matrix block of the second register group into the tenth register group;

[0150] The current historical global row maximum value of the eighth register group and the historical global row maximum value of the seventh register group are used to perform an exponential operation to obtain a correction value, which is then written into the eleventh register group. The correction value is used to correct the attention result matrix block obtained by the softmax operation performed on the query matrix block in the past.

[0151] The current row accumulation value is obtained by performing scaling operations using the correction value of the eleventh register group, the historical row accumulation value of the tenth register group, and the row accumulation value of the ninth register group, and then writing it into the twelfth register group.

[0152] The correction value of the eleventh register group is used to correct the current multiply-accumulate matrix block of the fifth register group, and then overwrites the fifth register group.

[0153] Matrix multiplication and addition operations are performed using the molecular matrix block after layout transformation of the fourth register group, the value matrix block pre-copied into the shared memory, and the multiply-add matrix block after correction of the fifth register group to obtain the multiply-add matrix block and overwrite it into the fifth register group.

[0154] The attention result matrix block is obtained by performing a division operation between the multiply-accumulate matrix block of the fifth register group and the current row accumulated value of the twelfth register group, and then overwrites and writes it into the fifth register group.

[0155] It is worth noting that the row accumulation value chunk_rsum is specifically the accumulation value of each row of the attention score matrix block, that is, the accumulation value of the j-th attention score matrix block in each row; where each row corresponds to one accumulation value.

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

[0157] The modifier is used to modify the historical row accumulation value accum_rsum to obtain the current row accumulation value accum_rsum_new.

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

[0159] In this embodiment, the row accumulation value chunk_rsum of the first product matrix block, the historical row accumulation value accum_rsum in the same row as the first product matrix block, the current row accumulation value accum_rsum_new, and the modification value Modifier are calculated. These values ​​are then used to perform a softmax operation with the second product matrix block to obtain the attention result matrix block, which is then written to the fifth register group regO.

[0160] Of course, these values ​​also need to be stored in registers, so the row accumulation values ​​of the first product matrix, chunk_rsum, historical row accumulation values, accum_rsum, current row accumulation values, and modifier are stored in different register groups: the ninth register group, the tenth register group, the eleventh register group, and the twelfth register group.

[0161] Therefore, in the first preset relationship S2 = the total number of registers in each register group, these register groups also include: the ninth register group, the tenth register group, the eleventh register group, and the twelfth register group. Similarly, the matrix elements held by one thread are distributed across two rows, and each row needs to calculate and store the row accumulation value chunk_rsum of the first product matrix block, the historical row accumulation value accum_rsum, the current row accumulation value accum_rsum_new, and the modifier. Therefore, the required ninth register group, tenth register group, eleventh register group, and twelfth register group each include two registers.

[0162] In an optional embodiment, the attention mechanism calculation method further includes:

[0163] A second thread group is controlled to copy the query matrix block, the key matrix block, and the value matrix block from global memory to the shared memory.

[0164] It is worth noting that this application provides two types of thread bundles: a first thread bundle and a second thread bundle. The first thread bundle, as a consumer thread bundle, mainly performs the attention mechanism calculation process, i.e., steps S11 to S16; the second thread bundle, as a producer thread bundle, mainly undertakes the cross-memory migration task of matrix blocks. Specifically, it controls a second thread bundle to copy the query matrix block, key matrix block, and value matrix block from global memory (GLM) to shared memory (GSM).

[0165] One computing unit can control one second thread group and at least two first thread groups to execute in parallel, creating a pipeline overlap effect between "data block migration" and "attention calculation process": while the second thread group supplies the data required for computation, the first thread group can simultaneously perform attention mechanism calculations. By parallelizing storage access and computational resource execution, the first thread group is prevented from being idle while waiting for data, significantly improving overall execution efficiency. Specifically, each first thread group executes steps S11~S18 for different query matrix blocks. For example, one first thread group executes query matrix block 1 to obtain an attention result matrix block using query matrix block 1, while another first thread group executes query matrix block 2 to obtain an attention result matrix block using query matrix block 2.

[0166] Further, the step of controlling a second thread group to copy the query matrix block, the key matrix block, and the value matrix block from global memory to the shared memory includes:

[0167] Control the first thread bundle to perform the following steps:

[0168] Create one stage to copy the query matrix blocks from the global memory to the shared memory; each stage executes at least one copying process of the query matrix blocks;

[0169] Create X1 stages to copy the key matrix blocks from the global memory to the shared memory; each stage performs one copy of the key matrix block, and the value of X1 is the same as the first quantity mentioned above;

[0170] Create X2 stages to copy the value matrix blocks from the global memory to the shared memory; each stage performs one copy of the value matrix block, and the value of X2 is the same as the second number mentioned above.

[0171] In the embodiments of this application, each stage can perform a copy of a matrix block. For key matrix blocks and value matrix blocks, multiple stages can be copied simultaneously, which can provide the data required for calculation for multiple first thread bundles, avoid the first thread bundles from being idle due to waiting for data, and significantly improve the overall execution efficiency.

[0172] In an optional embodiment, the attention mechanism calculation method further includes:

[0173] The position coordinates of the attention result matrix blocks in the shared memory are subjected to exponential and logarithmic operations to obtain exponential and logarithmic matrix blocks, which are then written into the shared memory.

[0174] In this embodiment, the first thread group can also be controlled to perform exponential and logarithmic operations (Log-Sum-Exp, lse) on the position coordinates of the first product matrix block in the second register group regS, obtaining an exponential and logarithmic matrix block and writing it to shared memory to provide data for subsequent use of the lse matrix in global memory. Alternatively, the obtained exponential and logarithmic matrix block can be written directly to global memory without going through shared memory. The exponential and logarithmic operation refers to summing the exponents of a set of values ​​and then taking the logarithm.

[0175] In a specific embodiment, for a certain artificial intelligence chip, its hardware resources are as follows: (1) One thread block includes one second thread bundle and two first thread bundles; each thread in each second thread bundle is limited to using a maximum of 32 registers, and each thread in each first thread bundle is limited to using a maximum of 240 registers; (2) The maximum capacity of shared memory is 228KB, the first storage area is 96KB, and the second storage area is 128KB. For this artificial intelligence chip, the attention head dimension is set to 128, the first objective function and the second objective function are solved, and the sequence length of the query matrix block is 64, the sequence length of the key matrix block is 128, and the sequence length of the value matrix block is 128. That is, the query matrix, key matrix and value matrix are divided according to the following dimensions: the size of the query matrix block: 64 128; Key matrix block size: 128 128; Size of the value matrix block: 128 128. Based on the above matrix partitioning, execute steps S11~S18, specifically according to the following instructions:

[0176] Initialize chunk_rmax, accum_rmax, accum_rmax_new to -inf;

[0177] Init chunk_rsum,accum_rsum,accum_rsum_new to 0 / / Initialization;

[0178] regO init as 0 / / Initialization;

[0179] Loop{ / Q outer loop /

[0180] ldmatrix regQ / / Writes the query matrix in chunks from shared memory to the first register set;

[0181] Loop { / KV internal circulation /

[0182] regS = mma(regQ, smemK, 0) / / Perform matrix multiplication on the query matrix block of the first register group and the key matrix block that has been copied into shared memory beforehand, and write the first product matrix block of the first precision type into the second register group;

[0183] calc chunk_rmax, accum_rmax, accum_rmax_new / / Calculate the maximum row value of the first product matrix block and write it to the sixth register group; write the historical global row value that is in the same row as the first product matrix block to the seventh register group; compare the maximum row value of the sixth register group with the historical global row value of the seventh register group, and write the larger value as the current global row value to the eighth register group;

[0184] exp((regS-accum_rmax_new) sf), reuse regS / / Perform exponential operation on the first product matrix block of the second register group and the current global row maximum value of the eighth register group to obtain the numerator matrix block of the softmax operation and overwrite it into the second register group; sf is the preset scaling factor;

[0185] calc chunk_rsum, accum_rsum / / Calculate the row accumulation value of the first product matrix block and write it to the ninth register group; write the historical row accumulation value that is in the same row as the first product matrix block to the tenth register group;

[0186] Modifier=exp((accum_rmax–accum_rmax_new) sf ) / / Perform an exponential operation using the current global row maximum value of the eighth register group and the historical global row maximum value of the seventh register group to obtain the correction value, which is then written to the eleventh register group;

[0187] accum_rsum_new= accum_rsum Modifier + chunk_rsum / / Perform scaling operations using the correction value of the eleventh register group, the historical row accumulation value of the tenth register group, and the row accumulation value of the ninth register group to obtain the current row accumulation value and write it to the twelfth register group;

[0188] Quantize regS / / Performs block quantization on the numerator matrix of the second register group;

[0189] Typecast regS to regTrans / / Divides the quantized molecular matrix into blocks and writes them to the third register group according to the second precision type;

[0190] Typecast regTrans to regP / / Writes the molecular matrix of the third register group into blocks according to the third precision type to the fourth register group;

[0191] stmatrix regP / / Writes the molecular matrix of the fourth register group into blocks and into shared memory;

[0192] ldmatrix regP / / Writes the molecular matrix blocks from shared memory to the fourth register group, so as to perform layout transformation on the molecular matrix blocks of the fourth register group through shared memory;

[0193] regO = modifier regO / / Use the correction value of the eleventh register group to correct the current multiply-accumulate matrix block of the fifth register group and overwrite it to the fifth register group;

[0194] regO = mma(regP, smemV, regO) / / Perform matrix multiplication and addition operations using the numerator matrix block after layout transformation of the fourth register group, the value matrix block pre-copied into shared memory, and the multiply-add matrix block corrected by the fifth register group, to obtain the multiply-add matrix block and overwrite it into the fifth register group.

[0195] regO = regO / accum_rsum_new / / Perform a division operation between the multiply-accumulate matrix block of the fifth register group and the current row accumulated value of the twelfth register group to obtain the attention result matrix block and overwrite it into the fifth register group;

[0196] }

[0197] Quantize regO / / Quantizes the attention result matrix of the fifth register group in blocks;

[0198] Typecast regO to regTrans / / Divide the quantized attention result matrix into blocks and write them into the third register group according to the second precision type;

[0199] Typecast regTrans to regP / / Writes the attention result matrix of the third register group into blocks according to the third precision type to the fourth register group;

[0200] Stmatrix regP / / Writes the attention result matrix of the fourth register group into blocks and into shared memory;

[0201] Gsm[coord]=lse / / Perform exponential and logarithmic operations on the position coordinates of the blocks of the attention result matrix in shared memory, and write the exponential and logarithmic matrices into shared memory in blocks;

[0202] Tma(cpb) lse / / Writes the exponent and logarithm matrices from shared memory to global memory in blocks;

[0203] Tma O / / Writes the attention result matrix from shared memory to global memory in chunks;

[0204] }

[0205] The register resources used in this scenario are shown in Table 1.

[0206] Table 1

[0207]

[0208] The shared memory resources used in this scenario are shown in Table 2, and their specific distribution is as follows: Figure 7 The query matrix block (Q block), value matrix block (V block), numerator matrix block (Trans block), and exponent and logarithm matrix block (lse block) are stored in the second storage area; the key matrix block (K block) and attention result matrix block (O block) are stored in the first storage area.

[0209] Table 2

[0210]

[0211] The attention mechanism provided in this application executes intermediate operations in a first-precision manner, avoiding the use of a single low-precision method. This reduces the cumulative error of intermediate operation results and ultimately improves the accuracy of the attention output. Furthermore, it precisely adapts to the register and shared memory resource characteristics of the AI ​​chip, without exceeding usage limits. In particular, the shared memory is close to its maximum capacity. While avoiding register overflow and shared memory contention, it mobilizes as many hardware resources as possible for computation, fully releasing the hardware's computing potential and achieving high computational performance based on efficient resource utilization.

[0212] See Figure 8This is a structural block diagram of an artificial intelligence chip provided in an embodiment of this application. The artificial intelligence chip includes a global memory (GLM) and one or more computing units; each computing unit includes shared memory (GSM) and one or more thread bundles. The thread bundles in each computing unit can be divided into multiple consumer thread bundles (such as consumer thread bundle 1, consumer thread bundle 2) and one producer thread bundle. Each consumer thread bundle can execute the attention mechanism calculation method described in any of the above embodiments. In the register space, regQ is the first register group, regS is the second register group, regTrans is the third register group, regP is the fourth register group, and regO is the fifth register group. The intermediate result register groups of Softmax include: the sixth register group, the seventh register group, the eighth register group, the ninth register group, the tenth register group, the eleventh register group, and the twelfth register group.

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

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

[0215] See Figure 9 , Figure 9 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.

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

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

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

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

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

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

[0222] 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 at least one first thread bundle to perform the following steps: Write the query matrix in blocks from shared memory into the first register group; The query matrix block of the first register group is multiplied with the key matrix block that has been copied into the shared memory in advance to obtain a first product matrix block of the first precision type and write it into the second register group. The first product matrix block of the second register group is used to perform exponential operation to obtain the numerator matrix block of the softmax operation and overwrite it into the second register group. The molecular matrix blocks of the second register group are quantized, and the quantized molecular matrix blocks are written into the third register group according to the second precision type. The molecular matrix blocks of the third register group are written into the fourth register group according to the third precision type, and the layout transformation of the molecular matrix blocks of the fourth register group is performed through the shared memory. The softmax operation is performed using the molecular matrix blocks after layout transformation of the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain the attention result matrix blocks of the first precision type and write them into the fifth register group; The attention result matrix of the fifth register group is quantized in blocks, and the quantized attention result matrix blocks are written into the third register group according to the second precision type. The attention result matrix of the third register group is divided into blocks according to the third precision type and written to the shared memory through the fourth register group; wherein the precision of the first precision type is greater than the precision of the second precision type, and the precision of the second precision type is greater than the precision of the third precision type.

2. The attention mechanism calculation method as described in claim 1, characterized in that, The step of writing the molecular matrix blocks of the third register group into the fourth register group according to the third precision type, and performing layout transformation on the molecular matrix blocks of the fourth register group through the shared memory, includes: The molecular matrix of the third register group is divided into blocks and written into the fourth register group according to the third precision type; The molecular matrix of the fourth register group is written into the shared memory in blocks; The molecular matrix of the shared memory is written into the fourth register group in blocks.

3. The attention mechanism calculation method as described in claim 1, characterized in that, The quantization process of dividing the molecular matrix of the second register group into blocks includes: The attention result matrix of the second register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type. The step of quantizing the attention result matrix of the fifth register group in blocks includes: The attention result matrix of the fifth register group is divided into blocks and mapped from a numerical representation of the first precision type to a numerical representation of the third precision type.

4. The attention mechanism calculation method as described in claim 1, characterized in that, The step of performing exponential operations using the first product matrix block of the second register group to obtain the numerator matrix block of the softmax operation and overwriting it into the second register group includes: Calculate the maximum row value of the first product matrix block in the second register group and write it to the sixth register group; Write the historical global row maximum value that is in the same row as the first product matrix block of the second register group into the seventh register group; Compare the maximum row value of the sixth register group with the historical global row maximum value of the seventh register group, and write the larger value as the current global row maximum value into the eighth register group; The first product matrix block of the second register group and the current global row maximum value of the eighth register group are subjected to exponential operation to obtain the numerator matrix block, which is then overwritten into the second register group.

5. The attention mechanism calculation method as described in claim 4, characterized in that, The step of performing a softmax operation on the molecular matrix blocks after layout transformation using the fourth register group and the value matrix blocks pre-copied into the shared memory to obtain attention result matrix blocks of the first precision type and writing them into the fifth register group includes: Calculate the row accumulation value of the first product matrix block in the second register group and write it to the ninth register group; Write the accumulated value of the historical row that is in the same row as the first product matrix block of the second register group into the tenth register group; The current historical global row maximum value of the eighth register group and the historical global row maximum value of the seventh register group are used to perform an exponential operation to obtain a correction value, which is then written into the eleventh register group. The correction value is used to correct the attention result matrix block obtained by the softmax operation performed on the query matrix block in the past. The current row accumulation value is obtained by performing scaling operations using the correction value of the eleventh register group, the historical row accumulation value of the tenth register group, and the row accumulation value of the ninth register group, and then writing it into the twelfth register group. The correction value of the eleventh register group is used to correct the current multiply-accumulate matrix block of the fifth register group, and then overwrites the fifth register group. Matrix multiplication and addition operations are performed using the molecular matrix block after layout transformation of the fourth register group, the value matrix block pre-copied into the shared memory, and the multiply-add matrix block after correction of the fifth register group to obtain the multiply-add matrix block and overwrite it into the fifth register group. The attention result matrix block is obtained by performing a division operation between the multiply-accumulate matrix block of the fifth register group and the current row accumulated value of the twelfth register group, and then overwrites and writes it into the fifth register group.

6. The attention mechanism calculation method as described in claim 1, characterized in that, The size of the query matrix block is the attention head dimension. The sequence length of the query matrix block; The size of the key matrix block is the attention head dimension. The sequence length of the key matrix block; The size of the value matrix block is the attention head dimension. The sequence length of the value matrix block; The sequence length of the key matrix block is equal to the sequence length of the value matrix block; The dimensions of the query matrix block, the key matrix block, and the value matrix block are determined through the following steps: Based on the first preset relationship, the first objective function is established as min(S) max -S1-S2); where S max The maximum number of registers used by each thread in the first thread bundle is limited, S1 is the preset number of reserved registers, and S2 is the actual number of registers that each thread in the first thread bundle needs to use when the attention mechanism is executed; the first preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block and the actual number of registers; Based on the second presupposed relationship, the second objective function is established as min(M) max -M); where M max M represents the maximum capacity of the shared memory, where M is the actual shared memory capacity required for all parallel execution of the attention mechanism by the first thread bundle in one computing unit; the second preset relationship includes at least the relationship between the attention head dimension, the sequence length of the query matrix block, the sequence length of the key matrix block, the sequence length of the value matrix block, and the actual shared memory capacity. Given a fixed attention head dimension, the first objective function and the second objective function are jointly solved to obtain the sequence length of the query matrix block, the sequence length of the key matrix block, and the sequence length of the value matrix block.

7. The attention mechanism calculation method as described in claim 6, characterized in that, The second preset relationship also includes the relationship between the first quantity, the second quantity, the third quantity, and the actual shared memory capacity; Wherein, the first quantity is the number of key matrix blocks that can be copied into the shared memory simultaneously, the second quantity is the number of value matrix blocks that can be copied into the shared memory simultaneously, and the third quantity is the number of parallel groups of the first thread bundle in one computing unit, and the first quantity and the second quantity are equal.

8. 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 7.

9. 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 7.

10. 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 7.

Images