Attention computation implementation method, device, medium, equipment and product

Abstract

The application discloses an attention calculation implementation method and device, medium, equipment and product, the method comprises the following steps: loading the first register group of the consumer thread group from the shared memory to the current to be calculated i query block, the internal split calculation is carried out to the query block and key block, and the corresponding attention score block is obtained; M K K The application can efficiently utilize existing hardware resources to improve the calculation performance, and is especially suitable for large-size query block and key block scenarios.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to an attention calculation method, apparatus, computer-readable storage medium, electronic device, and computer program product. Background Technology

[0002] With the rapid development of artificial intelligence technology, the attention mechanism has become a core operator in many fields such as natural language processing and computer vision. It achieves dynamic weighted aggregation of input data by calculating the interaction relationship between queries, keys, and values. However, the computation process of the attention mechanism usually involves large-scale matrix operations and complex data dependencies, resulting in extensive resource consumption. However, due to the constraints of artificial intelligence processor hardware resources, it is difficult to achieve good performance. Summary of the Invention

[0003] The purpose of this invention is to provide an attention calculation implementation method, apparatus, computer-readable storage medium, electronic device, and computer program product. By using a consumer thread group, the query block and key block are internally split and calculated to obtain the corresponding attention score block. The attention output block is written back to shared memory by reusing a specified register group. This can efficiently utilize existing hardware resources to improve computing performance, and is especially suitable for scenarios with large query blocks and key blocks.

[0004] A first aspect of the present invention provides an attention calculation implementation method, executed by a consumer thread group, the method comprising:

[0005] Load the i-th query block to be computed from shared memory into the first register group of the consumer thread group; where i ≥ 1;

[0006] The loaded query blocks are logically divided along the header dimension to obtain N1 query sub-blocks. Then, based on matrix multiplication and addition instructions, each query sub-block is compared with its corresponding key block in the j-th key block to obtain the j-th attention score block corresponding to the i-th query block; where N1≥1; 1≤j≤M. K M K The number of blocks in the key matrix;

[0007] Obtain M sequentially K Each key block is divided into attention score blocks, and mixed-precision attention fusion calculation is performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block; wherein, the key blocks and the value blocks are cached in batches in shared memory;

[0008] After dividing the attention output block logic into N2 batches, the target precision type conversion is performed by reusing the specified register group in each batch, and the output result after conversion of each batch is written back to the shared memory; wherein, the specified register group is the register group used to store the target precision type data during the attention fusion calculation process; N2≥1.

[0009] Optionally, the method further includes:

[0010] After obtaining the attention output blocks, M K The exponent and logarithm results corresponding to each attention score block are written to shared memory or directly to global memory.

[0011] Optionally, the sizes of the query sub-block and the key sub-block are adapted to the matrix multiplication and addition instructions.

[0012] Optionally, the step of calculating the corresponding key block in the j-th key block for each query sub-block based on matrix multiplication and addition instructions to obtain the j-th attention score block corresponding to the i-th query block includes:

[0013] Initialize the second register group of the consumer thread group;

[0014] Transpose the j-th key block to obtain the key transpose block; where the key transpose block is composed of N col The columns are composed of key blocks; N col ≥1;

[0015] The matrix multiplication and addition instruction multiplies N1 query sub-blocks with all key sub-blocks in the k-th column of the key transpose block, and then sequentially accumulates the corresponding sub-block product results into the k-th region of the second register group; where 1≤k≤N col ;

[0016] After traversing and calculating all key sub-blocks in the key transpose block, the N obtained from the second register group... col The summation results are combined to form the j-th attention score block.

[0017] Optionally, the step of sequentially obtaining M K Each attention score block is divided into several parts, and mixed-precision attention fusion calculations are performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block, including:

[0018] Initialize the historical maximum value tensor, the historical accumulation tensor, and the historical output block; wherein, the historical output block is cached in the third register group of the consumer thread group;

[0019] Based on the iteration rounds, M is sequentially cached in the second register group of the consumer thread group.K The attention score blocks are used for online softmax iterative calculation;

[0020] After each iteration, update the historical maximum tensor, historical summation tensor, and historical output blocks;

[0021] When the iterative calculation ends, the last obtained current output block is used as the attention output block.

[0022] Optionally, the M cached sequentially in the second register group of the consumer thread group K The online softmax calculation is performed on each attention score block, including:

[0023] Perform the following operations on the j-th attention score block processed in the current iteration round:

[0024] By comparing the current block maximum tensor with the historical maximum tensor, the current cumulative maximum tensor is obtained, and the corresponding update correction factor is calculated.

[0025] Based on the cumulative maximum value tensor and the scaling factor, the j-th attention score block is exponentially operated on, and the calculated exponential score block is cached by reusing the second register group;

[0026] The indexed score blocks are summed by accumulating the elements along the row dimension to obtain the block index and tensor. The current accumulated tensor is then calculated by combining the update correction factor and the historical accumulated tensor.

[0027] The historical output blocks are corrected by updating the correction factor and the historical accumulation and tensor, and the calculated historical correction numerator is cached by reusing the third register group;

[0028] After writing the indexed fraction blocks into the fourth register group for target precision type conversion, multiplying them with the corresponding value blocks and accumulating them into the third register group, the current numerator is obtained.

[0029] Divide the current numerator by the current sum tensor and cache the current output block obtained by the calculation by reusing the third register group.

[0030] Optionally, the designated register group is either the first register group or the fourth register group.

[0031] Optionally, the sizes of the query block, the key block, and the value block are determined by the upper limit of available register space, the upper limit of available shared memory space, and the matrix multiplication and addition instructions.

[0032] Optionally, when the size of the query block is fixed, the steps for obtaining the sequence length of the key block are as follows:

[0033] The primary optimization objective is to minimize the difference between the upper limit of register availability for a single thread within the consumer thread group and the number of registers required for attention fusion computation.

[0034] Based on the number of query blocks computed in parallel and the number of key blocks and value blocks cached in shared memory in each batch, the memory usage capacity is determined, and the second optimization objective is to minimize the difference between the upper limit of the available space in shared memory and the memory usage capacity.

[0035] The sequence length of the key block is obtained by jointly solving the first optimization objective and the second optimization objective.

[0036] A second aspect of the present invention provides an attention calculation implementation apparatus, comprising:

