Self-attention mechanism processing method and device, equipment and medium

By dividing tasks along the attention head dimension in the self-attention mechanism, the effective utilization of GPU parallel computing power is achieved, and the problem of excessive global memory access in the existing technology is solved, and the computing performance is improved.

CN119692405APending Publication Date: 2025-03-25SHANGHAI JIAOTONG UNIV +1
View PDF 0 Cites 0 Cited by

Patent Information

Application Number
CN202411614646.9
Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
Filing Date
2024-11-13
Publication Date
2025-03-25

AI Technical Summary

Technical Problem

The existing self-attention mechanism in the decoding calculation of the large language model causes each thread block to read all codebooks redundantly, increasing the number of global memory accesses and affecting the computing performance.

Method used

By dividing tasks along the attention head dimension of the self-attention mechanism, multiple query vectors and multiple sets of key-value vectors for each attention head are generated, and parallel operations are performed in the attention head dimension to reduce the number of accesses to global memory.

Benefits of technology

Effectively utilize the parallel computing power of the GPU, accelerate the calculation process of the self-attention mechanism, reduce the amount of access to global memory, and improve computing performance.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN119692405A_ABST
    Figure CN119692405A_ABST
Patent Text Reader

Abstract

The invention provides a processing method and device of a self-attention mechanism, equipment and a medium. Through a task division mode along the head dimension of the self-attention mechanism, the self-attention mechanism can effectively utilize the parallel computing capability of a GPU, the computing process of the self-attention mechanism is accelerated, the access amount of a global memory is reduced, and the computing performance is improved.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the field of data processing technology, and in particular to a processing method, device, equipment and medium for a self-attention mechanism. Background Art

[0002] The Self-Attention Mechanism plays a core role in the decoding calculation of large language model inference. It achieves effective use and dynamic attention of contextual information by calculating and weighted summing the correlation of different tokens in the input sequence, thereby improving the model's language understanding and generation capabilities.

[0003] At present, in the decoding calculation of large language model inference, the existing self-attention mechanism such as Flash Decoding (FD) is mainly divided into different thread blocks by the token dimension for parallel operation, and then reduced in the token dimension; but in the case of vector quantization, the storage of the codebook is divided by the attention head dimension. However, using the existing calculation method will cause each thread block to redundantly read all codebooks, increase the access volume of global memory, and thus affect the calculation performance. Therefore, an effective processing method of the self-attention mechanism is urgently needed to solve the above problems. Summary of the invention

[0004] In view of the above problems, the present invention is proposed to provide a processing method, device, equipment and medium for a self-attention mechanism that overcomes the above problems or at least partially solves the above problems.

[0005] To achieve the above objectives and other related objectives, the present invention provides a processing method for a self-attention mechanism, which is applied to a preset self-attention mechanism, and the method includes:

[0006] Obtaining an input sequence of the self-attention mechanism, and dividing the input sequence along the attention head dimension of the self-attention mechanism according to a preset division granularity, to generate multiple query vectors and multiple groups of key value vectors for each attention head;

[0007] For each query vector in a single attention head, determine each key vector and value vector corresponding to each query vector from multiple sets of key-value vectors, perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector;

[0008] The output vectors of all query vectors of each attention head are merged to calculate the target output of the self-attention mechanism.

[0009] Optionally, before the step of obtaining the input sequence of the self-attention mechanism, the step further includes:

[0010] Obtaining a codebook reading amount, an attention head dimension global specification memory reading amount, and a token dimension global specification memory reading amount from the attention mechanism, and determining a sum of the attention head dimension global specification memory reading amount and the token dimension global specification memory reading amount as a global specification memory reading amount;

[0011] According to a preset mean theorem, the preset division granularity of the self-attention mechanism is determined by utilizing the codebook reading amount and the global reduced memory reading amount.

[0012] Optionally, the step of dividing the input sequence along the attention head dimension of the self-attention mechanism according to a preset division granularity to generate multiple query vectors and multiple groups of key-value vectors for each attention head includes:

[0013] Performing embedding and position encoding processing on the input sequence to obtain an encoded input sequence;

