An NPU-based flash attention weight calculation stage pipeline parallel optimization method

Abstract

This invention discloses a pipeline parallel optimization method for the FlashAttention attention weight calculation stage based on an NPU. Addressing the serialization problem in the attention weight calculation pipeline of FP32 input scenarios, the method achieves initial parallelism between transport and calculation instructions by setting more precise synchronization control than existing solutions. Furthermore, it solves the problem of excessively long transport times in some rounds by readjusting the existing pipeline optimization technique, the Ping-Pong strategy, and enabling DB. For the pipeline blockage phenomenon caused by inappropriate use of cache space in the attention weight calculation stage of FP16 input scenarios, this invention proposes a DPP strategy by reasonably allocating cache space and redesigning the calculation pipeline based on the reallocated space, thus eliminating pipeline blockage caused by read-write conflicts.

Description

Technical Field

[0001] This invention belongs to the field of code optimization technology for core operator algorithms of large models, and relates to a pipelined parallel optimization method, specifically a pipelined parallel optimization method for the attention weight calculation stage of the FlashAttention forward operator (FA operator for short) based on NPU. Background Technology

[0002] Self-attention, a core component of the Transformer architecture, empowers deep learning models with the crucial ability to capture long-range dependencies in input sequences. However, the time and space complexity of classic self-attention operators increases quadratically with sequence length. This characteristic leads to a sharp increase in memory usage and computational overhead, posing a severe computational bottleneck for large models when processing long input sequences, significantly limiting the scale of model pre-training and the efficiency of inference deployment. Therefore, efficient acceleration optimization of self-attention operators has become a research hotspot in both academia and industry.

[0003] Represented by the FlashAttention series, these methods do not alter the model structure but delve into the storage hierarchy and instruction pipeline characteristics of the accelerator, reconstructing the computation order through block computation and operator fusion. Building upon the foundations laid by FlashAttention-1 / 2 in reducing global memory access and fully utilizing on-chip cache, FlashAttention-3 further introduces a producer-consumer asynchronous parallel model. Utilizing the asynchronous memory copying and thread bundle scheduling mechanisms of the underlying hardware, it achieves extreme latency hiding for matrix multiplication and Softmax operations. The latest FlashAttention-4 is deeply adapted to next-generation hardware architectures, proposing a "lazy" online Softmax rescaling and exponential operation simulation approximation technique. By eliminating up to 90% of redundant rescaling computations at the algorithm level and using high-throughput multiply-accumulate instructions to mask complex vector operation latency, it brings computational throughput close to the physical limits of the hardware. Because this precise acceleration scheme, evolving from memory-constrained to computationally pipelined, can significantly improve training and inference speeds without sacrificing model performance, it has become a core evolutionary direction for large model operator optimization.

[0004] In line with the trend of hardware-aware optimization, the domestic Ascend computing platform has also actively introduced and deepened the engineering adaptation of the FlashAttention series of algorithms. In the Ascend 910B AI processor and its supporting CANN heterogeneous computing architecture, the official high-performance fusion operator library is not limited to basic algorithm implementation, but deeply integrates the block computing and operator fusion architecture of FlashAttention-1 / 2, and further draws on the extreme optimization ideas of FlashAttention-3 / 4 for the underlying hardware pipeline.

[0005] Specifically, the existing Ascend FA operator adaptation scheme cleverly achieves parallel queueing of transport instructions and computation instructions by designing a Ping-Pong buffer strategy on the vector side, and uses an efficient memory handover mechanism to mask data transport latency. In terms of the collaboration between matrix and vector computation units, it introduces an asynchronous multi-stage pipelined parallel mechanism similar to FlashAttention-3 / 4. Through early firing and Cube dual firing strategies, the single-round serial computation link is deconstructed into four highly overlapping stages (attention score calculation, attention weight calculation, attention output calculation, attention output correction and rescaling).

[0006] The calculation process of the FA operator implemented based on Ascend hardware is as follows: Figure 1 As shown, it is mainly divided into four parts: attention score calculation, attention weight calculation, attention output calculation, and attention output correction and rescaling. Attention score calculation means performing Matmul calculation on the Query and the transposed Key. Attention weight calculation means performing position encoding, scaling, Mask, Softmax, and DropOut on the attention score. Attention output calculation means performing Matmul calculation on the result of the attention weight and the Value. Attention output correction and rescaling means multiplying the attention output of the previous round with the intermediate result exp value of the current round of Softmax, completing the correction, and then adding it to the attention output of the current round. The last inner loop iteration is divided by the sum value (the sum of rows) in the Softmax result to complete the rescaling operation of the attention output. The formula for calculating the exp value is shown in formula (1).

[0007]

[0008] To improve the utilization of the Ascend 910B computing unit, existing best practices have adopted a variety of parallel optimization strategies, including V-side pipelined parallelism, CV-side pipelined parallelism, and the IBShare strategy on the C-side.

[0009] The core computing unit in the Ascend chip is the AI ​​Core, which includes three basic computing units: Cube (matrix) computing unit, Vector (vector) computing unit, and Scalar (scalar) computing unit. It also includes storage units (including hardware storage and data transfer units) and data transfer units. The hardware architecture is divided into two types based on whether the Cube and Vector computing units are deployed on the same core: coupled architecture and discrete architecture. The Ascend 910B uses a discrete architecture, and its hardware structure diagram is shown below. Figure 2 As shown.

[0010] The Ascend 910B splits the AI ​​Core into two independent cores: Matrix Calculation (AI Cube, AIC) and Vector Calculation (AI Vector, AIV). Each core has its own Scalar unit, capable of independently loading its own code segments. This decouples matrix and vector calculations, allowing them to work together under unified system software scheduling to optimize computational efficiency. Data is transferred between AIV and AIC via off-chip memory (L2 Cache and GM).