[0037] The query data loading module is used to load the i-th query block to be computed from shared memory into the first register group of the consumer thread group; where i ≥ 1;

[0038] The attention score calculation module is used to logically partition the loaded query blocks along the header dimension, obtaining N1 query sub-blocks. Based on matrix multiplication and addition instructions, it calculates the attention score for each query sub-block and its corresponding key block in the j-th key block to obtain the j-th attention score block corresponding to the i-th query block; where N1≥1; 1≤j≤M K M K The number of blocks in the key matrix;

[0039] Attention output calculation module, used to sequentially obtain M K Each key block is divided into attention score blocks, and mixed-precision attention fusion calculation is performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block; wherein, the key blocks and the value blocks are cached in batches in shared memory;

[0040] The attention output unloading module is used to divide the attention output block logic into N2 batches, reuse the specified register group for target precision type conversion in each batch, and write the output result of each batch back to shared memory; wherein, the specified register group is the register group used to store target precision type data during the attention fusion calculation process; N2≥1.

[0041] A third aspect of the present invention provides a computer-readable storage medium comprising 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 calculation implementation method described in any of the first aspects of the present invention.

[0042] A fourth aspect of the present invention provides a computer program product, including computer instructions, which, when executed by a processor, implement the attention calculation method described in any of the first aspects of the present invention.

[0043] A fifth aspect of the present invention provides an electronic 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, when executing the computer program, implements the attention calculation implementation method described in any of the first aspects of the present invention.

[0044] Compared with existing technologies, embodiments of the present invention provide an attention calculation implementation method, apparatus, computer-readable storage medium, electronic device, and computer program product. The embodiments of the present invention first load the i-th query block to be calculated from shared memory into the first register group of the consumer thread group, and internally split and calculate the query block and key block to obtain the corresponding attention score block. This reduces the number of times the query block is copied between shared memory and registers, improving the calculation speed of the attention score, and is particularly suitable for scenarios with large query blocks and key blocks. Then, M is sequentially obtained... K The attention score blocks are divided into several blocks, and mixed-precision attention fusion calculations are performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block. Finally, when converting the attention output blocks into the target precision type in batches, the specified register group that stores the target precision type data during the fusion calculation is reused, which can efficiently utilize existing hardware resources. Attached Figure Description

[0045] Figure 1 This is a flowchart of an embodiment of thread block scheduling provided by the present invention;

[0046] Figure 2 This is a flowchart illustrating an embodiment of the attention calculation implementation method provided by the present invention;

[0047] Figure 3 This is a schematic diagram of an embodiment of the calculation matrix correlation size provided by the present invention;

[0048] Figure 4 This is a flowchart illustrating an embodiment of attention score block acquisition provided by the present invention;

[0049] Figure 5 This is a schematic diagram of an embodiment of data arrangement in a thread bundle provided by the present invention;

[0050] Figure 6 This is a schematic diagram of another embodiment of the data arrangement in the thread bundle provided by the present invention;

[0051] Figure 7 This is a schematic diagram of an embodiment of shared memory data storage provided by the present invention;

[0052] Figure 8 This is a schematic diagram of an embodiment of the attention calculation implementation device provided by the present invention;

[0053] Figure 9 This is a schematic diagram of the structure of an embodiment of the electronic device provided by the present invention;

[0054] Figure 10 This is a schematic diagram of the structure of an embodiment of the artificial intelligence processor provided by the present invention. Detailed Implementation

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

[0056] The artificial intelligence processor involved in this invention can be any one of CPU (Central Processing Unit), GPU (Graphics Processing Unit), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), DPU (Deep Learning Processing Unit), APU (Accelerated Processing Unit), and GPGPU (General-Purpose computing on Graphics Processing Unit), depending on its application to a specific product or technology in the embodiments of this invention.

[0057] This invention uses the execution of a thread block by a GPU's Streaming Multiprocessor (SM) as an example to illustrate the computational task relationships between consumer thread groups. This invention improves computational efficiency by assigning different roles or tasks to each thread group within a thread block through warp specialization scheduling.

[0058] See Figure 1 This is a flowchart of an embodiment of thread block scheduling provided by the present invention. Figure 1 In this code, a thread block consists of one producer warpgroup and two consumer warpgroups (Consumer Warpgroup 1 + Consumer Warpgroup 2). The producer warpgroup is responsible for copying data from Global Memory (GLM) to Group Shared Memory (GSM) to prepare data for subsequent computation. The two consumer warpgroups are responsible for executing attention computation methods (such as QK) on the loaded data within the GSM. T (Multiplication / softmax / PV multiplication and GSM read / write), to obtain the attention output block of the corresponding query block. It is worth noting that the embodiments of the present invention do not limit the number of consumer thread groups enabled, but determine it when the embodiments of the present invention are applied to specific products or technologies.

[0059] The following uses consumer warpgroup 1 as an example to illustrate the attention calculation implementation method, apparatus, computer-readable storage medium, electronic device, and computer program product provided in the embodiments of the present invention.

[0060] See Figure 2 This is a flowchart illustrating an embodiment of the attention calculation implementation method provided by the present invention.

[0061] A first aspect of the present invention provides a method for implementing attention calculation, executed by a consumer thread group, the method comprising steps S1 to S4, as follows:

[0062] Step S1: Load the i-th query block to be computed from shared memory into the first register group of the consumer thread group; where i ≥ 1;

[0063] Step S2: Logically partition the loaded query blocks according to the header dimension to obtain N1 query sub-blocks. Based on matrix multiplication and addition instructions, calculate the corresponding key sub-block in the j-th key block for each query sub-block to obtain the j-th attention score block corresponding to the ith query block; where N1≥1; 1≤j≤M K M K The number of blocks in the key matrix;

[0064] Step S3: Obtain M sequentially K Each key block is divided into attention score blocks, and mixed-precision attention fusion calculation is performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block; wherein, the key blocks and the value blocks are cached in batches in shared memory;

[0065] Step S4: After dividing the attention output block logic into N2 batches, the target precision type is converted by reusing the specified register group in each batch, and the output result after conversion of each batch is written back to the shared memory; wherein, the specified register group is the register group used to store the target precision type data during the attention fusion calculation process; N2≥1.

