Attention mechanism computation methods, electronic devices and storage media

Abstract

This invention relates to the field of artificial intelligence technology, providing a method, electronic device, and storage medium for calculating attention mechanisms. The method includes: configuring multiple thread bundles that share the same set of thread-local registers; breaking down the attention mechanism calculation process into multiple computational sub-steps and assigning these sub-steps to the multiple thread bundles, wherein each thread bundle is configured to execute its assigned sub-steps; and controlling the multiple thread bundles to execute the sub-steps in parallel in a pipelined manner by reading and writing to the shared thread-local registers to obtain the calculation result of the attention mechanism. This invention, by utilizing shared thread-local registers to construct a pipelined calculation mode for multi-threaded bundle collaboration, achieves efficient data copying and deep parallelism between different computational units during the attention mechanism calculation process, offering advantages such as low communication overhead, high hardware resource utilization, and strong operator execution performance.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to an attention mechanism computation method, electronic device, and storage medium. Background Technology

[0002] With the rapid development of artificial intelligence technology, large-scale attention mechanism computation accelerated by computing devices such as general-purpose graphics processing units (GPGPUs) and graphics processing units (GPUs) has become a common practice in the industry. Currently, multiple warp groups are typically used to process data collaboratively. For example, some warp groups are responsible for the interaction of data between global memory and shared memory, while other warp groups perform matrix multiplication and softmax operations to complete the forward computation of the attention mechanism.

[0003] However, the existing collaborative model is limited by the sequential dependencies between computational steps in the attention mechanism, making it difficult for hardware computing units responsible for different types of operations (such as matrix operations and special function operations) to achieve efficient parallel processing. Especially when processing feature data of specific dimensions, existing scheduling methods often fail to balance the load of different computing units, easily causing some computing units to be idle and waiting, thus restricting the overall computing performance. Summary of the Invention

[0004] This invention provides an attention mechanism calculation method, electronic device, and storage medium to address the shortcomings of related technologies in performing attention mechanism calculations, such as low computational pipeline parallelism efficiency and insufficient hardware resource utilization.

[0005] This invention provides a method for calculating attention mechanisms, comprising:

[0006] Configure multiple thread bundles, which share the same set of thread local registers;

[0007] The computation process of the attention mechanism is broken down into multiple computational sub-steps, and the multiple computational sub-steps are respectively assigned to the multiple thread bundles, wherein each thread bundle is configured to execute the corresponding assigned computational sub-step;

[0008] The multiple thread bundles are controlled to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread-local registers to obtain the computational result of the attention mechanism.

[0009] According to an attention mechanism calculation method provided by the present invention, the plurality of calculation sub-steps include a data interaction step and a plurality of arithmetic operation steps, and the step of assigning the plurality of calculation sub-steps to the plurality of thread bundles includes:

[0010] The data interaction step is assigned to the first thread bundle group, and the data interaction step is used to perform data interaction between global memory and shared memory;

[0011] The arithmetic operation steps are assigned to multiple second thread bundles, wherein the first thread bundle runs in parallel with the multiple second thread bundles, and intermediate calculation results are passed between the multiple second thread bundles in a pipeline manner.

[0012] According to an attention mechanism calculation method provided by the present invention, the data interaction step includes:

[0013] Record the query matrix block, key matrix block, and value matrix block required for computation from the global memory to the shared memory;

[0014] The output results calculated by the multiple second thread bundles are written back from the shared memory to the global memory.

[0015] According to the attention mechanism calculation method provided by the present invention, the plurality of arithmetic operation steps include matrix multiplication operation step, row maximum value operation step, exponentiation operation step, row summation operation step, weighted summation and output update step;

[0016] The step of assigning the plurality of arithmetic operation steps to a plurality of second thread bundles includes:

[0017] Assign the matrix multiplication operation steps of the query matrix block and the key matrix block and the row maximum value operation steps to the first operation thread bundle;

[0018] The exponentiation step is assigned to the second operation thread group;

[0019] The row summation operation step is assigned to the third operation thread bundle;

[0020] The weighted summation and output update steps are assigned to the fourth operation thread group;

[0021] The data flow of the multiple arithmetic operation steps sequentially passes through the first operation thread group, the second operation thread group, the third operation thread group, and the fourth operation thread group.

[0022] According to the attention mechanism calculation method provided by the present invention, the matrix multiplication operation step of the query matrix block and the key matrix block and the row maximum value operation step include:

[0023] The first intermediate result is obtained by performing matrix multiplication operations between the query matrix block and the key matrix block using the tensor core.

[0024] The maximum row value is obtained by performing a row maximum value operation using the vector core;

[0025] The matrix multiplication operation and the row maximum value operation are processed in parallel within the first operation thread group.

[0026] According to the attention mechanism calculation method provided by the present invention, the exponential operation step includes:

[0027] Obtain the first intermediate result and the maximum value of the row;

[0028] Based on a special function unit, the first intermediate result and the maximum value of the row are used to perform an exponential operation to obtain the second intermediate result;

[0029] The data type of the second intermediate result is converted to obtain the probability matrix.

[0030] According to an attention mechanism calculation method provided by the present invention, the row summation operation step includes:

[0031] The second intermediate result is summed row by row to obtain the normalization coefficient;

[0032] The weighted summation and output update steps include:

[0033] The matrix multiplication operation between the probability matrix and the value matrix block is performed using the tensor core to obtain the third intermediate result;

[0034] The third intermediate result is updated based on the normalization coefficient and the output result of the previous stage to obtain the output result of the current stage.

[0035] According to the attention mechanism calculation method provided by the present invention, the calculation process of the attention mechanism is divided into multiple calculation sub-steps, and the method further includes the following steps beforehand:

[0036] Divide the data tensor to be processed into multiple data blocks;

[0037] The control of the multiple thread bundles to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread-local registers includes:

[0038] Different thread bundles are controlled to process different computational sub-steps corresponding to different data blocks within the same time slice.

[0039] The present invention also provides an attention mechanism computing device, comprising:

[0040] A configuration unit is used to configure multiple thread bundles, which share the same set of thread local registers.

[0041] The allocation unit is used to break down the computation process of the attention mechanism into multiple computational sub-steps and allocate the multiple computational sub-steps to the multiple thread bundles respectively, wherein each thread bundle is configured to execute the corresponding allocated computational sub-step;

[0042] An execution unit is used to control the multiple thread bundles to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread-local registers, so as to obtain the computational result of the attention mechanism.

