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

Abstract

The application discloses an attention calculation implementation method and device, medium, equipment and product, and the method comprises the following steps: by multiplexing a first register group, sequentially loading N1 query subblocks of the ith query block from the shared memory to the first register group, and based on a matrix multiplication instruction, calculating the query subblock obtained each time and the corresponding key subblock in the jth key block to obtain the jth attention score block; sequentially obtaining M K attention score blocks, and performing mixed-precision attention fusion calculation on the corresponding value block to obtain the corresponding attention output block, and then performing target precision type conversion and writing back to the shared memory on the attention output block by batch multiplexing the specified register group storing the target precision type data in the fusion calculation process. The application can significantly reduce the register overhead, support large-scale attention calculation, reduce memory access delay and significantly improve the overall calculation performance.

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, which consumes hardware resources in a particularly inefficient manner, making it difficult to achieve ideal computational performance under the constraints of artificial intelligence processor hardware resources. Summary of the Invention

[0003] The purpose of this invention is to provide an attention computing implementation method, apparatus, computer-readable storage medium, electronic device, and computer program product. By cyclically reusing the first register group to load query sub-blocks, it can support large-scale attention score block computing while reducing register overhead. Furthermore, by directly reusing the specified register group corresponding to the target precision for data type conversion and unloading, it can avoid the additional resource overhead introduced by data flow, thereby achieving high-performance attention computing under limited hardware conditions.

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

[0005] By reusing the first register group of the consumer thread group, N1 query sub-blocks of the i-th query block are sequentially loaded from shared memory into the first register group; where i≥1; N1≥1;

[0006] Based on matrix multiplication and addition instructions, the query sub-block obtained in each loading is calculated with 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 1≤j≤M K M K The number of blocks in the key matrix;

[0007] Get 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] The attention output is multiplexed in blocks and batches to a designated register group for target precision type conversion, and the output results of each batch after conversion are written back to shared memory; wherein, the designated register group is a register group used to store target precision type data during the attention fusion calculation process.

