Operator execution method, electronic device, storage medium, and program product

Abstract

The present application relates to the technical field of artificial intelligence, and provides an operator execution method, an electronic device, a storage medium and a program product, the method comprising: determining the number of processing batches corresponding to each thread block based on the batch size of a query tensor; in any thread block, aggregating the query tensors of all attention heads corresponding to the number of processing batches, performing attention calculation with a key tensor and a value tensor to obtain a thread block calculation result; and integrating the thread block calculation results of each thread block to obtain an output tensor. The method provided by the present application determines the outer parallel loop by the batch size of the query tensor, changes the core operation of calculation from the originally low-efficiency vector x matrix operation to the matrix x matrix operation with higher calculation density, fully utilizes the hardware characteristics of the tensor core optimized for matrix multiplication, greatly improves the throughput and execution efficiency of the calculation unit, and further gives full play to the potential computing power of the GPU, thereby further improving the execution speed of the overall operator.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence, and in particular to an operator execution method, an electronic device, a storage medium and a program product. BACKGROUND

[0002] With the continuous evolution of deep learning technology, attention mechanisms, especially self-attention mechanisms in Transformer architecture, have become the core technology in the field of artificial intelligence such as natural language processing and computer vision. In existing optimization implementations for GPUs (Graphics Processing Unit), the mainstream solution usually adopts batch_size x head_num as the outermost parallel loop, which corresponds to the grid dimension in the GPU computing model. Under this mapping strategy, each thread block only processes attention calculation between a single query tensor and the selected key / value sequence.

[0003] However, due to the first dimension size of the query tensor being only 1, this extremely small dimension feature results in the inability of the computing mode to fully utilize the hardware characteristics of Tensor Core optimized for large-scale matrix multiplication. This mismatch between the computing load and the hardware characteristics directly leads to low utilization of the computing unit, which in turn prevents the full exploitation of the potential computing power of the GPU, limiting the improvement of overall computing efficiency. SUMMARY

[0004] The present application provides an operator execution method, an electronic device, a storage medium and a program product to solve the defect that the existing technology cannot fully utilize the hardware characteristics of Tensor Core optimized for matrix multiplication, resulting in low utilization of the computing unit and failure to fully exploit the potential computing power of the GPU.

[0005] The present application provides an operator execution method, comprising:

[0006] Obtaining an input tensor; the input tensor includes a query tensor, a key tensor and a value tensor;

[0007] Determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor;

[0008] In any thread block, aggregate the query tensors of all attention heads corresponding to the number of processing batches, and perform attention calculation with the key tensor and the value tensor to obtain the thread block calculation result;

[0009] Integrate the thread block calculation results of the thread blocks to obtain the output tensor of the input tensor.

[0010] According to the operator execution method provided by the present application, the key tensor and the value tensor are obtained by:

[0011] obtaining an initial key tensor, an initial value tensor and index information;

[0012] extracting a target subsequence corresponding to the index information from the initial key tensor and the initial value tensor to obtain a rearranged key tensor and a rearranged value tensor;

[0013] taking the rearranged key tensor and the rearranged value tensor as the key tensor and the value tensor.

[0014] According to the operator execution method provided by the application, the query tensor of all attention heads corresponding to the processing batch quantity is aggregated, and attention calculation is performed on the key tensor and the value tensor to obtain a thread block calculation result, which comprises:

[0015] The query tensor of all attention heads corresponding to the processing batch quantity is aggregated to construct a query matrix, and the number of rows of the query matrix is the total number of all attention heads.

[0016] Attention calculation is performed on the query matrix, the key tensor and the value tensor to obtain the thread block calculation result.

[0017] According to the operator execution method provided by the application, the attention calculation is performed on the query matrix, the key tensor and the value tensor to obtain the thread block calculation result, which comprises:

[0018] Matrix multiplication calculation is performed on the query matrix and the key tensor to obtain an attention score matrix.

[0019] Normalization operation is performed on the attention score matrix to obtain an attention weight matrix.

[0020] Calculation is performed on the attention weight matrix and the value tensor to obtain the thread block calculation result.

[0021] According to the operator execution method provided by the application, the number of rows of the attention score matrix is the total number of all attention heads corresponding to the processing batch quantity.

[0022] According to the operator execution method provided by the application, the thread block calculation results of the respective thread blocks are integrated to obtain an output tensor of the input tensor, which comprises:

[0023] The thread block calculation results of the respective thread blocks are integrated based on the batch index of the query tensor to obtain the output tensor.

[0024] According to the operator execution method provided by the application, the index information corresponds to the sequence dimension coordinates of the key tensor and the value tensor.

[0025] The obtaining step of the index information comprises:

[0026] calculating a correlation matrix between the query tensor and the initial key tensor;

[0027] based on the correlation matrix, selecting a preset number of positions with the highest correlation along the sequence dimension of the initial key tensor to obtain the index information.

[0028] The application further provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the operator execution method according to any one of the above when executing the program.

[0029] The application further provides a non-transitory computer-readable storage medium, which stores a computer program, wherein the computer program is executable on a processor to implement the operator execution method according to any one of the above.

[0030] The application further provides a computer program product, comprising a computer program, wherein the computer program is executable on a processor to implement the operator execution method according to any one of the above.

[0031] The operator execution method, the electronic device, the storage medium and the program product provided by the application determine the number of processing batches corresponding to each thread block through the batch size of the query tensor, aggregate the query tensors of all attention heads corresponding to the number of processing batches in the thread block, and perform attention calculation on the key tensor and the value tensor to obtain the thread block calculation result, integrate the thread block calculation results of each thread block to obtain the output tensor corresponding to the input tensor, and realize that the outer parallel loop is determined by the batch size of the query tensor, and the core operation of calculation is changed from the originally low-efficiency vector x matrix operation to the matrix x matrix operation with higher density, the hardware characteristics of Tensor Core optimized for matrix multiplication are fully utilized, the throughput and execution efficiency of the calculation unit are greatly improved, the potential computing power of the GPU is fully utilized, and the execution speed of the overall operator is further improved. BRIEF DESCRIPTION OF DRAWINGS

[0032] In order to more clearly illustrate the technical solutions in the application or prior art, the following will briefly introduce the drawings needed to be used in the embodiments or prior art description. Obviously, the drawings in the following description are some embodiments of the application, and for those skilled in the art, other drawings can also be obtained without creative labor.

[0033] Figure 1 is a structural schematic diagram of a graphics processor provided by the application;

[0034] Figure 2is a flowchart of an operator execution method provided by the present application;

[0035] Figure 3 is a flowchart of an operator execution method based on dynamic sequence selection provided by the present application;

[0036] Figure 4 is a comparison diagram between a traditional attention calculation process and an attention calculation process provided by the present application;