[0043] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and running on the processor, wherein the processor executes the computer program to implement the attention mechanism calculation method as described above.

[0044] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the attention mechanism computation method as described above.

[0045] The present invention also provides a computer program product, including a computer program that, when executed by a processor, implements the attention mechanism calculation method as described above.

[0046] The attention mechanism computation method, electronic device, and storage medium provided by this invention, by configuring multiple thread bundles to share the same set of thread-local registers, break the limitation of register resource isolation between thread bundles in traditional architectures. This allows different thread bundles to directly transfer intermediate computation data by accessing shared registers, thereby avoiding the read / write overhead and communication latency caused by frequent data interaction through shared memory or global memory. Furthermore, by decomposing the complex computation process of the attention mechanism into multiple functionally simple computational sub-steps and assigning them to different thread bundles, fine-grained division of computational tasks is achieved. This enables multiple thread bundles to use shared registers as data transfer media to execute their respective computational tasks in parallel in a pipelined manner. This allows hardware computing units responsible for different types of operations to work simultaneously, improving overall hardware resource utilization and thus enhancing the pipeline efficiency and overall throughput of the attention mechanism computation. Attached Figure Description

[0047] To more clearly illustrate the technical solutions in this invention or related technologies, the accompanying drawings used in the description of the embodiments or related technologies will be briefly introduced below. Obviously, the accompanying drawings described below are some embodiments of this invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0048] Figure 1 This is a schematic diagram of the structure of the general-purpose graphics processor provided by the present invention;

[0049] Figure 2 This is a flowchart illustrating the attention mechanism calculation method provided by the present invention;

[0050] Figure 3 This is a schematic diagram of a pipeline for parallel execution of multiple thread bundles provided by the present invention;

[0051] Figure 4 This is a schematic diagram of the structure of the attention mechanism computing device provided by the present invention;

[0052] Figure 5 This is a schematic diagram of the structure of the electronic device provided by the present invention. Detailed Implementation

[0053] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0054] With the rapid development of Artificial Intelligence (AI) technology, deep learning large-scale models have achieved remarkable results in fields such as natural language processing and computer vision. Among these large models, the Transformer architecture dominates, and the attention mechanism is a core component of the Transformer architecture. To meet the high-performance requirements of large model training and inference, general-purpose graphics processing units (GPGPUs) are typically used for large-scale parallel computing acceleration. In GPGPU operator development, how to fully utilize the various computing resources on the chip to improve operator performance is a key issue in the software development field.

[0055] In the implementation of attention mechanism forward operators based on GPGPU chips, multiple warp groups are typically used to process data collaboratively. For example, one warp group acts as a producer, responsible for the interaction of input and output tensors between global memory (GLM) and shared memory (GSM); two warp groups act as consumers, processing different query (Q) data respectively and sharing key (K) and value (V) data. Under the scheduling of the software programmer, these two consumer warp groups collaborate to complete the computation of the attention mechanism, including matrix multiplication and normalization calculations related to Softmax.

[0056] However, the above scheme suffers from performance bottlenecks when processing feature data with specific dimensions. Specifically, when the feature dimension processed by the attention head is small (e.g., headDim=64, where headDim refers to the size of the feature dimension processed by each attention head), the computational cost of matrix multiplications (such as QK and QKV) performed based on the Tensor Core (tcore) is significantly reduced. In contrast, the computational cost of Special Function Unit (SFU) operations (such as exponentiation and row summation) performed based on the Vector Core (vcore) is not proportionally reduced. This results in the performance bottleneck of the operator shifting from the Tensor Core to the Special Function Unit (SFU) computation of the Vector Core when processing such small-dimensional data. In the above-mentioned collaborative mode of two consumer thread bundles, due to the computational order dependency of the operator itself, i.e., executing QK matrix multiplication, Softmax calculation, QKV matrix multiplication, and output update in sequence, it is difficult for the two thread bundles to form an efficient pipeline to interleave the use of different computational units. This means that within the same time slice, the SFU computing unit is often not fully utilized, or the computing unit is idle due to waiting for data dependencies, resulting in low hardware utilization of the overall operator and failure to achieve ideal computing performance.

[0057] To address this, the present invention provides an attention mechanism computation method based on the hardware characteristics of shared registers. By configuring multiple thread bundles, and these multiple thread bundles can share the same set of thread local registers (TLRs), this hardware characteristic enables efficient data interaction between thread bundles. This allows the attention mechanism computation process to be further divided into multiple computational sub-steps, which are then assigned to different thread bundles for parallel execution. This enables different types of computational units to work simultaneously, improving the utilization of computational resources and overcoming the aforementioned shortcomings.

[0058] It should be noted that the attention mechanism computation method provided by this invention is mainly applied to model training or inference scenarios in the field of artificial intelligence deep learning, especially for large language models (LLM) based on the Transformer architecture. The execution entity of this method is typically an electronic device with parallel computing capabilities, specifically a general-purpose graphics processing unit (GPGPU), graphics processing unit (GPU), AI accelerator card, or similar application-specific integrated circuit (ASIC). The structure of the execution entity of this invention will be described below using a GPGPU as an example.

[0059] Figure 1 This is a schematic diagram of the structure of the general-purpose graphics processor provided by the present invention, as shown below. Figure 1 As shown, this general-purpose graphics processing unit (GPGPU) includes at least multiple streaming processor clusters (SPCs) and global memory. Each SPC handles one computational task, or multiple SPCs collaboratively process a computational task, with data shared between SPCs via global memory. Each SPC includes multiple compute units (CUs) and intermediate caches (such as L1 caches) for sharing data among the compute units. Each compute unit includes multiple compute cores, which are used to execute specific computational tasks. Computational cores are typically scheduled in warp units, and multiple warp units form warp groups to work collaboratively.

[0060] Each computing unit also includes registers, shared caches, or shared memory for hierarchical storage of data related to computing tasks. Shared caches or shared memory within a computing unit are used to share data between the computing cores of that unit. Here, a shared cache (also known as an on-chip cache) is a temporary memory, for example, it could be a GMB (Gemm MainBuffer), which has a smaller capacity than an L1 cache but a faster data exchange speed. In contrast, registers have a smaller capacity than shared caches but a faster data exchange speed. Specifically, registers can be thread-local registers (TLRs).