[0014] According to a preset division granularity, the encoded input sequence is divided along the head dimension of the self-attention mechanism to obtain one or more subsequences, and the one or more subsequences are assigned to each attention head of the self-attention mechanism to generate multiple query vectors and multiple groups of key-value vectors for each attention head; wherein each group of key-value vectors includes a key vector and a value vector corresponding to the key vector.

[0015] Optionally, performing matrix operations on each query vector and each key vector to obtain each attention score includes:

[0016] Transpose each key vector to obtain each transposed key vector;

[0017] Perform matrix multiplication on each query vector and each of the transposed key vectors to obtain each attention score.

[0018] Optionally, performing global reduction and normalization processing on each attention score to obtain the attention weight includes:

[0019] Each attention score is globally reduced to obtain an intermediate result, and the intermediate result is normalized to obtain an attention weight.

[0020] Optionally, performing a matrix operation on the attention weights and each value vector to obtain an output vector for each query vector includes:

[0021] The attention weights are weighted and summed for each value vector to obtain an output vector for each query vector.

[0022] Optionally, the merging of the output vectors of all query vectors of the attention heads to calculate the target output of the self-attention mechanism includes:

[0023] Merge the output vectors of each query vector in a single attention head to calculate the output representation of all query vectors of a single attention head;

[0024] The output representations of the attention heads are merged along the head dimension of the self-attention mechanism, and the target output of the self-attention mechanism is calculated.

[0025] In a second aspect, the present invention further provides a processing device for a self-attention mechanism, which is applied to a preset self-attention mechanism, and the device comprises:

[0026] A partitioning module is used to obtain an input sequence of the self-attention mechanism, and partition the input sequence along the attention head dimension of the self-attention mechanism according to a preset partitioning granularity to generate multiple query vectors and multiple groups of key value vectors for each attention head;

[0027] a processing module, configured to determine, for each query vector in a single attention head, each key vector and value vector corresponding to each query vector from multiple sets of key-value vectors, perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector;

[0028] The merging module is used to merge the output vectors of all query vectors of each attention head to calculate the target output of the self-attention mechanism.

[0029] In a third aspect, the present invention provides an electronic device, comprising: a memory and a processor; the memory is used to store a computer program; the processor is used to execute the computer program stored in the memory, so that the electronic device performs the steps of the processing method of the self-attention mechanism as described above.

[0030] In a fourth aspect, the present invention provides a computer-readable storage medium having a computer program stored thereon, which, when executed by an electronic device, implements the steps of the processing method of the self-attention mechanism as described above.

[0031] The above one or more technical solutions provided by the present invention may have the following advantages or at least achieve the following technical effects:

[0032] The processing method, device, equipment and medium of the self-attention mechanism of the present invention can effectively utilize the parallel computing capability of the GPU through the task division method along the head dimension of the self-attention mechanism, accelerate the calculation process of the self-attention mechanism, reduce the access amount of the global memory, and improve the computing performance. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] Figure 1 Shown is a flow chart of a method for processing a self-attention mechanism in one embodiment of the present invention;

[0034] Figure 2 It is a schematic diagram showing a flow chart of calculating a single query vector and its corresponding key vectors in one embodiment of the present invention;

[0035] Figure 3 Shown is a schematic diagram of a flow chart of calculating the attention weight of a single query vector and its corresponding value vectors in one embodiment of the present invention;

[0036] Figure 4 It is a schematic diagram of the functional modules of a processing device of a self-attention mechanism in one embodiment of the present invention;

[0037] Figure 5 Shown is a schematic diagram of an electronic device according to an embodiment of the present invention. DETAILED DESCRIPTION

[0038] The following describes the embodiments of the present invention by specific examples, and those skilled in the art can easily understand other advantages and effects of the present invention from the contents disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and the details in this specification can also be modified or changed in various ways based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that the following embodiments and features in the embodiments can be combined with each other without conflict.

[0039] It should be noted that the illustrations provided in the following embodiments are only schematic illustrations of the basic concept of the present invention, and thus the drawings only show components related to the present invention rather than being drawn according to the number, shape and size of components in actual implementation. In actual implementation, the type, quantity and proportion of each component may be changed arbitrarily, and the component layout may also be more complicated.

