Model reasoning method and device based on attention mechanism, equipment and storage medium

Abstract

The invention discloses a model reasoning method and device based on an attention mechanism, equipment and a storage medium, and the method comprises the steps: respectively segmenting a query block and a key block under the condition of determining that the query block and the key block which need to be operated meet a first preset condition; determining at least one attention weight block according to each query block obtained by segmentation and each key block obtained by segmentation; and determining an output result based on the at least one attention weight block and the value block corresponding to the key block. By segmenting the query block and the key block into the small blocks, the phenomenon that the intermediate result of the query block and the key block adopts a causal mask to shield unwanted results is avoided, at least one attention weight block determined according to the small blocks contains a result required by subsequent calculation, and only the attention weight blocks need to be calculated, so that the accuracy of subsequent calculation is improved. Therefore, the effect of shielding unwanted results by adopting the causal mask can be achieved, invalid calculation is reduced, occupation of calculation resources is reduced, and therefore the overall reasoning speed of the model can be increased.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of computer, in particular to a model inference method and device based on attention mechanism, equipment, storage medium and program product. BACKGROUND

[0002] In the field of deep learning, especially when dealing with large-scale data sets and complex tasks, large language models (LLM) have very wide applications. Among them, attention mechanism (Attention) has become an indispensable part of these models due to its advantages in long sequence understanding and processing.

[0003] As the sequence length grows, the proportion of attention mechanism calculation time in the entire model inference time gradually increases. From the structure of the GPU (Graphics Processing Unit) running the model, in order to optimize the calculation process of the attention mechanism, it is necessary to make the calculation process as much as possible in the SRAM (Static Random Access Memory). However, the sequence length is usually very long, and the space of SRAM is relatively small, which cannot store the complete query matrix, key matrix, value matrix and intermediate results, so it is necessary to store the query matrix, key matrix and value matrix on the HBM (High Bandwidth Memory), and then load them into the SRAM by block to calculate, in order to improve the calculation speed.

[0004] Although the block calculation method can improve the calculation speed, there is still room for optimization by analyzing the data flow of the attention calculation process. SUMMARY

[0005] The purpose of the present application is to propose a model inference method, device, equipment, storage medium and program product based on attention mechanism to solve the above problems of the prior art. The purpose is achieved by the following technical solutions.

[0006] The first aspect of the present application proposes a model inference method based on attention mechanism, which comprises:

[0007] When it is determined that the query block and the key block that need to be operated meet the first preset condition, the query block and the key block are respectively split; the first preset condition represents that the matrix multiplication operation result of the query block and the key block needs to use the causal mask;

[0008] Based on the segmented query blocks and key blocks, at least one attention weight block is determined; the attention weight block is the result calculated from one query block and one key block.

[0009] Based on the at least one attention weight block and the value block corresponding to the key block, determine the output result corresponding to the query block and the key block;

[0010] The query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0011] Based on the attention-based model inference method described in the first aspect above, this application has at least the following beneficial effects or advantages:

[0012] Considering the invalid computation of intermediate results requiring causal masking in the block-based computation process of the attention mechanism in the model, the query block and key block are divided into smaller blocks if the query block and key block to be computed meet the first preset condition. This avoids using causal masking to block unnecessary results in the intermediate results of the query block and key block. Then, at least one attention weight block is determined based on the smaller blocks. These attention weight blocks all contain the results required for subsequent computation. Based on the value blocks corresponding to at least one attention weight block and the key block, the output results corresponding to the query block and key block are determined. Since dividing the query block and key block into smaller blocks only requires computing the attention weight blocks containing the results required for subsequent computation, it can effectively block unnecessary results using causal masking. This not only reduces invalid computation of unnecessary intermediate results but also reduces the consumption of computing resources, thus improving the overall inference speed of the model.

[0013] A second aspect of this application proposes a model inference method based on an attention mechanism, wherein the attention mechanism needs to process data including a query matrix, a key matrix, and a value matrix, and the method includes:

[0014] For each query block included in the query matrix, the set of key blocks that need to be operated on is determined from each key block included in the key matrix according to the index of the query block;

[0015] For each key block in the key block set, if the query block and the key block meet the first preset condition, the query block and the key block are respectively segmented. A first thread block determines at least one attention weight block based on each segmented query block and each segmented key block, and determines the output result of the first thread block based on the at least one attention weight block and the value block corresponding to the key block; the attention weight block is the result calculated from one query block and one key block;

[0016] If the query block and the key block do not meet the first preset condition, a second thread block is used to determine the output result of the second thread block based on the query block and the key block; the first preset condition indicates that the matrix multiplication result of the query block and the key block needs to use a causal mask;

[0017] The output results of the first thread block and the second thread block are merged to obtain the output result of the query block;

[0018] The output results of each query block are concatenated to obtain the output result of the query matrix.

[0019] Wherein, the first thread block and the second thread block are both the smallest units executed by the graphics processor; the query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0020] Based on the attention-based model inference method described in the second aspect above, this application has at least the following beneficial effects or advantages:

[0021] For each query block, after determining the set of key blocks to be processed based on the query block index, the set of key blocks is used to distinguish between combinations of query blocks and key blocks that meet the first preset condition and combinations that do not meet the first preset condition. The combinations of query blocks and key blocks that meet the first preset condition are then divided into smaller blocks. Due to the change in the blocks, the divided blocks are not suitable for processing in the same thread block as the original whole block. These smaller blocks are processed by the first thread block, while the combinations of query blocks and key blocks that do not meet the first preset condition are processed in parallel by the second thread block to improve the calculation speed. Finally, the output results of the two thread blocks are merged to obtain the output result of the entire query block, thereby improving the overall inference speed of the model.

[0022] A third aspect of this application proposes a model inference device based on an attention mechanism, the device comprising:

[0023] The segmentation module is used to segment the query block and the key block respectively when the query block and the key block to be operated on meet the first preset condition; the first preset condition indicates that the matrix multiplication operation result of the query block and the key block needs to use a causal mask;

[0024] The weight calculation module is used to determine at least one attention weight block based on the various query blocks and key blocks obtained from the segmentation; the attention weight block is the result calculated from one query block and one key block.

[0025] An output determination module is used to determine the output result corresponding to the query block and the key block based on the at least one attention weight block and the value block corresponding to the key block;

[0026] The query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0027] A fourth aspect of this application proposes a model inference apparatus based on an attention mechanism, wherein the attention mechanism needs to process data including a query matrix, a key matrix, and a value matrix, and the apparatus includes:

[0028] The determination module is used to determine, based on the index of the query block included in the query matrix, the set of key blocks that need to be operated on in the key matrix.