[0061] GPGPUs may also include an L2 cache, a high-speed cache located between global memory and L1 cache, used to reduce global memory access latency. The L2 cache is typically a shared resource across the entire GPGPU chip, rather than being private to each SPC; multiple SPCs can share data through the L2 cache.

[0062] Furthermore, the computational units in SPC are divided into tensor computation units and vector computation units. Tensor computation units perform tensor computations, such as matrix multiplication and convolution. Within a tensor computation unit, the computational core, also called a tensor core (tcore), is used to perform specific computational tasks. Vector computation units primarily perform arithmetic and logical operations other than tensor computations such as matrix multiplication and convolution, such as accumulation, reduction, and regular addition, subtraction, multiplication, and division. Within a vector computation unit, the computational core, also called a vector core (vcore), is used to perform specific computational tasks.

[0063] In addition to the structures described above, the general-purpose graphics processor of this invention may also include other structures, and this invention does not specifically limit them.

[0064] Figure 2 This is a flowchart illustrating the attention mechanism calculation method provided by the present invention, as shown below. Figure 2 As shown, this method, through hardware-software co-design, fully utilizes hardware register resources and improves the utilization rate of hardware computing units, thereby enhancing the execution efficiency of the attention mechanism operator. The method includes the following steps:

[0065] Step S10: Configure multiple thread bundles, wherein the multiple thread bundles share the same set of thread local registers.

[0066] Specifically, in this step, the hardware resources for executing the computing task need to be configured first. Here, a warp refers to a set of logical or physical execution units consisting of several warps within a GPGPU computing unit. Each warp typically contains 32 or more threads that execute in parallel.

[0067] In traditional GPGPU architectures, thread-local registers (TLRs) are typically private, meaning that one thread bundle cannot directly access the register resources of another thread bundle. Cross-bundle communication usually requires the use of shared memory (GSM) or global memory (GLM), which introduces additional read / write overhead and latency.

[0068] The embodiments of this invention, based on a specific hardware architecture, configure a mechanism for sharing the same set of thread-local registers. This means that the configured multiple thread bundles are physically or logically mapped to the same register file region, or the hardware provides a specific access channel, allowing one thread bundle to directly read or write to the register space used by another thread bundle. This sharing feature is the foundation for achieving subsequent zero-copy communication or low-latency communication.

[0069] Step S20: The calculation process of the attention mechanism is divided into multiple calculation sub-steps, and the multiple calculation sub-steps are respectively assigned to the multiple thread bundles, wherein each thread bundle is configured to execute the corresponding assigned calculation sub-step.

[0070] Specifically, the computation process of the attention mechanism typically involves complex operations such as matrix multiplication between the query (Q) matrix, key (K) matrix, and value (V) matrix, scaling dot product, softmax normalization, and finally weighted summation.

[0071] In this step, to achieve fine-grained parallel processing, this complete long-chain computation process is broken down into several functionally single computational sub-steps. This breakdown is not arbitrary, but rather based on the type of operation (such as dense matrix computation, vector computation, special function computation, etc.) and the characteristics of the hardware units.

[0072] After the division is completed, these sub-steps are assigned to the multiple thread groups configured in step S10. Each thread group is specialized, meaning that a specific thread group is only responsible for executing a specific sub-step. For example, some thread groups are specifically responsible for data handling, some for matrix multiplication, and some for exponentiation. This allocation method makes the code logic within each thread group more compact and allows for instruction-level optimization for the specific type of operation it is responsible for.

[0073] Step S30: Control the multiple thread bundles to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread local registers to obtain the computational result of the attention mechanism.

[0074] Specifically, after tasks are assigned, multiple thread groups are started concurrently to execute multiple computational sub-steps in parallel in a pipeline manner. This parallel execution is based on the assignment logic of step S20, meaning each thread group only executes its assigned specific sub-step. Step S30 essentially drives the data blocks to be processed to flow sequentially between different thread groups according to the logical division of labor determined in S20.

[0075] The pipelined approach described above refers to the data stream to be processed. For example, due to memory limitations, a large data tensor is divided into multiple data blocks for processing. After the first thread group completes the first step of the calculation for the first data block, the result is passed to the second thread group for the second step of the calculation. At the same time, the first thread group can immediately start processing the first step of the calculation for the second data block.

[0076] In this embodiment of the invention, data transfer between different stages of the pipeline is achieved by reading and writing a shared Thread Local Register (TLR). Specifically, the intermediate results generated by the previous computation sub-step are directly stored in the shared TLR, and the thread bundle corresponding to the next computation sub-step directly reads data from the TLR for subsequent computation, or directly modifies the data at the original register location.

[0077] This approach avoids the cumbersome process of writing intermediate results back to shared memory (GSM) and then reading them out by the next group, reducing the number of memory access instructions and memory bandwidth pressure. As the pipeline fills up, multiple thread bundles process different data blocks at different stages simultaneously, and finally, after processing all sub-steps, the final calculation result of the attention mechanism (i.e., the output matrix) is obtained.

[0078] The method provided in this invention, by configuring multiple thread bundles to share the same set of thread-local registers, breaks the limitation of register resource isolation between thread bundles in traditional architectures. This allows different thread bundles to directly transfer intermediate computation data by accessing shared registers, thereby avoiding the read / write overhead and communication latency caused by frequent data interaction through shared memory or global memory. Furthermore, by decomposing the complex computation process of the attention mechanism into multiple functionally simple computational sub-steps and assigning them to different thread bundles, fine-grained division of computational tasks is achieved. This enables multiple thread bundles to use shared registers as data transfer media to execute their respective computational tasks in parallel in a pipelined manner. This allows hardware computing units responsible for different types of operations to work simultaneously, improving overall hardware resource utilization and thus enhancing the pipeline efficiency and overall throughput of the attention mechanism computation.

[0079] Based on any of the above embodiments, before step S20, the method further includes:

[0080] Divide the data tensor to be processed into multiple data blocks.

[0081] Specifically, in real-world large-scale model applications, the dimensions of the data tensors to be processed (typically including query tensors, key tensors, and value tensors) are often very large, far exceeding the capacity limitations of the GPGPU's on-chip high-speed memory (such as shared memory or registers). Therefore, a tiling strategy is required.

[0082] Specifically, the system divides the original data tensor into several smaller data blocks based on the sequence length or feature dimension, according to hardware resource constraints (such as the size of shared memory and the number of register files) and algorithm configuration (such as the dimension of the attention head, headDim). Each data block contains a small slice of data required for complete attention computation and serves as the basic unit for subsequent pipelined processing.

