Attention operator head dimension block calculation method applied to sea light DCU

Abstract

The invention relates to the field of heterogeneous parallel computing, in particular to an attention operator head dimension block computing method applied to a sea light DCU, which comprises the following steps of: expanding a sequence length block computing range in charge of each thread block in an attention operator and a matrix computing range processed by a single thread working group; a blocking parameter splitqk and a blocking parameter splitv are selected, wherein the blocking parameter splitqk and the blocking parameter splitv are selected; selecting a proper block cutting parameter splitqk for the head dimensions of the query tensor Q and the key tensor K for cutting, accumulating block calculation results, and calculating S block calculation results through a softmax function to obtain P block calculation results; selecting a proper block cutting parameter splitv for the head dimensions of the value tensor V and the output tensor O for cutting, wherein a block calculation result corresponds to a corresponding block of the output tensor O; and finally, segmenting the tensor O according to the splitv, and writing each partitioning result back to the global memory. The method is suitable for a heterogeneous parallel computing system composed of a CPU and a DCU, computing resources of a DCU computing unit can be saved, performance loss caused by resource overflow is reduced, and computing efficiency is improved.

Description

Technical Field

[0001] The present invention relates to the field of heterogeneous parallel computing, and in particular to a method for block-wise calculation of attention operator heads applied to Haiguang DCU. Background Art

[0002] With the deepening evolution of a new round of scientific and technological revolution and industrial transformation, the demand for computing resources in cutting-edge technology fields such as artificial intelligence, high-performance computing, and the Internet of Things is growing exponentially. In typical application scenarios such as deep learning model training, autonomous driving decision-making systems, and quantum chemistry simulations, the computing power provided by hardware resources has become a key factor restricting technological breakthroughs. While hardware resources provide the necessary computing power, traditional single-core computing architectures based on the CPU as the primary computing core are inefficient when processing parallel computing tasks, such as matrix multiplication, the core computing operation of neural networks, and cannot meet the high-throughput and high-efficiency computing requirements. Furthermore, with the slowdown of Moore's Law, the computing performance of a single CPU core has gradually slowed down.

[0003] Modern GPU architectures integrate thousands of processing cores and use a SIMT execution mode, enabling them to perform the same operations on massive amounts of data elements simultaneously. The CPU and GPU form a typical heterogeneous computing platform, where the CPU is responsible for handling complex control logic, task scheduling, and serialization operations, while the GPU is responsible for executing large-scale parallel computing tasks. Heterogeneous computing architectures provide better performance in specific tasks by dividing computing tasks into multiple subtasks and assigning them to different processing units. They can simultaneously support high-throughput computing and low-latency computing, and significantly reduce energy consumption by selecting the appropriate processor to perform specific tasks. As a representative of domestically produced high-performance GPGPUs, the Haiguang Deep Computing Processor (DCU) has become an important computing power carrier for AI training and inference scenarios with its large-scale parallel computing architecture and full-precision computing power support. However, with the surge in the number of parameters in large models, the computational complexity of traditional attention mechanisms has placed higher demands on hardware resources. As an important component of deep learning, the attention operator needs to ensure its efficiency. However, the attention operator requires a large amount of computing resources, and the amount of computing resources increases with the increase of head dimension. Excessive resource usage will lead to performance loss, making it impossible to fully utilize the computing power of the hardware. Therefore, it is necessary to improve the overall computing performance of the system by further subdividing the tasks and reducing resource usage. Summary of the Invention

[0004] In order to solve the technical problems existing in the prior art, the present invention provides a head-dimension block calculation method for the attention operator applied to Haiguang DCU, aiming to reduce resource usage during attention operator calculation and ensure that the operator maintains performance under larger head dimensions.

[0005] The first object of the present invention can be achieved by adopting the following technical solutions:

[0006] A method for calculating the block dimension of an attention operator head applied to Haiguang DCU, the method comprising:

[0007] S1. Based on the shared memory and registers of each computing unit on the Haiguang DCU and the matrix dimension requirements of the MMOP matrix calculation instructions, the sequence length block calculation range responsible for each thread block in the attention operator and the matrix calculation range processed by a single thread work group are expanded;