[0066] It should be noted that the query is divided into blocks. This involves partitioning the original query matrix into blocks according to preset dimensions (such as sequence length seqlen_q and head dimension headdim). Similarly, key blocks and value blocks are obtained by partitioning the original key matrix and original value matrix, respectively, and both have the same size. The first register group regQ is a register resource for the consumer thread group, used to temporarily store the query blocks to be computed. .

[0067] In step S2, the consumer thread group processes the loaded query blocks. Logical partitioning along the head dimension yields N1 query sub-blocks. This embodiment of the invention uses the mma (matrix multiply accumulate) instruction as the specified matrix multiplication and accumulation instruction. When the mma instruction is executed, it requires matrix A (e.g., query sub-blocks) to be stored in registers, and matrix B (e.g., key sub-blocks) to be stored in shared memory, effectively reducing register usage.

[0068] The granularity of each sub-block is determined based on the computational size of the mma instruction. For query and key sub-blocks, the computational size of the mma is m×n×K. B Among them, K B The number of bytes is m×n×K, ranging from 32B to 256B; when the matrix data type is FP16, the value is m×n×K. B =64×16×16k (k=1, 2,…,8), meaning the size of the query sub-block is m×K. B =64×16k, the size of the bond block is K B ×n=16k×16. Preferably, K is the number of sub-blocks to be defined in the query sub-block. B The dimension is set to the maximum value of 16k to speed up query chunking. The corresponding j-th attention score block The calculation. For example, when querying blocks. When the size of (FP16) is seqlen_q×headdim=64×256, it can be divided into two query sub-blocks of 64×128, and the corresponding key sub-block size is 128×16; when query blocks are used... When the size seqlen_q×headdim=64×192, it can be divided into two query sub-blocks of 64×96, and the corresponding key sub-block size is 96×16.

[0069] Based on the mma instruction, the product results between each query sub-block and the corresponding key sub-block are integrated to form the j-th attention score block. .

[0070] In step S3, M is traversed sequentially. K Each key is divided into blocks (i.e., the calculation process of step S2 is repeated) to obtain M. K Each attention score block is divided into several parts. After obtaining each attention score block, attention fusion calculation is performed with the corresponding value block (including performing an online softmax operation on the attention score block) to finally obtain the query block. Corresponding attention output blocks .

[0071] In this embodiment of the invention, key-based and value-based chunks are not loaded into shared memory all at once, but are cached in batches, such as... Figure 1 As shown in the diagram. Specifically, based on the shared memory capacity and computational requirements, only a portion of the key and value blocks are loaded at a time. After computation, the space is released promptly to cache the next batch of blocks, thereby supporting large-scale attention computation with limited shared memory resources.

[0072] In step S4, attention output blocks are formed. The system is logically divided into N² batches based on a preset batch size. The batch size is determined by a batch storage instruction (such as the Stmatrix instruction) to optimize memory access efficiency. For the attention output data of each batch, a specified register set is reused for target precision type conversion, avoiding the occupation of additional register resources. Furthermore, the converted output of each batch is efficiently written back to shared memory using the Stmatrix instruction.

[0073] This invention utilizes a consumer thread group to internally split and calculate query and key blocks to obtain corresponding attention score blocks. Furthermore, by reusing a specified register group, the attention output blocks are written back to shared memory, efficiently utilizing existing hardware resources to improve computational performance. In addition, this invention loads the complete query block from shared memory in one go. Accessing the first register group (regQ) reduces the number of shared memory accesses, thereby improving computational efficiency.

[0074] In an optional embodiment, the sizes of the query sub-block and the key sub-block are adapted to the matrix multiply-add instructions.

[0075] Furthermore, the sizes of the query block, the key block, and the value block are determined by the upper limit of available register space, the upper limit of available shared memory space, and the matrix multiplication and addition instructions.

[0076] It should be noted that, since the sizes of query subblocks and key subblocks are adapted to the mma instructions, and query subblocks are created by dividing the query into blocks... (seqlen_q×headdim) is obtained by logically dividing in the headdim direction, so seqlen_q also needs to be adapted to mma instructions.

[0077] Query blocks with data type FP16 Taking (64×256) and key segmentation as examples, this embodiment of the invention calculates attention score blocks. When using the mma instruction, only 128×16 key blocks can be processed at a time. However, this embodiment of the invention does not directly design the key block size to be 16×256 (i.e., sequence length seqlen_k=16). The reason is that if the seqlen_k of the key block is too small (e.g., only 16), it will cause frequent loading of small-sized key blocks from global memory to shared memory when calculating attention score blocks, making the video memory bandwidth a bottleneck and resulting in low register utilization and wasted resources. Instead, it is preferable to design seqlen_k to be as close as possible to the maximum value of the available register limit and / or the available shared memory space limit (i.e., using the largest key block size). This can make full use of register resources to store more intermediate calculation results and effectively reduce the number and frequency of memory access instructions, thereby improving the overall computational efficiency.

[0078] In an optional embodiment, step S2 above, which calculates the corresponding key block in the j-th key block based on matrix multiplication and addition instructions for each query sub-block to obtain the j-th attention score block corresponding to the i-th query block, includes:

[0079] Initialize the second register group of the consumer thread group;

[0080] Transpose the j-th key block to obtain the key transpose block; where the key transpose block is composed of N col The columns are composed of key blocks; N col ≥1;

[0081] The matrix multiplication and addition instruction multiplies N1 query sub-blocks with all key sub-blocks in the k-th column of the key transpose block, and then sequentially accumulates the corresponding sub-block product results into the k-th region of the second register group; where 1≤k≤N col ;

[0082] After traversing and calculating all key sub-blocks in the key transpose block, the N obtained from the second register group... col The summation results are combined to form the j-th attention score block.

[0083] like Figure 3 The diagram shown is a schematic representation of an embodiment of the computational matrix correlation dimension provided by the present invention. Figure 3 In this context, Q chunk, K chunk, and V chunk represent the query chunk, key chunk, and value chunk, respectively; S represents the attention score chunk calculated from the query chunk and key chunk; P represents the exponential score chunk of S after softmax; and O chunk is the attention output chunk corresponding to the query chunk.

[0084] like Figure 4 The diagram shown is a flowchart illustrating an embodiment of attention score block acquisition provided by the present invention. Figure 3 Taking a 64×256 query chunk (Q chunk) and a 32×256 key matrix chunk (K chunk), with data type FP16, as an example, the explanation is as follows: First, the consumer thread group divides the 64×256 query chunk (Q chunk) into two 64×128 query sub-chunks (Half Q) using the computation size adapted to the mma instruction; simultaneously, the key matrix chunk (K chunk) is transposed to obtain a 256×32 key transposed chunk. Next, the key transpose is divided into blocks. Divided into N col = 2 columns of key blocks, each column consists of 2 128×16 key blocks (Half k).