[0029] A first calculation module is used to segment each key block in the key block set, provided that the query block and the key block meet a first preset condition. A first thread block is used to determine at least one attention weight block based on each segmented query block and each segmented key block, and to determine the output result of the first thread block based on the at least one attention weight block and the value block corresponding to the key block. The first preset condition indicates that the matrix multiplication result between the query block and the key block needs to use a causal mask. The attention weight block is the result calculated from one query block and one key block.

[0030] The second calculation module is used to determine the output result of the second thread block based on the query block and the key block when the query block and the key block do not meet the first preset conditions.

[0031] The fusion module is used to fuse the output results of the first thread block and the output results of the second thread block to obtain the output result of the query block;

[0032] The output module is used to concatenate the output results of each query block to obtain the output result of the query matrix;

[0033] Wherein, the first thread block and the second thread block are both the smallest units executed by the graphics processor; the query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0034] The fifth aspect of this application provides an electronic device including a memory, a processor, and a computer program stored in the memory and executable on the processor, the processor executing the program to implement the method as described in the first or second aspect above.

[0035] The sixth aspect of this application provides a computer-readable storage medium having a computer program stored thereon, the program being executed by a processor to implement the method as described in the first or second aspect above.

[0036] An embodiment of the seventh aspect of this application provides a computer program product including a computer program that is executed by a processor to implement the method described in the first or second aspect above.

[0037] The above description is only an overview of the technical solution of this application. In order to better understand the technical means of this application, it can be implemented according to the contents of the specification. In order to make the above and other objects, features and advantages of this application more obvious and understandable, specific embodiments of this application are given below. Attached Figure Description

[0038] The accompanying drawings, which are included to provide a further understanding of this application and form part of this application, illustrate exemplary embodiments and are used to explain this application, but do not constitute an undue limitation of this application. In the drawings:

[0039] Figure 1 This is a schematic diagram of attention block computation provided in the prior art;

[0040] Figure 2 This is a schematic diagram illustrating the application of causal masking in an attention computation process provided in the prior art.

[0041] Figure 3 This is a flowchart illustrating an embodiment of an attention-based model inference method according to an exemplary embodiment;

[0042] Figure 4This is a schematic diagram illustrating the processing of a query block and a key block that meet a first preset condition according to an exemplary embodiment;

[0043] Figure 5 This is a schematic diagram illustrating the calculation output of a query block and a key block that meet a first preset condition, according to an exemplary embodiment.

[0044] Figure 6 This is a flowchart illustrating an embodiment of another attention-based model inference method according to an exemplary embodiment;

[0045] Figure 7 This is a schematic diagram illustrating the structure of an attention-based model inference device according to an exemplary embodiment.

[0046] Figure 8 This is a schematic diagram illustrating the structure of another attention-based model inference device according to an exemplary embodiment;

[0047] Figure 9 This is a schematic diagram of the hardware structure of an electronic device according to an exemplary embodiment;

[0048] Figure 10 This is a schematic diagram illustrating the structure of a storage medium according to an exemplary embodiment. Detailed Implementation

[0049] Exemplary embodiments will now be described in detail, examples of which are illustrated in the accompanying drawings. When the following description relates to the drawings, unless otherwise indicated, the same numbers in different drawings represent the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with this application. Rather, they are merely examples of systems and methods consistent with some aspects of this application as detailed in the appended claims.

[0050] The terminology used in this application is for the purpose of describing particular embodiments only and is not intended to be limiting of the application. The singular forms “a,” “the,” and “the” used in this application and the appended claims are also intended to include the plural forms unless the context clearly indicates otherwise. It should also be understood that the term “and / or” as used herein refers to and includes any or all possible combinations of one or more of the associated listed items.

[0051] It should be understood that although the terms first, second, third, etc., may be used in this application to describe various information, such information should not be limited to these terms. These terms are only used to distinguish information of the same type from one another. For example, without departing from the scope of this application, first information may also be referred to as second information, and similarly, second information may also be referred to as first information. Depending on the context, the word "if" as used herein may be interpreted as "when," "when," or "in response to determination."

[0052] Furthermore, it should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, stored data, displayed data, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties. Moreover, the collection, use and processing of the relevant data must comply with the relevant laws, regulations and standards of the relevant countries and regions, and corresponding operation portals are provided for users to choose to authorize or refuse.

[0053] As mentioned earlier, the sequence length of attention mechanisms is very long when dealing with complex tasks. Especially when dealing with complex application scenarios such as ultra-long text, high-resolution images, or large-scale time series, the demand for overall computing resources of attention mechanisms increases dramatically due to the significant increase in sequence length, which can even consume up to 50% of the total inference time of the model, which severely limits the inference efficiency of the model.

[0054] In existing technologies, block-based computation is used to improve the speed of attention calculation, such as... Figure 1 As shown, the calculation process for attention includes:

[0055] 1. Read the query matrix Q, key matrix K, and value matrix V of the attention mechanism into the GPU's HBM.

[0056] 2. Calculate the block size based on the size of the GPU's SRAM, and then divide the query matrix Q, key matrix K, and value matrix V into many blocks according to the block size.

[0057] 3. Execute the outer loop, querying each block Q of the matrix Q. i , divide Q into blocks i Load from HBM into SRAM.

[0058] 4. Execute the inner loop to traverse the blocks of the key matrix K and the value matrix V, and combine the blocks of the key matrix K and the value matrix V. j and V j Load from HBM into SRAM, and compute O on SRAM. i (j) =softmax(Q) i *Kj )*V j This formula represents Qi*K j The similarity blocks are obtained, and then the similarity blocks are normalized using the softmax function to obtain the attention weight blocks. Finally, the attention weight blocks and V are calculated. j The product of O is obtained i (j) .

[0059] 5. On the SRAM, for each O i (j) Fusion yields O i and O i Output to HBM, continue executing the outer loop until the last block of the query matrix Q is traversed.

[0060] 6. On HBM, each O i The concatenation result is the final output O of the attention mechanism.

[0061] The query matrix Q, key matrix K, and value matrix V are three matrices obtained by performing a linear transformation on the input sequence. Assuming the length of the input sequence is N, that is, the input sequence contains N tokens, if each token is represented by a d-dimensional vector, then the size of the query matrix Q is N×d, the size of the key matrix K is d×N, the size of the value matrix is ​​N×d, and the size of the corresponding final output O is N×d.