[0008] S2. Based on the header dimensions of the query tensor Q, key tensor K, value tensor V, and output tensor O, select the split parameters split_qk and split_v.

[0009] S3. Use the split parameter split_qk to split the header dimensions of the query tensor Q and the key tensor K respectively, obtaining multiple query tensor Q blocks and key tensor K blocks. Perform matrix multiplication on each query tensor Q block and key tensor K block to obtain block calculation results. Accumulate all block calculation results to obtain S block results. Calculate the S block results through the softmax function to obtain P block results.

[0010] S4. Use the block parameter split_v to split the head dimension of the value tensor V to obtain multiple value tensor V blocks, perform dot product calculations on the P block result and different value tensor V blocks, and calculate the output tensor O;

[0011] S5. Use the block parameter split_v to split the calculated output tensor O to obtain multiple output tensor O block results. Different output tensor O block results are written back to the global memory to obtain the final calculation result.

[0012] Specifically, the matrix dimension requirement of the MMOP matrix calculation instruction is 16×16×16. The matrix calculation executed by one MMOP matrix calculation instruction is performed by a thread work group, and each thread work group includes 64 threads.

[0013] Specifically, the split parameter split_qk is used to split the header dimension of the query tensor Q and the key tensor K into several sub-blocks, each block contains several consecutive headers;

[0014] The block parameter split_v is used to split the header dimension of the value tensor V and the output tensor O so that the calculation result of each block can be directly written into the corresponding sub-block of the output tensor O.

[0015] Specifically, the block parameter split_qk is selected based on the K dimension of the matrix calculation range processed by a single thread working group of the query tensor Q multiplied by the key tensor K, and the block parameter split_v is selected based on the N dimension of the matrix calculation range processed by a single thread working group of the probability distribution matrix P multiplied by the value tensor V. The header dimension of each block after slicing is an integer multiple of the dimension of the matrix calculation range processed by the single thread working group.

[0016] Specifically, the calculation formula of the S block result is as follows:

[0017]

[0018] Among them, S represents the attention score matrix, Q i Indicates the query tensor Q blocks after the query tensor Q is split, K i Indicates the key tensor K blocks after the key tensor K is split, d K Represents the head dimension of the key tensor K, split qk It is the block parameter split_qk.

[0019] Specifically, the calculation formula of the P block result is as follows:

[0020]

[0021] Specifically, the calculation formula of the output tensor O is as follows:

[0022] O=P·V=[PV0 PV1 … PV n-1 ]

[0023] Where P represents the block size of a sequence length block calculation range (block tile), n is the value of the block parameter split_v, V n-1 Represents the n-1th value tensor V block

[0024] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0025] This invention provides a block-based calculation method for the head dimension of an attention operator applied to a Haiguang DCU. This method modifies the original attention calculation method from calculating the entire head dimension in one go to multiple calculations. Based on the feature information of the attention operator input tensor, the head dimensions of the query tensor Q and key tensor K are segmented using the appropriate segmentation parameter split_qk. The results of the segmented calculations are then accumulated and softmax calculations performed. The head dimensions of the value tensor V and output tensor O are segmented using the appropriate segmentation parameter split_v. The results of the segmented calculations correspond to the corresponding blocks of the output tensor O. Finally, tensor O is segmented according to split_v, and each block is written back to global memory. This invention is applicable to heterogeneous parallel computing systems composed of CPUs and DCUs, saving computing resources in the DCU computing units, reducing performance losses caused by resource overflow, and improving computing efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the structures shown in these drawings without paying any creative work.

[0027] Figure 1 This is a flow chart of an attention calculation method based on head dimension segmentation in an embodiment of the present invention;

[0028] Figure 2 is a schematic diagram of the correspondence between block tiles and wavefront tiles in an embodiment of the present invention;

[0029] Figure 3 Schematic diagram of the corresponding relationship between the query tensor Q and the key tensor K block calculation in an embodiment of the present invention;

[0030] Figure 4 It is a schematic diagram of the correspondence between the P block result and the value tensor V block calculation in an embodiment of the present invention. DETAILED DESCRIPTION