[0085] During computation, the second register set regS (containing two Sacc regions) is first initialized to accumulate the sub-block product results. Then, the mma instruction (Sacc = mma(Half Q, Half k, Sacc)) is used to calculate the local attention score column-by-column, that is, multiplying the two query sub-blocks (Half Q) with the key sub-block in the k-th column of the key transpose block one by one, and accumulating the product results to the k-th region of the second register set regS. After completing N... col After the mma instruction of the 2 column key blocks is calculated, the summation results of the two regions in the second register group regS are combined as the corresponding attention score blocks.

[0086] In an optional embodiment, the step of sequentially obtaining M K Each attention score block is divided into several parts, and mixed-precision attention fusion calculations are performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block, including:

[0087] Initialize the historical maximum value tensor, the historical accumulation tensor, and the historical output block; wherein, the historical output block is cached in the third register group of the consumer thread group;

[0088] Based on the iteration rounds, M is sequentially cached in the second register group of the consumer thread group. K The attention score blocks are used for online softmax iterative calculation;

[0089] After each iteration, update the historical maximum tensor, historical summation tensor, and historical output blocks;

[0090] When the iterative calculation ends, the last obtained current output block is used as the attention output block.

[0091] It should be noted that in online softmax, the mathematical equivalence of block softmax is achieved by dynamically updating the historical maximum value tensor, historical accumulation, and tensor.

[0092] In this embodiment of the invention, the historical maximum value tensor (used to record the row maximum value of all calculated attention score blocks during the iteration process) is first initialized to -∞, the historical sum tensor (used to record the row sum of all calculated exponential score blocks) is initialized to 0, and a dedicated storage space (i.e., the third register group regO) is allocated in the consumer thread group to cache the historical output blocks, which is also initialized to 0.

[0093] In each iteration, the attention score blocks currently requiring processing in the second register group regS of the consumer thread group are subjected to online softmax iterative calculation. After each iteration, the historical maximum tensor, historical accumulation tensor, and historical output blocks are synchronously updated. In M K When the iterative calculation of each attention score block ends, the last obtained current output block is used as the query block. Corresponding attention output blocks .

[0094] In conjunction with the above embodiments, the step of sequentially caching M in the second register group of the consumer thread group K The online softmax calculation is performed on each attention score block, including:

[0095] Perform the following operations on the j-th attention score block processed in the current iteration round:

[0096] By comparing the current block maximum tensor with the historical maximum tensor, the current cumulative maximum tensor is obtained, and the corresponding update correction factor is calculated.

[0097] Based on the cumulative maximum value tensor and the scaling factor, the j-th attention score block is exponentially operated on, and the calculated exponential score block is cached by reusing the second register group;

[0098] The indexed score blocks are summed by accumulating the elements along the row dimension to obtain the block index and tensor. The current accumulated tensor is then calculated by combining the update correction factor and the historical accumulated tensor.

[0099] The historical output blocks are corrected by updating the correction factor and the historical accumulation and tensor, and the calculated historical correction numerator is cached by reusing the third register group;

[0100] After writing the indexed fraction blocks into the fourth register group for target precision type conversion, multiplying them with the corresponding value blocks and accumulating them into the third register group, the current numerator is obtained.

[0101] Divide the current numerator by the current sum tensor and cache the current output block obtained by the calculation by reusing the third register group.

[0102] It should be noted that in online softmax, the current output block is obtained. The calculation formulas involved are as follows:

[0103] (1) Query blocks The corresponding j-th attention score block : ;

[0104] (2) The current maximum value tensor of the chunk: chunk_rmax ;

[0105] (3) The historical maximum tensor accum_rmax (a tensor composed of the maximum values ​​of all rows in the 1st to j-1th attention score blocks): ;

[0106] (4) The current cumulative maximum value tensor accum_rmax_new: new ;

[0107] (5) Segmenting attention scores After performing the exponentiation, the resulting exponential fraction blocks :

[0108] (Formula ①); where sf is the scaling factor;

[0109] (6) Block exponent and tensor chunk_rsum: ;