[0037] Figure 5 is a workflow diagram of a single thread block provided by the present application;

[0038] Figure 6 is a structural diagram of an electronic device provided by the present application. DETAILED DESCRIPTION

[0039] In order to make the objects, technical solutions and advantages of the present application clearer, the technical solutions in the present application will be described clearly and completely below with reference to the drawings in the present application. Obviously, the described embodiments are part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative labor fall within the scope of protection of the present application.

[0040] It should be noted that various artificial intelligence devices, such as GPU, GPGPU (General-purpose computing on Graphics Processing Units), TPU (Tensor Processing Unit), etc., have become the key force to promote complex computing tasks such as deep learning by virtue of their powerful parallel computing capabilities. The execution subject of the operator execution method provided by the present application can be an artificial intelligence device such as GPU, GPGPU, TPU, etc. Hereinafter, taking GPU as an example, the structure of the execution subject of the present application will be briefly introduced.

[0041] Figure 1 is a structural diagram of a graphics processing unit provided by the present application, such as Figure 1As shown, the graphics processor 100 at least includes a plurality of stream processor clusters 101 and a display memory 102, wherein each stream processor cluster 101 includes a plurality of computing units 103, and each computing unit 103 can be a tensor computing unit or a vector computing unit. Here, the tensor computing unit and the vector computing unit are heterogeneous units. The tensor computing unit is mainly used to perform tensor computation, such as matrix multiplication operation, convolution operation, etc., and includes a plurality of tensor cores for performing specific computing tasks. The vector computing unit is mainly used to perform arithmetic logic operations, such as accumulation, reduction, general addition, subtraction, multiplication, division, etc. The vector computing unit includes a plurality of vector cores for performing specific computing tasks.

[0042] Each computing unit 103 at least includes an on-chip cache 104 and a register 105. Here, the display memory 102 can be an off-chip memory HBM (High Bandwidth Memory), or other types of memory.

[0043] The on-chip cache 104 is a temporary memory, for example, the on-chip cache 104 can be a GMB (GemmMain Buffer), which has a smaller capacity than the display memory 102, but a faster data exchange speed than the display memory 102. Compared with the on-chip cache 104, the register 105 has a smaller capacity than the on-chip cache 104, but a faster data exchange speed than the on-chip cache 104. The register 105 can be a TLR (Thread Local Register).

[0044] In addition to the above structure, the graphics processor 100 in the present application can also include other structures, which are not specifically limited by the present application. Figure 1 Based on the architecture diagram of the graphics processor shown, the present application provides an operator execution method, which can be applied to various scenarios, such as text processing, image processing, voice processing, etc. In different application scenarios, the input tensor has different physical meanings. For example, in the text processing scenario, the input tensor can be a word embedding vector or a context feature vector corresponding to each word element in the text sequence. For another example, in the image processing scenario, the input tensor can be a patch feature vector after image division or a pixel point feature vector of a feature map. For another example, in the voice processing scenario, the input tensor can be an audio frame feature or a spectrogram feature vector corresponding to a voice signal. For ease of understanding, the following describes the technical solutions provided by the present application by taking the execution of an Attention operator and the input tensor being a word embedding vector corresponding to each word element in a text sequence as an example.

[0045] Specifically, in the traditional operator implementation, each attention head is usually assigned to a different thread block for processing, resulting in a smaller query tensor dimension within each thread block, and the tensor core calculation degenerates into vector-by-matrix multiplication, which cannot fully utilize the powerful computing power of the tensor core in the graphics processor optimized for matrix multiplication.

[0046] To solve the above problems, the present application provides an operator execution method to improve the utilization of the computing unit and fully utilize the potential performance of the GPU. Figure 2 is a flowchart of the operator execution method provided by the present application, as Figure 2 shown, the method comprises:

[0047] Step 210, obtaining an input tensor; the input tensor comprises a query tensor, a key tensor and a value tensor.

[0048] Here, the input tensor comprises a query tensor, a key tensor and a value tensor. Among them, the query tensor Q, the key tensor K and the value tensor V refer to three core components generated after linear transformation for attention mechanism calculation, respectively for representing query information, queried information and actual content information. The query tensor is a multi-dimensional array stacked by data of multiple independent sample sequences in memory, and its shape usually contains batch dimension, attention head dimension, sequence length dimension, etc. The query tensor Q here can be represented as (batch_size, head_num, 1, head_dim_q). Among them, batch_size represents batch size; head_num represents the number of attention heads; 1 represents that each query head corresponds to a single position in the sequence; head_dim_q represents the feature dimension of each query head, i.e. the representation ability of the query tensor. In addition, the key tensor or the value tensor can be represented as (batch_size, 1, seq, head_dim_v). Among them, batch_size represents batch size; 1 represents that all heads share the same set of key and value matrices, reducing the storage and calculation overhead of K / V; seq represents sequence length; head_dim_v represents the feature dimension of the key / value, which is usually aligned with head_dim_q.

[0049] Specifically, first, data from the upper layer network or the preprocessing module is received, which constitutes the input tensor. For example, the input tensor can be a digital representation of a batch of natural language text sequences to be processed. For example, when the batch size is N, the input tensor contains N independent sentences or paragraphs, each of which corresponds to an initial word embedding vector of a word unit in the text, which carries the original semantic features of the word unit.

[0050] Here, the input tensor is logically divided into three parts, namely a query tensor, a key tensor, and a value tensor. For example, the query tensor can be used to represent the feature state of each word unit in the current text sequence when actively seeking context-related information. For example, when the word "line" in the sentence is used as a query object, the query tensor carries features to match the context for related clues such as "bank" or "walk" to determine its specific contextual meaning. In addition, the key tensor represents the index feature state of each word unit in the text sequence as a queried object. The value tensor represents the actual content semantic information of each word unit in the text sequence.

[0051] It should be noted that in the actual implementation of the deep learning framework, the input tensor can be obtained by reading the corresponding tensor data stored in the video memory. For example, for a natural language processing model, the shape of the input usually includes batch size, number of attention heads, sequence length, and dimension of each head.

[0052] In step 220, the number of processing batches corresponding to each thread block is determined based on the batch size of the query tensor.

[0053] Here, the batch size of the query tensor refers to the number of independent sample sequences processed simultaneously in one parallel computing process. The thread block here refers to the basic execution unit in a parallel computing device such as a graphics processing unit, which includes a group of threads that can work together and share memory. Here, the number of processing batches corresponding to each thread block refers to the number of batches of query tensors processed by a single thread block.