[0031] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It is obvious that the embodiments described are only some embodiments of the present invention, not all embodiments, and the implementation of the present invention is not limited to these. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of the present invention.

[0032] Example 1:

[0033] This paper proposes an attention calculation method based on head dimension block. By dynamically splitting the head dimensions of the query tensor (Q), key tensor (K), value tensor (V), and output tensor (O), combined with block accumulation and memory optimization strategies, it can save computing resources of the DCU computing unit, reduce performance losses caused by resource overflow, and significantly improve the computing efficiency of Haiguang DCU. The general attention formula includes:

[0034] The standard scaled dot product attention formula is as follows:

[0035]

[0036] Among them, d K represents the head dimension of the key tensor K, Represents a scaling factor that controls the dot product variance.

[0037] The attention score calculation formula is to input the query vector Q and the key vector K to generate the attention score matrix S through matrix multiplication. Each element S of the attention score matrix S is ij represents the relevance of the i-th query to the j-th key. The attention score matrix S is expressed as:

[0038] S=Q·K T ;

[0039] Scaled dot product attention formula, dividing the attention score by the scaling factor After scaling to prevent the gradient from disappearing due to excessive values, the attention score matrix is converted into a probability distribution matrix P through the softmax function:

[0040]

[0041] The weighted summation formula uses the probability distribution matrix P to perform weighted summation on the value vector V to obtain the output matrix O:

[0042] O=P·V;

[0043] According to the characteristic information of the attention operator input tensor, the query tensor Q and the key tensor K are split by selecting the appropriate block parameter split_qk, and the softmax calculation is performed after the results of the block calculation are accumulated; the value tensor V and the output tensor O are split by selecting the appropriate block parameter split_v, and the results of the block calculation correspond to the corresponding blocks of the output tensor O; finally, the tensor O is split according to split_v, and each block is written back to the global memory. The present invention is applicable to heterogeneous parallel computing systems composed of CPU and DCU, which can save the computing resources of the DCU computing unit, reduce the performance loss caused by resource overflow, and improve computing efficiency.

[0044] like Figure 1As shown in the flowchart of a method for calculating attention based on head dimension block, the method for calculating the head dimension block of the attention operator applied to Haiguang DCU described in the present invention includes the following steps:

[0045] S1. Based on the shared memory and registers owned by each computing unit on the Haiguang DCU, and in accordance with the matrix dimension requirements of the MMOP (Matrix Multiply-Operation) matrix calculation instruction, expand the sequence length block calculation range (block tile) responsible for each thread block in the attention operator and the matrix calculation range (wavefront tile) processed by a single thread work group (wavefront).

[0046] The Hygon DCU's computing core utilizes a GPU-like architecture. Each compute unit (CU) contains four SIMD cores, each equipped with 16 vector compute units and a shared scalar compute unit. Each thread can access both vector and scalar registers, with the total number of registers determined by the Wavefront size (64 threads / Wavefront). Shared memory utilizes a hierarchical memory architecture, with the L1 cache working in conjunction with shared memory to support data sharing within a thread block. Dual DMA engines enable parallel data transfer and computation between the CPU and DCU.

[0047] Specifically, the matrix dimensions required for an MMOP matrix calculation instruction are M×N×K, where M×N×K is 16×16×16. The matrix calculations executed by an MMOP matrix calculation instruction are performed by a wavefront (thread workgroup), and each wavefront contains 64 threads. The GPU groups a large number of threads into wavefronts according to a fixed number. Threads within the same wavefront execute exactly the same instruction sequence but process different data, maximizing computational throughput through hardware-level parallel pipelines. A wavefront tile refers to the range of matrix calculations handled by a single wavefront. Expanding the wavefront tile allows a wavefront to be more than limited to 16×16×16 matrix calculations. By fusing multiple MMOP instructions to expand the wavefront tile, expansion in M, N, and K dimensions is achieved.