[0083] It is understood that the method provided in this embodiment of the invention can be applied to various scenarios, such as text processing, image processing, and speech processing. In different application scenarios, the data tensor to be processed has different physical meanings. For example, in a text processing scenario, the data tensor to be processed can be text data related to tasks such as text generation and text recognition. As another example, in an image processing scenario, the data tensor to be processed can be image-related data used in tasks such as image preprocessing, image segmentation, and object detection. As yet another example, in a speech processing scenario, the data tensor to be processed can be speech-related data used in tasks such as speech recognition and speech synthesis. This embodiment of the invention does not specifically limit these aspects.

[0084] Accordingly, step S30 specifically involves controlling different thread bundles to process different computational sub-steps corresponding to different data blocks within the same time slice.

[0085] Specifically, this describes the parallel operation of the pipeline. When the system starts, multiple thread bundles share register resources in space, but process different data objects in time.

[0086] It should be noted that the time slice mentioned above refers to a logical processing cycle when the pipeline architecture is in a stable parallel state. During this cycle, each thread bundle is busy at the same time, using different hardware resources (such as tensor cores, special function units, etc.) to process tasks in different stages (i.e., different data blocks) of the pipeline, thereby maximizing the computational parallelism of the hardware.

[0087] For example, suppose the pipeline is divided into five computational sub-steps: data loading, matrix multiplication, exponentiation, normalization, and output update, corresponding to five thread groups, such as WarpGroup0, WarpGroup1, WarpGroup2, WarpGroup3, and WarpGroup4. Within the same time slice, WarpGroup0 is loading the (N+4)th data block from global memory; simultaneously, WarpGroup1 is performing matrix multiplication on the (N+3)th data block, WarpGroup2 is performing exponentiation on the (N+2)th data block, WarpGroup3 is performing normalization on the (N+1)th data block, and WarpGroup4 is performing weighted summation and output update on the Nth data block.

[0088] This mechanism ensures that at any given time, every thread bundle in the pipeline is busy, with data blocks flowing sequentially through each thread bundle. Intermediate results of the currently processed data blocks are seamlessly passed between bundles via shared registers without waiting for each other.

[0089] The method provided in this invention maximizes computational throughput by dividing large data tensors into blocks and combining them with pipelined parallel control of multi-threaded beamforming. This design utilizes the latency masking principle, which uses the execution time of computationally intensive steps (such as matrix multiplication) to mask the latency of memory-intensive steps (such as data loading). Simultaneously, it allows different types of hardware units on the chip (such as tcores and SFUs) to process tasks on different data blocks at the same time, thereby enabling the GPGPU to maintain extremely high hardware utilization and parallel efficiency when processing input sequences of arbitrary sizes.

[0090] Based on the above embodiments, the plurality of computational sub-steps include a data interaction step and a plurality of arithmetic operation steps. Accordingly, in step S20, assigning the plurality of computational sub-steps to the plurality of thread bundles respectively includes:

[0091] Step S21: Assign the data interaction step to the first thread group, the data interaction step being used to perform data interaction between global memory and shared memory;

[0092] Step S22: The plurality of arithmetic operation steps are assigned to a plurality of second thread bundles, wherein the first thread bundle and the plurality of second thread bundles run in parallel, and intermediate calculation results are passed between the plurality of second thread bundles in a pipeline manner.

[0093] Specifically, the entire computational process of the attention mechanism can be logically divided into two major categories of operations with completely different natures: one is the data interaction step, and the other is the arithmetic operation step.

[0094] For data interaction steps, they can be assigned to the first thread group (such as WarpGroup0). This first thread group is configured specifically for data transfer, and its core task is to perform data interaction between Global Memory (GLM) and Shared Memory (GSM). In the GPGPU architecture, global memory has a large capacity but high access latency, while shared memory has a small capacity but fast access speed and extremely high bandwidth. By setting up a dedicated thread group to handle this level of communication, the asynchronous copy characteristics of hardware, such as TMA (Tensor Memory Accelerator), can be utilized to ensure efficient issuance of data transfer instructions and efficient actual data transmission.

[0095] For multiple arithmetic operations, they can be assigned to multiple second thread groups (such as WarpGroup1~WarpGroup4). These second thread groups focus on performing intensive mathematical calculations, such as matrix multiplication or exponentiation. It's important to note that the first thread groups and the multiple second thread groups run in parallel on the timeline. This means that while a second thread group is calculating the current data block using data in registers, the first thread group is already in the background prefetching the next data block from global memory into shared memory, or writing the previously calculated data block back to global memory.

[0096] Meanwhile, intermediate computation results are passed between multiple second-thread bundles in a pipelined manner. That is, different second-thread bundles are responsible for different stages of arithmetic operations. The intermediate results of the previous second-thread bundle are directly passed to the next second-thread bundle through shared thread-local registers, rather than through shared memory, thus forming a compact computation pipeline.

[0097] The method provided in this invention decouples data transport from core computation and assigns it to different thread bundles for parallel execution. This effectively hides the high access latency of global memory, so that computational units (such as tcores) do not have to be idle while waiting for data loading. This achieves a high degree of overlap between computation and memory access, and significantly improves the throughput performance of the entire attention mechanism operator.

[0098] Based on any of the above embodiments, the data interaction step includes:

[0099] Record the query matrix block, key matrix block, and value matrix block required for computation from the global memory to the shared memory;

[0100] The output results calculated by the multiple second thread bundles are written back from the shared memory to the global memory.

[0101] Specifically, the data interaction steps include data loading and result writing back. During the data loading phase, the first thread group (e.g., WarpGroup0) is responsible for loading the query matrix blocks (i.e., Q blocks), key matrix blocks (i.e., K blocks), and value matrix blocks (i.e., V blocks) required for computation from global memory into shared memory. Since the query matrix, key matrix, and value matrix involved in the attention mechanism in large models are typically large in size and cannot all be placed into on-chip memory at once, a block-based strategy is usually adopted. The first thread group reads the Q, K, and V block data to be processed in the current time slice from global memory according to the preset block size and stores it in high-bandwidth shared memory for subsequent reading by the second thread group and computation in registers.