[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 size of the query sub-block and the key sub-block is adapted to the matrix multiplication and addition instruction, and the query sub-block is obtained by logically dividing the i-th query block by the head dimension.

[0012] Optionally, the step of calculating the corresponding key block in the j-th key block based on matrix multiplication and addition instructions for each loaded query sub-block 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; the key transpose block is composed of key sub-blocks in row N1;

[0015] The matrix multiply-add instruction is used to multiply the k-th query sub-block with all the key sub-blocks in the k-th row of the key transpose block in turn, and the corresponding sub-block product results are accumulated into the second register group; where 1≤k≤N1;

[0016] After traversing and calculating all the query sub-blocks, the accumulated result obtained from the second register group is used as 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 along the row dimension to obtain the block index and tensor. The current sum and tensor are then calculated by combining the update correction factor and the historical sum 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 N1 query sub-blocks of the i-th query block from shared memory into the first register group by reusing the first register group of the consumer thread group; where i≥1; N1≥1;

[0038] The attention score calculation module is used to calculate the corresponding key block in the j-th key block for each loaded query sub-block based on matrix multiplication and addition instructions, so as to obtain the j-th attention score block corresponding to the i-th query block; where 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 multiplex the attention output in blocks and batches to a designated register group, perform target precision type conversion, and write the output results of each batch back to shared memory; wherein, the designated register group is a register group used to store target precision type data during the attention fusion calculation process.

[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 the prior art, the 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 query sub-blocks by cyclically reusing the first register group, and after each query sub-block is loaded, it calculates with the corresponding key sub-block in the j-th key block in shared memory using matrix multiplication and addition instructions to obtain the j-th attention score block; then, it sequentially obtains M... K The attention score blocks are divided into several segments, 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 to the target precision type in batches, the specified register group storing the target precision type data during the fusion calculation is reused. This invention can significantly reduce register overhead, thereby supporting attention calculations for large query blocks and key blocks, reducing memory access latency, fully exploiting the hardware computing potential, and thus significantly improving overall computing performance. 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 computation matrix correlation dimension 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: By reusing the first register group of the consumer thread group, load N1 query sub-blocks of the i-th query block from shared memory into the first register group in sequence; where i≥1; N1≥1;

[0063] Step S2: Based on matrix multiplication and addition instructions, calculate the corresponding key block in the j-th key block and the query sub-block obtained each time it is loaded, to obtain the j-th attention score block corresponding to the i-th query block; where 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: The attention output is multiplexed into a designated register group in blocks and batches, the target precision type is converted, and the output results of each batch after conversion are written back to the shared memory; wherein, the designated register group is the register group used to store the target precision type data during the attention fusion calculation process.

[0066] It should be noted that the query is divided into blocks. It is a matrix block obtained by dividing the original query matrix according to a preset dimension (such as sequence length seqlen_q and head dimension headdim); similarly, key blocks and value blocks are matrix blocks obtained by dividing the original key matrix and the original value matrix respectively, and both have the same size.

[0067] In steps S1 and S2, the query is divided into blocks. The process is broken down into smaller query sub-blocks, and the corresponding query sub-blocks are loaded sequentially from shared memory into the first register set regQ. After each query sub-block is loaded, it is immediately multiplied with the corresponding key sub-block in the j-th key block in shared memory using a specified matrix multiply-add instruction. The calculation results of all query sub-blocks are then accumulated to obtain the complete attention score block. .

[0068] This invention calculates attention score blocks in batches by repeatedly reusing the first register group regQ. This reduces the register requirements for a single operation. Furthermore, this embodiment of the invention uses the mma (matrix multiply-accumulate) instruction as the designated matrix multiplication-accumulation instruction. When the mma instruction is executed, it requires matrix A (such as a lookup sub-block) to be stored in a register, and matrix B (such as a key sub-block) to be stored in shared memory, further reducing register usage.

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

[0070] In this embodiment of the invention, the mixed-precision computation uses low-precision (e.g., FP16) data types for multiplication operations of query sub-blocks, key sub-blocks, and value blocks to improve computation speed and reduce resource consumption. The remaining intermediate computation processes in the fusion computation use high-precision (e.g., FP32) data types to reduce accumulation errors and ensure computational accuracy. Furthermore, value blocks and key blocks employ a synchronous batch caching strategy; that is, key blocks and value blocks are moved in batches from global memory to shared memory for caching. Figure 1 As shown, based on the shared memory capacity and computational requirements, only 2 key blocks and 2 value blocks are moved in each batch, and the space is released in a timely manner after the computation is completed to cache the next batch of blocks, thereby supporting large-scale attention computation with limited shared memory resources.

[0071] In step S4, the attention output is divided into blocks. The data is loaded into a designated register set in batches for target precision (e.g., FP16) conversion. The output of each batch is efficiently written back to shared memory using batch store instructions (e.g., Stmatrix instructions). Therefore, the size of each batch conversion is determined by the batch store instructions. The designated register set is the register space used to store target precision type data during the attention fusion calculation in step S3. This embodiment of the invention avoids additional register overhead introduced by data type conversion and prevents register overflow by reusing the designated register set in batches.

[0072] In summary, the embodiments of the present invention can support large-scale attention score block computation while reducing register overhead by cyclically reusing the first register group regQ to load query sub-blocks. Furthermore, by directly reusing the specified register group corresponding to the target precision for data type conversion and unloading, it can avoid the additional resource overhead introduced by data flow, thereby achieving high-performance attention computation under limited hardware conditions. It is especially suitable for large-size query blocks and key blocks.

[0073] In an optional embodiment, the sizes of the query sub-block and the key sub-block are adapted to the matrix multiplication and addition instructions, and the query sub-block is obtained by logically dividing the i-th query block by the head dimension.

[0074] It should be noted that the embodiments of the present invention use the mma instruction as the specified matrix multiplication and addition instruction. When performing "sub-block level" calculations, the granularity of the sub-block division is determined by the calculation size of the mma instruction.

[0075] Specifically, for query subblocks and key subblocks, the computational size of mma is m×n×K. B Among them, K BThe 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-blocks. B The dimension is set to the maximum value of 16k to minimize the value of N1, thereby reducing the number of times the consumer thread group moves query subblocks from shared memory to the first register set RegQ, while enabling a single instruction to process more data, thus further accelerating query block processing. The corresponding j-th attention score block The calculation process. For example, when querying blocks. When the size of (FP16) is seqlen_q×headdim=64×256, it is logically partitioned along the head dimension (headdim), resulting in two 64×128 query sub-blocks, with corresponding key sub-blocks of size 128×16. Based on the mma instruction, only one query sub-block is calculated at a time, and the product results of each query sub-block and its corresponding key block are gradually integrated through accumulation, ultimately generating a complete attention score block. .

[0076] In an optional embodiment, 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 multiply-accumulate instruction.

[0077] It is worth noting 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 along the headdim direction, so the value of seqlen_q must also match the calculation dimension of the mma instruction.

[0078] Queries of type FP16 are divided into blocks. Taking (64×256) and key segmentation as examples, this embodiment of the invention calculates attention score blocks. While the `mma` instruction processes 128×16 key blocks in a single operation, the key block size is not directly designed to be 16×256 (i.e., sequence length `seqlen_k` = 16). This is because if the `seqlen_k` of the key block is too small (e.g., only 16), it will require frequent loading of small key blocks from global memory to shared memory when calculating attention score blocks. Similarly, the value blocks, due to the same size constraint as the key blocks, will also experience increased data movement, i.e., increased inner loop iterations. This not only makes memory access bandwidth a performance bottleneck but also leads to low register utilization due to insufficient data, resulting in wasted hardware resources. Conversely, 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 maximum available shared memory space (i.e., using the largest feasible key block size). This allows for full utilization of register resources to store more intermediate computation results and effectively reduces the number and frequency of memory access instructions, thereby significantly improving overall computational efficiency.

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

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

[0081] Transpose the j-th key block to obtain the key transpose block; the key transpose block is composed of key sub-blocks in row N1;

[0082] The matrix multiply-add instruction is used to multiply the k-th query sub-block with all the key sub-blocks in the k-th row of the key transpose block in turn, and the corresponding sub-block product results are accumulated into the second register group; where 1≤k≤N1;

[0083] After traversing and calculating all the query sub-blocks, the accumulated result obtained from the second register group is used as the j-th attention score block.

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

[0085] like Figure 4The diagram shown is a flowchart illustrating an embodiment of attention score block acquisition provided by the present invention. Figure 3 The following explanation uses a 64×256 query chunk (Q chunk) and a 48×256 key matrix chunk (Kchunk) of type FP16 as examples. First, the consumer thread group, using the computation size adapted to the mma instruction, divides the 64×256 query chunk (Q chunk) into two 64×128 query sub-chunks (Half Q). Simultaneously, the key matrix chunk (Kchunk) is transposed to obtain a 256×48 key transposed chunk. Next, the key transpose is divided into blocks. Divided into 2 rows of key blocks, each row consisting of 3 128×16 key blocks (Half k); among which, the key transpose blocks... The number of rows divided is the same as the number of sub-blocks in the query.

[0086] This invention employs an iterative calculation strategy. First, the second register group regS (containing three Sacc regions) is initialized to accumulate the sub-block product results. Then, the k=1th query sub-block Half Q is loaded into the first register group, and the mma instruction (Sacc=mma(Half Q, Half k, Sacc)) is used to perform the following operations: ① The current query sub-block is divided into blocks with the key transpose. ① Multiply all key sub-blocks in the k-th row sequentially; ② Accumulate the resulting sub-block product into the corresponding Sacc region of the second register group regS. Process N1 query sub-blocks in a loop, and use the final accumulated result in the second register group regS as the corresponding attention score block.

[0087] In calculating attention score blocks, this embodiment of the invention avoids the resource constraints caused by processing large-scale matrices at once by sequentially processing and accumulating the product of query sub-blocks and key sub-blocks, and supports matrix calculation between query blocks of large size headdim and key blocks of large size seqlen_k.

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

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

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

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

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

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

[0094] In this embodiment of the invention, in the consumer thread group, 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.

[0095] In each iteration, online softmax calculation is performed on the currently pending attention score block in the second register group regS of the consumer thread group. After each iteration, the historical maximum tensor, historical accumulation tensor, and historical output block are synchronously updated. In M K When all iterative calculations of the attention score blocks are completed, the last obtained current output block is used as the query block. Corresponding attention output blocks .

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

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

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

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

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

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

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

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

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

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

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

[0107] (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): ;

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

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

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

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

[0112] (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) ;

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

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

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

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

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

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

[0119] 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):

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

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

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

[0123] Loop { / Q outer loop /

[0124] Loop{ / KV internal circulation /

[0125] loop{ / / Load Q in batches

[0126] Ldmatrix Q ->regQ / / Load query subblocks from shared memory into the first register group;

[0127] regS=mma(regQ,smemK,regS) / / Corresponding to step S2, the k-th query sub-block is multiplied sequentially by all the key sub-blocks in the k-th row of the key transpose block, and the corresponding sub-block product results are accumulated into the second register group regS, so that the second register group regS finally obtains the attention score block. ;

[0128] }

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

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

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

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

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

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

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

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

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

[0138] }

[0139] Loop{

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

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

[0142] }

[0143] Glm[coord] = LSE / / Write the exponent and logarithm result (LogSumExp, LSE) directly into global memory to save shared memory space;

[0144] }

[0145] It is worth noting that the embodiments of the present invention compress the use of registers by reusing register groups multiple times; under the same hardware conditions, it can support attention calculation for large query blocks (large headdim) and key blocks (large seqlen_k), and can effectively reduce memory access latency, fully tap the hardware computing potential, and thus significantly improve overall computing performance.

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

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

[0148] It should be noted that in the above implementation instruction example, the remaining space size of the shared memory (GSM) needs to be determined. After the shared memory has completed storing the query block, key block, value block, and the calculated attention output block, if the remaining space is insufficient to accommodate the exponentiation and logarithm (LogSumExp, LSE) results, the LSE results are directly written to the global memory; otherwise, if the remaining space is sufficient, the LSE results are written to the shared memory first.

[0149] In an optional embodiment, the method further includes: the designated register group is either the first register group or the fourth register group.

[0150] 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 type conversion and data unloading operations in batches, with the number of unloaded batches being less than that of register group regP (the fourth register group). The specific register group used for type conversion is determined based on actual needs.

[0151] It is worth noting that if the first register group regQ is reused, although attention output can be divided into blocks... Type conversion and unloading operations are performed, but this will block the next query block from being loaded from shared memory into the register, meaning we need to wait for the attention output block. Only after processing is complete can the query sub-blocks in the next query block 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... The number of batches processed is greater than that of the first register group regQ (i.e., regP conversion is slower), but this will not cause the data loading of the next stage to be blocked. Preferably, the designated register group is the fourth register group regP.

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

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

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

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

[0156] 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 the "detail" of each word (or position) represented by the model when calculating attention. The larger the headdim value, the stronger the model's ability to capture complex relationships, but the corresponding resource overhead is also greater. This embodiment of the invention achieves large-scale attention score calculation by reusing the first register group regQ multiple times. Furthermore, by reusing multiple registers in the online softmax iteration calculation, the overhead of registers can be reduced to support attention calculation for larger query blocks (large headdim) and key blocks (large seqlen_k). In addition, headdim is usually preset according to actual needs.

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

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

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

[0160] 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), each thread has a register overhead of four. Based on this, when the consumer thread group holds a 64×(headdim / N1) query subblock, the number of registers per thread, q_tlr, is 4×(64×headdim / N1) / (64×16).

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

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