[0048] After expanding the block tile size, each thread block will calculate matrix tiles of corresponding sizes, and the size of the block tile directly affects shared memory usage. For example, if the block tile size is m_block*n_block, the query tensor Q will use tiles of m_block*headDim, and shared memory of the same size will be allocated for Q to handle data. The K and V tensors will also use tiles of n_block*headDim, and shared memory of the same size will also be allocated. Ultimately, this thread block will calculate an O tensor of m_block*n_block.

[0049] Expanding the wavefront tile size causes each wavefront to compute a correspondingly sized matrix block. Typically, four wavefronts are allocated to each thread block, and each wavefront computes an O-tensor of size (m_block / 4)*n_block. Wavefront tiles affect register usage. Larger tiles can reduce loop count and the number of registers used for loop control, but they also require more registers for computation. Smaller tiles increase loop count, and loop unrolling adds additional registers, but reduces the number of registers used for each computation.

[0050] like Figure 2 The following is a schematic diagram of the corresponding relationship between block tile and wavefront tile. T When the wavefront tile is 16×64×16 and m_block and n_block are 64, one block tile computes a 64×64 S block, and each wavefront computes a 16×64 S block. When the P·V wavefront tile is 16×32×16, one wavefront computation is performed twice, computing a 16×32 O block, ultimately resulting in a 16×64 O block. When m_block and n_block are 64, one block tile computes a 64×64 O block.

[0051] S2. Based on the query tensor Q, key tensor K, value tensor V, and the header dimension of the output tensor O, select the split parameters split_qk and split_v.

[0052] Specifically, the split_qk parameter is used to split the header dimension of the query tensor Q and key tensor K into several sub-blocks, each containing several consecutive headers. For example, if the total number of headers is 8, split_qk = 2 will be divided into 4 blocks, each processing 2 headers. The split_v parameter is used to split the header dimension of the value tensor V and the output tensor O, so that the calculation results of each block can be directly written to the corresponding sub-block of the output tensor O, reducing memory conflicts.

[0053] Each compute unit in the Hygon DCU can only host a limited number of resident thread blocks. Choosing a smaller tile size requires more thread blocks to complete the computation. At the same time, each thread block uses fewer computational resources, so the limited number of resident thread blocks cannot fully utilize resources, resulting in wasted excess resources. Furthermore, each compute unit has limited shared memory and register resources. Overly large tiles can lead to resource overflow and performance loss.

[0054] Specifically, the block parameter split_qk is selected based on the K dimension of the wavefront tile of the query tensor Q multiplied by the key tensor K, and the block parameter split_v is selected based on the N dimension of the wavefront tile of the probability distribution matrix P multiplied by the value tensor V. The header dimension of each block after slicing is an integer multiple of the wavefront tile dimension.

[0055] The split parameter, split_qk, is selected based on the size of the wavefront tile. For example, if the wavefront tile is m*n*k, headDimQK / n is selected, while split_v is headDimV / n. The Q times K wavefront tile may not be the same as the P times V wavefront tile, and the head dimensions may also be different. If headDim does not divide n evenly, for example, 160 / 64 is not an integer, then split is selected as 3. Because each data read is a fixed-size tile, the last split tile requires boundary verification. In this case, the wavefront tile can be reduced to a divisible size. The performance impact needs to be weighed in actual use.

[0056] S3. Use the block parameter split_qk to split the head dimensions of the query tensor Q and the key tensor K respectively to obtain multiple query tensor Q blocks and key tensor K blocks. Perform matrix multiplication on each query tensor Q block and key tensor K block to obtain the block calculation results. Accumulate all the block calculation results to obtain the S block result. Calculate the S block results through the softmax function to obtain the P block result.

[0057] Specifically, the calculation formula for the S block result is as follows:

[0058]

[0059] Among them, S represents the S block result, that is, the attention score matrix, Q i Indicates the query tensor Q blocks after the query tensor Q is split, K i Indicates the key tensor K blocks after the key tensor K is split, d K Represents the head dimension of the key tensor K, split qk It is the block parameter split_qk.