[0011] In the matrix operation module, the SLU controls the scalar instruction flow, the L1 buffer register stores the input matrix data, the L0A and L0B registers store the two input matrix data after data format preprocessing, the CUBE (the unit that actually performs matrix operations) performs various matrix operations, and the L0C register stores the result matrix. The matrix calculation module is primarily responsible for performing matrix multiplication-related operations. In the vector operation module, there is only one general-purpose register, UB (Unified Buffer), used to store all vector data, and the VEC (the unit that actually performs vector calculations) is used to perform all vector operations. The vector calculation module is primarily responsible for performing all vector operations.

[0012] In addition to storage and computing units, AI Core also contains data transfer units. For example... Figure 2 As shown, MTE2 represents a load instruction, with its load path from GM to the L1 Buffer, L0A / L0B of the AIC, and UB of the AIV. MTE1 represents a load instruction within the AIC, with its load path from the L1 Buffer to L0A / L0B. MTE3 represents a load instruction within the AIV, with its load path from UB to GM. Fixp represents a load instruction within the AIC, with its load path from L0C to the L1 Buffer and GM.

[0013] After performing performance analysis on the Ascend 910B AI processor on the FA operator (the corresponding operator in this fusion operator library is aclnnFlashAttentionScore) in the Ascend official fusion operator library (v0.8-8.0RC3), it was found that there is a pipeline blockage phenomenon in the computational pipeline of attention weight calculation.

[0014] The specific analysis is as follows:

[0015] The first step is a pipeline analysis of the attention weight calculation stage in the FP32 scenario.

[0016] By analyzing the profiling pipeline diagram of the baseline operator in the FP32 scenario, such as... Figure 3 As shown, the execution flow of the attention weight calculation stage exhibits a clear serialization characteristic, meaning that the input of data in round i+1 must wait until the output of data in round i is completely completed before it can start. This phenomenon causes the expected parallel benefits of the Ping-Pong strategy to completely fail.

[0017] The root cause lies in the fact that existing implementations fail to differentiate between FP32 and FP16 scenarios. To mitigate the risk of memory overflow in FP16 scenarios, existing implementations forcibly insert a WaitFlag(MTE3ToMTE2) synchronization wait instruction at the beginning of the attention weight calculation code loop. However, in FP32 scenarios, since data does not require additional type conversion, this coarse-grained synchronization mechanism actually creates excessive constraints, resulting in severe pipeline congestion (serialization).

[0018] Secondly, we analyze the pipeline of attention weight calculation in the FP16 scenario. Compared to the FP32 scenario, the FP16 scenario has an additional type conversion operation at the end of the attention weight calculation stage. This is because the type of attention weight is FP32 regardless of the input scenario. In order to ensure that the two input types of attention output calculation are consistent (the value type is FP16), it is necessary to convert the attention weight from FP32 type to FP16 type.

[0019] In the existing baseline implementation, to accommodate data storage after the Cast operation, a large, idle buffer (used to store the maskBuf of the DropMask matrix) was chosen as the target address for the Cast operation and the source address for the dropout. The initial intention of choosing this buffer was to delay pipeline blocking as much as possible (Dropout is the penultimate operation in the attention weight calculation stage). However, this strategy based on idle reuse ignores the single-channel characteristic of the Ascend MTE2 (dropout) pipeline.

[0020] like Figure 4As shown, VEC2 represents the Dropout and Cast vector calculations, and VEC1 represents the remaining vector calculations after excluding the vector calculations in VEC2. In actual execution, the DropMask matrix loading operation in round i must wait for the calculation results of round i-1 to be loaded out of maskBuf before it can be executed. Since the MTE2 channel is serial, this dependency not only blocks the DropMask matrix loading in the current round but also causes a head-of-line blocking effect, meaning that all subsequent loading instructions (including the input data in round i+1) are blocked. Although logically the input data in round i+1 can start loading immediately after the data handover between PingBuffer and PongBuffer in round i, the head-of-line blocking effect forces the pipeline to degenerate into a serial loading-out mode, resulting in idle computation units and pipeline obstruction. Summary of the Invention