[0040] In the following description, numerous details are discussed to provide a more thorough explanation of the embodiments of the present invention. However, it is obvious to those skilled in the art that the embodiments of the present invention can be implemented without these specific details. In other embodiments, well-known structures and devices are shown in the form of block diagrams rather than in detail to avoid making the embodiments of the present invention difficult to understand.

[0041] The terms "first", "second", etc. in the specification and claims of the disclosed embodiments and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the terms used in this way can be interchanged where appropriate, so that the embodiments of the disclosed embodiments described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions.

[0042] Unless otherwise stated, the term "plurality" means two or more.

[0043] In the embodiment of the present disclosure, the character " / " indicates that the preceding and following objects are in an "or" relationship. For example, A / B indicates: A or B.

[0044] The term "and / or" is a description of the association relationship between objects, indicating that three relationships can exist. For example, A and / or B means: A or B, or, A and B.

[0045] The technical solutions in the embodiments of the present invention will be described in detail below in conjunction with the accompanying drawings in the embodiments of the present invention.

[0046] See also Figure 1 An embodiment of the present invention provides a method for processing a self-attention mechanism, which is applied to a preset self-attention mechanism. The method may include the following steps S10 to S30:

[0047] Step S10, obtaining an input sequence of the self-attention mechanism, and dividing and encoding the input sequence along the head dimension of the self-attention mechanism according to a preset division granularity, to generate multiple query vectors and multiple groups of key-value vectors for each attention head.

[0048] The input sequence can be composed of a series of elements, which can be words in text, pixels in images, time segments in audio, etc.

[0049] In natural language processing, the input sequence is generally a sentence or document consisting of a series of words or characters.

[0050] The partition granularity refers to how many small blocks the input sequence is divided into for parallel processing; it can be divided by sequence length (i.e. each head processes a different part of the sequence) or feature dimension (i.e. each head processes a different feature subset).

[0051] It should be noted that in the self-attention mechanism, the determination of the partition granularity usually needs to consider the architecture of the model and the available computing resources. The partition granularity affects the parallelism and memory efficiency of the model. Smaller partition granularity can improve parallelism and memory efficiency, but may increase computational overhead because more calculations are required to merge structures. Larger partition granularity may reduce parallelism, but can reduce computational and memory requirements. In practical applications, the choice of partition granularity usually needs to be determined through experiments to achieve the best balance between performance and efficiency.

[0052] Head Dimension refers to the dimension of the feature vector processed by each attention head in the self-attention mechanism.

[0053] If the total feature dimension of the self-attention mechanism is d model , and the number of heads is h, then the feature dimension of each head is d head Usually d model h. This means that each head processes a subset of the input sequence, thereby capturing information from different perspectives in parallel.

[0054] As an example, if d model =512 and h=8, then the head dimension of each head is d head =d mo d e l h=512 / 8=64.

[0055] In parallel computing, the parallel capability of GPU can be used to assign different attention heads to different thread blocks for calculation. Each thread block processes the calculation of one head, so that multiple heads can be processed at the same time, thereby improving the calculation efficiency.

[0056] Thread block refers to the basic unit for executing parallel computing on GPU. In the self-attention mechanism, each thread block can be responsible for the calculation of one or more heads.

[0057] The query vector can be composed of a token dimension and a head dimension. The key vector and value vector are composed of a token dimension and a head dimension respectively.

[0058] If each thread block (TB) processes the calculation of an attention head, each thread block (TB) holds a portion of the head dimension of all tokens of the key vector and calculates with the query vector after transposition. According to the properties of matrix multiplication, each TB will hold a partial result (i.e., attention score) of the same shape (1, Tocken). Before proceeding to the next step, the partial result needs to be reduced and Softmax normalized before being calculated with the value vector.

[0059] In the specific implementation, first, the input sequence of the self-attention mechanism is obtained, and the input sequence can be encoded (including embedding and position encoding) to obtain an encoded input sequence; then the encoded input sequence can be divided along the head dimension of the self-attention mechanism according to the preset division granularity to obtain one or more subsequences; the one or more subsequences encoded by position can be assigned to each attention head of the self-attention mechanism, and then each attention head can independently process its own subsequence, thereby generating multiple query vectors (Q) and multiple groups of key-value vectors for each attention head (where each group of key-value vectors includes a key vector (K) and its corresponding value vector (V)). Since the calculations of each attention head are independent, these calculations can be performed in parallel, thereby improving processing efficiency, especially when using parallel computing hardware such as GPU.