[0060] In this embodiment, the block parameters are selected in combination with hardware characterization. The block parameter split_qk depends on the K dimension of the wavefront tile, and the split_v parameter depends on the N dimension of the wavefront tile. The value of the split parameter allows each block to reduce the shared memory and register requirements when processing K-dimensional blocks. The block parameter split_qk is used to split the header dimensions of the query tensor Q and the key tensor K to obtain smaller sub-blocks. After blocking, the size of each sub-block must be an integer multiple of the dimension of the wavefront tile. When the block size is an integer multiple of the warp processing width, the boundary condition judgment during calculation can be eliminated, thereby achieving fully merged memory access.

[0061] Specifically, to perform matrix multiplication on the query tensor Q and the key tensor K, the softmax calculation can only be performed after the calculation of a sequence block is completed. If the head dimension is divided into blocks, all blocks need to be calculated separately and then accumulated before the softmax calculation can be performed. When calculating the product of the query tensor Q and the key tensor K, due to the global normalization characteristics of the softmax operation in the attention mechanism, any block calculation strategy must ensure that no local softmax calculation can be performed before the complete S matrix is obtained. When the head dimension block strategy is adopted, the result of each sub-block calculation is actually a component of the complete attention matrix. The softmax calculation can only be performed after all blocks have completed the calculation and the complete result is obtained through the accumulation operation.

[0062] Perform softmax calculation on the S block results, and convert the attention score matrix into a probability distribution matrix P through the softmax function to obtain P blocks. Figure 3 The following figure shows the block-wise calculation relationship between the query tensor Q and the key tensor K. When the header dimension is 192 and split_qk is 3, Q and K are each divided into 3 blocks. Each block corresponds to a dot product calculation, and the results of each block are added together to obtain an S-block result. The S-block result is then subjected to a softmax calculation to obtain the P-block result. The calculation formula for the P-block result is as follows:

[0063]

[0064] S4. Use the block parameter split_v to split the head dimension of the value tensor V into multiple blocks of the value tensor V. The original one-time calculation of the entire n_block × headDim block is split into split_v times. The dot product calculation is performed on the P block result and the different value tensor V blocks to calculate the output tensor O. The calculated output tensor O is written to shared memory.

[0065] Assume that the value of the block parameter split_v is equal to n, then the head dimension of the value tensor V is divided into n sub-blocks, each sub-block participates in the calculation independently, and the calculation formula of the output tensor O is as follows:

[0066] O=P·V=[PV0 PV1 … PV n-1 ]

[0067] Among them, P represents the block size of a sequence length block calculation range (block tile), V n-1 Represents the n-1th value tensor V block.

[0068] like Figure 4 The figure below shows the relationship between the block results of P and the block calculation of the value tensor V. When the block parameter split_v is 3 and the header dimension is 192, the value tensor V and the output tensor O are divided into 3 blocks. Each block of V is dot-producted with P, and the block result at the corresponding position in the output tensor O is finally obtained. The block result at the corresponding position of the output tensor is written to shared memory.

[0069] Since each calculation after block splitting does not need to wait for the entire sequence block to be read, calculation can be performed after each header dimension block is read. This allows for better masking of reading and calculation. Furthermore, data does not need to be reused, so shared memory can be reduced to 1 / split. Furthermore, registers can be released immediately after each block is read and calculated, reducing the risk of register overflow.

[0070] S5. Use the block parameter split_v to split the calculated output tensor O to obtain multiple output tensor O block results. Different output tensor O block results are written back to the global memory to obtain the final calculation result.

[0071] After performing a softmax calculation on the results of the S blocks to obtain the P blocks, the final result O can be calculated. The results of each block of the value tensor V are independent of each other; the calculated result corresponds to the block corresponding to the output tensor O. Therefore, different result blocks can be written directly to global memory without waiting for other blocks, thus obtaining the final Attention calculation result. Since the calculated output tensor O needs to be written back to global memory, that is, when writing back the tensor O, the result must be written to shared memory and then to global memory. According to global memory, this means writing back m_block * headdim O, so this global memory block must be split in the code.