[0021] To address the pipeline congestion problem in baseline algorithms, this invention provides a pipeline parallel optimization method for the FlashAttention attention weight calculation stage based on an NPU. The baseline algorithm uses the aclnnFlashAttentionScore implemented in the existing cann-ops-adv (v0.8-8.0RC3) repository (Huawei's official open-source repository for fusion operators), which is the forward implementation of FlashAttention (referred to as the FA operator). For the attention weight calculation stage, this invention proposes a precision-aware differentiated optimization strategy: for the FP32 scenario, deep pipeline optimization through delayed handover and double buffering strategies effectively masks data transfer latency; for the FP16 scenario, a DPP storage rearrangement scheme is proposed, which eliminates pipeline congestion caused by read-write conflicts by optimizing the spatial layout of the unified buffer.

[0022] The objective of this invention is achieved through the following technical solution:

[0023] A pipelined parallel optimization method for the FlashAttention attention weight calculation stage based on NPU includes the following steps:

[0024] Step 1: For the FP32 scenario, a deep optimization strategy combining delayed handover and double buffering is adopted. The specific steps are as follows:

[0025] Step 11: Delay the data handover from PongBuffer to PingBuffer from the scaling operation to the blanking operation added after Dropout. This blanking operation is a scalar multiplication (Muls), which multiplies each number in the data to be calculated by 1.

[0026] Step 12: Logically split PongBuffer and CommonBuffer into two sub-blocks, denoted as PongBuffer[0], PongBuffer[half], CommonBuffer[0], and CommonBuffer[half] respectively;

[0027] Step 13: Move half of the position encoding matrix, the entire AttnMask matrix (the mask matrix required for the attn_mask operation), and half of the attention score matrix into CommonBuffer[0], MaskPingBuffer, and PongBuffer[0] in sequence. The premise for moving half of the position encoding matrix is ​​that the position encoding matrix part corresponding to CommonBuffer[0] in the previous round has been calculated. The premise for moving the AttnMask matrix is ​​that the attn_mask calculation in the previous round has been completed. Due to the delayed handover in Step 1, the premise for moving half of the attention score matrix is ​​that the scalar multiplication operation in the previous round has been completed.

[0028] Step 14: After inserting the synchronization wait instruction of Set / WaiFlag (Mte2toV), move the other half of the position encoding matrix and the other half of the attention score matrix into CommonBuffer[half] and PongBuffer[half] in sequence. After issuing the move instruction, start the addition of position encoding and scaling calculations corresponding to CommonBuffer[0] and PongBuffer[0].

[0029] Step 15: After inserting the Set / WaiFlag (Mte2toV) synchronization wait instruction, start the addition of position encoding and scaling calculations for CommonBuffer[half] and PongBuffer[half]. After completing this calculation, reset PongBuffer to the full view, meaning that subsequent calculations will be performed on the full PongBuffer.

[0030] Step 16: After completing the calculation of attn_mask, first issue the instruction to move the DropMask matrix into MaskPongBuffer, and then start the calculation of Softmax and Dropout in sequence. Before the calculation of Dropout, a synchronization wait instruction of Set / WaiFlag (Mte2toV) needs to be inserted.

[0031] Step 17: Finally, complete the scalar multiplication (Muls) operation, transfer the data from PongBuffer to PingBuffer, and let PingBuffer complete the subsequent data transfer. The prerequisite for starting the scalar multiplication (Muls) calculation is that the previous round of data transfer is completed, to avoid PingBuffer sending data in a chaotic manner.

[0032] Step 2: For FP16 scenarios, adopt the DPP (Double-Ping-Pong) strategy. The specific steps are as follows:

[0033] Step 21: Split PingBuffer into two sub-blocks, denoted as PingBuffer[0] and PingBuffer[half] respectively. The same applies to PongBuffer, PseBuffer, and CommonBuffer.

[0034] Step 22: Decompose and rearrange the instructions of the calculation process in the attention weight calculation stage. Divide the calculation process into VEC1 (from the beginning to before FlashSoftmax) and VEC2 (from FlashSoftmax to the end of Cast). In order to distinguish the calculations of different Buffers, the corresponding Buffer will be added before the corresponding calculation process. For example, PongBuffer[0]-VEC1 means VEC1 calculation with PongBuffer[0] enabled.

[0035] Step 23: Move half of the position encoding matrix (FP16) into PseBuffer[0], complete the type conversion operation and save the result into CommonBuffer[0]. The premise of moving the position encoding matrix is ​​that the previous round of type conversion operation is completed, and the current round of type conversion operation needs to insert the Set / WaiFlag (Mte2toV) synchronous wait instruction before execution.

[0036] Step 24: Move half of the attention score matrix and the complete AttnMask matrix into PonBuffer[0] and MaskPongBuffer in sequence. The prerequisite for moving the attention score matrix is ​​that the previous round of PongBuffer[0]-VEC1 calculation is completed. The prerequisite for moving the AttnMask matrix is ​​that the previous round of PingBuffer[half]-VEC1 calculation is completed.

[0037] Step 25: After inserting the synchronization wait instruction of Set / WaiFlag (Mte2toV), first move the other half of the attention score matrix and the position encoding matrix into PingBuffer[half] and PseBuffer[half] respectively, then start the calculation of PongBuffer[0]-VEC1, and dump the final calculation result into PingBuffer[0].

[0038] Step 26: After inserting the Set / WaiFlag (Mte2toV) synchronization wait instruction, first complete the type conversion operation of PseBuffer[half], convert this part of the position encoding matrix to FP32 type and dump it into CommonBuffer[half]. Then start the calculation of PingBuffer[half]-VEC1. After PingBuffer[half]-VEC1 is completed, the views of PingBuffer need to be merged, that is, the subsequent calculations will be performed on the complete PingBuffer.

[0039] Step 27: Start PingBuffer-VEC2 calculation. Note that the timing of importing and calculating the DropMask matrix here is consistent with the description in the FP32 scenario, so it will not be repeated here.

[0040] Step 28: After the calculation is completed, the result (FP16) is dumped into PongBuffer[half], and PongBuffer[half] is responsible for moving the data out. The prerequisite for Cast execution in PingBuffer-VEC2 is that the previous round of data movement is completed, so as to avoid data scrambling in PongBuffer[half].

[0041] Compared with the prior art, the present invention has the following advantages:

[0042] 1. To address the issue of serialization in the attention weight calculation stage of FP32 input scenarios, this invention achieves initial parallelism between transport instructions and calculation instructions by setting more precise synchronization control than existing solutions. Furthermore, this invention solves the problem of excessively long transport times in some rounds by readjusting the existing pipeline optimization technique, the Ping-Pong strategy, and enabling DB.

[0043] 2. To address the issue of pipeline blockage caused by inappropriate use of cache space during the attention weight calculation stage in FP16 input scenarios, this invention proposes the DPP strategy by reasonably allocating cache space and redesigning the computation pipeline based on the reallocated space. Attached Figure Description

[0044] Figure 1This is a flowchart of the FA operator calculation based on Ascend hardware.

[0045] Figure 2 It uses the Ascend 910B hardware architecture.

[0046] Figure 3 This is a computational pipeline for the attention weight calculation stage of the FA operator in an FP32 scene.

[0047] Figure 4 This is a computational pipeline for the attention weight calculation stage of the FA operator in an FP16 scene.

[0048] Figure 5 This is a computational pipeline diagram of the attention weight calculation stage of the FA operator after the initial optimization scheme based on precise synchronous control in the FP32 scenario.

[0049] Figure 6 This is a computational pipeline diagram of the attention weight calculation stage of the FA operator after optimization of the deep pipeline optimization scheme based on delayed handover and double buffering mechanism in the FP32 scenario.

[0050] Figure 7 This is the calculation process for the attention weight calculation stage after the FA operator instructions are decomposed and rearranged in the FP16 scenario.

[0051] Figure 8 This is a computational pipeline of the attention weight calculation stage of the FA operator in the FP16 scene after optimization using the DPP strategy.

[0052] Figure 9 This is an operator correctness verification.

[0053] Figure 10 This is a diagram showing the pipeline optimization verification results for the attention weight calculation stage of the FA operator in the FP32 scenario.

[0054] Figure 11 This is a diagram showing the pipeline optimization verification results for the attention weight calculation stage of the FA operator in the FP16 scenario. Detailed Implementation

[0055] The technical solution of the present invention will be further described below with reference to the accompanying drawings, but it is not limited thereto. Any modifications or equivalent substitutions to the technical solution of the present invention that do not depart from the spirit and scope of the technical solution of the present invention should be covered within the protection scope of the present invention.

[0056] Given that existing FA operator implementations suffer from pipeline blocking during the attention weight calculation stage, this invention provides a pipeline parallel optimization method for the FlashAttention attention weight calculation stage based on an NPU. The following sections detail optimization schemes for pipeline blocking in different scenarios.

[0057] I. Pipeline Parallel Optimization Scheme for Attention Weight Calculation Stage in FP32 Scenarios:

[0058] To address the flow congestion problem in this scenario, this invention proposes a preliminary optimization strategy based on precise synchronization control. The core idea is to relax the coarse-grained MTE3ToMTE2 dependency into a finer-grained MTE3ToVector dependency.

[0059] The data loading and unloading are handled independently by PongBuffer and PingBuffer, respectively. Before the unloading is complete, PingBuffer, not PongBuffer, is occupied. At this point, PongBuffer can load the next round of data and perform some computations. The key is that data cannot be handed over to PingBuffer before the unloading is finished. Therefore, the original WaitFlag (MTE3ToMTE2) synchronization instruction is changed to the WaitFlag (MTE3ToV) synchronization instruction and placed before the data handover operation between PongBuffer and PingBuffer (usually a scaling operation). This means that the loading of the next round of data is no longer constrained by the unloading of the previous round; it only needs to ensure that the data in PingBuffer is unloaded before being overwritten by the new round of computation.

[0060] like Figure 5 As shown, VEC1 computation represents the computation enabling PongBuffer, and VEC2 computation represents the computation enabling PingBuffer. After synchronization logic correction, PongBuffer can immediately start the prefetching operation of the next round of data after completing the data handover of the current round, realizing instruction parallelism for data loading, unloading, and partial vector computation. However, this strategy still has certain limitations. Because PongBuffer only undertakes a small amount of vector computation tasks before data handover, the PongBuffer enable computation time is too short. In rounds where data unloading takes a long time, such as... Figure 5 The time taken for the VEC1 calculation in round i+1 cannot cover the out in round i-i, thus blocking the calculation of VEC2 in round i+1 and the in in round i+2.

[0061] To further address the issue of computing units idling due to long turnaround cycles, this invention proposes a deep optimization strategy that combines delayed handover with double buffering.

[0062] To address the pipeline congestion caused by the long-tail delay in data transfer from PongBuffer to PingBuffer, this invention postpones the data handover from PongBuffer to PingBuffer from the scaling operation to a new blanking operation added after Dropout. This blanking operation is a scalar multiplication (Muls), which multiplies each number in the data to be computed by 1. This strategy aims to extend the computational lifetime of PongBuffer, enabling it to handle more vector operation load, thereby alleviating the pipeline congestion problem caused by the long-tail delay in data transfer.

[0063] After the delayed handover, PongBuffer completes the data loading and all vector calculations, while PingBuffer only enables data loading out. This means that PongBuffer forms a serial pipeline of loading-calculation-loading between different calculation rounds. To solve the serial pipeline problem caused by the delayed handover, i.e., the inability to perform the next round of loading during PongBuffer calculation, this invention introduces a double buffering mechanism on PongBuffer.

[0064] This invention logically splits the PongBuffer into two sub-blocks, performing only half of the data loading and computation at a time. This fine-grained pipeline design enables a parallel processing mode within the PongBuffer, where the computation of the current half of the data is performed while the loading of the next half is completed.

[0065] The optimized flow chart, such as Figure 6 As shown in the diagram, VEC1 and VEC2 represent the computations of the first and second halves of the PongBuffer up to Softmax, respectively, while VEC3 represents the computation after integrating the logical view of the PongBuffer. Since the data handover is delayed until the blank operation after the Dropout operation, the PingBuffer only performs the role of data transfer at this point; as can be seen from the diagram, the computation layer is almost unaware of the data transfer. Furthermore, a double-buffering mechanism is enabled for the PongBuffer, allowing each round of computation to start earlier without waiting for the entire transfer to complete, thereby reducing the idle time of the computation unit.

[0066] II. Pipeline optimization scheme for the attention weight calculation stage in FP16 scenarios:

[0067] To address the pipeline congestion issue in this scenario, this invention proposes a DPP strategy that utilizes buffer idle time windows for efficient reuse. By analyzing the lifecycle of each storage space during the attention weight calculation phase, it was found that after completing the initial scaling operation and handing over data to the PingBuffer, the PongBuffer enters a long idle state until the entire attention weight calculation process ends. Combined with the characteristic that the data volume is halved after the Cast operation, the remaining space in the PongBuffer is sufficient to accommodate the transformed FP16 results.

[0068] Based on this observation, the DPP strategy logically restructured the UB space:

[0069] (1) Input splitting: PingBuffer and PongBuffer are each split into two halves. The input data is not loaded all at once, but loaded into the first half of PongBuffer and the second half of PingBuffer in two separate operations.

[0070] (2) Calculation set: After PongBuffer completes the preprocessing calculation, it hands over the data to the first half of PingBuffer, and PingBuffer completes the remaining vector calculation.

[0071] (3) Loopback output: The target address of the Cast operation no longer points to the storage space MaskPongBuf of DropMask, but points to the latter half of PongBuffer.

[0072] This strategy actually constructs two nested Ping-Pong processes: one consisting of half a PongBuffer and half a PingBuffer, and the other independently handled by half a PingBuffer. Its core advantage lies in completely decoupling the address dependency between the intermediate variable storage address and the output result.

[0073] To accommodate the DPP's storage structure, this invention performs fine-grained instruction decomposition and rearrangement of the attention weight calculation phase. The invention divides the calculation process into VEC1 (from start to before FlashSoftmax) and VEC2 (from FlashSoftmax to the end of Cast), and schedules instructions according to the following timing sequence: Figure 7 As shown.

[0074] In the first stage, taking advantage of the high bandwidth of MTE2, three data transfers are initiated simultaneously: the input matrix mm1Res[0] is transferred into PongBuffer[0], half of the position encoding is transferred into PseBuffer, and the complete AttnMask matrix is ​​transferred in. This parallel transfer maximizes bandwidth utilization. Only half of the position encoding matrix is ​​transferred in to facilitate the rapid startup of VEC1-PongBuffer[0] with mm1Res[0], while reserving buffer space for temporary use by AttnMask calculation. After VEC1-PongBuffer[0] completes its calculation, the data is handed over to PingBuffer[0].

[0075] In the second stage, while PongBuffer[0] is executing VEC1, the transfer of the input matrix mm1Res[half] to PingBuffer[half] and the transfer of the remaining position encoded data are immediately started (it is necessary to wait for the AttnMask calculation of PongBuffer[0] to be completed to avoid trampling). After the transfer is completed, the calculation of VEC1-PingBuffer[half] is started.

[0076] In the third stage, after VEC1-PingBuffer[half] is completed, the logical views of PingBuffer are merged. After the views are merged, the calculation of VEC2-PingBuffer begins. The DropMask matrix required for Dropout calculation in VEC2 can be performed concurrently with the view merging operation, using the Softmax calculation time to mask the loading of DropMask. It is important to note that merging views does not incur additional overhead because of the continuity of address allocation; only the buffer size flag needs to be modified.

[0077] In the fourth stage, the Cast operation of VEC2-PingBuffer converts the FP32 calculation result into FP16 and writes it to the idle PongBuffer[half] for retrieval. For example... Figure 8 As shown, the adjusted pipeline eliminates the dependency on maskBuf: while VectorCompute2 is executing in the i-th round, MTE2 can start the data loading in the (i+1)-th round in advance; when data is loaded out in the i-th round, VectorCompute1 in the (i+1)-th round is already executing. This extreme parallel scheduling significantly improves the instruction issue density and effectively alleviates the potential pipeline bubble problem in FP16 scenarios.

[0078] Example:

[0079] This experiment was conducted on Huawei Cloud AI Development Platform ModelArts, and the instance specifications selected for the cloud environment were... When compiling and deploying the FA operator, the instance specification needs to be temporarily changed to... To speed up compilation (estimated to take about 30 minutes). ascend-snt9b3 corresponds to the Ascend 910B2 AI processor, which has 24 AICs and 48 AIVs. The kernel function execution mode of the FA operator is KERNEL_TYPE_MIX_AIC_1_2, meaning that 2 AIVs share one AIC.

[0080] The baseline FA operator uses the aclnnFlash-AttentionScoreV2 implementation from the cann-ops-adv (v0.8-8.0RC3) repository (although this repository is iterating, the core computational functionality of aclnnFlashAttentionScoreV2 remains unchanged). The corresponding software environment is cann-toolkit_8.0.0, cann-kernels-910b__8.0.0, torch 2.4.0, and torch-npu 2.4.0.post2. The operator is developed in AscendC, while the test code and model are developed in Python.

[0081] Before starting the experiment, the compilation and deployment of the FA operator need to be completed. The FA operator should be deployed to the CANN package, and then the eager mode should be registered using the cppextension provided by torch. This allows the baseline or optimized FA operator to be directly called within torch-npu for subsequent testing. Assume the corresponding operator name in torch-npu is npu_flash_attention_v2.

[0082] 1. Baseline implementation of attention weight calculation

[0083] The baseline attention weight calculation primarily employs a Ping-Pong pipeline strategy. The buffers used for this strategy are PingBuffer and PongBuffer. PongBuffer mainly serves as an inbound buffer and enables a small amount of computation (adding positional encoding and scaling) during the attention weight calculation phase. After PongBuffer completes the scaling operation, it hands the data over to PingBuffer, which then performs subsequent calculations (AttnMask, softmax, and dropout) and performs outbound processing (acting as an outbound buffer).

[0084] In this process, the buffer used to store the AttnMask matrix is ​​denoted as MaskPingBuffer, and the buffer used to store the dropout matrix is ​​denoted as MaskPongBuffer. Since the AttnMask matrix and the dropout matrix are of data type uint8_t, the theoretical size of these two buffers is only one-quarter of that of PingBuffer (data type FP32). However, in the baseline implementation, MaskPongBuffer needs to be used as a dropout buffer (data type FP16) in FP16 input scenarios; therefore, the size of MaskPongBuffer is half that of PingBuffer.

[0085] The buffer used to store the position encoding matrix is ​​relatively complex. In the FP32 input scenario, the external input position encoding matrix is ​​of type FP32. In this case, the baseline directly stores the position encoding matrix into the CommonBuffer. The CommonBuffer is a general-purpose buffer allocated by the baseline on the UB. Its function is to serve as temporary space. Calculations such as softmax and dropout require a certain amount of temporary space during computation, and this temporary space is provided by the CommonBuffer. Its size is the same as the PingBuffer. In the FP16 input scenario, the external input position encoding matrix is ​​of type FP16. In this case, the baseline first moves the position encoding matrix into the PseBuffer, and then performs a type conversion to dump the FP32 type position encoding matrix into the CommonBuffer. The PseBuffer serves as a transition buffer here, and its size is only half that of the PingBuffer.

[0086] Based on the above buffer, pseudocode for the baseline implementation of attention weight calculation is given.

[0087] Pseudocode 1 Implementation of the ProcessVec1() function

[0088]

[0089] As shown in pseudocode 1, this invention uses sizeof(Input_T) to represent the input scenario. sizeof(Input_T)=2 indicates an input scenario of FP16, and sizeof(Input_T)=4 indicates an input scenario of FP32. eventId represents a synchronous event, V represents vector computation, Mte2 represents loading, and Mte3 represents loading out. Taking eventIdVtoMte2B as an example, the code shows that this synchronous event sets its flag after the attn_mask computation is complete. After the attn_mask computation is complete, PseBuffer (FP16 scenario), PongBuffer, and MaskPingBuffer are all in an idle state. At this time, the position encoding matrix, AttnMask matrix, and attention score matrix for the next round can be loaded into the corresponding buffer space in advance. Therefore, in lines 2-3 of the code, that is, at the beginning of the next loop, the Waitflag of eventIdVtoMte2B is set. This means that the loading of the position encoding matrix, AttnMask matrix, and attention score matrix for the next round needs to wait for the attn_mask computation of the previous round to finish.

[0090] Attention weight calculation first issues instructions sequentially to load the attention score matrix, AttnMask matrix, and positional encoding matrix (these three are loaded simultaneously in the pipeline). After loading is complete, in FP16 scenarios, a type conversion operation is performed to convert the FP16 type positional encoding matrix to an FP32 type positional encoding matrix and dump it into the CommonBuffer. Then, ComputePse, ComputeScale, and Compute-AttnMask are executed sequentially. Before executing Softmax, instructions to load the dropout matrix are issued first, the purpose of which is to use the long computation time of Softmax to mask the loading time of the dropout matrix. ComputeSoftmax is executed simultaneously with the dropout loading. After the softmax and dropout matrices are loaded, ComputeDropout is executed.

[0091] When moving data out, the scenario needs to be differentiated. In an FP32 scenario, the data in the PingBuffer can be moved to the GM after the dropout calculation is complete. In an FP16 scenario, a type conversion is required first. The FP32 type result stored in the PingBuffer is converted to an FP16 type result and dumped to the MaskPongBuffer, which then handles the data movement. As can be seen from the code, in the FP32 scenario, the next round of dropout matrix loading only needs to wait for the previous round of dropout calculation to complete, while in the FP16 scenario, it also needs to wait for the previous round of data movement to finish.

[0092] Before the next round of loading, FP32, compared to FP16, directly sets an instruction to wait for the previous round of loading to finish (lines 6-7). This is mainly because, firstly, in the FP32 scenario, the position encoding matrix is ​​directly loaded into the CommonBuffer. As can be seen from the code, ComputeAttnMask, ComputeSoftmax, and ComputeDropout all require the CommonBuffer as a temporary buffer for input. Therefore, the loading of the position encoding matrix must wait until the dropout calculation is finished. Secondly, although the last line of the code sets a synchronization event to wait for the vector calculation to start after loading, the baseline code sets the wait instruction to before the loading starts on line 30. This means that this synchronization event only serves the FP16 scenario (type conversion on line 32).

[0093] 2. Pipeline optimization implementation in the attention weight calculation stage of FP32 scenario

[0094] This invention first removes the synchronization instructions Mte3toMte2 in lines 4-7. After removal, the loading of the next round's position encoding matrix, AttnMask, and attention score matrix will occur immediately after the current round's attn_mask calculation is complete. However, this causes a conflict between loading the next round's position encoding matrix and subsequent calculations in the current round. Therefore, it is necessary to adjust the temporary space required for subsequent calculations. In the FP32 scenario, PseBuffer is completely idle, so it can be used to replace CommonBuffer as the temporary space required for subsequent calculations.

[0095] Next, the handover operation is adjusted. This invention adds an operation called Muls after the dropout calculation, which represents scalar multiplication. This multiplies the result in PingBuffer / PongBuffer by 1. This operation is mainly used for data handover. It multiplies the data in PongBuffer by 1 and transfers the result to PingBuffer, where PingBuffer completes the data transfer. To ensure that data scrambling does not occur in the next round, the next round needs to wait for the previous round's data transfer to finish before calculating Muls.

[0096] Furthermore, DB is enabled for PongBuffer and CommonBuffer. At this time, the attention weight calculation first loads half of the attention score matrix and position encoding matrix into PongBuffer[0] and CommonBuffer[0], and simultaneously loads the AttnMask matrix into MaskPingBuffer. Then, it sets the wait for the Mte2toV synchronization instruction. After setting the synchronization instruction, it directly sends the other half of the attention score matrix and position encoding matrix into PongBuffer[half] and CommonBuffer[half], and simultaneously enables the addition of position encoding and scaling calculations in PongBuffer[0] and CommonBuffer[0]. After this calculation and the loading of the other half of the attention score matrix and position encoding matrix are completed, the addition of position encoding and scaling calculations in PongBuffer[half] and CommonBuffer[half] are completed. The subsequent part is consistent with the existing implementation. The adjusted implementation is shown in pseudocode 2.

[0097] Pseudocode 2 Implementation of the ProcessVec1 WithFP32() function

[0098]

[0099] 3. Pipeline optimization implementation for the attention weight calculation stage in FP16 scenarios

[0100] PingBuffer, PongBuffer, and PseBuffer are divided into two logical regions, denoted as PingBuffer[0] and PingBuffer[half], PongBuffer[0] and PongBuffer[half], and PseBuffer[0] and PseBuffer[half], respectively.

[0101] First, move half of the attention score matrix, half of the position encoding matrix, and the complete AttnMask matrix into PongBuffer[0], PseBuffer[0], and MaskPingBuffer in sequence. After setting the synchronous Mte2toV instruction, first issue the instruction to move the other half of the attention score matrix and position encoding matrix into PingBuffer[half] and PseBuffer[half]. Then, enable the addition of position encoding, scaling, and attn_mask calculation enabled by PongBuffer[0], PseBuffer[0], and MaskPingBuffer. Note that only half of the AttnMask matrix is ​​used at this time (moving into AttnMask means moving the whole matrix, while other data is half). After completing this part of the attn_mask calculation, the data in PongBuffer[0] needs to be handed over to PingBuffer[0].

[0102] After completing the above calculations and loading the other half of the attention score matrix and the position encoding matrix, start the calculation of adding position encoding, scaling, and attn_mask enabled by PingBuffer[half], PseBuffer[half], and MaskPingBuffer.

[0103] After completing the above operations, PingBuffer performs the subsequent calculations, following the same calculation process as the baseline. Finally, during type conversion, the data in PingBuffer is converted from FP32 to FP16 and dumped into PongBuffer[half]. PongBuffer[half] then handles the subsequent data transfer. With this adjustment, the Mte3toMte2 synchronization instruction previously added to MaskPongBuffer needs to be removed; the remaining synchronization instructions remain unchanged. The adjusted code is shown in pseudocode 3.

[0104] Pseudocode 3: Implementation of the ProcessVec1 WithFP16() function

[0105]

[0106] 4. Operator correctness testing and performance testing

[0107] (1) Operator correctness test

[0108] The correct test of the operator is to compare the execution result of the optimized FA operator with the result of the FlashAttention forward process calculated by pure torch_npu according to the same calculation process. The comparison interface is the assertRtolEqual() interface in the TestCase class of torch_npu, and prec is set to 1e-3.

[0109] Operator correctness test results are as follows Figure 9 As shown, the operator's execution result is correct.

[0110] (2) Operator performance test

[0111] While ensuring correctness, this invention verifies that three pipeline optimizations improve single-operator performance compared to the baseline. The experimental verification test cases uniformly adopted a batch size of 4, an attention head size of 8, a hidden layer dimension size of 128, and the same length for Query and Kev / Valued sequences, selected from the set [2048, 4096, 6144, 8192].

[0112] Figure 10 This study demonstrates the relative speedup of the FA operator, after incorporating deep pipeline optimization based on delayed handover and double buffering mechanisms, compared to the official baseline scheme in single-operator tests at different sequence lengths (S=2048 to 8192), under FP32 data precision. Experimental results show that the optimized scheme achieves stable positive gains across all test sequence lengths, with speedup ratios ranging from 1.030x to 1.037x, resulting in an average performance improvement of over 3%. This result proves that by refining the coarse-grained WaitFlag (MTE3ToMTE2) synchronization dependency to MTE3ToV, and combining delayed handover data with the internal double buffering mechanism of PongBuffer, the serialization deadlock of data ingress and egress in the original baseline scheme is alleviated to some extent, enabling the vector computation unit to achieve efficient parallel pipeline overlap with the MTE transport unit.

[0113] Figure 11This paper demonstrates the relative speedup of the FA operator after implementing the DPP strategy optimization, compared to the official baseline, under different sequence lengths (S=2048 to 8192) at FP16 data precision. Experimental results show that the optimized scheme achieves stable positive gains across all test sequence lengths, with speedup ratios ranging from 1.015x to 1.030x, resulting in an average performance improvement of over 2%. Compared to pipelining optimization in FP32 scenarios, the gains from the DPP strategy are smaller. This is because the DPP strategy primarily addresses the head-of-line blocking effect caused by excessively long outgoing times in some rounds. While the DPP strategy eliminates unified buffer read / write address conflicts caused by data type conversion to some extent, head-of-line blocking does not occur in all rounds. Therefore, compared to the over 3% performance improvement from serialized to parallelized pipelining in FP32 scenarios, the performance improvement in FP16 scenarios is relatively lower, as it only optimizes certain rounds.

Claims

1. A method for optimizing the flow parallel of the attention weight calculation phase of FlashAttention based on NPU, characterized in that The method includes the following steps: Step 1: For the FP32 scenario, a deep optimization strategy combining delayed handover and double buffering is adopted. The specific steps are as follows: Step 11: Delay the data handover from PongBuffer to PingBuffer from the scaling operation to the blanking operation added after Dropout. This blanking operation is a scalar multiplication, which multiplies each number in the data to be calculated by 1. Step 12: Logically split PongBuffer and CommonBuffer into two sub-blocks, denoted as PongBuffer[0], PongBuffer[half], CommonBuffer[0], and CommonBuffer[half] respectively; Step 13: Sequentially move half of the position encoding matrix, the entire AttnMask matrix, and half of the attention score matrix into CommonBuffer[0], MaskPingBuffer, and PongBuffer[0]. Step 14: After inserting the Set / WaiFlag synchronization wait instruction, move the other half of the position encoding matrix and the other half of the attention score matrix into CommonBuffer[half] and PongBuffer[half] in sequence. After issuing the move instruction, start the addition of position encoding and scaling calculations corresponding to CommonBuffer[0] and PongBuffer[0]. Step 15: After inserting the Set / WaiFlag synchronization wait instruction, start the addition of position encoding and scaling calculations for CommonBuffer[half] and PongBuffer[half]. After completing this calculation, reset PongBuffer to the full view, meaning that subsequent calculations will be performed on the full PongBuffer. Step 16: After completing the calculation of attn_mask, first issue the instruction to move the DropMask matrix into MaskPongBuffer, and then start the calculation of Softmax and Dropout in sequence. Before the calculation of Dropout, a synchronization wait instruction of Set / WaiFlag needs to be inserted. Step 17: Finally, complete the scalar multiplication operation, transfer the data from PongBuffer to PingBuffer, and let PingBuffer handle the subsequent data transfer. Step 2: For FP16 scenarios, the DPP strategy is adopted. The specific steps are as follows: Step 21: Split PingBuffer into two sub-blocks, denoted as PingBuffer[0] and PingBuffer[half] respectively. The same applies to PongBuffer, PseBuffer, and CommonBuffer. Step 22: Decompose and rearrange the instructions for the calculation process of the attention weight calculation stage, divide the calculation process into VEC1 and VEC2, and add the corresponding buffer before the corresponding calculation process in order to distinguish the calculations of different buffers enabled. Step 23: Move half of the position encoding matrix into PseBuffer[0], complete the type conversion operation, and save the result into CommonBuffer[0]. Step 24: Move half of the attention score matrix and the complete AttnMask matrix into PonBuffer[0] and MaskPongBuffer in sequence; Step 25: After inserting the Set / WaiFlag synchronization wait instruction, first move the other half of the attention score matrix and the position encoding matrix into PingBuffer[half] and PseBuffer[half] respectively, then start the calculation of PongBuffer[0]-VEC1, and dump the final calculation result into PingBuffer[0]. Step 26: After inserting the Set / WaiFlag synchronization wait instruction, first complete the type conversion operation of PseBuffer[half], convert this part of the position encoding matrix to FP32 type and dump it into CommonBuffer[half]. Then start the calculation of PingBuffer[half]-VEC1. After PingBuffer[half]-VEC1 is completed, the view of PingBuffer needs to be merged, that is, the subsequent calculation will be performed on the complete PingBuffer. Step 27: Start PingBuffer-VEC2 computation. The timing of loading and computation of the DropMask matrix is ​​consistent with the description in the FP32 scenario. Step 28: After completing the calculation, the result is transferred to PongBuffer[half], and PongBuffer[half] is used to move the data out.

2. The NPU-based FlashAttention attention weight calculation stage pipeline parallel optimization method according to claim 1, characterized in that In step 13, the premise for moving half of the position encoding matrix is ​​that the position encoding matrix part corresponding to CommonBuffer[0] in the previous round has been calculated, the premise for moving the AttnMask matrix is ​​that the attn_mask in the previous round has been calculated, and due to the delayed handover in step one, the premise for moving half of the attention score matrix is ​​that the scalar multiplication operation in the previous round has been completed.

3. The NPU-based FlashAttention attention weight calculation stage pipeline parallelism optimization method according to claim 1, characterized in that In step 17, the scalar multiplication calculation is initiated only after the previous round of data transfer has ended, in order to avoid data loss during PingBuffer transmission.

4. The NPU-based FlashAttention attention weight calculation stage pipeline parallelism optimization method according to claim 1, characterized in that In step 23, the prerequisite for moving the position encoding matrix is ​​that the previous round of type conversion operation is completed, and the current round of type conversion operation needs to insert a Set / WaiFlag synchronous wait instruction before it is executed.

5. The NPU-based FlashAttention attention weight calculation stage pipeline parallelism optimization method according to claim 1, characterized in that In step 24, the prerequisite for moving the attention score matrix is ​​that the previous round of PongBuffer[0]-VEC1 calculation is completed, and the prerequisite for moving the AttnMask matrix is ​​that the previous round of PingBuffer[half]-VEC1 calculation is completed.

6. The NPU-based FlashAttention attention weight calculation stage pipeline parallelism optimization method according to claim 1, characterized in that In step 28, the prerequisite for executing Cast in PingBuffer-VEC2 is that the previous round of data transfer has ended, in order to prevent data scrambling in PongBuffer[half].

Images