[0062] Typically, in autoregressive language modeling, the attention mechanism requires applying a causal mask to the matrix operation result S = Q*K of the query matrix Q and the key matrix K. This means setting the result where the column index is greater than the row index in S to negative infinity, so that the result set to negative infinity becomes 0 during normalization. This masks the result that is not needed in subsequent calculations, ensuring that no token is affected by subsequent tokens.

[0063] The application of causal masks in the above block-based computation process, such as Figure 2 As shown, the entire matrix is ​​the result S of matrix operations Q*K (i.e., the intermediate result of attention calculation), located on the diagonal ( Figure 2 The white blocks below the dashed line are all blocks that need to be calculated. The parts above the diagonal are parts that are skipped in the calculation. The white blocks on the diagonal need to be masked by using a causal mask, that is, setting the column index of the white block to negative infinity to block out the part that exceeds the diagonal.

[0064] Depend on Figure 2It can be seen that although setting the result of the column index being greater than the row index in the white block to negative infinity has a masking effect, since attention calculation is performed by matrix operations, this part of the result always participates in the entire attention calculation process, and therefore this part of the result is an invalid calculation.

[0065] To address the aforementioned technical issues, this application proposes an attention-based model inference method. By determining that the query block and key block requiring computation meet a first preset condition, the query block and key block are divided into smaller segments. This avoids using causal masks to mask unnecessary intermediate results of the query block and key block. Then, at least one attention weight block is determined based on these smaller segments. These attention weight blocks all contain the results needed for subsequent computation. Based on the value blocks corresponding to at least one attention weight block and the key block, the output results corresponding to the query block and key block are determined. Since dividing the query block and key block into smaller segments only requires computation of the attention weight blocks containing the results needed for subsequent computation, the causal mask effectively masks unnecessary results. This not only reduces invalid computation of unnecessary intermediate results but also reduces computational resource consumption, thus improving the overall inference speed of the model.

[0066] The technical solution of this application and how it solves the aforementioned technical problems are described in detail below with specific embodiments. The listed specific embodiments can be combined with each other, and the same or similar concepts or processes may not be described again in some embodiments. The embodiments of this application will be described in detail below with reference to the accompanying drawings.

[0067] Example 1

[0068] Figure 3 This is a flowchart illustrating an embodiment of an attention-based model inference method according to an exemplary embodiment, such as... Figure 3 As shown, it includes the following steps 310-330:

[0069] Step 310: If the query block and key block to be operated on meet the first preset condition, the query block and the key block are split separately. The first preset condition indicates that the matrix multiplication result of the query block and the key block needs to be masked by a causal mask.

[0070] Step 320: Based on the query blocks and key blocks obtained from the segmentation, determine at least one attention weight block.

[0071] Step 330: Based on at least one attention weight block and the value block corresponding to the key block, determine the output results corresponding to the query block and the key block.

[0072] In this embodiment, the query block, key block, and value block refer to the query matrix, key matrix, and value matrix of the attention mechanism, respectively. These three matrices are intermediate data obtained based on the model's input data, specifically through a linear transformation of the input data.

[0073] In other words, after splitting the query matrix, multiple query blocks are obtained; after splitting the key matrix, multiple key blocks are obtained; and after splitting the value matrix, multiple value blocks are obtained, with the number of key blocks being the same as the number of value blocks.

[0074] It should be noted that, based on the existing partitioning principles of query matrix, key matrix, and value matrix, the sequence length of the key block is the same as the sequence length of the value block, while the sequence length of the query block can be different from the sequence length of the key block.

[0075] The query blocks and key blocks that need to be processed can be understood as combinations of query blocks and key blocks that participate in attention calculations, as described above. Figure 2 As shown, the combination of query block Qi and key block Kj corresponding to the white blocks located below the diagonal and on the diagonal both need to be calculated.

[0076] The first precondition can be understood as a condition that requires further segmentation of the query block and the key block, and it is used to characterize the matrix multiplication result of the query block and the key block, which requires the use of a causal mask.

[0077] In other words, when setting this condition, it is necessary to ensure that the column index and row index of the matrix obtained by multiplying the query block and the key block are equal. As mentioned above... Figure 2 In the middle, the query block and key block corresponding to the white block located on the diagonal meet the first preset condition.

[0078] Query partitioning is the result of splitting query blocks. The splitting direction of the query blocks is consistent with the splitting direction of the query matrix, that is, splitting by row.

[0079] Key block partitioning is the result of splitting key blocks. The splitting direction of the key blocks is consistent with the splitting direction of the key matrix, that is, splitting along the column direction.

[0080] Attention weight blocks can be understood as the result calculated from a query block and a key block, representing the dependencies between partial elements of the input sequence contained in the query block.

[0081] Attention weight blocks are determined based on the query blocks and key blocks obtained from further segmentation. For combinations of query blocks and key blocks corresponding to intermediate results that need to be blocked, the calculation is skipped. Therefore, the determined attention weight blocks all contain the results required for subsequent calculations.

[0082] For example, such as Figure 4As shown, the query block Q3 and key block K3 that need to be processed meet the first preset condition. The query block Q3 is divided into two query blocks Q31 and Q32, and the key block K3 is divided into two key blocks K31 and K32. Based on these query blocks and key blocks, three attention weight blocks can be determined.

[0083] The value block corresponding to the key block can be understood as a key block and a value block with the same index. The index of the key block indicates the position of the key block on the key matrix, and the index of the value block indicates the position of the value block on the value matrix.

[0084] For example, the above Figure 1 The inner loop execution flow in the code will repeatedly retrieve the K at the same index. j and V j The computation is performed by loading from HBM into SRAM, where "j" is the index of the key block and the value block.

[0085] The output can be understood as the result of the operation corresponding to the combination of the query block and the key block that need to be operated on.

[0086] This completes the above. Figure 3 The attention-based model inference process illustrated involves dividing the query and key blocks into smaller segments if the query and key blocks meet a first preset condition. This avoids using causal masks to mask unnecessary intermediate results. Then, at least one attention weight block is determined based on these smaller segments. These attention weight blocks contain the results needed for subsequent computations. Finally, the output result corresponding to the combination of query and key blocks is determined based on the value blocks corresponding to at least one attention weight block and the key block. Since dividing the query and key blocks into smaller segments only requires computing the attention weight blocks containing the results needed for subsequent computations, causal masks effectively mask unnecessary results. This not only reduces invalid computations of unnecessary intermediate results but also reduces computational resource consumption, thus improving the overall inference speed of the model.

[0087] In one optional embodiment of this application, the determination in step 310 that the query block and key block to be calculated meet the first preset condition includes the following specific steps:

[0088] For each query block included in the query matrix, based on the index of the query block, determine the set of key blocks to be operated on in each key block included in the key matrix. For each key block in the key block set, if the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block, determine that the query block and the key block meet the first preset condition.

[0089] The index of a query block can be understood as the number of the query block in the query matrix, which indicates its position in the query matrix.

[0090] A key block set can be understood as a set of key blocks that need to be operated on with the query block to be processed, and the set contains at least one key block.

[0091] Based on the index of the query block, determine the set of key blocks that need to be operated on from among the key blocks included in the key matrix. Specifically, determine the maximum index of the key blocks that need to be operated on based on the index of the query block, and then determine the set of key blocks that need to be operated on based on the maximum index.

[0092] The specific formula for calculating the maximum index is as follows:

[0093]

[0094] Among them, T c ' represents the maximum index of the key block to be operated on, i represents the index of the query block, and B r B represents the sequence length of the query block. c This represents the sequence length of the key blocks. Therefore, the set of key blocks to be operated on includes key blocks with indices j≤T. c ',Right now

[0095] If the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block, it means that the column index and row index of the matrix obtained by multiplying the query matrix and the key matrix are equal, indicating that the matrix multiplication result of the query block and the key block needs to use a causal mask.

[0096] As mentioned above Figure 2 As shown, if the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block, then the white block corresponding to the matrix multiplication result of the query block and the key block is located on the diagonal. In other words, by comparing the product of the index of the query block and the sequence length of the query block with the product of the index of the key block and the sequence length of the key block, the white block located on the diagonal can be identified.

[0097] It is understood that the above-described method for determining the set of key blocks that need to be operated on is merely an illustrative example and does not constitute a limitation on the solution of this application. For example, the set of key blocks that need to be operated on can be pre-defined for each query block.

[0098] In this embodiment, by determining the set of key blocks that need to be operated on for the query block to be processed, and then when traversing each key block in the set, the product of the index and sequence length of the query block and the currently traversed key block is used to determine whether the query block and the currently traversed key block meet the conditions. Thus, the combination of query block and key block can be determined through simple product calculation, so as to divide the query block and key block into smaller blocks and speed up the inference speed of the model.

[0099] In one optional embodiment of this application, the above step 310, which involves dividing the query block and the key block, includes: dividing the query block and the key block evenly, and ensuring that the number of query blocks obtained from the division is the same as the number of key blocks obtained from the division.

[0100] In other words, both the query block and the key block are divided equally to obtain query blocks and key blocks of the same size, and the number of query blocks and key blocks is the same. This makes it easier to filter the combination of query blocks and key blocks that need to be processed later.

[0101] For example, as described above Figure 4 As shown, the query block Q3 and key block K3 that need to be calculated meet the first preset condition. The query block Q3 is evenly divided into two query blocks Q31 and Q32, and the key block K3 is also evenly divided into two key blocks K31 and K32.

[0102] Depend on Figure 4 It can be seen that the intersection of the extensions of the dividing lines of the two query blocks and the dividing lines of the two key blocks is located on the diagonal of the matrix multiplication block of query block Q3 and key block K3, which is beneficial for filtering the combination of query blocks and key blocks to be used in subsequent operations.

[0103] It is understood that the above-described method of dividing query blocks and key blocks is only a preferred embodiment. This application can also divide key blocks and query blocks unevenly, or divide key blocks and query blocks evenly but with different numbers of blocks. As long as the key blocks and query blocks are divided into smaller blocks, it will help reduce invalid calculations and thus speed up model inference.

[0104] In one optional embodiment of this application, step 320 above, which determines at least one attention weight block based on the segmented query blocks and the segmented key blocks, includes the following specific steps:

[0105] For each query block obtained from the segmentation, the set of key blocks to be operated on is determined in each key block obtained from the segmentation based on the index of the query block. The index of the query block indicates the position of the query block in its respective query block. For each key block in the set of key blocks, the attention weight block corresponding to the query block and the key block is determined based on the query block and the key block.

[0106] The index of a query block can be understood as the number of the query block within the query block, indicating its position within the query block.

[0107] A key block set can be understood as a set of key blocks that need to be operated on with the query block to be processed, and the set contains at least one key block.

[0108] Based on the index of the query block, the principle of determining the set of key blocks that need to be operated on in each key block obtained by splitting can be the same as the principle of determining the set of key blocks that need to be operated on, and this application will not elaborate on this further.

[0109] Attention weight blocks can be understood as the normalized result of the product of a query block and a key block, which represents the dependency between partial elements of the input sequence contained in the query block.

[0110] It should be noted that by combining the query block with each key block in the corresponding key block set and calculating the attention weight block for each combination, these attention weight blocks all contain the results required for subsequent calculations, that is, the attention weight blocks all contain the valid results for subsequent calculations.

[0111] For example, as described above Figure 4 As shown, the set of key blocks that need to be processed for query block Q31 includes key block K31, and the set of key blocks that need to be processed for query block Q32 includes two key blocks K31 and K32. Based on query block Q31 and the corresponding set of key blocks, we can determine the weight attention block A corresponding to Q31 and K31. Based on query block Q32 and the corresponding set of key blocks, we can determine the weight attention block B corresponding to Q32 and K31, and the weight attention block C corresponding to Q32 and K32.

[0112] From the above Figure 4 It can be seen that the weight attention blocks corresponding to Q32 and K32 do not need to be calculated, because they are results that need to be masked. This saves the computational resources required to perform causal masking and reduces unnecessary invalid calculations.

[0113] In this embodiment, by determining the set of key blocks that need to be calculated for the query block, and then traversing each key block in the set to calculate the attention weight block of the query block and the currently traversed key block, it is not necessary to calculate the attention weight block of the query block and all key blocks. This reduces unnecessary invalid calculations and speeds up model inference.

[0114] In one optional embodiment of this application, the process of determining the attention weight blocks corresponding to the query block and the key block based on the query block and the key block includes the following specific steps:

[0115] Matrix multiplication is performed on the query block and the key block to obtain a similarity block. If the query block and the key block meet the second preset condition, causal masking is performed on the similarity block, and the processed similarity block is normalized to obtain an attention weight block. If the query block and the key block do not meet the second preset condition, the similarity block is normalized to obtain an attention weight block.

[0116] Similarity segmentation can be understood as the product of query segmentation and key segmentation, where each element represents the attention score between the two sequences.

[0117] The second precondition can be understood as requiring the use of a causal mask in the matrix multiplication result of the query block and the key block. In other words, when setting this condition, it is necessary to ensure that the column and row indices of the matrix obtained by multiplying the query block and the key block are equal.

[0118] Similar to the first preset condition mentioned above, the second preset condition can be that the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block.

[0119] It should be noted that the above description of the second preset condition is merely an example. This application does not limit the specific setting method, as long as it ensures that the column index and row index on the matrix obtained by multiplying the query block and the key block are equal.

[0120] Causal masking can be understood as a function operation that fills a specified region in a matrix with a specific value. In this embodiment, causal masking means setting any item in the matrix with a column index greater than the row index to negative infinity, that is, setting the items located in the upper right part of the matrix diagonal to negative infinity, so that the items set to negative infinity become 0 during normalization. These masked items will not affect subsequent calculations.

[0121] The similarity block is processed by causal masking. Specifically, a mask block of the same size as the similarity block is created, and then the elements above the diagonal of the similarity block are filled with negative infinity using this mask block.

[0122] Normalization can be understood as scaling the elements of a matrix in a certain way so that the entire matrix meets specific requirements. For example, normalization can be performed using methods commonly employed by attention mechanisms, such as softmax or safe softmax, and this application does not specifically limit the methods used.

[0123] For example, as described above Figure 4 As shown, after calculating the similarity block between query block Q31 and key block K31, since query block Q31 and key block K31 meet the second preset condition, it is necessary to perform causal masking on this similarity block and then normalize it to obtain the weight attention block A corresponding to query block Q31 and key block K31; after calculating the similarity block between query block Q32 and key block K31, since query block Q32 and key block K31 do not meet the second preset condition, it is possible to directly perform normalization on this similarity block to obtain the weight attention block B corresponding to query block Q32 and key block K31.

[0124] In this embodiment, by determining whether the combination of query block and key block meets the second preset condition, two different calculation processes are used to obtain the corresponding attention weight block, so as to accurately ensure that any item with a column index greater than the row index in the intermediate results of query block and key block is blocked, thereby speeding up the model inference.

[0125] In one optional embodiment of this application, step 330 above, which determines the output result corresponding to the query block and the key block based on at least one attention weight block and the value block corresponding to the key block, includes the following specific steps:

[0126] The value block is segmented to obtain individual value blocks. For each query block, the value blocks to be processed are obtained from each value block based on the attention weight block corresponding to the query block. Matrix multiplication is performed on the attention weight block corresponding to the query block and the obtained value blocks to obtain the output block corresponding to the query block. The output blocks corresponding to each query block are concatenated to obtain the output results corresponding to the query block and the key block.

[0127] Value partitioning is the result of splitting value blocks. The splitting direction of the value blocks is the same as the splitting direction of the value matrix, that is, splitting by rows. It is important to note that the sequence length of the value blocks is the same as the sequence length of the key blocks.

[0128] The value blocks that need to be processed can be understood as value blocks that need to be processed in conjunction with the attention weight blocks.

[0129] Based on the attention weight block corresponding to the query block, the value block that needs to be operated on is obtained from each value block. Specifically, the value block that needs to be operated on can be obtained from each value block based on the key block corresponding to the attention weight block.

[0130] For example, such as Figure 5 As shown, the value block V3 corresponding to key block K3 is divided into two value blocks V31 and V32, which are the same number of key blocks. The attention weight block A corresponding to query block Q31 corresponds to key block K31, so the value block that needs to be processed is V31. The attention weight blocks B and C corresponding to query block Q32 correspond to key blocks K31 and K32 respectively, so the value blocks that need to be processed are V31 and V32.

[0131] Furthermore, matrix multiplication is performed on the attention weight block A and value block V31 corresponding to query block Q31 to obtain the output block corresponding to query block Q31. Similarly, matrix multiplication is performed on the attention weight blocks B and C and value blocks V31 and V32 corresponding to query block Q32 to obtain the output block corresponding to query block Q32. These two output blocks are then concatenated to obtain the output results of query block Q3 and key block K3.

[0132] In this embodiment, by dividing the query block according to the attention weight corresponding to each query block, the value blocks obtained by the segmentation are traversed to obtain the output blocks corresponding to the query block dimensions. Then, the final output results corresponding to the query block and the key block can be obtained by simply concatenating the output blocks, so as to simplify the calculation and speed up the model inference.

[0133] In an optional embodiment of this application, the above-described attention-based model inference method further includes the following specific steps:

[0134] If the query block and key block to be operated on do not meet the first preset condition, perform matrix multiplication on the query block and the key block to obtain a similarity block; normalize the similarity block to obtain an attention weight block; perform matrix multiplication on the attention weight block and the value block corresponding to the key block to obtain the output result.

[0135] The fact that the query block and key block that need to be operated on do not meet the first preset condition can be understood as the matrix multiplication result of the query block and key block not needing to use a causal mask. That is, the column index of the matrix obtained by multiplying the query matrix and key matrix is ​​less than the row index. Therefore, there is no need to further divide the query block and key block, and attention calculation can be performed directly on the query block and key block.

[0136] As mentioned above Figure 2 In the middle, the query block and key block corresponding to the white block located below the diagonal do not meet the first preset condition.

[0137] Example 2

[0138] Based on the above embodiment one, Figure 6 The flowchart illustrates an embodiment of another model inference method based on an attention mechanism according to an exemplary embodiment. It should be noted that the data that the attention mechanism needs to process includes a query matrix, a key matrix, and a value matrix. The query matrix includes multiple query blocks, the key matrix includes multiple key blocks, and the value matrix contains multiple value blocks. The query matrix, key matrix, and value matrix are intermediate data obtained based on the input data of the model.

[0139] like Figure 6 As shown, it includes the following steps 610-650:

[0140] Step 610: For each query block included in the query matrix, determine the set of key blocks that need to be operated on in each key block included in the key matrix according to the index of the query block.

[0141] Step 620: For each key block in the key block set, if the query block and the key block meet the first preset condition, the query block and the key block are segmented respectively. The first thread block determines at least one attention weight block based on each query block and each key block obtained by the segmentation, and determines the output result of the first thread block based on at least one attention weight block and the value block corresponding to the key block.

[0142] Step 630: If the query block and the key block do not meet the first preset condition, the second thread block is used to determine the output result of the second thread block based on the query block and the key block. The first preset condition indicates that the matrix multiplication result of the query block and the key block needs to use a causal mask.

[0143] Step 640: Merge the output results of the first thread block and the second thread block to obtain the output result of the query block.

[0144] Step 650: Concatenate the output results of each query block to obtain the output result of the query matrix.

[0145] In this embodiment, both the first thread block and the second thread block are the smallest units executed by the graphics processor. In a graphics processor, a thread block contains multiple threads, and the threads within a thread block can share data through shared memory. Each thread can execute instructions using different data.

[0146] The first thread block is used to process query blocks and key blocks that have undergone block changes. Since the query blocks and key blocks are divided into smaller blocks, the calculations performed on these smaller blocks are necessarily different from the calculations performed on the original whole block. Therefore, a separate thread block is needed to perform the calculations on these smaller blocks.

[0147] The second thread block is used to process query blocks and key blocks that have not undergone block changes; that is, the second thread block performs the calculation of the original whole block.

[0148] It should be noted that the first thread block and the second thread block are independent of each other, so they can be executed in parallel.

[0149] For example, as described above Figure 2 As shown, for each query block Q contained in the query matrix Q i In the corresponding row of white blocks, only the combination of the key block and query block corresponding to the last white block needs to be calculated using the first thread block; the combinations of the key block and query block corresponding to the remaining white blocks are all calculated using the second thread block.

[0150] The output results of the first thread block and the second thread block are merged to obtain the output result of the query block. Specifically, the output results of the first thread block and the second thread block are summed to obtain the output result of the query block.

[0151] It should be noted that the preferred implementation of this embodiment can be found in the relevant description in Embodiment 1, and will not be repeated here.

[0152] The scheme of this application embodiment, for each query block, after determining the set of key blocks to be operated on based on the index of the query block, distinguishes between combinations of query blocks and key blocks that meet the first preset condition and combinations of query blocks and key blocks that do not meet the first preset condition based on the key block set, and divides the combinations of query blocks and key blocks that meet the first preset condition into smaller blocks. Due to the change in the blocks, the divided blocks are not suitable to be processed in the same thread block as the original whole block. The first thread block processes these smaller blocks, and in parallel, the second thread block processes the combinations of query blocks and key blocks that do not meet the first preset condition to improve the calculation speed. Finally, the output results of the two thread blocks are merged to obtain the output result of the entire query block, thereby improving the overall inference speed of the model.

[0153] In one optional embodiment of this application, step 630 above uses a second thread block to determine the output result of the second thread block based on the query block and the key block, including the following specific steps:

[0154] The second thread block performs matrix multiplication on the query block and the key block to obtain the similarity block; the second thread block normalizes the similarity block to obtain the attention weight block; the second thread block performs matrix multiplication on the attention weight block and the value block corresponding to the key block to obtain the output of the second thread block.

[0155] It should be noted that the preferred implementation of this embodiment can be found in the relevant description in Embodiment 1, and will not be repeated here.

[0156] The execution entity in this application embodiment can be an application, service, instance, functional module in software form, virtual machine (VM), container, or cloud server, or hardware device with data processing function (such as server or terminal device) or hardware chip (such as CPU, GPU, FPGA, NPU, AI accelerator card, or DPU). The device for implementing attention-based model inference can be deployed on the computing device of the application providing the corresponding service or on a cloud computing platform providing computing power, storage, and network resources. The cloud computing platform can provide services in the following modes: IaaS (Infrastructure as a Service), PaaS (Platform as a Service), SaaS (Software as a Service), or DaaS (Data as a Service). Taking the platform providing SaaS (Software as a Service) as an example, the cloud computing platform can utilize its own computing resources to provide training for attention-based model inference models or execution of attention-based model inference modules. The specific application architecture can be built according to service requirements. For example, the platform can provide building services based on the above model to application parties or individuals using platform resources, and further invoke the above model and realize online or offline model inference based on attention mechanism based on model inference requests submitted by relevant client or server devices.

[0157] Corresponding to the aforementioned embodiments of the attention-based model inference method, this application also provides embodiments of the attention-based model inference apparatus.

[0158] Figure 7 This is a schematic diagram illustrating the structure of an attention-based model inference apparatus according to an exemplary embodiment. The apparatus is used to execute the attention-based model inference method provided in Embodiment 1 above, such as... Figure 7 As shown, the attention-based model inference device includes:

[0159] The segmentation module 710 is used to segment the query block and the key block respectively when the query block and the key block to be operated on meet the first preset condition; the first preset condition indicates that the matrix multiplication operation result of the query block and the key block needs to use a causal mask.

[0160] The weight calculation module 720 is used to determine at least one attention weight block based on each query block and each key block obtained from the segmentation; the attention weight block is the result calculated from one query block and one key block.

[0161] The first output determination module 730 is used to determine the output result corresponding to the query block and the key block based on the at least one attention weight block and the value block corresponding to the key block.

[0162] The query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0163] In an optional implementation, the segmentation module 710 is specifically used to, during the process of determining that the query blocks and key blocks to be operated on meet the first preset condition, determine the set of key blocks to be operated on in the key blocks included in the key matrix according to the index of the query block for each query block included in the query matrix; and for each key block in the key block set, determine that the query block and the key block meet the first preset condition if the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block; wherein, the index of the query block indicates the position of the query block on the query matrix, and the index of the key block indicates the position of the key block on the key matrix.