[0054] Specifically, the batch size value in the query tensor can be extracted by analyzing the dimension information of the query tensor. Then, based on the batch size of the query tensor, the number of processing batches corresponding to each thread block can be determined. For example, the number of processing batches corresponding to each thread block can be obtained by calculating the ratio of the batch size of the query tensor to the number of thread blocks launched. In a specific embodiment, if the batch size is 6, the number of thread blocks is fixed at 3, i.e., the grid dimension is 3, then the number of processing batches corresponding to each thread block is 2.

[0055] It should be noted that the number of thread blocks launched can be pre-set or flexibly configured. In the case of flexible configuration of the number of thread blocks launched, the number of thread blocks launched can be configured to be equal to the batch size value or in a specific proportional relationship. For example, if the batch size is N, the dimension of the parallel computing grid is configured to be N rows and 1 column, i.e., N thread blocks are launched.

[0056] It should be noted that determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor means that each thread block will be responsible for processing all the computing tasks of the number of batch samples. Compared with the product of the batch size and the number of attention heads as the start-up number in the traditional technology, this step greatly changes the granularity of task allocation, raises the level of parallelism to the batch dimension, to realize higher density matrix operations within the thread block, and fully exert the computing power potential of the computing unit.

[0057] At step 230, within any thread block, the query tensors of all attention heads corresponding to the number of processing batches are aggregated, and the attention calculation is performed with the key tensor and the value tensor to obtain the thread block calculation result.

[0058] Here, the attention head refers to an independent attention calculation channel in the multi-head attention mechanism, which is used to capture semantic associations in different subspaces. In addition, the thread block calculation result here refers to the intermediate data or local result generated after completing a specific computing task within a single thread block.

[0059] It should be noted that within each started thread block, the corresponding number of processing batches of sample data are processed. At this time, a single thread block no longer processes a single attention head, but loads the query tensors of all attention heads under the number of processing batches as a whole, and uses multiple threads within the thread block to load the data slices of the query tensors corresponding to all attention heads under the number of processing batches in parallel.

[0060] Specifically, the query tensors of all attention heads corresponding to the number of processing batches can be aggregated, that is, the shape of the query tensor can be regarded as a two-dimensional matrix composed of the number of attention heads and the head dimension, such as a matrix with a shape of number of processing batches It can be understood that the aggregated query tensor can be considered to contain the query feature set of all word pieces in the sentence in all attention heads under the number of processing batches.

[0061] Subsequently, the aggregated query tensor matrix and the corresponding key tensor matrix are subjected to matrix multiplication to obtain an attention score matrix. The attention score matrix is then subjected to normalization and matrix multiplication with the value tensor matrix. In the text processing scenario, the thread block calculation result can be used to represent the local context semantic representation of the batch number of sentences processed by the attention mechanism. Specifically, the thread block calculation result can be considered as a vector set of all word units in the batch number of sentences after fusing the information of other word units in the sentence, eliminating the ambiguity of polysemous words, and enhancing the semantic association. For example, for the sentence "apple is rich in vitamins", the vector representation of "apple" after the thread block calculation has integrated the fruit semantic features brought by "vitamins", rather than the semantic features of technology companies.

[0062] It should be noted that, since all attention heads are aggregated in the thread block, the core operation in the calculation process is changed to matrix multiplication with a larger scale, rather than vector multiplication with a matrix. After the calculation is completed, the thread block obtains the calculation results of all attention heads for the batch number of samples, i.e., the thread block calculation result.

[0063] It should be further noted that, by aggregating all attention heads in the thread block, the core operation of the calculation is changed from the originally low-efficiency vector x matrix operation to the higher-density matrix x matrix operation. The matrix multiplication of this scale can make more full use of the hardware characteristics of Tensor Core optimized for matrix multiplication, improve the utilization of the calculation unit, and thus fully exploit the potential computing power of the GPU. In addition, the matrix multiplication of this scale can also perfectly match the instruction characteristics of the tensor acceleration unit Tensor Core in the hardware such as GPU, significantly improving the utilization rate and instruction emission efficiency of the hardware calculation unit.

[0064] Step 240, integrating the thread block calculation results of the respective thread blocks to obtain an output tensor of the input tensor.

[0065] Here, the output tensor refers to the final data representation processed by the attention mechanism after integrating all calculation results, which is used as the input of the subsequent neural network layer.

[0066] Specifically, when each thread block completes the attention calculation of the batch quantity of samples it is responsible for, the threads in the thread block write the thread block calculation result stored in the on-chip register or cache to the specified offset position in the memory according to the memory address mapping rule of the output tensor. It can be understood that, since each thread block is responsible for different batch indexes, the memory areas they write are mutually non-overlapping, so the write-back operation can be performed in parallel without complex atomic lock mechanism. After all the started thread blocks are executed and the write-back is completed, the complete output tensor is formed in the memory. In the text processing scenario, the output tensor can be the complete semantic matrix of all text sequences in the batch after self-attention feature extraction, which can be considered as the result of physically splicing the context semantic representation calculated for all sentences in the original batch order in the memory. It can be understood that the output tensor here will be used as the input of the next layer of neural network to provide feature data with deep context dependency for subsequent sentiment analysis, machine translation or text generation tasks.

[0067] The operator execution method provided by the embodiment of the present application determines the processing batch quantity corresponding to each thread block through the batch size of the query tensor, aggregates all attention heads of the query tensor corresponding to the processing batch quantity in the thread block, performs attention calculation with the key tensor and the value tensor to obtain thread block calculation results, integrates the thread block calculation results of each thread block to obtain an output tensor corresponding to the input tensor, and realizes that the outer parallel loop is determined by the batch size of the query tensor, and the core operation of calculation is changed from the originally low-efficiency vector x matrix operation to the matrix x matrix operation with higher calculation density, fully utilizes the hardware characteristics of Tensor Core optimized for matrix multiplication, greatly improves the throughput and execution efficiency of the calculation unit, and further fully develops the potential computing power of the GPU, and further improves the execution speed of the overall operator.

[0068] It should be noted that the standard self-attention mechanism faces a serious calculation bottleneck when processing long sequences due to the complexity of O(n 2 ). In order to alleviate this problem, the current sparse attention method such as DSA (Dynamic Sparse Attention) aims to break through this limitation, and its core idea is to calculate on demand, not to process the full sequence seq. However, in the existing GPU implementation, a binary mask corresponding to the complete sequence length seq is usually generated according to the correlation index, that is, the invalid data is not removed, and the matrix operation is directly performed on the complete sequence length. Then, the weights of irrelevant positions are set to zero by using the mask after the calculation is completed.

[0069] It can be understood that although this way realizes sparsity in mathematical logic, it is dense in physical calculation. That is, the hardware actually performs operations on full-length sequences, resulting in a large amount of computing power being wasted on invalid data that will eventually be shielded, failing to truly exert the performance advantage of sparse algorithms. To solve this problem, based on any of the above embodiments, the obtaining step of the key tensor and the value tensor comprises:

[0070] obtaining an initial key tensor, an initial value tensor, and index information;

[0071] extracting a target sub-sequence corresponding to the index information from the initial key tensor and the initial value tensor to obtain a rearranged key tensor and a rearranged value tensor;

[0072] using the rearranged key tensor and the rearranged value tensor as the key tensor and the value tensor.

[0073] Here, the initial key tensor and the initial value tensor refer to original data that has not been screened and contains complete sequence length information. The dimensions thereof are usually related to the complete length of the input text and contain a large amount of redundant information that is not relevant to the current attention calculation.

[0074] In addition, the index information here is identification data indicating which key-value pairs at which positions are truly required to participate in the calculation in the current calculation task. The index information corresponds to the logical address or physical offset of the key tensor and the value tensor in the sequence dimension and is usually generated by a preposed correlation evaluation module. The target sub-sequence here refers to an effective data segment with high correlation selected from complete data according to the index information.

[0075] The rearranged key tensor and the rearranged value tensor here refer to compact tensors formed by reorganizing the extracted target sub-sequence according to a specific rule. The sequence dimension thereof is significantly reduced from the original complete length seq to the selected length seq_select. Thus, accurate data acquisition before calculation can be performed to avoid subsequent invalid calculation on invalid data.

[0076] Specifically, first, the original data, i.e., the initial key tensor and the initial value tensor, can be read from the video memory, which is pre-stored or output by the previous layer network. It should be noted that the two tensors correspond to the complete context window size in the sequence dimension, for example, when processing long documents, the dimension can be very large. At the same time, an index tensor is received as index information, which indicates which key-value pairs at which positions in the complete sequence are relevant to the current query tensor. The index information is usually generated based on a certain sparsity strategy and contains the physical address or logical coordinates of the key-value pairs required for each query in the original sequence.

[0077] Then, the parallel computing device can access the memory address space of the initial key tensor and the initial value tensor discontinuously according to the coordinates recorded in the index information. The specific vector data pointed to by the index is copied out to form a target subsequence. Subsequently, the target subsequences extracted in a scattered manner are stored continuously in the memory to construct a new rearranged key tensor and a rearranged value tensor.

[0078] It should be noted that the shapes of the rearranged key tensor and the rearranged value tensor obtained at this time have undergone substantial changes, and the sequence dimension is no longer the original complete length, but is compressed to the selected length defined by the index information, which can be denoted as seq select. For example, if the original sequence length is 4096 and the index selects only the first 32 most relevant tokens, the rearranged tensor has only 32 in the sequence dimension.

[0079] Finally, the compressed and rearranged tensor can be formally defined as the key tensor and the value tensor required for subsequent attention calculation steps. That is, the scale of the input data of the matrix multiplication operation performed inside the subsequent thread block has been greatly reduced. The subsequent attention mechanism calculation will directly act on this compact rearranged tensor, rather than acting on the original large tensor and relying on a mask to mask invalid values.

[0080] Compared with the prior art, which usually performs matrix operations on the complete sequence length and then masks the irrelevant part through a mask, this results in a large amount of computing power being consumed by the computing unit to process invalid data that will eventually be set to zero, and a huge memory bandwidth being occupied. The key tensor and value tensor acquisition method provided in the embodiments of the present application extracts the target subsequence corresponding to the index information from the initial key tensor and the initial value tensor to obtain a rearranged key tensor and a rearranged value tensor. The rearranged key tensor and the rearranged value tensor are used as the key tensor and the value tensor, the amount of data participating in the operation is reduced from the O(seq) level to the O(seq select) level, the amount of calculation of the subsequent matrix multiplication operation is reduced by an order of magnitude, and the execution speed of the operator is greatly improved. In particular, when processing an ultra-long text sequence, this performance improvement is particularly obvious, and the technical problem of wasting T computing power in the traditional sparse attention implementation is fundamentally solved.

[0081] Based on any of the above embodiments, in step 230, the query tensors of all attention heads corresponding to the processing batch number are aggregated, and attention calculation is performed on the key tensor and the value tensor to obtain a thread block calculation result, including:

[0082] The query tensors of all attention heads corresponding to the processing batch number are aggregated to construct a query matrix, and the number of rows of the query matrix is the total number of the all attention heads.

[0083] Attention is calculated based on the query matrix, the key tensor, and the value tensor to obtain the thread block calculation result.

[0084] Here, the query matrix refers to a two-dimensional data structure constructed during the computation process, which logically stacks the query tensors of the various attention heads that were originally processed separately or independently. The number of rows in the query matrix corresponds to the total number of all attention heads for the number of processing batches.

[0085] Specifically, within each thread block responsible for processing a batch of samples, the first step is to read the query tensors of all attention heads for that batch size. Physically, these query tensors are typically stored contiguously or accessible with a specific stride. Logically, instead of treating each attention head's query tensor as an independent vector, a tensor-view transformation operation is used to concatenate all query tensors for that batch size along the row dimension, thus constructing a tensor with a shape equal to the batch size. The query matrix is ​​(head_num, head_dim_q). This embodiment of the invention uses a batch size of 1 as an example. Subsequent embodiments will use the batch size of the query vector and the number of thread blocks to be started, i.e., each thread block processes one batch of data for illustration. For example, if the model has 12 attention heads, each with a dimension of 64, and the batch size is 1, the constructed query matrix will be a 12-row, 64-column two-dimensional matrix.

[0086] Then, matrix multiplication can be performed using the constructed query matrix and the key tensor. Here, the key tensor is typically considered as a matrix with a shape equal to the number of processing batches. A matrix of [head dimension, sequence length] is generated. By performing the operation of query matrix × key tensor, the attention score matrix for all attention heads on the sequence is obtained at once, with the shape of [total number of all attention heads corresponding to the batch size, sequence length]. Subsequently, this score matrix is ​​Softmax normalized, and then the result is multiplied by the value tensor a second time to obtain the output vector containing all attention heads in the batch after attention weighting, which is the thread block computation result. Thus, the entire computation process for obtaining the computation result of any thread block mainly consists of two large-scale general matrix multiplications.

[0087] It should be noted that in modern parallel computing hardware such as a graphics processing unit (GPU), there are hardware units dedicated to accelerating matrix multiplication, such as Tensor Core, which have certain shape requirements for the dimensions of the input matrix, for example, the M, N, and K dimensions must be multiples of 16 or 8. It should be noted that if each head is calculated separately, the M dimension is usually 1, and only a less efficient vector matrix multiplication instruction can be used, and the high-performance matrix calculation unit cannot be activated.