[0072] In summary, the present invention provides a block calculation method for the head dimension of the attention operator applied to Haiguang DCU. By modifying the original attention calculation method from calculating the entire head dimension in one go to multiple calculations, the head dimension of the query tensor Q and the key tensor K is divided by the appropriate block parameter split_qk according to the feature information of the attention operator input tensor, and the softmax calculation is performed after the results of the block calculation are accumulated; the head dimension of the value tensor V and the output tensor O is divided by the appropriate block parameter split_v, and the results of the block calculation correspond to the corresponding blocks of the output tensor O; finally, the tensor O is divided according to split_v, and each block is written back to the global memory separately. The present invention is applicable to heterogeneous parallel computing systems composed of CPUs and DCUs, which can save computing resources of DCU computing units, reduce performance losses caused by resource overflow, and improve computing efficiency.

[0073] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be considered as equivalent replacement methods and are included in the scope of protection of the present invention.

Claims

1. A block-wise calculation method for attention operator head dimensions applied to Haiguang DCU, characterized in that: The following steps are involved: S1. Based on the shared memory and registers of each computing unit on the Haiguang DCU and the matrix dimension requirements of the MMOP matrix calculation instructions, the sequence length block calculation range responsible for each thread block in the attention operator and the matrix calculation range processed by a single thread work group are expanded; S2. Based on the header dimensions of the query tensor Q, key tensor K, value tensor V, and output tensor O, select the split parameters split_qk and split_v. S3. Use the split parameter split_qk to split the header dimensions of the query tensor Q and the key tensor K respectively, obtaining multiple query tensor Q blocks and key tensor K blocks. Perform matrix multiplication on each query tensor Q block and key tensor K block to obtain block calculation results. Accumulate all block calculation results to obtain S block results. Calculate the S block results through the softmax function to obtain P block results. S4. Use the block parameter split_v to split the head dimension of the value tensor V to obtain multiple value tensor V blocks, perform dot product calculations on the P block result and different value tensor V blocks, and calculate the output tensor O; S5. Use the block parameter split_v to split the calculated output tensor O to obtain multiple output tensor O block results. Different output tensor O block results are written back to the global memory to obtain the final calculation result.

2. The method for calculating the attention operator head dimension block applied to Haiguang DCU according to claim 1 is characterized by: The matrix dimension requirement of the MMOP matrix calculation instruction is 16×16×16. The matrix calculation executed by one MMOP matrix calculation instruction is performed by a thread work group, and each thread work group includes 64 threads.

3. The attention operator head dimension block calculation method applied to Haiguang DCU according to claim 1 is characterized by: The split parameter split_qk is used to split the header dimension of the query tensor Q and the key tensor K into several sub-blocks, each containing several consecutive headers; The block parameter split_v is used to split the header dimension of the value tensor V and the output tensor O so that the calculation result of each block can be directly written into the corresponding sub-block of the output tensor O.

4. The attention operator head dimension block calculation method applied to Haiguang DCU according to claim 3 is characterized by: The block parameter split_qk is selected based on the K dimension of the matrix calculation range processed by a single thread working group of the query tensor Q multiplied by the key tensor K, and the block parameter split_v is selected based on the N dimension of the matrix calculation range processed by a single thread working group of the probability distribution matrix P multiplied by the value tensor V. The header dimension of each block after slicing is an integer multiple of the dimension of the matrix calculation range processed by the single thread working group.

5. The method for calculating the attention operator head dimension block applied to Haiguang DCU according to claim 4 is characterized in that: The calculation formula of the S block result is as follows: Among them, S represents the attention score matrix, Q i Indicates the query tensor Q blocks after the query tensor Q is split, K i Indicates the key tensor K blocks after the key tensor K is split, d K Represents the head dimension of the key tensor K, split qk It is the block parameter split_qk.

6. The attention operator head dimension block calculation method applied to Haiguang DCU according to claim 5 is characterized by: The calculation formula of the P block result is as follows:

7. The attention operator head dimension block calculation method applied to Haiguang DCU according to claim 5 is characterized by: The calculation formula of the output tensor O is as follows: O=P·V=[PV0 PV1 … PV n-1 ] Where P represents the block size of a sequence length block calculation range (block tile), n is the value of the block parameter split_v, V n-1 Represents the n-1th value tensor V block.

Images