[0102] In the result write-back phase, after multiple second-thread bundles complete a series of complex pipeline calculations, the final calculation results are temporarily stored in shared memory or registers. At this time, the first-thread bundle is also responsible for writing the final output results calculated by the multiple second-thread bundles from shared memory back to global memory. This step ensures that the final output of the attention mechanism is correctly persisted in global memory for use by subsequent neural network layers. In this process, the first-thread bundle acts as a storage interface, shielding the computational bundles from direct access to slow memory.

[0103] In this embodiment of the invention, by clearly defining the specific matrix block loading and result write-back mechanism, the on-chip shared memory resources are efficiently reused, enabling the method to process input data tensors of arbitrary scale without being limited by on-chip storage capacity. At the same time, all read and write operations involving global memory are centrally managed by a single thread bundle, reducing memory access conflicts and facilitating the use of hardware batch data transfer instructions (such as TMA) to maximize the utilization of global memory bandwidth.

[0104] Based on any of the above embodiments, the plurality of arithmetic operation steps include matrix multiplication, row maximum value calculation, exponentiation, row summation, weighted summation, and output update steps. Accordingly, step S22 specifically includes:

[0105] Assign the matrix multiplication operation steps of the query matrix block and the key matrix block and the row maximum value operation steps to the first operation thread bundle;

[0106] The exponentiation step is assigned to the second operation thread group;

[0107] The row summation operation step is assigned to the third operation thread bundle;

[0108] The weighted summation and output update steps are assigned to the fourth operation thread group;

[0109] The data flow of the multiple arithmetic operation steps sequentially passes through the first operation thread group, the second operation thread group, the third operation thread group, and the fourth operation thread group.

[0110] Specifically, when the computation process of the attention mechanism is broken down into data interaction steps and multiple arithmetic operation steps, these arithmetic operation steps can include matrix multiplication, row maximum value calculation, exponentiation, row summation, weighted summation, and output update. These steps constitute the core computational chain of the standard attention mechanism operator.

[0111] To adapt to the pipeline characteristics of the hardware, multiple arithmetic operation steps can be assigned to multiple second thread bundles (such as WarpGroup1~WarpGroup4) as follows: First, the matrix multiplication operation steps of the query (Q) matrix block and the key (K) matrix block, as well as the row maximum value operation steps, are assigned to the first operation thread bundle (such as WarpGroup1). This thread bundle is located at the beginning of the arithmetic pipeline and is responsible for calculating the product of Q and the transposed K, i.e., the attention score matrix, and simultaneously calculating the maximum value (Row Max) of each row in the score matrix. This maximum value will be used for the subsequent numerical stability processing of Softmax.

[0112] Secondly, the exponentiation steps are assigned to the second operation thread group (such as WarpGroup2). This thread group receives data from upstream and is responsible for performing non-linear exponentiation operations, which is the core non-linear transformation part of the Softmax function.

[0113] Next, the row summation step is assigned to the third operation thread group (such as WarpGroup3). This thread group is responsible for performing row summation on the result after the exponentiation operation to obtain the normalized coefficients (i.e., the denominator) of Softmax.

[0114] Finally, the weighted summation and output update steps are assigned to the fourth operation thread group (e.g., WarpGroup4). This thread group is at the end of the pipeline and is responsible for multiplying the probability matrix with the value (V) matrix block and updating the final output in conjunction with the current normalization coefficients.

[0115] Understandably, the data flow of multiple arithmetic operations passes sequentially through the first, second, third, and fourth operation thread groups. This means that the data flows like a factory assembly line, with each thread group only responsible for its assigned operation. After processing, the data remains in a shared register for direct use by the next group.

[0116] In this embodiment of the invention, load balancing of computational tasks is achieved by decomposing the complex attention computation logic into four independent pipeline stages and assigning them to different thread bundles. In particular, this allocation method fully considers the differences in computational power requirements of different computational steps, spatially decoupling computationally intensive matrix multiplication (handled by WarpGroup1 and WarpGroup4) from exponential operations that rely on Special Function Units (SFUs) (handled by WarpGroup2). This avoids pipeline congestion caused by a single thread bundle processing a certain step (such as a special nonlinear operation), allowing the special function units of the tensor core (tcore) and vector core (vcore) to work in parallel to the maximum extent, thereby improving the overall throughput.

[0117] Based on any of the above embodiments, the matrix multiplication operation step between the query matrix block and the key matrix block, and the row maximum value operation step, include:

[0118] The first intermediate result is obtained by performing matrix multiplication operations between the query matrix block and the key matrix block using the tensor core.

[0119] The maximum row value is obtained by performing a row maximum value operation using the vector core;

[0120] The matrix multiplication operation and the row maximum value operation are processed in parallel within the first operation thread group.

[0121] Specifically, when the first computation thread group (such as WarpGroup1) executes its task, it uses the tensor core to perform matrix multiplication operations on the query matrix block and the key matrix block to obtain the first intermediate result. The tensor core is a hardware unit in the GPGPU specifically designed to accelerate matrix multiply accumulating (MMA) operations. It can perform MMA calculations of Q and K extremely quickly to generate the attention score matrix (i.e., the first intermediate result).

[0122] At the same time, this thread bundle also utilizes vector cores (vcores) to perform row maximum value calculations, obtaining the row maximum value of the first intermediate result. Vector cores typically refer to general-purpose computing units in GPGPUs, which excel at handling logical, comparison, and arithmetic operations on scalars or vectors.

[0123] Matrix multiplication and row maximum value calculation are processed in parallel within the first operation thread bundle. This means that within the same instruction stream or time slice, the software can interleave matrix multiplication instructions (primarily using tcore) and row maximum value reduction instructions (primarily using vcore). Through fine-grained instruction scheduling, while the tensor core executes matrix multiplication instructions, the vector core is not idle but utilizes instruction-level parallelism to execute row maximum value lookup instructions. For example, when the tcore performs MMA calculations on Q and K to generate the attention score matrix, the vcore can be used to find the maximum value of each row generated, achieving parallelism. It should be understood that because these two types of instructions use different hardware execution units on the chip, they can be physically executed in parallel within the same time slice.

[0124] This invention effectively hides the overhead of row maximum value calculation by implementing matrix multiplication (using tcore) and row maximum value calculation (using vcore) in parallel within the same thread bundle. In related technologies, it is usually necessary to wait for the matrix multiplication to be completely completed before starting the reduction operation to calculate the maximum value separately. However, by allocating these two steps to the same thread bundle, this invention not only reduces the communication frequency between different thread bundles but also significantly reduces the overall execution time of this stage by utilizing the parallel cooperation of hardware units.