[0088] The method provided by the embodiments of the present application directly changes the arithmetic intensity of the underlying calculation by aggregating the query tensors of multiple attention heads into a query matrix, significantly increases the M dimension in matrix multiplication, i.e., the total number of attention heads, so that it can perfectly match the input requirements of the high-performance calculation unit, thereby greatly improving the floating-point operations per second of the operator, significantly reducing the calculation delay, and reducing the calculation waste.

[0089] Based on any of the above embodiments, the attention calculation based on the query matrix, the key tensor, and the value tensor is performed to obtain the thread block calculation result, including:

[0090] The matrix multiplication calculation based on the query matrix and the key tensor is performed to obtain an attention score matrix.

[0091] The normalization operation is performed on the attention score matrix to obtain an attention weight matrix.

[0092] The calculation based on the attention weight matrix and the value tensor is performed to obtain the thread block calculation result.

[0093] Based on any of the above embodiments, the number of rows of the attention score matrix is the total number of all attention heads corresponding to the processing batch.

[0094] Here, the attention score matrix refers to an intermediate matrix generated by calculating the similarity between the query tensor and the key vector in the attention mechanism, and the numerical value represents the attention degree of the query to different position information. The attention weight matrix here is the output result of the normalization operation, which directly determines the proportion of the value vector of each position in the final information aggregation.

[0095] Specifically, first, a first matrix multiplication operation can be performed by taking the query matrix of all attention heads as the multiplicand and the key tensor as the multiplier. It should be noted that since each row of the query matrix represents the query tensor of an attention head and each column of the key tensor represents the key vector of a position in the sequence, the inner product results of all attention heads for each position in the sequence in the processing batch quantity are calculated in parallel through this matrix multiplication. Thus, the number of rows of the generated attention score matrix corresponds to the total number of all attention heads corresponding to the processing batch quantity, and the number of columns corresponds to the sequence length. It can be understood that in the embodiment of the present application, the number of columns can be the length of the target subsequence.

[0096] It should be noted that this matrix multiplication operation fully utilizes the matrix operation instructions in the GPU and completes the similarity calculation of all heads and all word positions corresponding to the processing batch quantity at one time.

[0097] Then, each row of the attention score matrix is processed. Since each row corresponds to the scoring of a specific attention head on the entire sequence, a nonlinear activation function such as the Softmax function is applied in this dimension. In specific calculations, the maximum value in the row is usually subtracted to ensure numerical stability, then the exponent is calculated, and finally the sum of all exponents in the row is divided. A scaling factor can also be introduced in this step as needed, which is usually the inverse of the square root of the head dimension, to prevent the gradient from disappearing due to the large dot product result. After this normalization processing, the original score matrix is converted into an attention weight matrix with values between 0 and 1, and the shape remains unchanged.

[0098] Finally, a second matrix multiplication operation is performed by taking the generated attention weight matrix as the left multiplier and the value tensor as the right multiplier. It can be understood that the value tensor stores the actual semantic information of each position in the sequence. Through matrix multiplication, a weighted sum operation is actually performed, that is, the vectors in the value tensor are linearly combined according to the attention weight matrix. The number of rows of the output matrix obtained by calculation is still the total number of all attention heads corresponding to the processing batch quantity, and the number of columns is the dimension of the head. This output matrix constitutes the final calculation result of the processing batch quantity of samples responsible by the thread block, where each row of data is the context vector extracted by the attention mechanism for the corresponding attention head.

[0099] The method provided by the embodiment of the application can utilize the large-dimension matrix form constructed to maintain a high calculation density and pipeline saturation at the hardware level, thereby improving the utilization of the calculation unit and fully exerting the potential computing power of the image processor after optimization.

[0100] Based on any of the above embodiments, step 240 comprises:

[0101] Based on the batch index of the query tensor, the thread block calculation results of the thread blocks are integrated to obtain the output tensor.

[0102] Specifically, after obtaining the thread block calculation results of the thread blocks, the threads in the thread blocks write the thread block calculation results stored in the on-chip register or cache to the specified offset position in the display memory according to the memory address mapping rule of the output tensor. It can be understood that, since each thread block is responsible for a batch index corresponding to a batch quantity, the memory areas written by them are mutually non-overlapping. That is, the thread block calculation results of each thread block are written to the corresponding memory area according to the batch index corresponding to the batch quantity responsible by each thread block. After all the started thread blocks are executed and the write-back is completed, the complete output tensor is formed in the display memory.

[0103] Based on any of the above embodiments, the index information corresponds to the sequence dimension coordinates of the key tensor and the value tensor;

[0104] The obtaining step of the index information comprises:

[0105] The correlation degree matrix between the query tensor and the initial key tensor is calculated.

[0106] Based on the correlation degree matrix, a preset number of positions with the highest correlation degree are selected along the sequence dimension of the initial key tensor to obtain the index information.

[0107] Here, the correlation degree matrix refers to a data table used to measure the rough similarity between the query and the key, which is calculated in advance before the formal weighted summation calculation, and its function is similar to that of a filter. Here, the preset number refers to the maximum number of key-value pairs that each query tensor is allowed to focus on according to the system computing power budget or model accuracy requirement, which can be regarded as a sparsity parameter.

[0108] Specifically, first, an initial similarity measurement calculation can be performed on the query tensor and the initial key tensor without rearrangement. In actual operator implementation, in order to avoid consuming too much computing power in this step, the initial key tensor here can be represented in low precision, or the calculation is only performed at a coarse granularity after down-sampling, or a standard dot product operation is used. It should be noted that the core of the calculation of the correlation matrix here is to obtain the matching score between each query tensor in the query tensor and each key vector in the initial key tensor. The correlation matrix generated by the operation usually covers the entire sequence length, and each element value of the correlation matrix represents the importance of the key-value pair at the corresponding position to the current query.

[0109] Then, the correlation matrix generated can perform a Top-k sorting or selection algorithm in the sequence dimension direction. For example, for each row corresponding to a query tensor, the correlation scores of all positions in the row are compared, and the positions with the highest scores are selected, the number of which is strictly limited by the preset number. For example, if the preset number is 64, only the indices of the top 64 positions are retained.

[0110] It can be understood that the selected index set constitutes the index information, which records the specific coordinates of the context fragments most critical to the current semantic understanding in the original long sequence, and provides an address basis for accurately extracting the target sub-sequence from the huge initial tensor subsequently.

[0111] It should be noted that the operator execution method provided by the embodiments of the present application is not only suitable for the sparse attention mechanism DSA, but also can be applied to any sparse attention model based on index selection or gating mechanism, such as BigBird, Longformer, etc., and has wide applicability.