[0163] When the attention score blocks calculated by the consumer thread group are 64×seqlen_k in size, with data type float (i.e., FP32), their 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).

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

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

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

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

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

[0169] p_tlr = max(s_tlr / 2, stmatrix_num); where stmatrix_num is the minimum number of registers (4) required to execute the stmatrix instruction.

[0170] The fourth register set, regP, was used for the first time for exponential fractional block decomposition. The first type conversion (FP32 → FP16) requires s_tlr / 2 of the register space; the second type conversion and batch storage (Stmatrix) operation is performed on the attention output in batches, requiring stmatrix_num of the register space. Therefore, taking the maximum value between s_tlr / 2 and stmatrix_num as p_tlr ensures that the above operations can be executed successfully.

[0171] Of course, if the first register group regQ is used for type conversion of the attention output block, then p_tlr = s_tlr / 2, and regQ needs to be multiplexed by N1 to complete the type conversion of the attention output block. In addition, when s_tlr / 2 ≥ stmatrix_num, the total register overhead (total_tlr) of a single thread remains the same regardless of whether regQ or regP is used for type conversion. This is because p_tlr is equal to s_tlr / 2 in both cases, and since the smallest feasible N1 value is selected as much as possible in the embodiments of this invention, the space allocated by regQ when storing query sub-blocks is also greater than stmatrix_num, which is sufficient to meet the batch type conversion and writing of attention output blocks.

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

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

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

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