[0060] Step S20, for each query vector in a single attention head, determine the key vectors and value vectors corresponding to each query vector from multiple groups of key-value vectors, and perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector.

[0061] Among them, the output vector can be represented as the vector representation calculated for each query vector in each attention head; it can be obtained by weighted summing the value vector (V) using the attention weight (calculated by the query vector (Q) and the key vector (K)).

[0062] It should be noted that for a single attention head, each query vector will interact with its corresponding multiple sets of key-value vectors; each set of key-value vectors includes a key vector (K) and its corresponding value vector (V).

[0063] In the specific implementation, for each query vector in a single attention head, first, it is necessary to determine the key vectors corresponding to the query vector and the value vectors corresponding to the key vectors from multiple groups of key-value vectors; then, the query vector and each key vector are subjected to matrix operations (such as calculating dot products or matrix multiplications) to obtain attention scores; these attention scores are converted into attention weights after global reduction and Softmax normalization; and then, the attention weights are weightedly summed with each value vector to obtain the output vector of the query vector.

[0064] Step S30, merging the output vectors of all query vectors of each attention head to calculate the target output of the self-attention mechanism.

[0065] Among them, the target output can be the combination of the outputs of all attention heads in the self-attention mechanism to form the final output of the self-attention mechanism.

[0066] In the specific implementation, after obtaining the output vector of each query vector in a single attention head, the output vectors of all query vectors in a single attention head can be calculated; then the output vectors of each attention head are concatenated along the head dimension of the self-attention mechanism and merged into the target output of the self-attention mechanism.

[0067] For further information, see Figure 2 In one embodiment, step S20 may include the following sub-steps S201 to S202:

[0068] Sub-step S201, transposing each key vector to obtain each transposed key vector;

[0069] Sub-step S202, performing matrix multiplication operation on each query vector and each of the transposed key vectors to obtain each attention score.

[0070] In a specific implementation, for each query vector in a single attention head, the key vectors corresponding to the query vector are transposed to obtain the transposed key vectors; then, the query vector and the transposed key vectors can be matrix multiplied to obtain the attention scores of the query vector.

[0071] For further information, see Figure 2 In one embodiment, step S20 may further include the following sub-step S203:

[0072] Sub-step S203, performing global reduction on each attention score to obtain an intermediate result, and normalizing the intermediate result to obtain an attention weight.

[0073] The intermediate result may be a calculation result obtained by performing global reduction on each attention score of each query vector.

[0074] In the specific implementation, for each query vector in a single attention head, the attention scores of each query vector are globally reduced to calculate the intermediate result of each query vector; then the intermediate result can be normalized by softmax and converted into the attention weight of the query vector.

[0075] As an example, see Figure 2 , if each thread block in the self-attention mechanism processes the calculation of an attention head, and the key cache corresponding to each query vector includes 4 sets of key-value vectors (each set of key-value vectors includes a key vector and its corresponding value vector); after matrix multiplication of the query vector and the transpose of the 4 key vectors, 4 attention scores are obtained, and then the 4 attention scores are global reduced and Softmax normalized to calculate the weight matrix of the query vector, and then the attention weight can be calculated with the 4 value vectors corresponding to the query vector.

[0076] For further information, see Figure 3 In one embodiment, step S20 may further include the following sub-step S204:

[0077] Sub-step S204, performing weighted summation of the attention weights on each value vector to obtain an output vector for each query vector.

[0078] In a specific implementation, after calculating the attention weight of each query vector, the attention weight can be used to perform weighted summation on the corresponding value vectors to obtain the output representation of the query vector (ie, the output vector).

[0079] In this embodiment, by dividing the tasks along the dimension of the self-attention mechanism head, the self-attention mechanism can effectively utilize the parallel computing capability of the GPU, accelerate the calculation process of the self-attention mechanism, reduce the access to the global memory, and improve the computing performance.