[0112] The index information acquisition method provided by the embodiments of the present application calculates the correlation matrix between the query tensor and the initial key tensor; based on the correlation matrix, a preset number of positions with the highest correlation are selected along the sequence dimension of the initial key tensor to obtain the index information, which realizes adaptive judgment of which context information is relevant and which is redundant noise according to the specific semantic content input, dynamic selection based on content, and ensures that the subsequent rearranged key tensor and rearranged value tensor have extremely small data volume but contain the part with the highest information density in the original sequence. Therefore, not only the computational complexity is greatly reduced, breaking the quadratic limitation of sequence length on computing power, but also the long-distance dependence capturing ability of the model is maximized, realizing the best balance between performance and accuracy.

[0113] Based on any of the above embodiments, Figure 3 is a flowchart of the operator execution method based on dynamic sequence selection provided by the present application, as Figure 3 shown, the method comprises:

[0114] At the beginning, the query tensor Q, the initial key tensor K and the initial value tensor V, and the index information are input. Among them, the shape of the query tensor is usually (batch_size, head_num, 1, head_dim_q), and the initial shape of the key tensor and the value tensor is (batch_size, 1, seq, head_dim_v). At this time, the sequence length seq represents the complete context window, for example, 8192 or more. At the same time, the dynamic selection index information indices is obtained, which can be calculated by the pre-posed lightweight algorithm, indicating which position of the key-value pair is related for each batch data.

[0115] Then, based on the index information, dynamic sequence selection is performed to obtain the key tensor K_select and the value tensor V_select used for attention calculation. In detail, according to the input dynamic selection index information indices, the parallel processor performs a non-continuous memory reading operation to extract the data of the specified position from the complete initial key tensor K and the initial value tensor V. This process greatly compresses the originally large sequence dimension seq to seq_select, for example, from 8192 to 256. Thus, through this step, the rearranged key sub-sequence K_select and value sub-sequence V_select are generated. At this time, the shape of the two tensors changes to (batch_size, 1, seq_select, head_dim_v). It should be noted that through dynamic sequence selection, irrelevant data is physically removed, ensuring that subsequent calculations are only performed on valid information, fundamentally eliminating invalid calculations, and avoiding the huge waste of computing power caused by calculating and masking on the complete sequence.

[0116] Next, the GPU parallel dimension is configured. In detail, according to the batch size batch_size of the query tensor, the grid dimension of the GPU is configured. Specifically, the grid Grid dimension is set to (batch_size, 1), which means that the number of thread blocks Block equal to the batch size will be started. Each thread block will be allocated to process all data in a complete batch Batch, that is, the number of batches corresponding to each thread block is 1. It should be noted that this configuration strategy is the basis for implementing subsequent multi-head aggregation calculation, aiming to provide each thread block with sufficient computing load to feed the Tensor Core and fully utilize the potential computing power of the Tensor Core.

[0117] Then, the batch processing and core calculation are executed in a loop, and it is determined whether all batch processing is completed. In the case where the batch processing is not completed, for each batch, that is, within each thread block, first, the query matrix constructed by loading the query tensor corresponding to the batch from the global memory and the generated key tensor K_select and value tensor V_select are loaded. At this time, the thread block aggregates all the query tensors of attention heads head_num in the batch into a query matrix with a shape of (head_num, head_dim_q). Then, the matrix multiplication is performed by using the Tensor Core unit inside the GPU. Among them, the matrix multiplication is performed on the query matrix and the key tensor, and the calculation result is normalized by Softmax. The normalized weight matrix is multiplied by the value tensor, and the calculation result of all attention heads in the batch is output, that is, the thread block calculation result corresponding to the batch is obtained.

[0118] After it is determined that all batches are executed, the output tensor can be obtained by integrating the thread block calculation results corresponding to all thread blocks. For example, the thread block calculation results calculated by each thread block can be written back to the predetermined address of the global memory. Then, these scattered thread block calculation results can be logically integrated into a complete output tensor, and the shape of the output tensor can be (batch_size, head_num, 1, head_dim_v). The output tensor is the calculation result of the attention mechanism after dynamic sparse acceleration, which can be directly transmitted to the next layer of the neural network.

[0119] The method provided by the embodiment of the application first reduces the data size from O(seq) to O(seq_select) based on the dynamic selection step, directly reducing the floating-point operation demand by several times to several tens of times. Secondly, the grid dimension configuration and the aggregation calculation inside the thread block ingeniously reorganize the originally fine-grained vector operation into large-scale matrix operation, so that the core calculation link can fully utilize the hardware advantages of the Tensor Core in the GPU, and the whole process is closely linked, which not only solves the waste of long sequence calculation, but also solves the problem of low hardware utilization rate in sparse calculation. Therefore, by combining the dynamic sequence selection with the coarse-grained parallel strategy based on the batch size, the invalid calculation can be significantly reduced and the GPU parallel computing capability can be fully utilized.

[0120] In an embodiment, the operator execution method provided by the present application is described in detail in combination with the NVIDIA GPU hardware architecture and specific model parameters. In order to quantitatively illustrate the execution process and effect of the present scheme, the following model parameters are set as examples: the batch size (batch_size) is 4, the number of attention heads (head_num) is 16, the dimension of each head of the query tensor Q and the key and value tensors K and V (head_dim_q, head_dim_v) is 64. In addition, the complete sequence length seq of the input text is set to 8192, and the effective sequence length (seq_select) selected by the sparse algorithm is 256. In this context, the input tensors Q, K and V carry the data information of the model queries, keys and values, and the index Indices indicates the effective data position after sparse selection.

[0121] First, the data preparation and rearrangement step is performed. The input tensor Q, the initial key tensor K, the initial value tensor V and the index information Indices are prepared in the GPU memory. At this time, the initial K and V correspond to the complete length of 8192 in the sequence dimension, and contain a large amount of redundant information that is not needed for current calculation. A dedicated GPU kernel function is called by the system to collect data from the initial K and V according to the coordinate information indicated by Indices, extract the target subsequence, and generate the rearranged key tensor K_select and the rearranged value tensor V_select with a shape of (4, 1, 256, 64). Among them, “4” represents the batch size, “256” represents the effective sequence length selected, and “64” represents the feature dimension. This step effectively compresses the calculation data from the sparse space of 8192 dimensions to the dense space of 256 dimensions, laying a data foundation for subsequent efficient calculation.

[0122] Next, the number of thread blocks to be started is determined and the parallel computation grid is configured. Based on the key dimension information that the batch size of the query tensor Q is 4, the grid dimension of the GPU is set to (4, 1). This means that the hardware scheduler will start 4 independent thread blocks to execute the subsequent attention calculation task. Each thread block is allocated to be responsible for processing the data of one batch. Unlike the prior art, which may start a thread block for each attention head, i.e., 4x16=64 blocks, the present embodiment converges the parallel granularity to the batch level, thereby significantly increasing the calculation load density within a single thread block and reducing the overhead of thread block scheduling.