[0164] In an optional implementation, the segmentation module 710 is specifically used to uniformly segment the query block and the key block during the segmentation process; wherein the number of query blocks obtained by segmentation is consistent with the number of key blocks obtained by segmentation.

[0165] In an optional implementation, the weight calculation module 720 is specifically used to determine, for each segmented query block, the set of key blocks to be operated on in each segmented key block according to the index of the query block; the index of the query block indicates the position of the query block on the query block; and for each key block in the key block set, the attention weight block corresponding to the query block and the key block is determined based on the query block and the key block.

[0166] In an optional implementation, the weight calculation module 720 is specifically used to perform matrix multiplication on the query block and the key block to obtain a similarity block during the process of determining the attention weight block corresponding to the query block and the key block based on the query block and the key block; if the query block and the key block meet a second preset condition, perform causal masking processing on the similarity block and normalize the processed similarity block to obtain the attention weight block; if the query block and the key block do not meet the second preset condition, perform normalization processing on the similarity block to obtain the attention weight block; wherein, the second preset condition is that the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block.

[0167] In one optional implementation, the first output determining module 730 is specifically used to segment the value block to obtain various value blocks; for each query block, according to the attention weight block corresponding to the query block, to obtain the value block to be operated on from the segmented value blocks; to perform matrix multiplication operation according to the attention weight block corresponding to the query block and the obtained value block to obtain the output block corresponding to the query block; and to concatenate the output blocks corresponding to each query block to obtain the output result corresponding to the query block and the key block.