[0110] (7) Historical accumulation and tensor accum_rsum (attention score blocks from 1 to j-1 in (Total row sum calculated under the condition of maximum tensor) ;

[0111] (8) Update the Modifier (for scale alignment): (Formula ②);

[0112] (9) Current accumulated sum tensor accum_rsum_new: (Formula ③);

[0113] (10) Historical revision section: (Formula ④);

[0114] Current calculation of the molecular part: =Historical Correction of Molecular Part+ (Formula ⑤);

[0115] Current output blocks ,Right now =Currently calculated numerator / Current accumulated sum tensor (Formula ⑥).

[0116] The chunk_rmax, accum_rmax, ..., accum_rsum_new mentioned above are the names of each calculation parameter in the implementation instructions (code), and each occupies the corresponding register resources.

[0117] Specifically, the consumer thread group implements the above calculation formula to obtain query chunks using the following example implementation instructions. The corresponding attention output blocks; among them, the query block, key block, and value block are all of FP16 data type, while the intermediate results of the online softmax calculation are of FP32 data type (to avoid increasing the accumulation error):

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

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

[0120] regO / regS init as 0 / / Initialize the third register group regO / the second register group regS;

[0121] Loop { / Q outer loop /

[0122] Ldmatrix Q -> regQ / / Divide the query into chunks Load from shared memory into the first register set regQ;

[0123] Loop{ / KV internal circulation /

[0124] regS=mma(regQ,smemK,regS) / / Corresponding to step S2, multiply the N1 query sub-blocks with all the key sub-blocks in the k-th column of the key transpose block, and sequentially accumulate the corresponding sub-block product results to the k-th region in the second register group, so that the second register group regS obtains the attention score block. ;

[0125] calc chunk_rmax, accum_rmax, accum_rmax_new / / Calculate the current chunk maximum value tensor, the historical maximum value tensor, and the cumulative maximum value tensor;

[0126] exp((S-accum_rmax_new) sf), reuse regS / / Corresponds to formula ①, by reusing the second register group regS, the calculated exponential fraction blocks are cached. ;

[0127] calc chunk_rsum, accum_rsum, accum_rsum_new / / Calculate the chunk exponent and tensor, the historical accumulated sum and tensor, and the current accumulated sum and tensor;

[0128] Modifier=exp((accum_rmax–accum_rmax_new) sf) / / Corresponds to formula ②;

[0129] accum_rsum_new=accum_rsum Modifier +chunk_rsum / / Corresponds to formula ③;

[0130] regO = modifier acc_rsum regO / / Corresponds to formula ④, by reusing the third register group regO, the calculated historical correction numerator is cached;

[0131] Typecast regS to regP / / Writes the exponential fraction blocks to the fourth register group regP for target precision type conversion (FP16), i.e., FP32 → FP16;

[0132] regO=mma(regP,smemV,regO) / / Corresponds to formula ⑤, multiply the precision-converted exponential fraction blocks by the corresponding value blocks and accumulate them into the third register group regO;

[0133] regO = regO / accum_rsum_new / / Corresponds to formula ⑥, by reusing the third register group regO, the calculated current output block is cached. ;

[0134] }

[0135] Loop{

[0136] Typecast regO to regP / / After the KV inner loop ends, the last cached data in regO (i.e. (For query blocks) Corresponding attention output blocks Batch type conversion is achieved by reusing the fourth register group regP;

[0137] Stmatrix regP / / Writes the output of each batch of transformations in regP back to shared memory;

[0138] }

[0139] Gsm[coord] = LSE / / Write the exponent and logarithm result (LogSumExp, LSE) to shared memory;

[0140] Tma(cpb) LSE / / Transfer LSE via tensor memory accessor;

[0141] Tma O / / Transfer attention output chunks via tensor memory accessor

[0142] }

[0143] It is worth noting that, by reusing register groups multiple times and compressing register usage, this embodiment of the invention supports attention calculation for larger query blocks (large headdim) and key blocks (large seqlen_k) under the same hardware conditions, thereby reducing memory access latency and improving computational performance.

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

[0145] After obtaining the attention output blocks, M K The exponent and logarithm results corresponding to each attention score block are written to shared memory or directly to global memory.

[0146] It should be noted that in the above implementation instruction example, if the shared memory (GSM) has insufficient remaining space to accommodate the LSE result after storing the query block, key block, value block, and the calculated attention output block required for computation, it will directly write the result to the global memory; otherwise, if there is sufficient remaining space, the LSE result will be written to the shared memory first.

[0147] In an optional embodiment, the designated register group is either the first register group or the fourth register group.

[0148] It should be noted that in the above implementation instruction example, the attention output block is divided using the fourth register set regP. Perform batch type conversions and data unloading / writing operations. Since the data type stored in the first register group regQ is the same as that in regP, it can also be used for block attention output. Perform a one-time type conversion and unloading operation. The specific register set used for the type conversion is determined based on actual needs.

[0149] It is worth noting that if the first register group regQ is reused, although attention output can be divided into blocks... A one-time type conversion and unloading operation is performed, but this will block the next query block from being loaded from shared memory into the register, meaning it needs to wait for the attention output block. Only after processing is complete can the next query chunk be loaded into regQ for QK processing. T Multiplication, which impacts computational performance to some extent. While reusing the fourth register set regP requires block-based attention output, this approach... Performing batch type conversion and data unloading operations (i.e., regP conversion is slower) will not block the data loading of the next stage. Preferably, the designated register group is the fourth register group regP.

[0150] In an optional embodiment, when the size of the query block is fixed, the sequence length of the key block is obtained as follows:

[0151] The primary optimization objective is to minimize the difference between the upper limit of register availability for a single thread within the consumer thread group and the number of registers required for attention fusion computation.

[0152] Based on the number of query blocks computed in parallel and the number of key blocks and value blocks cached in shared memory in each batch, the memory usage capacity is determined, and the second optimization objective is to minimize the difference between the upper limit of the available space in shared memory and the memory usage capacity.

[0153] The sequence length of the key block is obtained by jointly solving the first optimization objective and the second optimization objective.

[0154] It should be noted that the size of the query block is seqlen_q × headdim. When the data type is FP16, seqlen_q is set to 64 to adapt to the mma instruction. headdim determines how "detailed" the representation of each word (or position) is during attention computation; a larger value allows the model to capture more complex relationships, but also increases resource overhead. This embodiment of the invention reduces register overhead by reusing multiple registers in online softmax iteration computation, thus supporting attention computation for larger query blocks (large headdim) and key blocks (large seqlen_k). Furthermore, headdim is typically pre-set according to actual needs.

[0155] When the query block (FP16) size is 64 × headdim, the key block (FP16) size is seqlen_k × headdim, and the value block (FP16) is the same size as the key block; where seqlen_k is a multiple of 16 to accommodate the mma instruction. When executing the above implementation example, the register overhead for each thread in the consumer thread group is calculated as follows:

[0156] (1) The number of registers q_tlr allocated to a single thread from the first register set regQ is:

[0157] q_tlr=4×(64×headdim) / (64×16).

[0158] When a consumer thread group (consisting of 4 thread bundles) holds a 64×16 matrix (FP16), each thread bundle is responsible for processing 16×16 matrix blocks. For example... Figure 5 The diagram shown is a schematic representation of an embodiment of data arrangement in a thread bundle provided by the present invention. Figure 5In the diagram, t0 to t31 represent thread symbols, and r0 to r3 represent register symbols (each register can store two FP16 data points). In a 16×16 matrix block (FP16), thread t0 holds eight data points in columns 1, 2, 9, and 10 of rows 1 and 9, occupying four private registers (r0 to r3) of thread t0. In other words, when the consumer thread group holds a 64×16 matrix (FP16), the register overhead for each thread is four. Based on this, when the consumer thread group holds a 64×headdim query block, the number of registers per thread, q_tlr, is 4×(64×headdim) / (64×16).

[0159] (2) The number of registers s_tlr allocated to a single thread from the second register set regS is:

[0160] s_tlr=8×(64×seqlen_k) / (64×16).

[0161] When the attention score is calculated by the consumer thread group The size is 64×seqlen_k, the data type is float (i.e., FP32), and its storage layout is interleaved. When the consumer thread group holds a 64×16 matrix (FP32), each thread bundle is responsible for processing 16×16 matrix blocks. For example... Figure 6 The diagram illustrates another embodiment of the data arrangement in a thread bundle provided by this invention. In a 16×16 matrix block (FP32), thread t0 holds eight data points in columns 1, 2, 9, and 10 of rows 1 and 9, occupying eight private registers of thread t0 (i.e., r0 to r7, each register storing one FP32 data point). In other words, when the consumer thread group holds a 64×16 matrix (FP32), the register overhead for each thread is eight. Based on this, when the consumer thread group holds a 64×seqlen_k attention score block, the number of registers per thread is s_tlr = 8×(64×seqlen_k) / (64×16).

[0162] (3) chunk_rmax, accum_rmax, accum_rmax_new, and Modifier each occupy 2 registers.

[0163] As mentioned earlier, when the consumer thread group holds a 64×16 matrix (FP32), a single thread is responsible for processing two rows of data in the matrix. Therefore, when the consumer thread group holds 64×seqlen_k attention score blocks... In (FP32) mode, a single thread is also responsible for processing. The two rows of data in the dataset. Therefore, each thread needs to maintain chunk_rmax, accum_rmax, accum_rmax_new, and Modifier corresponding to these two rows of data to support the dynamic update and computational needs in online softmax iterative calculations.

[0164] (4) chunk_rsum, accum_rsum, and accum_rsum_new each occupy 2 registers.

[0165] Similarly, when the consumer thread group holds 64×seqlen_k exponentially fractional blocks... In (FP32) mode, a single thread is also responsible for processing. The two rows of data. Therefore, each thread maintains the chunk_rsum, accum_rsum, and accum_rsum_new corresponding to the two rows of data.

[0166] (5) The number of registers p_tlr allocated to a single thread from the fourth register group regP is:

[0167] p_tlr = max(s_tlr / 2, stmatrix_num); where stmatrix_num is the minimum number of registers (4) used when executing the stmatrix instruction.

[0168] The fourth register set, regP, was used for the first time for exponential fractional block decomposition. The type conversion (FP32 → FP16) requires s_tlr / 2 of space; the second operation, which involves batch type conversion and batch storage (stmatrix) of attention output blocks, requires stmatrix_num of space. Only by taking the maximum value between s_tlr / 2 and stmatrix_num can these operations be achieved.

[0169] Of course, if the first register group regQ is used for type conversion of the attention output block, then p_tlr = s_tlr / 2; in addition, when s_tlr / 2 ≥ stmatrix_num, the total_tlr used by a single thread is consistent regardless of whether regQ or regP is used for type conversion, because p_tlr is always equal to s_tlr / 2, and the space of regQ is sufficient to perform type conversion and write out the attention output block at once.

[0170] (6) The number of registers o_tlr allocated to a single thread from the third register group regO is:

[0171] o_tlr=8×(64×headdim) / (64×16).

[0172] The size of the attention output block (FP32) calculated by the consumer thread group is 64×headdim.

[0173] In summary, the total register overhead for a single thread is total_tlr = q_tlr + s_tlr + 7 × 2 + p_tlr + o_tlr.

[0174] It should be noted that the first optimization objective is ;in, This represents the upper limit of register availability. The number of registers required for attention fusion computation (i.e., the total register overhead for a single thread). The second optimization objective is... ;in, This represents the maximum available space for shared memory. This refers to the amount of memory used.

[0175] When jointly solving for the first and second optimization objectives, situations may arise where the number of registers used is close to the maximum available register limit, but there is still ample shared memory space; or, the memory usage is close to the maximum available shared memory space, but there is still ample register space. The choice of which solution to use depends on the specific requirements.

[0176] For example, when a thread block enables two consumer thread groups, the number of query blocks computed in parallel in shared memory, Num1, is 2, and the number of key blocks and value blocks cached in shared memory in each batch, Num2, is at least 1. When Num2=1, although pipelined processing is possible, ping-pong operations are not possible, leading to a decrease in computational performance. Therefore, when jointly solving for the first and second optimization objectives, we can use Num2=Num1 as a basis, first solve for the maximum seqlen_k, and then determine whether to increase the value of Num2 based on the remaining space in shared memory.

[0177] In a specific embodiment, for a certain artificial intelligence chip, when two consumer thread groups are enabled in the thread block, a total of 256 registers are configured in a single thread, and the upper limit of the available registers is set to 240 (with some margin reserved); the upper limit of the available shared memory space is 224KB, with a high address space of 96KB and a low address space of 128KB. High memory access bandwidth can be obtained by accessing the high bank and the low bank simultaneously.

[0178] When headdim=256, the size of the query block (FP16) is 64×256. By jointly solving the first and second optimization objectives, the seqlen_k of the key block (FP16) is 32 (i.e., the size of the key block is 32×256), and Num2=2. The total overhead of the single-threaded register is shown in Table 1, and the memory usage is shown in Table 2. Among them, sizeof(X) is used to characterize the number of bytes occupied by a single data when the data type is X.

[0179] Table 1. Total Register Overhead for Single-Threaded Threads with headdim=256

[0180]

[0181] Table 2. Memory Usage When headdim=256

[0182]

[0183] When headdim=192, the size of the query block (FP16) is 64×192. By jointly solving the first and second optimization objectives, the seqlen_k of the key block (FP16) is obtained as 64 (i.e., the size of the key block is 64×192), and Num2=2. The total overhead of the single-threaded register is shown in Table 3, and the memory usage is shown in Table 4.

[0184] Table 3. Total Register Overhead for Single-Threaded Programming with headdim=192

[0185]

[0186] Table 4. Memory Usage When headdim=192

[0187]

[0188] like Figure 7 The diagram shown is a schematic representation of an embodiment of the shared memory data storage provided by the present invention. The data storage spaces in Tables 2 and 4 are as follows: Figure 7 As shown, the different blocks of data in combinations such as query blocks and key blocks, key blocks and value blocks, and query blocks and attention output blocks are all located in different address regions. This allows for simultaneous access to the high-address storage region (corresponding to...) when reading data. Figure 7 The "HighAddr" and low bank memory regions (corresponding to) Figure 7 Use "LowAddr" to obtain high memory access bandwidth.

[0189] As shown in Tables 1 to 4, when headdim=256, seqlen_k=32, and when headdim=192, seqlen_k=64 are the maximum consumption / limit of computing resources used by this artificial intelligence chip. If the above seqlen_k is increased by the minimum amount (i.e. +16), it will directly lead to the overflow of register and / or shared memory resources, causing a sharp drop in performance.

[0190] The attention calculation method provided in this invention, under the premise of not overflowing register and shared memory resources, firstly, loads the i-th query block to be calculated from shared memory into the first register group of the consumer thread group in one go, and performs internal split calculation on the loaded query block and key block to obtain the corresponding attention score block. This reduces the number of times the query block is copied between shared memory and registers, improving the calculation speed of attention score. Secondly, in the iterative calculation of softmax, register groups are reused multiple times, which can compress register usage. Thirdly, by reusing the designated register group that stores target precision type data during the execution of fusion calculation, the attention output block is written back to shared memory, which can avoid the additional resource overhead introduced by data flow. Therefore, this invention can significantly improve the efficiency of register usage, thereby supporting the calculation of larger query blocks (large headdim) and key blocks (large seqlen_k), fully tapping the hardware computing potential, and ultimately maximizing memory access efficiency and computing performance.

[0191] See Figure 8 This is a schematic diagram of an embodiment of the attention calculation implementation device provided by the present invention.

[0192] A second aspect of the present invention provides an attention calculation implementation apparatus, comprising:

[0193] The query data loading module 11 is used to load the i-th query block to be calculated from shared memory into the first register group of the consumer thread group; where i ≥ 1;

[0194] The attention score calculation module 12 is used to logically divide the loaded query blocks into N1 query sub-blocks based on the header dimension. Then, based on matrix multiplication and addition instructions, it calculates the attention score for each query sub-block in relation to the corresponding key block in the j-th key block to obtain the j-th attention score block corresponding to the i-th query block; where N1 ≥ 1; 1 ≤ j ≤ M. K M K The number of blocks in the key matrix;

[0195] Attention output calculation module 13 is used to sequentially obtain M KEach key block is divided into attention score blocks, and mixed-precision attention fusion calculation is performed with the corresponding value blocks to obtain the attention output block corresponding to the i-th query block; wherein, the key blocks and the value blocks are cached in batches in shared memory;

[0196] Attention output unloading module 14 is used to divide the attention output block logic into N2 batches, reuse the specified register group for target precision type conversion in batches, and write the output result after conversion of each batch back to shared memory; wherein, the specified register group is the register group used to store target precision type data during the attention fusion calculation process; N2≥1.

[0197] It should be noted that the attention calculation implementation device provided in the second aspect embodiment of the present invention can implement all the processes of the attention calculation implementation method described in any of the first aspect embodiments. The functions and technical effects of each module and unit in the device are the same as the functions and technical effects of the attention calculation implementation method described in any of the first aspect embodiments, and will not be repeated here.

[0198] A third aspect of the present invention provides a computer-readable storage medium comprising 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 calculation implementation method described in any of the first aspects of the present invention.

[0199] A fourth aspect of the present invention provides a computer program product, including computer instructions, which, when executed by a processor, implement the attention calculation method described in any of the first aspects of the present invention.

[0200] See Figure 9 This is a schematic diagram of an embodiment of the electronic device provided by the present invention.

[0201] A fifth aspect of the present invention provides an electronic device including a processor 21, a memory 22, and a computer program stored in the memory 22 and configured to be executed by the processor 21, wherein the processor, when executing the computer program, implements the attention calculation implementation method described in any of the first aspects of the present invention.

[0202] Preferably, the computer program can be divided into one or more modules / units (such as computer program one, computer program two, ...), and the one or more modules / units are stored in the memory 22 and executed by the processor 21 to complete the present invention. The one or more modules / units can be a series of computer program instruction segments capable of performing specific functions, and the instruction segments are used to describe the execution process of the computer program in the electronic device.

[0203] The processor 21 can be any one of a CPU (Central Processing Unit), GPU (Graphics Processing Unit), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), DPU (Deep Learning Processing Unit), APU (Accelerated Processing Unit), and GPGPU (General-Purpose Computing on Graphics Processing Unit). The processor 21 is the control center of the electronic device, connecting various parts of the electronic device via various interfaces and lines.

[0204] The memory 22 mainly includes a program storage area and a data storage area. The program storage area can store the operating system, applications required for at least one function, etc., and the data storage area can store related data, etc. In addition, the memory 22 can be a high-speed random access memory, or a non-volatile memory, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, and a flash card, etc., or the memory 22 can also be other volatile solid-state storage devices.

[0205] It should be noted that the aforementioned electronic devices may include, but are not limited to, processors and memory, as will be understood by those skilled in the art. Figure 9 The structural block diagram shown is merely a structural example of the above-described electronic device and does not constitute a limitation on the structure of the above-described electronic device. The above-described electronic device may include more or fewer components than shown, or combine certain components, or different components.

[0206] See Figure 10This is a schematic diagram of an embodiment of the artificial intelligence processor provided by the present invention.

[0207] The artificial intelligence processor provided in this embodiment of the invention includes multiple computing units (CUs). The CU is the core component for performing computing tasks and can run multiple threads in parallel. The threads can be divided into multiple consumer thread groups (such as consumer thread group 1 and consumer thread group 2) and one producer thread group.

[0208] Each consumer thread group can execute the attention computation implementation method described in any of the above embodiments. The CU also includes other hardware resources, such as Group Shared Memory (GSM) and Thread Level Registers (TLRs). The data storage method of GSM is as follows: Figure 7 As shown, the GLM is used to store query blocks, key blocks, value blocks, and attention output blocks. If the GSM space is large enough, the LSE is stored; otherwise, the LSE is stored directly in the GLM. In the register space, regQ is the first register group, regS is the second register group, regO is the third register group, and regP is the fourth register group. The intermediate result register group of Softmax includes: chunk_rmax, accum_rmax, accum_rmax_new, chunk_rsum, accum_rsum, accum_rsum_new, and Modifier.

[0209] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the technical principles of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. An attention computation implementation method, characterized by, The method is executed by a consumer thread group, and the method comprises: loading a current i-th query sub-block to be calculated from shared memory to a first register group of the consumer thread group; wherein i≥1; Logical division is performed on the obtained query block in the head dimension to obtain N1 query sub-blocks, and based on the matrix multiplication and addition instruction, each query sub-block is calculated with the corresponding key sub-block in the jth key block to obtain the jth attention score block corresponding to the ith query block; wherein N1≥1; 1≤j≤M K ; M K is the number of key matrix blocks; M K key score blocks and value score blocks are obtained in sequence, and mixed-precision attention fusion calculation is performed on the corresponding value blocks to obtain an attention output block corresponding to the i-th query block; wherein, the key blocks and the value blocks are cached in batches in the shared memory; after dividing the attention output sub-block logic into N2 batches, multiplexing a specified register group by batch to convert the target precision type, and writing the output result of each batch after conversion back to the shared memory; wherein the specified register group is a register group used to store data of the target precision type during the execution of the attention fusion calculation process; N2≥1.

2. The attention computation implementation method of claim 1, wherein, The method further comprises: After obtaining the attention output blocks, the sum of the indices corresponding to the M K attention score blocks is taken and the logarithm of the result is written to shared memory or directly to global memory.

3. The attention computation implementation method of claim 1, wherein, The sizes of the query sub-block and the key sub-block are adapted to the matrix multiplication and addition instruction.

4. The attention computation implementation method of claim 1, wherein, Based on the matrix multiplication and addition instruction, the calculation of each query sub-block and the corresponding key sub-block in the j-th key sub-block is performed to obtain the j-th attention score sub-block corresponding to the i-th query sub-block, comprising: initializing a second register group of the consumer thread group; The jth key block is transposed to obtain a key transpose block; wherein the key transpose block is composed of N col key sub-blocks of the column; N col ≥ 1. By means of a matrix multiplication instruction, multiply N1 query sub-blocks with all the key sub-blocks in the kth column of the key transpose block, and sequentially accumulate the corresponding sub-block multiplication results into the kth area in the second register group; wherein, 1≤k≤N col ; After traversing all the key sub-blocks in the key transpose tile, the N col accumulation results obtained by the second register bank are combined as the jth attention score tile.

5. The attention computation implementation method of claim 1, wherein, The M K attention scores are obtained in sequence, and mixed-precision attention fusion calculation is performed on the corresponding value blocks to obtain the attention output block corresponding to the i-th query block, including: initializing a historical maximum value tensor, a historical cumulative sum tensor and a historical output sub-block; wherein the historical output sub-block is cached in a third register group of the consumer thread group; based on an iteration round, M K attention score tiles sequentially buffered in a second register group of the consumer thread group are subjected to an iterative computation of online softmax; after each iteration, updating the historical maximum value tensor, the historical cumulative sum tensor and the historical output sub-block; when the iterative calculation is completed, the current output sub-block obtained last is taken as the attention output sub-block.

6. The attention computation implementation method of claim 5, wherein, The pair of M K The online softmax iterative calculation of the attention score blocks includes: The following operations are performed on the j-th attention score sub-block processed in the current iteration round: by comparing the current sub-block maximum value tensor with the historical maximum value tensor, obtaining the current cumulative maximum value tensor, and calculating the corresponding update correction factor; based on the cumulative maximum value tensor and the scaling factor, performing exponential operation on the j-th attention score sub-block, and caching the calculated exponential score sub-block by multiplexing the second register group; performing element accumulation sum calculation on the row dimension of the exponential score sub-block to obtain a sub-block exponential sum tensor, and combining the update correction factor and the historical cumulative sum tensor to calculate the current cumulative sum tensor; modifying the historical output sub-block by the update correction factor and the historical cumulative sum tensor, and caching the calculated historical modified numerator by multiplexing the third register group; writing the exponential score sub-block into the fourth register group for conversion of the target precision type, multiplying it with the corresponding value sub-block and accumulating it into the third register group to obtain the current calculation numerator; dividing the current calculation numerator by the current cumulative sum tensor, and caching the calculated current output sub-block by multiplexing the third register group.

7. The attention computation implementation method of claim 6, wherein, The specified register group is the first register group or the fourth register group.

8. The attention computation implementation method of claim 1, wherein, The sizes of the query sub-block, the key sub-block and the value sub-block are determined by the upper limit of the available registers, the upper limit of the available shared memory space and the matrix multiplication and addition instruction.

9. The attention computation implementation method of claim 1, wherein, When the size of the query sub-block is fixed, the sequence length of the key sub-block is obtained as follows: minimizing the difference between the upper limit of the available registers of a single thread in the consumer thread group and the required number of registers for attention fusion calculation as the first optimization target; determining a memory occupation capacity based on a number of the query chunks calculated in parallel and a number of each batch of the key chunks and the value chunks cached in the shared memory, and taking a difference between an upper limit of available space of the shared memory and the memory occupation capacity as a second optimization target; obtaining a sequence length of the key chunks by jointly solving the first optimization target and the second optimization target.

10. An attention computation implementing apparatus characterized by comprising: The method comprises: a query data loading module configured to load an ith query chunk to be calculated from the shared memory to a first register group of the consumer thread group; wherein i≥1; An attention score calculation module is configured to perform logical division on the head dimension of the loaded query block to obtain N1 query sub-blocks, and perform calculation on each query sub-block and a corresponding key sub-block in the jth key block based on a matrix multiplication and addition instruction to obtain a jth attention score block corresponding to the ith query block; wherein N1≥1; 1≤j≤M K ; M K is the number of key matrix blocks. An attention output calculation module is configured to sequentially acquire M K key score blocks, and perform mixed-precision attention fusion calculation on the corresponding value blocks to obtain an i-th query block corresponding attention output block; wherein, the key blocks and the value blocks are batch cached in a shared memory. an attention output unloading module configured to, after dividing the attention output chunk logic into N2 batches, multiplex a specified register group to perform conversion of a target precision type according to batches, and write an output result of each batch after conversion back to the shared memory; wherein the specified register group is a register group used to store data of the target precision type in the process of performing attention fusion calculation; and N2≥1.

11. A computer readable storage medium, characterized in that, The computer readable storage medium comprises a stored computer program; wherein the computer program, when executed, controls a device where the computer readable storage medium is located to perform the attention calculation implementation method according to any one of claims 1 to 9.

12. A computer program product, characterised in that, The computer program product comprises computer instructions, which, when executed by a processor, implement the attention calculation implementation method according to any one of claims 1 to 9.

13. An electronic device, comprising: The device comprises a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, and the processor, when executing the computer program, implements the attention calculation implementation method according to any one of claims 1 to 9.

Images