[0125] Based on any of the above embodiments, the exponentiation calculation step includes:

[0126] Obtain the first intermediate result and the maximum value of the row;

[0127] Based on a special function unit, the first intermediate result and the maximum value of the row are used to perform an exponential operation to obtain the second intermediate result;

[0128] The data type of the second intermediate result is converted to obtain the probability matrix.

[0129] Specifically, when the second operation thread group (such as WarpGroup2) performs the exponentiation step, it first obtains the first intermediate result and the maximum row value. As described in the previous embodiment, since multiple thread groups share the same set of thread local registers (TLRs), the acquisition action here does not require the execution of actual data copy instructions at the hardware level. Instead, it directly reads the corresponding address written by the first operation thread group (such as WarpGroup1) in the register file. Here, the first intermediate result is the score matrix (denoted as S) obtained by multiplying the query matrix block and the key matrix block, and the maximum row value is the maximum element (denoted as m) of each row of the score matrix.

[0130] Next, based on the special function unit, the first intermediate result and the maximum value of the row are used to perform exponential operations to obtain the second intermediate result. In the specific calculation logic, to prevent numerical overflow and ensure the stability of the calculation, the Sm method is usually used for processing. The second operation thread group calls the special function unit (SFU) in the GPGPU to perform this operation. The SFU is a hardware module specifically designed for calculating transcendental functions (such as sine, cosine, reciprocal, exponent, etc.). In this step, the SFU calculates exp(Sm), and the value obtained is the second intermediate result. It should be understood that exp here represents the natural exponential function.

[0131] Finally, the second intermediate result is converted to a different data type to obtain the probability matrix. In practical deep learning deployments, to balance accuracy and speed, exponentiation is typically performed at higher precision (e.g., FP32) to accumulate values, while subsequent matrix multiplications, to leverage the high performance of tcores, often require lower precision (e.g., FP16 or BF16). Therefore, after completing the exponentiation calculation, the second computation thread bundle executes a Typecast instruction to convert the data from FP32 to FP16 format, generating the unnormalized probability matrix (denoted as P).

[0132] In this embodiment of the invention, by separating and allocating the specific nonlinear operation of exponentiation to a second computation thread bundle, the typically scarce and independent SFU resources in the chip can be fully utilized. This design avoids resource contention between general-purpose computing units (vcores) and special-purpose computing units (SFUs) within a single thread bundle, solving the pipeline stall problem caused by mixed instruction streams in related technologies. Especially for scenarios with a small headDim (e.g., 64), it can significantly improve the throughput of SFUs and ensure that data stored in shared registers can quickly flow to the next stage after nonlinear transformation.

[0133] Based on any of the above embodiments, the row summation operation step includes:

[0134] The second intermediate result is summed row by row to obtain the normalization coefficient;

[0135] The weighted summation and output update steps include:

[0136] The matrix multiplication operation between the probability matrix and the value matrix block is performed using the tensor core to obtain the third intermediate result;

[0137] The third intermediate result is updated based on the normalization coefficient and the output result of the previous stage to obtain the output result of the current stage.

[0138] Specifically, when the third computation thread group (such as WarpGroup3) performs row summation, it first reads the second intermediate result calculated in the previous stage from the shared register, and then performs row summation on this second intermediate result to obtain the normalized coefficients. This is a reduction operation, and the calculated row sums will be used as the denominator of the Softmax formula for subsequent normalization of the probability matrix. This simple statistical step is separated into a pipeline stage to avoid blocking subsequent matrix multiplication operations.

[0139] When the fourth computation thread group (e.g., WarpGroup3) performs the weighted summation and output update steps, it first uses the tensor core to perform matrix multiplication of the probability matrix (P) and the value matrix (V) block, obtaining the third intermediate result. Here, the high-performance tensor core (tcore) is used again to perform the PV calculation. Since the probability matrix P has been prepared and its type converted by the second computation thread group (e.g., WarpGroup2), and the value matrix block V has been preloaded into shared memory by the first thread group (e.g., WarpGroup0), this step can be performed efficiently, and the calculated third intermediate result is the attention output part contributed by the current data block.

[0140] Subsequently, the third intermediate result is updated based on the normalization coefficients and the output of the previous stage to obtain the output of the current stage. In scenarios using block-based computation, the final output is updated incrementally. The fourth computation thread group uses the normalization coefficients calculated by the third computation thread group, as well as the accumulated output and normalization statistics from the previous stage (i.e., when processing the previous data block), to rescale and accumulate the current calculation result according to the online Softmax update formula. The updated result will serve as the final output of the current stage, awaiting writing back to global memory by the first thread group (such as WarpGroup0).

[0141] In this embodiment of the invention, fine-grained parallelism of computation is achieved by separating row summation (obtaining normalization coefficients) from the final weighted summation (matrix multiplication) and output update. The fourth operation thread bundle can focus on performing intensive matrix multiplication using tcore, while the scalar or vector operations required to update the output are interspersed at the end of the matrix operation pipeline. This design not only ensures the numerical correctness of the computation results (through the online update mechanism) but also ensures that the throughput of the computation pipeline is not limited by the latency of the normalization operation when processing long sequence data, thereby achieving high-performance execution of the overall attention mechanism operator.

[0142] Based on any of the above embodiments Figure 3 This is a schematic diagram of a pipeline for parallel execution of multiple thread bundles provided by the present invention, as shown below. Figure 3As shown, it demonstrates how multiple arithmetic operation steps are assigned to four operation thread bundles (i.e., WarpGroup1~WarpGroup4), which work together to achieve pipelined parallel computing. Figure 3 The horizontal axis represents different stages of the pipeline, and the vertical axis represents different thread bundles.

[0143] In the row containing WarpGroup1, we can see that the steps of multiplying the Q and K matrices overlap vertically with the steps of calculating the row maximum value. This is because WarpGroup1 uses the tensor core and vector core to process these two steps in parallel. While calculating the matrix multiplication of Q and K, the idle vector core is used to extract the row maximum value synchronously, thus saving the time of calculating the row maximum value separately.

[0144] from Figure 3 From the overall view shown, after WarpGroup1 completes the QK matrix multiplication and row maximum value operation for the first data block, it passes the result to WarpGroup2 via a shared TLR. At this point, the second time slice begins, and WarpGroup2 starts processing the exponentiation operation for the first data block; meanwhile, WarpGroup1 is not idle, but immediately begins processing the QK matrix multiplication and row maximum value operation for the second data block.