[0123] Subsequently, the aggregated attention calculation is performed within any thread block. Taking the first thread block responsible for processing the 0th batch as an example, the thread block first loads the corresponding query tensor Q slice with the original shape of (16, 1, 64) and the rearranged key tensor K select slice with the shape of (1, 256, 64) from the video memory. During the loading and calculation process, the thread block aggregates the data of all 16 attention heads within the same batch: the query tensor Q is regarded as a two-dimensional matrix with the shape of (16, 64), where the 16 rows respectively represent the 16 attention heads; K select is regarded as a two-dimensional matrix with the shape of (256, 64). Efficient matrix multiplication is performed by using the internal Tensor Core of the GPU: Scores = Q_slice@K_select_slice.T. At this time, the operation core is represented as the matrix multiplication of (16, 64) and (64, 256), and the output is an attention score matrix with the shape of (16, 256). Then, the Softmax normalization operation is performed on the last dimension of the score matrix to obtain the weight matrix Weights, and the second matrix multiplication is performed between Weights and the rearranged value tensor V select slice with the shape of (256, 64): Output = Weights@V_select_slice, and finally the thread block calculation result with the shape of (16, 64) is obtained.

[0124] Finally, the calculation results are integrated to obtain the output tensor. After the above four thread blocks complete the calculation tasks of their respective batches in parallel, the thread block calculation results with the shape of (16, 64) obtained by each thread block are written back to the corresponding address position of the global memory. After all the thread blocks are written, the complete output tensor is formed.

[0125] Through the above process, on the one hand, the calculation amount is greatly reduced from the original calculation on the complete sequence length seq = 8192 to only on seq select = 256, avoiding the waste of computing power on invalid data; on the other hand, by aggregating all the attention heads within the thread block, the core calculation task is reshaped into a regular large-scale matrix multiplication (such as (16, 64)@ (64, 256)), which can efficiently fill the calculation pipeline of the Tensor Core, greatly improving the instruction execution efficiency of the hardware and the overall operator throughput.

[0126] Based on any of the above embodiments, Figure 4 is a comparison diagram between the conventional attention calculation process and the attention calculation process provided by the present application, as shown in Figure 4As shown, the upper half shows the traditional inefficient task mapping method flow, i.e. the traditional attention calculation flow, and the lower half shows the efficient task mapping method flow provided by the application, i.e. the attention calculation flow provided by the application. The two are significantly distinguished by two dimensions of key / value input comparison and core calculation operation comparison in the middle.

[0127] In detail, in the traditional implementation flow, the parallel granularity of the calculation task is too fine, resulting in waste of hardware resources. In detail, in the data input stage, the traditional method uses the complete sequence, i.e. the key tensor (K) and the value tensor (V) keep the original shape (seq, dim_v), without screening, which means that the calculation disk is extremely large and contains a large amount of invalid information calculation.

[0128] In the grid dimension configuration stage, the traditional method sets the grid dimension to (batch_size x head_num), which means that the number of thread blocks started is equal to the product of the batch size and the number of attention heads.

[0129] In the specific thread block internal processing, such as thread block 0, thread block 1 to thread block N-1, each thread block is only responsible for processing a single query tensor Q, whose shape is only (1, dim_q), and K / V: (seq, dim_v). Since the query tensor is single, the calculation degenerates to vector-matrix system operation, i.e. (1, dim_q)@(dim_v, seq). This calculation mode cannot effectively fill the Tensor Core in the GPU, resulting in the problem of low hardware utilization rate as indicated in the figure.

[0130] In the attention calculation flow provided by the application, the calculation efficiency is qualitatively changed by restructuring the task mapping and introducing dynamic sequence selection. In detail, in the data input stage, a dynamic selection mechanism is introduced. According to the pre-calculated index information, a key subset is extracted from the complete sequence to obtain K / V_select. At this time, the data shape becomes (seq_select, dim_v). As noted in the figure, seq_select << seq, which is much smaller than the original length, so the calculation disk is small.

[0131] In the grid dimension configuration stage, the grid dimension is configured to (batch_size, 1), i.e. only according to the batch size to start the thread block. Compared with the traditional method, the number of thread blocks started is greatly reduced.

[0132] In the thread block internal processing, in any thread block, instead of processing a single vector, a matrix combined operation is performed. For example Figure 4As shown, the query tensors Q(1, dim_q) of "head 0", "head 1" to "head H-1" in the same batch are aggregated together to construct a "Q matrix", i.e. a query matrix, which has a shape of (head_num, dim_q). The core calculation mode is that since the Q matrix (query matrix) contains all attention heads and the K / V has been filtered, the final operation is changed to a matrix-matrix multiplication operation, i.e. (head_num, dim_q)@(dim_v, seq_select). It should be noted that this large-scale matrix multiplication can fully utilize the parallel computing capability of Tensor Core, significantly improving the overall throughput and hardware utilization.

[0133] Based on any of the above embodiments, Figure 5 is a workflow schematic diagram of a single thread block provided by the application, as Figure 5 shown, the flow includes:

[0134] First, the calculation flow starts from the data loading stage, loading data to the thread block shared memory, where the data includes the query tensor corresponding to the batch, and the key tensor and value tensor selected according to the index information. In order to adapt to the dimension requirement of matrix multiplication, the key tensor is transposed to change its shape to [dim_v, seq_select], where it is assumed that dim_q and dim_v have the same dimension. In addition, by aggregating the query tensors of all attention heads in the batch, a query matrix with a shape of [head_num, dim_q] is constructed.

[0135] Then enter stage one: attention score calculation. Use Tensor Core to accelerate the execution of the first step of matrix multiplication to multiply the query matrix with a shape of [head_num, dim_q] and the transposed key matrix. This operation calculates the similarity of all attention heads for the selected sequence at one time, generating an attention score matrix Scores with a shape of [head_num, seq_select], i.e. obtaining the attention score matrix.

[0136] Then comes stage two: Softmax normalization. For the generated attention score matrix Scores, calculate the maximum value and sum along the seq_select dimension. In detail, process along the seq_select dimension. In order to ensure numerical stability, usually first calculate the maximum value of each row and perform subtraction operation, then execute exponential operation and sum operation, i.e. stable Softmax operation. After this nonlinear transformation, the original score is mapped to a probability distribution, and an attention weight matrix with the same shape [head_num, seq_select] is obtained.

[0137] Finally, the third stage is entered: context vector calculation, which uses Tensor Core to accelerate the second step of matrix multiplication. At this time, the value tensor V_select_slice is introduced, which is also loaded into shared memory and maintains the shape [seq_select, dim_v]. The attention weight matrix obtained by the above calculation is subjected to the second step of matrix multiplication with the value tensor V_select_slice, and this process also uses Tensor Core for acceleration. Through weighted summation, the output matrix Output with the shape [head_num, dim_v] is finally generated, which is the thread block calculation result and contains the final calculation result of all attention heads in this batch. After the calculation is completed, the thread block writes the Output matrix back to the corresponding position in the global memory, completing the operator execution task for this batch. The entire process completes data interaction in the on-chip shared memory and continuously uses Tensor Core for matrix calculation, greatly improving the calculation efficiency.