[0176] It should be noted that the first optimization objective is ;in, This represents the upper limit of available registers. 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.

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

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

[0179] 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 a certain margin reserved); the upper limit of the available shared memory space is 224KB, with its high address space (high bank) being 96KB and its low address space (low bank) being 128KB. High memory access bandwidth can be obtained by accessing the high bank and low bank simultaneously.

[0180] When headdim=256, the query block size of FP16 type is 64×256. By jointly solving the first and second optimization objectives, the seqlen_k of FP16 type key block is obtained as 48 (i.e., the key block size is 48×256), and Num2=2. The total overhead of single-threaded registers 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.

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

[0182]

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

[0184]

[0185] 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 space in Table 2 is 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 7The "HighAddr" and low bank memory regions (corresponding to) Figure 7 The "LowAddr" option is used to obtain high memory access bandwidth. Furthermore, the query blocks, key blocks, value blocks, and corresponding attention output blocks required by the two consumer thread groups in the pipeline occupy the entire shared memory space (e.g., ...). Figure 7 As shown in the figure, there is no remaining space to store the LSE, so the shared memory is skipped and the LSE is written directly to the global memory.

[0186] As shown in the table above, when headdim=256, seqlen_k=48 is the maximum consumption / limit of computing resources used by this artificial intelligence chip. If the seqlen_k is increased by the minimum amount (i.e. +16), it will directly cause the overflow of register and shared memory resources, resulting in a sharp drop in performance.