[0145] When the process enters the third time slice, the pipeline reaches full capacity. WarpGroup3 and WarpGroup4 begin processing the subsequent steps of the first data block (row summation and PV matrix multiplication and output update, respectively); meanwhile, WarpGroup2 is processing the second data block, and WarpGroup1 has already begun processing the third data block.

[0146] Figure 3 This clearly demonstrates that by breaking down the long computational chain of the attention mechanism into different sub-steps and assigning them to different thread bundles, temporal overlap of computational tasks is achieved. For example, at any given moment, all four thread bundles are working simultaneously, but they are processing different computational stages of different data blocks. This design maximizes the utilization of various computing units (tcore, vcore, SFU) on the GPGPU, avoiding pipeline stalls caused by single resource bottlenecks in traditional serial execution.

[0147] Furthermore, for operators in other scenarios, such as those with a relatively large headDim (e.g., 128), the performance bottleneck lies in matrix multiplication. In other words, when processing large-dimensional data, the performance bottleneck is the computation of the tensor core. For this scenario, this invention can still fully utilize the tcore through WarpGroup1 and WarpGroup4. Therefore, this invention is not dependent on a specific input shape; whether dealing with computationally intensive large-dimensional data or bandwidth-constrained small-dimensional data, it can leverage hardware performance through pipelined division of labor.

[0148] The attention mechanism computing device provided by the present invention is described below. The attention mechanism computing device described below and the attention mechanism computing method described above can be referred to in correspondence.

[0149] Based on any of the above embodiments Figure 4 This is a schematic diagram of the attention mechanism computing device provided by the present invention, as shown below. Figure 4 As shown, the device includes:

[0150] Configuration unit 410 is used to configure multiple thread bundles, wherein the multiple thread bundles share the same set of thread local registers;

[0151] The allocation unit 420 is used to divide the calculation process of the attention mechanism into multiple calculation sub-steps and allocate the multiple calculation sub-steps to the multiple thread bundles respectively, wherein each thread bundle is configured to execute the corresponding allocated calculation sub-step;

[0152] The execution unit 430 is used to control the multiple thread bundles to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread local registers, so as to obtain the computational result of the attention mechanism.

[0153] The apparatus provided in this invention, by configuring multiple thread bundles to share the same set of thread-local registers, breaks the limitation of register resource isolation between thread bundles in traditional architectures. This allows different thread bundles to directly transfer intermediate computation data by accessing shared registers, thereby avoiding the read / write overhead and communication latency caused by frequent data interaction through shared memory or global memory. Furthermore, by decomposing the complex computation process of the attention mechanism into multiple functionally simple computational sub-steps and assigning them to different thread bundles, fine-grained division of computational tasks is achieved. This enables multiple thread bundles to use shared registers as data transfer media to execute their respective computational tasks in parallel in a pipelined manner. Consequently, hardware computing units responsible for different types of operations can work simultaneously, improving overall hardware resource utilization and thus enhancing the pipeline efficiency and overall throughput of the attention mechanism computation.

[0154] Based on any of the above embodiments, the plurality of calculation sub-steps include a data interaction step and a plurality of arithmetic operation steps, and the allocation unit includes:

[0155] The first allocation subunit is used to allocate the data interaction step to the first thread bundle group, wherein the data interaction step is used to perform data interaction between global memory and shared memory;

[0156] The second allocation subunit is used to allocate the plurality of arithmetic operation steps to a plurality of second thread bundles, wherein the first thread bundle and the plurality of second thread bundles run in parallel, and the plurality of second thread bundles transfer intermediate calculation results in a pipeline manner.

[0157] Based on any of the above embodiments, the data interaction step includes:

[0158] Record the query matrix block, key matrix block, and value matrix block required for computation from the global memory to the shared memory;

[0159] The output results calculated by the multiple second thread bundles are written back from the shared memory to the global memory.

[0160] Based on any of the above embodiments, the plurality of arithmetic operation steps include matrix multiplication operation steps, row maximum value operation steps, exponentiation operation steps, row summation operation steps, weighted summation and output update steps; correspondingly, the second allocation subunit is specifically used for:

[0161] Assign the matrix multiplication operation steps of the query matrix block and the key matrix block and the row maximum value operation steps to the first operation thread bundle;

[0162] The exponentiation step is assigned to the second operation thread group;

[0163] The row summation operation step is assigned to the third operation thread bundle;

[0164] The weighted summation and output update steps are assigned to the fourth operation thread group;

[0165] The data flow of the multiple arithmetic operation steps sequentially passes through the first operation thread group, the second operation thread group, the third operation thread group, and the fourth operation thread group.

[0166] Based on any of the above embodiments, the matrix multiplication operation step between the query matrix block and the key matrix block, and the row maximum value operation step, include:

[0167] The first intermediate result is obtained by performing matrix multiplication operations between the query matrix block and the key matrix block using the tensor core.

[0168] The maximum row value is obtained by performing a row maximum value operation using the vector core;

[0169] The matrix multiplication operation and the row maximum value operation are processed in parallel within the first operation thread group.

[0170] Based on any of the above embodiments, the exponentiation calculation step includes:

[0171] Obtain the first intermediate result and the maximum value of the row;

[0172] Based on a special function unit, the first intermediate result and the maximum value of the row are used to perform an exponential operation to obtain the second intermediate result;

[0173] The data type of the second intermediate result is converted to obtain the probability matrix.

[0174] Based on any of the above embodiments, the row summation operation step includes:

[0175] The second intermediate result is summed row by row to obtain the normalization coefficient;

[0176] The weighted summation and output update steps include:

[0177] The matrix multiplication operation between the probability matrix and the value matrix block is performed using the tensor core to obtain a third intermediate result;

[0178] The third intermediate result is updated based on the normalization coefficient and the output result of the previous stage to obtain the output result of the current stage.

[0179] Based on any of the above embodiments, the device further includes a data partitioning unit, the data partitioning unit being used for:

[0180] Divide the data tensor to be processed into multiple data blocks;

[0181] Accordingly, the execution unit is specifically used for:

[0182] Different thread bundles are controlled to process different computational sub-steps corresponding to different data blocks within the same time slice.