[0138] Figure 6 An example of an entity structure diagram of an electronic device is shown in Figure 6 The electronic device can include a processor 610, a communication interface 620, a memory 630, and a communication bus 640, wherein the processor 610, the communication interface 620, and the memory 630 communicate with each other through the communication bus 640. The processor 610 can invoke the logical instructions in the memory 630 to execute the operator execution method, which includes: obtaining an input tensor; the input tensor includes a query tensor, a key tensor, and a value tensor; determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor; in any thread block, aggregate the query tensors of all attention heads corresponding to the number of processing batches, and perform attention calculation with the key tensor and the value tensor to obtain the thread block calculation result; and integrate the thread block calculation results of the thread blocks to obtain the output tensor of the input tensor.

[0139] In addition, the logic instructions in the memory 630 described above can be implemented in the form of software functional units and sold or used as independent products, and can be stored in a computer readable storage medium. Based on such understanding, the technical solutions of the present application essentially or the parts that contribute to the prior art or parts of the technical solutions can be embodied in the form of a software product. The computer software product is stored in a storage medium, and includes several instructions for making a computer device (which can be a personal computer, a server, or a network device, etc.) execute all or part of the steps of the methods described in the various embodiments of the present application. The aforementioned storage medium includes: a U disk, a mobile hard disk, a read-only memory (ROM, Read-Only Memory), a random access memory (RAM, Random Access Memory), a magnetic disk or an optical disk, and various media that can store program codes.

[0140] In another aspect, the present application also provides a computer program product, which comprises a computer program, the computer program can be stored on a non-transitory computer readable storage medium, and the computer program can be executed by a processor to enable a computer to execute the operator execution method provided by the above-mentioned methods. The method comprises: obtaining an input tensor; the input tensor comprises a query tensor, a key tensor and a value tensor; determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor; in any thread block, aggregating the query tensors of all attention heads corresponding to the number of processing batches to perform attention calculation with the key tensor and the value tensor to obtain a thread block calculation result; and integrating the thread block calculation results of the thread blocks to obtain an output tensor of the input tensor.

[0141] In another aspect, the present application also provides a computer program product, which comprises a computer program, the computer program can be stored on a non-transitory computer readable storage medium, and the computer program can be executed by a processor to enable a computer to execute the operator execution method provided by the above-mentioned methods. The method comprises: obtaining an input tensor; the input tensor comprises a query tensor, a key tensor and a value tensor; determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor; in any thread block, aggregating the query tensors of all attention heads corresponding to the number of processing batches to perform attention calculation with the key tensor and the value tensor to obtain a thread block calculation result; and integrating the thread block calculation results of the thread blocks to obtain an output tensor of the input tensor.

[0142] The device embodiments described above are merely illustrative, wherein the units described as separate components can or can not be physically separate, and the components displayed as units can or can not be physical units, i.e., can be located in one place, or can be distributed to multiple network units. Part or all of the modules can be selected to achieve the purposes of the embodiments according to actual needs. Those skilled in the art can understand and implement without creative labor.

[0143] Through the description of the above embodiments, those skilled in the art can clearly understand that the embodiments can be realized by means of software and the necessary general hardware platform, and of course can also be realized by hardware. Based on such understanding, the above technical solutions can be embodied in the form of a software product, which can be stored in a computer readable storage medium, such as a ROM / RAM, a magnetic disk, an optical disk, etc., and includes a number of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute the methods described in each embodiment or some parts of the embodiments.

[0144] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement for part of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. An operator execution method characterized by, The method comprises the following steps: obtaining an input tensor; the input tensor comprises a query tensor, a key tensor and a value tensor; determining the number of processing batches corresponding to each thread block based on the batch size of the query tensor; in any thread block, aggregating the query tensors of all attention heads corresponding to the number of processing batches, performing attention calculation with the key tensor and the value tensor to obtain a thread block calculation result; integrating the thread block calculation results of the thread blocks to obtain an output tensor of the input tensor; the step of aggregating the query tensors of all attention heads corresponding to the number of processing batches, performing attention calculation with the key tensor and the value tensor to obtain a thread block calculation result comprises: aggregating the query tensors of all attention heads corresponding to the number of processing batches to obtain a query matrix, the number of rows of the query matrix being the total number of all attention heads; performing attention calculation based on the query matrix, the key tensor and the value tensor to obtain the thread block calculation result.

2. The operator execution method of claim 1, wherein, The step of obtaining the key tensor and the value tensor comprises: obtaining an initial key tensor, an initial value tensor and index information; extracting target subsequences corresponding to the index information from the initial key tensor and the initial value tensor to obtain a rearranged key tensor and a rearranged value tensor; taking the rearranged key tensor and the rearranged value tensor as the key tensor and the value tensor.

3. The operator execution method of claim 1, wherein, The step of performing attention calculation based on the query matrix, the key tensor and the value tensor to obtain the thread block calculation result comprises: performing matrix multiplication calculation based on the query matrix and the key tensor to obtain an attention score matrix; performing normalization operation on the attention score matrix to obtain an attention weight matrix; performing calculation based on the attention weight matrix and the value tensor to obtain the thread block calculation result.

4. The operator execution method of claim 3, wherein, The number of rows of the attention score matrix is the total number of all attention heads corresponding to the number of processing batches.

5. The operator execution method according to any one of claims 1 to 4, characterized by, The step of integrating the thread block calculation results of the thread blocks to obtain the output tensor of the input tensor comprises: integrating the thread block calculation results of the thread blocks based on the batch index of the query tensor to obtain the output tensor.

6. The operator execution method of claim 2, wherein, The index information corresponds to the sequence dimension coordinates of the key tensor and the value tensor. The step of obtaining the index information comprises: calculating a correlation matrix between the query tensor and the initial key tensor; based on the correlation matrix, selecting a preset number of positions with the highest correlation along the sequence dimension of the initial key tensor to obtain the index information.

7. An electronic device comprising a memory, a processor, and a computer program stored on the memory and executable on the processor, characterized in that, The processor executes the computer program to realize the operator execution method according to any one of claims 1 to 6.

8. A non-transitory computer-readable storage medium having stored thereon a computer program, characterized in that, The computer program is executed by the processor to realize the operator execution method according to any one of claims 1 to 6.

9. A computer program product comprising a computer program, characterized in that, The computer program is executed by the processor to realize the operator execution method according to any one of claims 1 to 6.

Images