[0187] The attention calculation method provided in this invention, without overflowing register and shared memory resources, achieves large-scale attention score calculation by repeatedly reusing the first register group (regQ). Secondly, it reduces register overhead through reusing multiple registers in online softmax iteration calculation. Thirdly, it avoids additional resource overhead introduced by data transfer by reusing a specified register group to write attention output blocks back to shared memory. Therefore, under the same hardware conditions, this invention can support attention calculation for large query blocks (large headdim) and key blocks (large seqlen_k), effectively reduce memory access latency, fully exploit the hardware's computing potential, and significantly improve overall computing performance.

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

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

[0190] The query data loading module 11 is used to load N1 query sub-blocks of the i-th query block from shared memory into the first register group by reusing the first register group of the consumer thread group; where i≥1; N1≥1;

[0191] Attention score calculation module 12 is used to calculate the corresponding key block in the j-th key block and the query sub-block obtained each time it is loaded, based on matrix multiplication and addition instructions, so as to obtain the j-th attention score block corresponding to the i-th query block; where 1≤j≤M K M K The number of blocks in the key matrix;

[0192] Attention output calculation module 13 is 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;

[0193] Attention output unloading module 14 is used to multiplex the attention output in blocks and batches to a designated register group, perform target precision type conversion, and write the output results of each batch after conversion back to shared memory; wherein, the designated register group is a register group used to store target precision type data during the attention fusion calculation process.

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

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

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

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

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

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

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

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

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

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

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

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