[0183] Figure 5 An example is a schematic diagram of the physical structure of an electronic device, such as... Figure 5As shown, the electronic device may include a processor 510, a communication interface 520, a memory 530, and a communication bus 540, wherein the processor 510, the communication interface 520, and the memory 530 communicate with each other through the communication bus 540. The processor 510 can call logical instructions in the memory 530 to execute an attention mechanism calculation method. The method includes: configuring multiple thread bundles, the multiple thread bundles sharing the same set of thread local registers; dividing the attention mechanism calculation process into multiple calculation sub-steps, and assigning the multiple calculation sub-steps to the multiple thread bundles respectively, wherein each thread bundle is configured to execute the corresponding assigned calculation sub-steps; controlling the multiple thread bundles to execute the multiple calculation sub-steps in parallel in a pipelined manner by reading and writing the shared thread local registers to obtain the calculation result of the attention mechanism.

[0184] Furthermore, the logical instructions in the aforementioned memory 530 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to related technologies, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0185] On the other hand, the present invention also provides a computer program product, which includes a computer program that can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the attention mechanism calculation method provided by the above methods. The method includes: configuring multiple thread bundles that share the same set of thread local registers; dividing the attention mechanism calculation process into multiple calculation sub-steps and assigning the multiple calculation sub-steps to the multiple thread bundles respectively, wherein each thread bundle is configured to execute the corresponding assigned calculation sub-steps; controlling the multiple thread bundles to execute the multiple calculation sub-steps in parallel in a pipelined manner by reading and writing the shared thread local registers to obtain the calculation result of the attention mechanism.

[0186] In another aspect, the present invention also provides a non-transitory computer-readable storage medium storing a computer program thereon. When executed by a processor, the computer program implements the attention mechanism calculation method provided by the above methods. The method includes: configuring multiple thread bundles, the multiple thread bundles sharing the same set of thread local registers; dividing the attention mechanism calculation process into multiple calculation sub-steps and assigning the multiple calculation sub-steps to the multiple thread bundles respectively, wherein each thread bundle is configured to execute the corresponding assigned calculation sub-steps; controlling the multiple thread bundles to execute the multiple calculation sub-steps in parallel in a pipelined manner by reading and writing the shared thread local registers to obtain the calculation result of the attention mechanism.

[0187] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without any creative effort.

[0188] Through the above description of the embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus necessary general-purpose hardware platforms, and of course, it can also be implemented by hardware. Based on this understanding, the above technical solutions, in essence or the parts that contribute to the related technology, can be embodied in the form of software products. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods described in the various embodiments or some parts of the embodiments.

[0189] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A method for calculating an attention mechanism, characterized in that, include: Configure multiple thread bundles to be mapped to the same register file region so that the multiple thread bundles share the same set of thread local registers; The computation process of the attention mechanism is broken down into multiple computational sub-steps, and the multiple computational sub-steps are respectively assigned to the multiple thread bundles, wherein each thread bundle is configured to execute the corresponding assigned computational sub-step; The multiple thread bundles are controlled to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread-local registers to obtain the computational result of the attention mechanism; The plurality of computational sub-steps include a data interaction step and a plurality of arithmetic operation steps. The step of assigning the plurality of computational sub-steps to the plurality of thread bundles includes: The data interaction step is assigned to the first thread bundle group, and the data interaction step is used to perform data interaction between global memory and shared memory; The arithmetic operation steps are assigned to multiple second thread bundles, wherein the first thread bundle and the multiple second thread bundles run in parallel, the data stream of the multiple arithmetic operation steps passes through different second thread bundles in sequence, and the multiple second thread bundles transfer intermediate calculation results in a pipeline manner by reading and writing the shared thread local registers.

2. The attention mechanism calculation method according to claim 1, characterized in that, The data interaction steps include: Record the query matrix block, key matrix block, and value matrix block required for computation from the global memory to the shared memory; The output results calculated by the multiple second thread bundles are written back from the shared memory to the global memory.

3. The attention mechanism calculation method according to claim 1, characterized in that, The multiple arithmetic operation steps include matrix multiplication, row maximum value operation, exponentiation, row summation, weighted summation, and output update. The step of assigning the plurality of arithmetic operation steps to a plurality of second thread bundles includes: Assign the matrix multiplication operation steps of the query matrix block and the key matrix block and the row maximum value operation steps to the first operation thread bundle; The exponentiation step is assigned to the second operation thread group; The row summation operation step is assigned to the third operation thread bundle; The weighted summation and output update steps are assigned to the fourth operation thread group; The data flow of the multiple arithmetic operation steps sequentially passes through the first operation thread group, the second operation thread group, the third operation thread group, and the fourth operation thread group.

4. The attention mechanism calculation method according to claim 3, characterized in that, The matrix multiplication operation steps of the query matrix block and the key matrix block, and the row maximum value operation steps include: The first intermediate result is obtained by performing matrix multiplication operations between the query matrix block and the key matrix block using the tensor core. The maximum row value is obtained by performing a row maximum value operation using the vector core; The matrix multiplication operation and the row maximum value operation are processed in parallel within the first operation thread group.

5. The attention mechanism calculation method according to claim 4, characterized in that, The exponentiation calculation steps include: Obtain the first intermediate result and the maximum value of the row; Based on a special function unit, the first intermediate result and the maximum value of the row are used to perform an exponential operation to obtain the second intermediate result; The data type of the second intermediate result is converted to obtain the probability matrix.

6. The attention mechanism calculation method according to claim 5, characterized in that, The row summation operation steps include: The second intermediate result is summed row by row to obtain the normalization coefficient; The weighted summation and output update steps include: The matrix multiplication operation between the probability matrix and the value matrix block is performed using the tensor core to obtain the third intermediate result; The third intermediate result is updated based on the normalization coefficient and the output result of the previous stage to obtain the output result of the current stage.

7. The attention mechanism calculation method according to any one of claims 1 to 6, characterized in that, The process of breaking down the attention mechanism computation into multiple computational sub-steps also includes, prior to: Divide the data tensor to be processed into multiple data blocks; The control of the multiple thread bundles to execute the multiple computational sub-steps in parallel in a pipelined manner by reading and writing the shared thread-local registers includes: Different thread bundles are controlled to process different computational sub-steps corresponding to different data blocks within the same time slice.

8. An electronic device comprising a memory, a processor, and a computer program stored in the memory and running on the processor, characterized in that, When the processor executes the computer program, it implements the attention mechanism calculation method as described in any one of claims 1 to 7.

9. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by the processor, it implements the attention mechanism calculation method as described in any one of claims 1 to 7.

Images