[0080] Based on the above embodiments, a second embodiment of a method for processing a self-attention mechanism of the present invention is proposed. In this embodiment, before step S10, the method may further include the following sub-steps A10 to A20:

[0081] Step A10, obtain the codebook reading amount of the self-attention mechanism, the global reduced memory reading amount of the attention head dimension, and the global reduced memory reading amount of the token dimension, and determine the sum of the global reduced memory reading amount of the attention head dimension and the global reduced memory reading amount of the token dimension as the global reduced memory reading amount.

[0082] The codebook read amount refers to the amount of video memory read from a storage medium (such as a hard disk or memory).

[0083] The global reduction in the amount of video memory read for the attention head dimension may depend on the number and size of attention heads.

[0084] The token dimension globally reduces the memory read, depending on the sequence length and the embedding dimension.

[0085] The global reduction memory read amount refers to the memory read amount of the global reduction operation on different dimensions in the self-attention mechanism; it may include the global reduction memory read amount of the attention head dimension and the global reduction memory read amount of the token dimension.

[0086] In a specific implementation, a dedicated memory analysis tool (such as NVIDIA's Nsight tool or other GPU performance analysis tools) can be used to track and analyze the memory access pattern (including the usage of memory and video memory) of the self-attention mechanism during execution, and then the codebook reading amount, the global reduced video memory reading amount of the attention head dimension, and the global reduced video memory reading amount of the token dimension can be obtained; then, the global reduced video memory reading amount of the attention head dimension and the global reduced video memory reading amount of the token dimension can be added together to calculate the global reduced video memory reading amount.

[0087] Step A20, according to a preset mean theorem, using the codebook reading amount and the global reduced memory reading amount, determine the preset division granularity of the self-attention mechanism.

[0088] The mean value theorem (also called the fundamental inequality) states that for two positive real numbers A and B, The equality holds if and only if A=B.

[0089] When determining the partition granularity, the codebook reading amount and the global reduced memory reading amount can be regarded as two related quantities, and a relatively balanced partition granularity can be found through the mean value theorem.

[0090] As an example, if the total number of thread blocks of the self-attention mechanism is N, the split granularity (split_factor) is the number of thread blocks split in the attention head dimension;

[0091] Among them, the codebook reading amount can be expressed as The increase in the global reduction memory read of the attention head dimension can be expressed as Traffic Reduce_head =headreduce size×split_factor, the token dimension global reduction memory read amount can be expressed as The total read volume can be expressed as Traffic Total =Traffic CodeBook +Traffic Reduce_head +Traffic Reduce_tockenAccording to the preset mean theorem, the minimum reading amount of the two parts is achieved when the codebook reading amount is equal to the global reduced memory reading amount; therefore, before executing the self-attention mechanism, the mean theorem can be used to determine the optimal division granularity using the codebook reading amount and the global reduced memory reading amount.

[0092] It should be noted that the preset partitioning granularity is determined by combining the mean value theorem with the codebook reading amount and the global reduced memory reading amount. It needs to be tested and adjusted many times and optimized according to the actual situation to find the partitioning granularity that best suits the self-attention mechanism and improve the performance and efficiency of the self-attention mechanism.

[0093] In the specific implementation, the codebook reading amount and the global reduced memory reading amount reflect the memory usage of the self-attention mechanism in different operations. Through the mean theorem, a division granularity can be found so that at this granularity, the codebook reading and the global reduced memory reading can be balanced to a certain extent, avoiding excessive memory usage in one aspect and affecting the performance and efficiency of the model.

[0094] In this embodiment, the codebook reading amount, the global reduced memory reading amount of the attention head dimension, and the global reduced memory reading amount of the token dimension are obtained, and the sum of the global reduced memory reading amount of the attention head dimension and the global reduced memory reading amount of the token dimension is determined as the global reduced memory reading amount; according to the preset mean theorem, the codebook reading amount and the global reduced memory reading amount are used to determine the preset partitioning granularity. Thus, through the mean theorem, a partitioning granularity can be found, so that under this partitioning granularity, the codebook reading and the global reduced reading can be balanced to a certain extent, avoiding excessive memory usage in one aspect and affecting the performance and efficiency of the self-attention mechanism, so as to reduce unnecessary memory access and improve efficiency.

[0095] Based on the above embodiments, a third embodiment of a method for processing a self-attention mechanism of the present invention is proposed. In this embodiment, step S10 may include the following sub-steps S101 to S102:

[0096] Sub-step S101, embedding and position encoding the input sequence to obtain an encoded input sequence.

[0097] The encoded input sequence may be an input sequence that has undergone processing including embedding and position encoding.

[0098] It should be noted that in the self-attention mechanism, the input sequence usually needs to be encoded before being divided. The purpose of encoding is to enable the model (self-attention mechanism) to process and understand the input data.

[0099] In the specific implementation, an embedding layer can be used to convert the input sequence into a numerical vector (i.e., an embedding vector) that the model can process. In order to enable the self-attention mechanism to understand the sequential information in the input sequence, position encoding can be added to the embedding vector to obtain an encoded input sequence.

[0100] Sub-step S102, according to a preset division granularity, divides the encoded input sequence along the head dimension of the self-attention mechanism to obtain one or more sub-sequences, and assigns the one or more sub-sequences to each attention head of the self-attention mechanism to generate multiple query vectors and multiple groups of key-value vectors for each attention head.

[0101] Among them, the subsequence can be a subset (or one or more continuous segments) of the input sequence after embedding and position encoding. They are processed independently in the self-attention mechanism to improve computational efficiency and the parallel processing capability of the model.

[0102] In the specific implementation, according to the preset division granularity, the encoded input sequence can be divided along the head dimension of the self-attention mechanism (e.g., equal length division or equal structure division) to obtain one or more subsequences; each subsequence after division is assigned to a different attention head of the self-attention mechanism, and each attention head independently generates its own multiple query (Q), key (K) and value (V) vectors. Since the calculation of each attention head is independent, these calculations can be performed in parallel, thereby improving processing efficiency, especially when using parallel computing hardware such as GPU.

[0103] In this embodiment, the input sequence is embedded and positionally encoded to obtain an encoded input sequence; the encoded input sequence is divided along the head dimension of the self-attention mechanism according to a preset division granularity to obtain one or more subsequences, and the one or more subsequences are assigned to each attention head of the self-attention mechanism to generate multiple query vectors and multiple groups of key value vectors for each attention head. Thus, by dividing the input sequence along the self-attention head dimension for parallel computing, not only can the efficiency of parallel computing be improved, but also the amount of global memory access can be reduced, thereby improving the computing performance of the self-attention mechanism.

[0104] Based on the above embodiments, a fourth embodiment of a method for processing a self-attention mechanism of the present invention is proposed. In this embodiment, step S30 may include the following sub-steps S301 to S302:

[0105] Sub-step S301, merging the output vectors of each query vector in a single attention head, and calculating the output representation of all query vectors of a single attention head.

[0106] In a specific implementation, after obtaining the output vector of each query vector in a single attention head, for each attention head, the output vectors of all query vectors of the single attention head can be merged.

[0107] Sub-step S302, merging the output representations of each attention head along the head dimension of the self-attention mechanism, and calculating the target output of the self-attention mechanism.

[0108] In the specific implementation, after calculating the output vectors of all query vectors of each attention head, the output vectors of each attention head can be concatenated along the attention head dimension. The concatenated vectors are then passed through a linear layer (usually called the output projection layer), which can further integrate the information and convert it into the target output of the self-attention mechanism.

[0109] In this embodiment, the output vectors of each query vector in a single attention head are merged to calculate the output representation of all query vectors of the single attention head; the output representations of each attention head are merged along the head dimension of the self-attention mechanism to calculate the target output of the self-attention mechanism. Thus, the computing time can be reduced and the computing efficiency can be improved through parallel processing and optimized memory access mode.

[0110] Based on the same inventive concept, the fifth embodiment of the present invention also provides a processing device for the self-attention mechanism corresponding to the processing method of the self-attention mechanism in the previous embodiment. Since the principle of solving the problem by the device in the fifth embodiment of the present invention is similar to the processing method of the self-attention mechanism in the previous embodiment of the present invention, the implementation of the device can refer to the implementation of the method, and the repeated parts will not be repeated. Figure 4 , the processing device of the self-attention mechanism of the present invention may include:

[0111] A partitioning module 10 is used to obtain an input sequence of the self-attention mechanism, and partition the input sequence along the attention head dimension of the self-attention mechanism according to a preset partitioning granularity to generate multiple query vectors and multiple groups of key value vectors for each attention head;

[0112] A processing module 20 is used to determine, for each query vector in a single attention head, each key vector and value vector corresponding to each query vector from multiple sets of key-value vectors, perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector;

[0113] The merging module 30 is used to merge the output vectors of all query vectors of each attention head to calculate the target output of the self-attention mechanism.

[0114] In addition, the present invention also provides a computer-readable storage medium on which a computer program is stored, and when the program is executed by a processor, the processing method of the above-mentioned self-attention mechanism is implemented.

[0115] Figure 5 is a schematic block diagram of an electronic device provided in an embodiment of the present application. Figure 5 As shown, the electronic device includes: at least one processor 401, a memory 402, at least one network interface 403 and a user interface 405. The various components in the electronic device are coupled together through a bus system 404. It can be understood that the bus system 404 is used to realize the connection and communication between these components. In addition to the data bus, the bus system 404 also includes a power bus, a control bus and a status signal bus. However, for the sake of clarity, Figure 5 In the specification, various buses are labeled as bus systems.

[0116] The user interface 405 may include a display, a keyboard, a mouse, a trackball, a click gun, keys, buttons, a touch pad or a touch screen.

[0117] It is understood that the memory 402 can be a volatile memory or a non-volatile memory, and can also include both volatile and non-volatile memories. Among them, the non-volatile memory can be a read-only memory (ROM), a programmable read-only memory (PROM), which is used as an external cache. By way of example but not limitation, many forms of RAM are available, such as static random access memory (SRAM), synchronous static random access memory (SSRAM). The memory described in the embodiments of the present invention is intended to include but is not limited to these and any other suitable categories of memory.

[0118] The memory 402 in the embodiment of the present invention is used to store various categories of data to support the operation of the electronic device 400. Examples of these data include: any executable program for operating on the electronic device 400, such as an operating system 4021 and an application 4022; the operating system 4021 includes various system programs, such as a framework layer, a core library layer, a driver layer, etc., for implementing various basic services and processing hardware-based tasks. The application 4022 may include various applications, such as a media player (Media Player), a browser (Browser), etc., for implementing various application services. The processing method for implementing the self-attention mechanism provided in the embodiment of the present invention may be included in the application 4022.

[0119] The method disclosed in the above embodiment of the present invention can be applied to the processor 401, or implemented by the processor 401. The processor 401 may be an integrated circuit chip with signal processing capabilities. In the implementation process, each step of the above method can be completed by the hardware integrated logic circuit in the processor 401 or the instruction in the form of software. The above processor 401 may be a general-purpose processor, a digital signal processor (DSP, Digital Signal Processor), or other programmable logic devices, discrete gates or transistor logic devices, discrete hardware components, etc. The processor 401 can implement or execute the methods, steps and logic block diagrams disclosed in the embodiment of the present invention. The general-purpose processor 401 may be a microprocessor or any conventional processor, etc. The steps of the processing method of the self-attention mechanism provided in combination with the embodiment of the present invention can be directly embodied as a hardware decoding processor to execute, or a combination of hardware and software modules in the decoding processor to execute. The software module can be located in a storage medium, which is located in a memory, and the processor reads the information in the memory and completes the steps of the above method in combination with its hardware.

[0120] In an exemplary embodiment, the electronic device 400 may be implemented by one or more application specific integrated circuits (ASIC), DSP, programmable logic device (PLD), complex programmable logic device (CPLD) to execute the aforementioned method.

[0121] In summary, through the task division method along the head dimension of the self-attention mechanism, the self-attention mechanism can effectively utilize the parallel computing capability of the GPU, accelerate the calculation process of the self-attention mechanism, reduce the access amount of the global memory, and improve the computing performance.

[0122] The above embodiments are merely illustrative of the principles and effects of the present invention, and are not intended to limit the present invention. Anyone familiar with the art may modify or alter the above embodiments without departing from the spirit and scope of the present invention. Therefore, all equivalent modifications or alterations made by a person of ordinary skill in the art without departing from the spirit and technical concept disclosed by the present invention shall still be covered by the claims of the present invention.

Claims

1. A processing method of a self-attention mechanism, characterized in that: Applied to a preset self-attention mechanism, the method includes: Obtaining an input sequence of the self-attention mechanism, and dividing the input sequence along the attention head dimension of the self-attention mechanism according to a preset division granularity, to generate multiple query vectors and multiple groups of key value vectors for each attention head; For each query vector in a single attention head, determine each key vector and value vector corresponding to each query vector from multiple sets of key-value vectors, perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector; The output vectors of all query vectors of each attention head are merged to calculate the target output of the self-attention mechanism.

2. The method according to claim 1, characterized in that Before the step of obtaining the input sequence of the self-attention mechanism, it also includes: Obtaining a codebook reading amount, an attention head dimension global specification memory reading amount, and a token dimension global specification memory reading amount from the attention mechanism, and determining a sum of the attention head dimension global specification memory reading amount and the token dimension global specification memory reading amount as a global specification memory reading amount; According to a preset mean theorem, the preset division granularity of the self-attention mechanism is determined by utilizing the codebook reading amount and the global reduced memory reading amount.

3. The method according to claim 1, characterized in that The step of dividing the input sequence along the attention head dimension of the self-attention mechanism according to the preset division granularity to generate multiple query vectors and multiple groups of key-value vectors for each attention head includes: Performing embedding and position encoding processing on the input sequence to obtain an encoded input sequence; According to a preset division granularity, the encoded input sequence is divided along the head dimension of the self-attention mechanism to obtain one or more subsequences, and the one or more subsequences are assigned to each attention head of the self-attention mechanism to generate multiple query vectors and multiple groups of key-value vectors for each attention head; wherein each group of key-value vectors includes a key vector and a value vector corresponding to the key vector.

4. The method according to claim 1, characterized in that: The matrix operation of each query vector and each key vector to obtain each attention score includes: Transpose each key vector to obtain each transposed key vector; Perform matrix multiplication on each query vector and each of the transposed key vectors to obtain each attention score.

5. The method according to claim 1, characterized in that The global reduction and normalization of each attention score to obtain the attention weight includes: Each attention score is globally reduced to obtain an intermediate result, and the intermediate result is normalized to obtain an attention weight.

6. The method according to claim 1, characterized in that The matrix operation of the attention weights and each value vector to obtain an output vector for each query vector includes: The attention weights are weighted and summed for each value vector to obtain an output vector for each query vector.

7. The method according to claim 1, characterized in that The output vector of all query vectors of each attention head is merged to calculate the target output of the self-attention mechanism, including: Merge the output vectors of each query vector in a single attention head to calculate the output representation of all query vectors of a single attention head; The output representations of the attention heads are merged along the head dimension of the self-attention mechanism, and the target output of the self-attention mechanism is calculated.

8. A processing device for a self-attention mechanism, characterized in that: Applied in a preset self-attention mechanism, the device comprises: A partitioning module is used to obtain an input sequence of the self-attention mechanism, and partition the input sequence along the attention head dimension of the self-attention mechanism according to a preset partitioning granularity to generate multiple query vectors and multiple groups of key value vectors for each attention head; a processing module, configured to determine, for each query vector in a single attention head, each key vector and value vector corresponding to each query vector from multiple sets of key-value vectors, perform matrix operations on each query vector and each key vector to obtain each attention score, perform global reduction and normalization processing on each attention score to obtain an attention weight, and perform matrix operations on the attention weight and each value vector to obtain an output vector for each query vector; The merging module is used to merge the output vectors of all query vectors of each attention head to calculate the target output of the self-attention mechanism.

9. An electronic device, characterized in that: The electronic device comprises: a memory and a processor, wherein the memory is used to store a computer program, and the processor is used to execute the computer program stored in the memory, so that the electronic device performs the steps of the method as claimed in any one of claims 1 to 7.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a program, and when the program is run, it is used to implement the steps of the method according to any one of claims 1 to 7.