[0168] In an alternative implementation, the apparatus further includes ( Figure 7 (Not shown in the image):

[0169] The second output determination module is used to determine that if the query block and key block to be operated on do not meet the first preset condition, perform matrix multiplication on the query block and the key block to obtain a similarity block; perform normalization on the similarity block to obtain an attention weight block; and perform matrix multiplication on the attention weight block and the value block corresponding to the key block to obtain an output result.

[0170] The specific implementation process of the functions and roles of each unit in the above device can be found in the implementation process of the corresponding steps in the above method, and will not be repeated here.

[0171] Figure 8 This is a schematic diagram illustrating the structure of an attention-based model inference device according to an exemplary embodiment. This device is used to execute the attention-based model inference method provided in Embodiment 2 above, such as... Figure 8 As shown, the attention-based model inference device includes:

[0172] The determining module 810 is used to determine, based on the index of the query block, the set of key blocks that need to be operated on in the key matrix, for each query block included in the query matrix.

[0173] The first calculation module 820 is used to segment each key block in the key block set, provided that the query block and the key block meet a first preset condition. A first thread block is used to determine at least one attention weight block based on each segmented query block and each segmented key block, and to determine the output result of the first thread block based on the at least one attention weight block and the value block corresponding to the key block. The first preset condition indicates that the matrix multiplication result between the query block and the key block needs to use a causal mask. The attention weight block is the result calculated from one query block and one key block.

[0174] The second calculation module 830 is used to determine the output result of the second thread block based on the query block and the key block when the query block and the key block do not meet the first preset conditions.

[0175] The fusion module 840 is used to fuse the output results of the first thread block and the output results of the second thread block to obtain the output results of the query block.

[0176] The output module 850 is used to concatenate the output results of each query block to obtain the output result of the query matrix.

[0177] Wherein, the first thread block and the second thread block are both the smallest units executed by the graphics processor; the query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

[0178] In an optional implementation, the second calculation module 820 is specifically used to, during the process of determining the output result of the second thread block based on the query block and the key block using the second thread block, perform matrix multiplication on the query block and the key block to obtain a similarity block; perform normalization processing on the similarity block to obtain an attention weight block; and perform matrix multiplication on the attention weight block and the value block corresponding to the key block to obtain the output result of the second thread block.

[0179] The specific implementation process of the functions and roles of each unit in the above device can be found in the implementation process of the corresponding steps in the above method, and will not be repeated here.

[0180] For the device embodiments, since they basically correspond to the method embodiments, the relevant parts can be referred to in the description of the method embodiments. The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this application according to actual needs. Those skilled in the art can understand and implement this without creative effort.

[0181] This application also provides an electronic device corresponding to the attention-based model reasoning method provided in the foregoing embodiments, for executing the aforementioned attention-based model reasoning method.

[0182] Figure 9 The diagram illustrates a hardware structure of an electronic device according to an exemplary embodiment. The electronic device includes a communication interface 601, a processor 602, a memory 603, and a bus 604. The communication interface 601, processor 602, and memory 603 communicate with each other via the bus 604. The processor 602 can execute the attention-based model inference method described above by reading and executing machine-executable instructions in the memory 603 corresponding to the control logic of the attention-based model inference method. The specific details of this method are described in the above embodiments and will not be repeated here.

[0183] The memory 603 mentioned in this application can be any electronic, magnetic, optical, or other physical storage system, and can contain stored information such as executable instructions, data, etc. Specifically, the memory 603 can be RAM (Random Access Memory), flash memory, storage drive (such as hard disk drive), any type of storage disk (such as optical disc, DVD, etc.), or similar storage media, or combinations thereof. Communication between this system network element and at least one other network element is achieved through at least one communication interface 601 (which can be wired or wireless), and the Internet, wide area network, local area network, metropolitan area network, etc., can be used.

[0184] Bus 604 can be an ISA bus, PCI bus, or EISA bus, etc. The bus can be divided into an address bus, a data bus, a control bus, etc. The memory 603 is used to store programs, and the processor 602 executes the programs after receiving execution instructions.

[0185] Processor 602 may be an integrated circuit chip with signal processing capabilities. In implementation, each step of the above method can be completed by the integrated logic circuitry in the hardware of processor 602 or by instructions in software form. The processor 602 can be a general-purpose processor, including a network processor (NP), digital signal processor (DSP), application-specific integrated circuit (ASIC), field-programmable gate array (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. It can implement or execute the methods, steps, and logic block diagrams disclosed in the embodiments of this application. The general-purpose processor can be a microprocessor or any conventional processor. The steps of the methods disclosed in the embodiments of this application can be directly manifested as execution by a hardware decoding processor, or execution by a combination of hardware and software modules in the decoding processor.

[0186] The electronic device provided in this application embodiment and the model reasoning method based on the attention mechanism provided in this application embodiment are based on the same inventive concept and have the same beneficial effects as the methods they adopt, operate or implement.

[0187] This application also provides a computer-readable storage medium corresponding to the attention-based model inference method provided in the foregoing embodiments. Please refer to [link / reference]. Figure 10 As shown, the computer-readable storage medium is an optical disc 30, on which a computer program (i.e., a program product) is stored. When the computer program is run by a processor, it executes the model inference method based on the attention mechanism provided in any of the aforementioned embodiments.

[0188] It should be noted that examples of the computer-readable storage medium may also include, but are not limited to, phase-change memory (PRAM), static random access memory (SRAM), dynamic random access memory (DRAM), other types of random access memory (RAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other optical and magnetic storage media, which will not be elaborated here.

[0189] The computer-readable storage medium provided in the above embodiments of this application and the model reasoning method based on attention mechanism provided in the embodiments of this application are based on the same inventive concept and have the same beneficial effects as the methods adopted, run or implemented by the applications stored therein.

[0190] This application also provides a computer program product corresponding to the attention-based model reasoning method provided in the foregoing embodiments. The computer program product includes a computer program that is executed by a processor to implement the attention-based model reasoning method provided in the foregoing embodiments.

[0191] The computer program products provided in the above embodiments of this application and the model reasoning method based on attention mechanism provided in the embodiments of this application are based on the same inventive concept and have the same beneficial effects as the methods adopted, run or implemented by the applications stored therein.

[0192] Other embodiments of this application will readily occur to those skilled in the art upon consideration of the specification and practice of the invention disclosed herein. This application is intended to cover any variations, uses, or adaptations of this application that follow the general principles of this application and include common knowledge or customary techniques in the art not disclosed herein. The specification and examples are to be considered exemplary only, and the true scope and spirit of this application are indicated by the following claims.

[0193] It should also be noted that the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitation, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element.

[0194] The above description is merely a preferred embodiment of this application and is not intended to limit this application. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of this application should be included within the scope of protection of this application.

Claims

1. A model inference method based on an attention mechanism, characterized in that, The method includes: If the query block and key block to be operated on meet the first preset condition, the query block and the key block are respectively segmented; the first preset condition indicates that the matrix multiplication result of the query block and the key block needs to be masked by a causal mask. Based on the segmented query blocks and key blocks, at least one attention weight block is determined; the attention weight block is the result calculated from one query block and one key block. Based on the at least one attention weight block and the value block corresponding to the key block, determine the output result corresponding to the query block and the key block; The query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

2. The method according to claim 1, characterized in that, The determination that the query block and key block to be processed meet the first preset condition includes: For each query block included in the query matrix, the set of key blocks that need to be operated on is determined from each key block included in the key matrix according to the index of the query block; For each key block in the key block set, if the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block, then the query block and the key block are determined to meet the first preset condition. The index of the query block indicates the position of the query block on the query matrix, and the index of the key block indicates the position of the key block on the key matrix.

3. The method according to claim 1, characterized in that, The step of splitting the query block and the key block respectively includes: The query block and the key block are each uniformly divided; The number of query blocks obtained from the split is consistent with the number of key blocks obtained from the split.

4. The method according to claim 1, characterized in that, The step of determining at least one attention weight block based on the segmented query blocks and key blocks includes: For each query block obtained from the segmentation, the set of key blocks that need to be operated on is determined in each key block obtained from the segmentation based on the index of the query block; the index of the query block indicates the position of the query block in the query block; For each key block in the key block set, based on the query block and the key block, determine the attention weight block corresponding to the query block and the key block.

5. The method according to claim 4, characterized in that, The step of determining the attention weight blocks corresponding to the query block and the key block based on the query block and the key block includes: Perform matrix multiplication on the query block and the key block to obtain the similarity block; If the query block and the key block meet the second preset condition, the similarity block is subjected to causal masking processing, and the processed similarity block is normalized to obtain the attention weight block. If the query block and the key block do not meet the second preset condition, the similarity block is normalized to obtain the attention weight block; The second preset condition is that the product of the index of the query block and the sequence length of the query block is equal to the product of the index of the key block and the sequence length of the key block.

6. The method according to claim 1, characterized in that, The step of determining the output result corresponding to the query block and the key block based on the at least one attention weight block and the value block corresponding to the key block includes: The value block is divided into individual value sub-blocks; For each query block, the value block to be processed is obtained from each value block obtained by segmentation, based on the attention weight block corresponding to the query block; Based on the attention weight block and the obtained value block corresponding to the query block, perform matrix multiplication to obtain the output block corresponding to the query block; The output blocks corresponding to each query block are concatenated to obtain the output results corresponding to the query block and the key block.

7. The method according to any one of claims 1-6, characterized in that, The method further includes: If the query block and key block to be operated on do not meet the first preset condition, matrix multiplication is performed on the query block and key block to obtain the similarity block; The similarity block is normalized to obtain the attention weight block; Perform matrix multiplication on the value blocks corresponding to the attention weight blocks and the key blocks to obtain the output results.

8. A model inference method based on an attention mechanism, characterized in that, The attention mechanism needs to process data including a query matrix, a key matrix, and a value matrix, and the method includes: For each query block included in the query matrix, the set of key blocks that need to be operated on is determined from each key block included in the key matrix according to the index of the query block; For each key block in the key block set, if the query block and the key block meet the first preset condition, the query block and the key block are respectively segmented. A first thread block determines at least one attention weight block based on each segmented query block and each segmented key block, and determines the output result of the first thread block based on the at least one attention weight block and the value block corresponding to the key block; the attention weight block is the result calculated from one query block and one key block; If the query block and the key block do not meet the first preset condition, a second thread block is used to determine the output result of the second thread block based on the query block and the key block; the first preset condition indicates that the matrix multiplication result of the query block and the key block needs to use a causal mask; The output results of the first thread block and the second thread block are merged to obtain the output result of the query block; The output results of each query block are concatenated to obtain the output result of the query matrix. Wherein, the first thread block and the second thread block are both the smallest units executed by the graphics processor; the query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

9. The method according to claim 8, characterized in that, The step of using a second thread block to determine the output result of the second thread block based on the query block and the key block includes: The second thread block performs matrix multiplication on the query block and the key block to obtain a similarity block; the second thread block then performs normalization on the similarity block to obtain an attention weight block. The second thread block performs matrix multiplication on the value blocks corresponding to the attention weight block and the key block to obtain the output result of the second thread block.

10. A model inference device based on an attention mechanism, characterized in that, The device includes: The segmentation module is used to segment the query block and the key block respectively when the query block and the key block to be operated on meet the first preset condition; the first preset condition indicates that the matrix multiplication operation result of the query block and the key block needs to use a causal mask; The weight calculation module is used to determine at least one attention weight block based on the various query blocks and key blocks obtained from the segmentation; the attention weight block is the result calculated from one query block and one key block. An output determination module is used to determine the output result corresponding to the query block and the key block based on the at least one attention weight block and the value block corresponding to the key block; The query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

11. A model inference device based on an attention mechanism, characterized in that, The attention mechanism needs to process data including a query matrix, a key matrix, and a value matrix, and the device includes: The determination module is used to determine, based on the index of the query block included in the query matrix, the set of key blocks that need to be operated on in the key matrix. A first calculation module is used to segment each key block in the key block set, provided that the query block and the key block meet a first preset condition. A first thread block is used to determine at least one attention weight block based on each segmented query block and each segmented key block, and to determine the output result of the first thread block based on the at least one attention weight block and the value block corresponding to the key block. The first preset condition indicates that the matrix multiplication result between the query block and the key block needs to use a causal mask. The attention weight block is the result calculated from one query block and one key block. The second calculation module is used to determine the output result of the second thread block based on the query block and the key block when the query block and the key block do not meet the first preset conditions. The fusion module is used to fuse the output results of the first thread block and the output results of the second thread block to obtain the output result of the query block; The output module is used to concatenate the output results of each query block to obtain the output result of the query matrix; Wherein, the first thread block and the second thread block are both the smallest units executed by the graphics processor; the query block belongs to the query matrix of the attention mechanism, the key block belongs to the key matrix of the attention mechanism, and the value block belongs to the value matrix of the attention mechanism. The query matrix, the key matrix, and the value matrix are intermediate data obtained based on the input data of the model.

12. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, The processor executes the program to implement the method as described in any one of claims 1-9.

13. A computer-readable storage medium having a computer program stored thereon, characterized in that, The program is executed by a processor to implement the method as described in any one of claims 1-9.

14. A computer program product, comprising a computer program, characterized in that, The computer program is executed by a processor to implement the method according to any one of claims 1-9.

Images