[0206] 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: N1 query sub-blocks of an i-th query block are sequentially loaded from the shared memory to a first register group of the consumer thread group by multiplexing the first register group; wherein i≥1; N1≥1; Based on the matrix multiplication and addition instruction, the query sub-block obtained each time 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, 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; The attention output block is batch-multiplexed to a specified register group, a target precision type conversion is performed, and each batch-converted output result is written back to the shared memory; wherein the specified register group is a register group used for storing data of the target precision type in the process of performing attention fusion calculation.

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-blocks and the key sub-blocks are adapted to the matrix multiplication and addition instruction, and the query sub-blocks are obtained by logically dividing the i-th query block in the head dimension.

4. The attention computation implementation method of claim 3, wherein, Based on the matrix multiplication and addition instruction, the query sub-block obtained each time is calculated with the corresponding key sub-block in the j-th key block to obtain the j-th attention score block corresponding to the i-th query block, comprising: A second register group of the consumer thread group is initialized; The j-th key block is transposed to obtain a key transposed block; wherein the key transposed block is composed of N1 key sub-blocks; The k-th query sub-block is multiplied with all key sub-blocks in the k-th row of the key transposed block in sequence by a matrix multiplication and addition instruction, and the corresponding sub-block product result is accumulated to the second register group; wherein 1≤k≤N1; After all the query sub-blocks are calculated, the accumulated result of the second register group is taken as the j-th attention score block.

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 ith query block, including: A history maximum value tensor, a history accumulated sum tensor and a history output block are initialized; wherein the history output 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, the history maximum value tensor, the history accumulated sum tensor and the history output block are updated; When the iterative calculation is completed, the current output block finally obtained is taken as the attention output block.

6. The attention computation implementation method of claim 5, wherein, The pair of M K attention score patches are sequentially buffered in the second register group of the consumer thread group, and the iteration calculation of online softmax includes: The following operations are performed on the j-th attention score block processed in the current iteration round: A current accumulated maximum value tensor is obtained by comparing the current block maximum value tensor with the history maximum value tensor, and a corresponding update correction factor is calculated; Based on the accumulated maximum value tensor and the scaling factor, an exponential operation is performed on the j-th attention score block, and the calculated exponentialized score block is cached by multiplexing the second register group; A block exponential sum tensor is obtained by performing row-dimension accumulated sum calculation on the exponentialized score block, and a current accumulated sum tensor is calculated in combination with the update correction factor and the history accumulated sum tensor; The history output block is corrected by the update correction factor and the history accumulated sum tensor, and the calculated history corrected numerator part is cached by multiplexing the third register group; The exponentialized score block is written into a fourth register group for target precision type conversion, multiplied with a corresponding value block and accumulated to the third register group to obtain a current calculation numerator part; The current calculation numerator part is divided by the current accumulated sum tensor, and the calculated current output block is cached 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 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.

9. The attention computation implementation method of claim 1, wherein, When the size of the query block is fixed, the steps for obtaining the sequence length of the key block are as follows: 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. 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. The sequence length of the key block is obtained by jointly solving the first optimization objective and the second optimization objective.

10. An attention computation implementing apparatus characterized by comprising: include: The query data loading module is used to load N1 query sub-blocks of the i-th query block from shared memory into the first register group by reusing the first register group of the consumer thread group; where i≥1; N1≥1; An attention score calculation module is configured to, based on the matrix multiplication-addition instruction, calculate each obtained query sub-block with a corresponding key sub-block in the jth key sub-block to obtain a jth attention score sub-block corresponding to the ith query sub-block, where 1≤j≤M K ; M K is the number of key matrix sub-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. The attention output unloading module is used to multiplex the attention output in blocks and batches to a designated register group, perform target precision type conversion, and write the output results of each batch back to shared memory; wherein, the designated register group is a register group used to store target precision type data during the attention fusion calculation process.

11. 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 computation implementation method as described in any one of claims 1 to 9.

12. A computer program product, characterised in that, It includes computer instructions that, when executed by a processor, implement the attention calculation implementation method as described in any one of claims 1 to 9.

13. An electronic device, comprising: The method 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 computation implementation method as described in any one of claims 1 to 9.

Images