Multi-head attention calculation task allocation method and related products
By dynamically splitting multi-head attention computation tasks and allocating them according to actual and target computational loads, the problem of unbalanced computational load caused by causal masking is solved, thereby improving system performance and hardware resource utilization.
Patent Information
- Application Number
- CN202511846881.3
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-12-09
- Publication Date
- 2026-03-13
AI Technical Summary
In autoregressive language models based on the Transformer architecture, the unbalanced computational load caused by causal masking results in the overall system performance being constrained by the computationally intensive unit, forming a significant time bottleneck.
By dynamically splitting multi-head attention computation tasks, tasks are allocated based on the actual computational load of the current processing core and the target computational load of each processing core, ensuring that the computational load of each processing core is roughly the same and avoiding overload.
It effectively solves the problem of unbalanced computational load caused by causal masking, and improves the overall throughput and hardware resource utilization of parallel computing systems.
Smart Images

Figure CN121658233A_ABST
Abstract
Description
Technical Field
[0001] This disclosure generally relates to the field of artificial intelligence technology. More specifically, this disclosure relates to a method for multi-head attention mechanism computational tasks, an artificial intelligence processor, a chip, and a board. Background Technology
[0002] In the field of attention-based neural network computation, especially in autoregressive language models based on the Transformer architecture, causal masks are typically used to ensure the autoregressive properties during sequence generation. The role of a causal mask is to restrict attention at each position to only that position and its preceding positions, thus forming a strictly lower triangular attention matrix.
[0003] In existing parallel computing systems, the total query sequence (total_seq_q) is typically distributed evenly across the processing units. After calculating the attention score matrix, due to the effect of causal masking, the actual computational load of subsequent operations based on the effective region (such as mask application, Softmax normalization, and weighted summation with V) increases linearly as the query position moves forward. This results in an uneven load distribution among processing units handling different query sequence segments of the same length. This uneven distribution of computational load causes the overall system performance to be constrained by the processing unit with the highest computational load. Even if other processing units have completed their calculations earlier, they must wait for the slowest processing unit to complete its task, thus forming a significant time bottleneck.
[0004] Therefore, there is an urgent need in this field for a novel computational allocation method that can effectively solve the problem of unbalanced computational load under causal masking. Summary of the Invention
[0005] In order to at least solve one or more technical problems described in the background section above, this disclosure proposes the following technical solutions and several embodiments thereof.
[0006] In the first aspect, this disclosure proposes a multi-head attention computation task allocation method, which is applied to an artificial intelligence processor including multiple processing cores. The method is executed by each processing core and includes: loading a sequence interval computation task of multi-head attention allocated to the processing core, wherein the sequence interval computation task is dynamically split based on a comparison between the actual computational load of the current processing core and the target computational load of each processing core, wherein the actual computational load of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask, and the target computational load refers to the average computational load of each processing core participating in the multi-head attention computation task; and executing the sequence interval computation task.
[0007] In a second aspect, this disclosure discloses an artificial intelligence processor comprising multiple processing cores, each processing core comprising a control unit and at least two computational units, wherein the control unit is configured to perform the step of loading a sequence interval computation task of multi-head attention allocated to the processing core in the multi-head attention computation task allocation method described in the first aspect; and the computational units are configured to perform the sequence interval computation task.
[0008] In the third aspect, this disclosure discloses a chip configured to include an artificial intelligence processor as described in the second aspect.
[0009] In the fourth aspect, this disclosure discloses a board that includes the chip described in the third aspect.
[0010] Through the multi-head attention computation task allocation method, artificial intelligence processor, chip, and board proposed above, this disclosed embodiment loads a sequence interval computation task of multi-head attention, allocated to a single processing core of the artificial intelligence processor. The sequence interval computation task is dynamically split based on a comparison between the actual computational load of the current processing core and the target computational load of each processing core. The actual computational load of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask. The target computational load refers to the average computational load of each processing core participating in the multi-head attention computation task. The sequence interval computation task is then executed. The loading of the sequence interval computation task, dynamically split based on the comparison between the actual computational load of the current processing core and the target computational load of each processing core, effectively prevents processing core overload and avoids any processing core being assigned too many tasks, thus avoiding becoming a system bottleneck. This significantly improves the overall throughput and hardware resource utilization of the parallel computing system. Attached Figure Description
[0011] The above and other objects, features, and advantages of exemplary embodiments of this disclosure will become readily apparent upon reading the following detailed description with reference to the accompanying drawings. In the drawings, several embodiments of this disclosure are illustrated by way of example and not limitation, and like or corresponding reference numerals denote like or corresponding parts, wherein: Figure 1 This diagram illustrates the structure of the board card according to an embodiment of this disclosure. Figure 2 This diagram illustrates the structure of the combined processing apparatus in the chip according to an embodiment of the present disclosure. Figure 3 An exemplary internal structure diagram of an artificial intelligence processor in some embodiments of this disclosure is shown; Figure 4 An exemplary structural diagram of a processor core in some embodiments of this disclosure is shown; Figure 5 An exemplary schematic diagram is shown illustrating how one processor core writes data to another processor core in one of the embodiments of this disclosure; Figure 6 An exemplary schematic diagram of a two-layer, three-stage pipeline in some embodiments of this disclosure is shown; Figure 7 An exemplary flowchart of a multi-head attention computation task allocation method in some embodiments of this disclosure is shown; Figure 8(a) shows a schematic diagram of the actual computational load when the query sequence and key sequence lengths are the same in some embodiments of this disclosure; Figure 8(b) illustrates the actual computational load when the query sequence and key sequence lengths are different in some embodiments of this disclosure; Figure 9(a) illustrates a schematic diagram of the top-to-bottom alignment in some embodiments of this disclosure; Figure 9(b) illustrates a schematic diagram of bottom-up alignment in some embodiments of this disclosure; Figure 10 An exemplary flowchart of a multi-head attention computation task allocation method in some embodiments of this disclosure is shown; Figure 11 This disclosure shows a partial exemplary flowchart of a computational task splitting method in some embodiments; Figure 12 This disclosure shows a partial exemplary flowchart of a computational task splitting method in some embodiments; Figure 13 A partial exemplary flowchart of a computational task splitting method in some embodiments of this disclosure is shown. Detailed Implementation
[0012] The technical solutions in the embodiments of this disclosure will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of this disclosure, not all of them. Based on the embodiments in this disclosure, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this disclosure.
[0013] It should be understood that the terms “comprising” and “including” used in this disclosure and claims indicate the presence of the described features, integrals, steps, operations, elements and / or components, but do not exclude the presence or addition of one or more other features, integrals, steps, operations, elements, components and / or collections thereof.
[0014] It should also be understood that the terminology used herein is for the purpose of describing particular embodiments only and is not intended to limit the scope of this disclosure. As used in this disclosure and claims, the singular forms “a,” “an,” and “the” are intended to include the plural forms unless the context clearly indicates otherwise. It should also be understood that the term “and / or” as used in this disclosure and claims refers to any combination and all possible combinations of one or more of the associated listed items, and includes such combinations.
[0015] As used in this specification and claims, the term "if" may be interpreted, depending on the context, as "when," "once," "in response to determination," or "in response to detection." Similarly, the phrase "if determined" or "if [described condition or event] is detected" may be interpreted, depending on the context, as "once determined," "in response to determination," "once [described condition or event] is detected," or "in response to detection of [described condition or event]."
[0016] The specific embodiments disclosed herein will now be described in detail with reference to the accompanying drawings.
[0017] During the parallel execution of multi-head attention computation tasks by multiple processing cores, due to the existence of causal masks, the query sequence segments allocated to different processing cores face different computational loads in the actual computation process. As shown in Figure 8(b), the area of the trapezoid in the blank part represents the effective numerical region of the attention score matrix, and the shaded part of the upper right triangle represents the numerical region of the attention score matrix that becomes invalid due to the existence of causal masks. This part has no effect on subsequent computations. If the query sequence (seq_q) is directly and evenly distributed across the processing cores (core0, core1, core2), meaning the query sequences obtained from the splits in core0, core1, and core2 all have the same number of queries (seq_q / 3), then the query sequence segment assigned to core2, after calculating the attention score matrix, undergoes the application of a causal mask, followed by computational operations based on the effective numerical region of the attention score matrix (such as Softmax normalization and weighted summation with the value vector V). These operations result in the highest actual computational cost for core2. Conversely, the query sequence segment assigned to core0, after calculating the attention score matrix, undergoes the application of a causal mask, followed by computational operations based on the effective numerical region of the attention score matrix (such as Softmax normalization and weighted summation with the value vector V). Due to the effect of the causal mask on the attention score matrix, these operations result in the lowest actual computational cost. Therefore, core2 > core1 > core0. Obviously, the actual amount of computation corresponding to each processing core is different. This uneven distribution of computational load causes the overall performance of the system to be constrained by the processing core with the largest amount of computation. Even if other processing cores have completed their calculations in advance, they must wait for the slowest processing core to complete its task, thus forming a significant time bottleneck.
[0018] Taking the multi-head attention algorithm in a large language model as an example, it includes the following steps: a) Calculate the attention score matrix: Calculate the dot product qk of the two tensors (i.e., the query sequence seq_q and the key sequence seq_k); b) Apply causal masking: add causal_mask to qk; c) Perform Softmax on each row; d) Perform a dot product between the result of Softmax and the value vector V to obtain the output result.
[0019] When the query sequence (seq_q) is evenly divided into several segments and distributed to the computational units of different processing cores, each processing core's computational unit is responsible for processing a continuous segment of query positions. Because the computational cost of each row is different (the computational cost of Softmax and the weighted summation with the value vector V is proportional to the number of valid elements in that row), the total computational cost of each processing core's computational unit depends on the range of rows (query positions) it processes. If each processing core's computational unit processes query sequence segments containing the same number of query positions, during attention computation, the application of causal masks results in later rows requiring more computation. This places a heavier computational load on the computational units of the processing cores processing later sequence segments, thus becoming a computational bottleneck.
[0020] Therefore, this disclosure proposes a multi-attention head computation task allocation method to effectively solve the problem of unbalanced computational load caused by causal masking.
[0021] Figure 1 A schematic diagram of the structure of a board 10 according to an embodiment of the present invention is shown. Figure 1 As shown, board 10 includes chip 101, which is a system-on-chip (SoC) integrating one or more combined processing units. These combined processing units are artificial intelligence computing units used to support various deep learning and machine learning algorithms, meeting the intelligent processing needs of complex scenarios in fields such as computer vision, speech, natural language processing, and data mining. In particular, deep learning technology is widely used in cloud intelligence. A significant characteristic of cloud intelligence applications is the large volume of input data, placing high demands on the platform's storage and computing capabilities. Board 10 in this embodiment is suitable for cloud intelligence applications, possessing massive off-chip storage, on-chip storage, and substantial computing power.
[0022] Chip 101 is connected to external device 103 via external interface device 102. External device 103 may be, for example, a server, computer, camera, monitor, mouse, keyboard, network card, or Wi-Fi interface. Data to be processed can be transmitted from external device 103 to chip 101 via external interface device 102. The calculation results from chip 101 can be transmitted back to external device 103 via external interface device 102. Depending on the application scenario, external interface device 102 may have different interface forms, such as a PCIe interface.
[0023] The board 10 also includes a storage device 104 for storing data, which includes one or more memory cells 105. The storage device 104 is connected to and transmits data with the controller 106 and the chip 101 via a bus. The controller 106 in the board 10 is configured to regulate the state of the chip 101. Therefore, in one application scenario, the controller 106 may include a microcontroller (MCU).
[0024] Figure 2 This is a structural diagram illustrating the combined processing device in chip 101 of this embodiment. (As shown) Figure 2 As shown, the combined processing device 20 includes a computing device 201, an interface device 202, a processing device 203, and a DRAM 204.
[0025] The computing device 201 is configured to perform user-specified operations. It is mainly implemented as a single-core intelligent processor or a multi-core intelligent processor to perform deep learning or machine learning calculations. It can interact with the processing device 203 through the interface device 202 to jointly complete the user-specified operations.
[0026] Interface device 202 is used to transmit data and control commands between computing device 201 and processing device 203. For example, computing device 201 can obtain input data from processing device 203 via interface device 202 and write it to on-chip storage device of computing device 201. Further, computing device 201 can obtain control commands from processing device 203 via interface device 202 and write them to on-chip control cache of computing device 201. Alternatively or optionally, interface device 202 can also read data from storage device of computing device 201 and transmit it to processing device 203.
[0027] The processing device 203, as a general-purpose processing device, performs basic controls including but not limited to data transfer and starting and / or stopping the computing device 201. Depending on the implementation, the processing device 203 may be one or more types of processors, such as a central processing unit (CPU), a graphics processing unit (GPU), or other general-purpose and / or special-purpose processors. These processors include, but are not limited to, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc., and their number can be determined according to actual needs. As mentioned above, the computing device 201 of this invention can be considered as having a single-core structure or a homogeneous multi-core structure. However, when the computing device 201 and the processing device 203 are considered together, they are considered to form a heterogeneous multi-core structure.
[0028] DRAM 204 is an off-chip memory used to store data to be processed. It is DDR and typically 16G or larger in size. It is used to store data in computing device 201 and / or processing device 203.
[0029] Figure 3 A schematic diagram of the internal structure of computing device 201 is shown. Computing device 201 is used to process input data such as computer vision, speech, natural language, and data mining. The computing device 201 in the figure adopts a multi-core hierarchical structure design. As a system-on-a-chip, computing device 201 includes multiple clusters, and each cluster includes multiple processor cores. In other words, computing device 201 is constructed in a hierarchical structure of system-on-a-chip, cluster, and processor core.
[0030] From the perspective of system-on-a-chip hierarchy, such as Figure 3 As shown, the computing device 201 includes an external storage controller 301, a peripheral communication module 302, an on-chip interconnect module 303, a synchronization module 304, and multiple clusters 305.
[0031] There can be multiple external storage controllers 301; two are shown as an example in the figure. These controllers are used to access external storage devices, such as those issued by the processor core, in response to access requests from the processor core. Figure 2The DRAM 204 in the chip allows data to be read from or written to external devices. The peripheral communication module 302 receives control signals from the processing device 203 via the interface device 202, initiating the computing device 201 to execute tasks. The on-chip interconnect module 303 connects the external storage controller 301, the peripheral communication module 302, and multiple clusters 305 to transmit data and control signals between modules. The synchronization module 304 is a global barrier controller (GBC) used to coordinate the working progress of each cluster and ensure information synchronization. The multiple clusters 305 are the computing core of the computing device 201; four are shown exemplary in the figure. With hardware development, the computing device 201 of this invention can also include eight, sixteen, sixty-four, or even more clusters 305. The clusters 305 are used to efficiently execute deep learning algorithms.
[0032] From the perspective of cluster hierarchy, such as Figure 3 As shown, each cluster 305 includes multiple processor cores (IPU cores) 306 and one memory core (MEM core) 307.
[0033] Four processor cores 306 are shown exemplarily in the figure, but the present invention does not limit the number of processor cores 306. Its internal architecture is as follows: Figure 4 As shown. Each processor core 306 includes three main modules: a control module 41, an arithmetic module 42, and a storage module 43.
[0034] The control module 41 coordinates and controls the operation of the computation module 42 and the storage module 43 to complete the deep learning task. It includes an instruction fetch unit (IFU) 411 and an instruction decode unit (IDU) 412. The instruction fetch unit 411 fetches instructions from the processing device 203, and the instruction decode unit 412 decodes the fetched instructions and sends the decoding result as control information to the computation module 42 and the storage module 43.
[0035] The computation module 42 includes a vector operation unit 421 and a matrix operation unit 422. The vector operation unit 421 is used to perform vector operations and can support complex operations such as vector multiplication, addition, and nonlinear transformations; the matrix operation unit 422 is responsible for the core computations of deep learning algorithms, namely matrix multiplication and convolution.
[0036] Storage module 43 is used to store or move related data, including neuron RAM (NRAM) 431, weight RAM (WRAM) 432, input / output direct memory access (IODMA) 433, and move direct memory access (MVDMA) 434. NRAM 431 is used to store feature maps for computation by processor core 306 and intermediate results after computation; WRAM 432 is used to store the weights of the deep learning network; IODMA 433 controls the memory access of NRAM 431 / WRAM 432 and DRAM 204 through broadcast bus 309; MVDMA 434 controls the memory access of NRAM 431 / WRAM 432 and SRAM 308.
[0037] Back Figure 3 The storage core 307 is primarily used for storage and communication, namely storing shared data or intermediate results among processor cores 306, and performing communication between cluster 305 and DRAM 204, communication between clusters 305, and communication between processor cores 306. In other embodiments, the storage core 307 has scalar operation capabilities and is used to perform scalar operations.
[0038] Storage core 307 includes a shared memory unit (SRAM) 308, a broadcast bus 309, a cluster direct memory access (CDMA) module 310, and a global direct memory access (GDMA) module 311. SRAM 308 acts as a high-performance data relay station. Data multiplexed between different processor cores 306 within the same cluster 305 does not need to be obtained from DRAM 204 by each processor core 306 individually. Instead, it is relayed between processor cores 306 via SRAM 308. Storage core 307 only needs to quickly distribute the multiplexed data from SRAM 308 to multiple processor cores 306, thereby improving inter-core communication efficiency and significantly reducing on-chip and off-chip I / O access.
[0039] Broadcast bus 309, CDMA 310, and GDMA 311 are used to perform communication between processor cores 306, communication between clusters 305, and data transfer between cluster 305 and DRAM 204, respectively. These will be explained below.
[0040] The broadcast bus 309 is used to complete high-speed communication between the processor cores 306 within the cluster 305. In this embodiment, the broadcast bus 309 supports inter-core communication methods including unicast, multicast, and broadcast. Unicast refers to point-to-point (i.e., data transmission from one processor core to another) data transmission. Multicast is a communication method that transmits a piece of data from SRAM 308 to several specific processor cores 306. Broadcast is a communication method that transmits a piece of data from SRAM 308 to all processor cores 306, and is a special case of multicast.
[0041] CDMA 310 is used to control SRAM 308 access between different clusters 305 within the same computing device 201. Figure 5 This diagram illustrates the operation of CDMA310 when one processor core attempts to write data to another processor core in a different cluster. In this application scenario, the same computing device comprises multiple clusters. For simplicity, only clusters 0 and 1 are shown in the diagram. Both clusters 0 and 1 contain multiple processor cores; similarly, for ease of explanation, only processor core 0 is shown in cluster 0, and only processor core 1 is shown in cluster 1. Processor core 0 intends to write data to processor core 1.
[0042] First, processor core 0 sends a unicast write request to write data into its local SRAM 0. CDMA 0 acts as the master and CDMA 1 acts as the slave. The master pushes the write request to the slave, that is, the master sends the write address AW and the write data W to transmit the data to SRAM 1 of cluster 1. Then the slave sends a write response B as a reply. Finally, processor core 1 of cluster 1 sends a unicast read request to read the data from SRAM 1.
[0043] Back Figure 3The GDMA 311, in conjunction with the external memory controller 301, controls memory access from the SRAM 308 to the DRAM 204 of the cluster 305, or reads data from the DRAM 204 into the SRAM 308. As previously described, communication between the DRAM 204 and the NRAM 431 or WRAM 432 can be achieved through two channels. The first channel is a direct connection between the DRAM 204 and the NRAM 431 or WRAM 432 via the IODAM 433; the second channel involves first transmitting data between the DRAM 204 and SRAM 308 via the GDMA 311, and then transmitting data between the SRAM 308 and the NRAM 431 or WRAM 432 via the MVDMA 434. Although the second channel appears to require more components and has a longer data flow, in some embodiments, the bandwidth of the second channel is actually much greater than that of the first channel. Therefore, communication between the DRAM 204 and the NRAM 431 or WRAM 432 may be more efficient via the second channel. The embodiments of the present invention can select the data transmission channel according to their own hardware conditions.
[0044] In other embodiments, the functions of GDMA 311 and IODMA 433 can be integrated into the same component. For ease of description, GDMA 311 and IODMA 433 are considered different components. For those skilled in the art, any component whose functions and technical effects are similar to those of this invention falls within the scope of protection of this invention. Furthermore, the functions of GDMA 311, IODMA 433, CDMA 310, and MVDMA 434 can also be implemented by the same component. Similarly, any component whose functions and technical effects are similar to those of this invention falls within the scope of protection of this invention.
[0045] In the aforementioned hardware architecture, the components inside the computing device 201 are referred to as on-chip, and those outside the computing device 201 are referred to as off-chip. During the execution of various layers of deep learning computation, numerous off-chip and on-chip accesses are required, particularly reading data from DRAM 204 into the computing device 201 and then storing the computation results from the computing device 201 back into DRAM 204. Such frequent accesses consume significant hardware resources. To address this issue, one embodiment of the present invention significantly reduces off-chip and on-chip accesses by planning a weight transfer mechanism, while simultaneously lowering the on-chip weight storage requirements.
[0046] One of the key reasons why the computing device 201 has strong computing power is its three-level operation hierarchy of system-on-chip, cluster, and processor core, combined with a three-level memory design of DRAM, SRAM, and NRAM / WRAM, which allows data to be cached and computed at appropriate levels, forming a sufficient pipeline.
[0047] The computing device 201 performs calculations in three main stages: Load stage: loading data; Compute stage: transferring data, performing calculations, and transferring intermediate results; Store stage: storing the results.
[0048] More specifically, this embodiment employs a two-layer, three-stage production line, such as... Figure 6 As shown, the first-layer loading stage 601, calculation stage 602, and write-back stage 603 occur at the cluster level. In the first-layer loading stage 601, the GDMA 330 loads data from DRAM 204 into SRAM 308. In the first-layer calculation stage 602, the cluster 305 performs calculations on the loaded on-chip cell diagram and generates calculation results. In the first-layer write-back stage 603, the GDMA 330 writes the calculation results back from SRAM 308 into DRAM 204.
[0049] Since cluster 305 includes multiple processor cores 306, the first-level computation stage 602 actually divides the on-chip cell graph into corresponding subgraphs through storage core 307 and broadcasts them to at least one processor core 306 for computation. Therefore, the second-level tertiary watershed occurs in processor core 306. More specifically, the second-level loading stage 604 loads the subgraph from SRAM 308 into NRAM 431 using MVDMA 434. The second-level computation stage 605 moves the subgraph and subweights to the arithmetic module 42 for computation, and then moves the intermediate results back to NRAM 431. The second-level write-back stage 606 is when MVDMA 434 writes the intermediate results from NRAM 431 back to SRAM 308.
[0050] The first layer of the pipeline refers to the fact that the first layer load stage 601, the first layer computation stage 602, and the first layer store-back stage 603 can be performed in parallel. Taking the same cluster 305 processing the j-th on-chip cell graph, the (j+1)-th on-chip cell graph, and the (j+2)-th on-chip cell graph as an example, firstly, the j-th on-chip cell graph is loaded into SRAM 308 in the first layer load stage 601. Then, the j-th on-chip cell graph is computed in the first layer computation stage 602, and the first computation result is transferred back to SRAM 308. Simultaneously, while the j-th on-chip cell graph is being computed, the (j+1)-th on-chip cell graph is loaded into SRAM 308 in the first layer load stage 607. When the first calculation result is stored back into DRAM 204 in the first layer store-back stage 603, the (j+1)th on-chip cell diagram is calculated in the first layer calculation stage 608, and the second calculation result is transferred back into SRAM 308. At the same time, the (j+2)th on-chip cell diagram is loaded into SRAM 308 in the first layer load stage 610. The first layer pipeline proceeds in this manner.
[0051] To facilitate the aforementioned pipelined operation, the SRAM 308 in this embodiment includes two storage spaces: ping-pong and pong-pong. Data pipelined according to the ping-pong attribute of the SRAM 308 is divided into three types: input / output ping-pong (IOparity), input ping-pong (input parity), and no ping-pong (no parity). IOparity supports parallel loading, computation, and write-back. To achieve IOparity, the ping-pong and pong-pong storage cells need to be exactly equal, used separately for loading and write-back. Input ping-pong only supports parallel write-back and computation, adding extra time for data transfer on the SRAM 308. Compared to IOparity, the ping-pong and pong-pong storage cells do not need to be exactly equal, but an additional cache of the same size as the write-back storage space needs to be allocated. No ping-pong means that loading / write-back and computation are serial, and no additional space is required.
[0052] To implement the aforementioned first-level pipeline, the SRAM 308 in this embodiment has ping-pong and pong memory cells of the same size to achieve an input / output ping-pong effect. Continuing with... Figure 6To clarify, the storage area involved in the first-level loading stage 601, the first-level computation stage 602, and the first-level write-back stage 603 of the j-th on-chip cell diagram is limited to the ping storage unit, while the storage area involved in the first-level loading stage 607, the first-level computation stage 608, and the first-level write-back stage 609 of the (j+1)-th on-chip cell diagram is limited to the pong storage unit, and the storage area involved in the first-level loading stage 610, the first-level computation stage 611, and the first-level write-back stage 612 of the (j+2)-th on-chip cell diagram is again limited to the ping storage unit. In this manner, ping storage units and pong storage units are used alternately for storage.
[0053] The second-level pipeline refers to the parallel operation of the second-level load stage 604, the second-level computation stage 605, and the second-level store-back stage 606. Consider an example where the same processor core 306 wants to process the i-th, (i+1)-th, and (i+2)-th subgraphs in the j-th on-chip unit graph. First, the i-th subgraph is broadcast to NRAM 431 in the second-level load stage 604. Then, the i-th subgraph is computed in the second-level computation stage 605 to produce the i-th intermediate result, which is then moved back to NRAM 431. Simultaneously, the (i+1)-th subgraph is broadcast to NRAM 431 in the second-level load stage 613. The i-th intermediate result is stored back into SRAM 308 in the second layer storage-back stage 606. At the same time, the (i+1)-th subgraph is calculated in the second layer calculation stage 614 to generate the (i+1)-th intermediate result, and the (i+1)-th intermediate result is moved back into NRAM 431. The (i+2)-th subgraph is loaded into NRAM 431 in the second layer loading stage 615.
[0054] Considering that each cluster 305 has different tasks and their completion times are naturally different, the synchronization module 304 in this embodiment will use the synchronization barrier instruction to synchronize the task completion time to avoid timing errors.
[0055] The above describes the relevant hardware implementation of the embodiments disclosed herein. Based on the above hardware environment, this disclosure proposes a scheme for multi-head attention computing task allocation. This scheme dynamically splits the computing tasks by comparing the estimated actual computing load of the current processing core with the target computing load corresponding to each processing core participating in the multi-head attention computing task, thereby effectively preventing processing core overload and avoiding the problem of any processing core being assigned too many tasks and becoming a system bottleneck.
[0056] Figure 7 An exemplary flowchart of a multi-head attention computation task allocation method in some embodiments of this disclosure is shown. Figure 7 As shown, the method includes: Step S701: Load the sequence interval computation task of the multi-head attention allocated to the processing core. The sequence interval computation task is dynamically split according to the comparison result of the actual computation amount of the current processing core and the target computation amount of each processing core. The actual computation amount of the current processing core is estimated according to the input sequence of the multi-head attention computation task and the scope of the causal mask. The target computation amount refers to the average computation amount of each processing core participating in the multi-head attention computation task.
[0057] In the above steps, the input sequence includes query sequences, key sequences, value sequences, etc. The input sequence is stored in the global memory area of the AI processor. Each processing core of the AI processor retrieves NRAM / WRAM on demand for computation via Move-DMA.
[0058] Step S702: Perform the sequence interval calculation task.
[0059] The artificial intelligence processor includes a control unit and at least two of the aforementioned computing units. Step S701 can, for example, be performed by... Figure 4 The control module 41 in the processor core 306 of the AI processor shown executes this, which can also be called a control unit. Step S702 is executed by the various processing units, corresponding to... Figure 4 The processor core 306 shown contains a computation module 42. As previously described, the computation module 42 may include multiple computation units, each of which can be configured as a multi-stage computation pipeline. Multiple computation units constitute multiple multi-stage computation pipelines, and the multiple multi-stage computation pipelines can execute their respective unit computation tasks in parallel.
[0060] In step S701, the input sequence may include at least a query sequence and a key sequence. The query sequence includes multiple query sequence positions, for example, seq_q{query position 0, query position 1, query position 2, query position 3, ..., query position m}. The key sequence includes multiple key sequence positions, for example, seq_k{key position 0, key position 1, key position 2, key position 3, ..., key position n}. The attention score matrix can be represented by a regular shape such as a rectangle or square, where the rows of the attention score matrix correspond to the query positions of the query sequence, and the columns correspond to the key positions of the key sequence.
[0061] Following the multi-head attention algorithm, after obtaining the attention score matrix, a causal mask is applied. Due to the causal mask, the attention score matrix is divided into effective and ineffective regions. The effective region, obtained by applying the causal mask to the attention score matrix during attention calculation, corresponds to the lower triangular portion (including the diagonal) of the matrix. Under the causal mask, the effective region only focuses on the current and previous positions. The ineffective region, on the other hand, contains matrix elements that are forcibly set to invalid values under the causal mask and will not participate in subsequent calculations. The ineffective region is the upper triangular portion of the attention score matrix (masked to negative infinity). Under the causal mask, the ineffective region cannot focus on future positions; in the subsequent Softmax calculation, these negative infinity values are converted to zero weights.
[0062] An attention score matrix can be obtained from the input sequence of a multi-head attention computation task. Based on the attention score matrix and the effective range of the causal mask, the actual computational load of each processing core can be estimated. Then, based on a comparison between the actual computational load of each processing core and its average computational load, the input sequence is dynamically split into tasks, generating task descriptors corresponding to the split results. This disclosed embodiment ensures that the computational load of each processing core is approximately the same by comparing the actual computational load of each processing core with its average computational load, avoiding the problem of one processing core being overloaded while others are idle, thereby improving hardware utilization and overall computational efficiency.
[0063] Back Figure 4 The execution of the multi-head attention computation task by the computing device 201 can be divided into two parallel pipelines. The control unit handles instruction fetching, decoding, and issuing, while the arithmetic units execute the instructions to perform the computation. For example, the control module 41 sequentially fetches compiled instructions from the I-Cache; after decoding, it simultaneously issues multiple operations within the same instruction to the back-end queue for execution by the arithmetic units in the arithmetic module 42. For instance, after receiving an operation, Move-DMA moves the next standard data block from SRAM into the current processing core's NRAM. Each arithmetic unit in the computing module 42 is responsible for performing operations on the data in the NRAM to complete the multi-head attention computation task.
[0064] In the above steps, the sequence interval computation task refers to the sequence segments that have already been split. This task can be processed by functional modules of a deep learning framework or kernel function library, or by a host machine in a heterogeneous computing system. The sequence interval computation task can be a task descriptor or a list of task descriptors. When a processing core executes a computation task, it can obtain a task descriptor from the task pool or move the computation task to on-chip storage and then read it from the on-chip storage. Each processing core is responsible for obtaining its own computation task from a specified storage location, then loading data from memory according to the pointers and parameters in the computation task, performing the computation, and writing back the result.
[0065] In some embodiments, the actual computational cost of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask. This may include determining the effective region in the attention score matrix based on the scope of the causal mask; estimating the computational cost required to generate the effective region as the actual computational cost of the current processing core, wherein the attention score matrix is obtained by performing a dot product operation on the query sequence and the key sequence.
[0066] The embodiments disclosed herein estimate the actual computational load of the current processing core by using an effective region, which can distribute the computational tasks more evenly to each processing core, thereby reducing waiting time and improving overall computational efficiency.
[0067] In some embodiments, the computational cost required to generate the effective region is estimated as the actual computational cost of the current processing core. This includes calculating a first computational cost based on the area of the geometry of the effective region, and using the first computational cost as the actual computational cost of the current processing core. The geometry can be a triangle or a trapezoid, etc.
[0068] Figure 8(a) illustrates the actual computational complexity when the query sequence and key sequence lengths are the same in some embodiments of this disclosure. As shown in Figure 8(a), when the sequence length seq_q of the query sequence and the sequence length seq_k of the key sequence are the same, the geometry of the effective region can be a triangle. The first computational complexity is calculated according to the triangle area formula, where the side length and height of the triangle can both be represented by the sequence length of the query sequence.
[0069] In some embodiments, Figure 8(b) illustrates a schematic diagram of the actual computational cost when the query sequence and key sequence lengths are different in some embodiments of this disclosure. As shown in Figure 8(b), when the sequence length of the query sequence is less than the sequence length of the key sequence, the geometry of the effective region can be a trapezoid. The first computational cost is calculated according to the trapezoid area formula, wherein the length of the upper base of the trapezoid is determined based on the absolute value of the difference between the sequence length of the key sequence and the sequence length of the query sequence and the starting index of the query sequence; the length of the lower base of the trapezoid is determined based on the absolute value of the difference between the sequence length of the key sequence and the sequence length of the query sequence and the ending position of the query sequence; and the height of the trapezoid is determined based on the absolute value of the difference between the starting index and the ending position of the query sequence.
[0070] Taking Figure 8(a) as an example, assuming the sequence length seq_q of the query sequence is equal to the sequence length seq_k of the key sequence, theoretically, the triangular region in the lower left corner is valid, while the triangular region in the upper right corner is actually invalid. The computational cost corresponding to the valid region can be calculated using the triangle area formula.
[0071] For example, when the length of the key sequence seq_k is greater than the length of the query sequence seq_q, i.e., delta_qk is greater than 0, where delta_qk is the difference between the length of the key sequence seq_k and the length of the query sequence seq_q, in this case, after performing a dot product operation on the key sequence and the query sequence, and adding a casual_mask, a trapezoidal region is obtained due to the effective range of the causal mask. This trapezoidal region is the effective region. The actual computational cost of each processing core can be represented by calculating the area of the trapezoid. As shown in Figure 8(b), the actual computational cost of each processing core can be represented by the blank trapezoidal part. When deriving the formula for calculating the area of the trapezoid, both multiplication and addition operations need to be considered; therefore, it is necessary to multiply by 2 when representing the actual computational cost.
[0072] Taking the trapezoidal area of core1 in Figure 8(b) as an example, the length of the upper base is begin_seq_q + delta_qk, the length of the lower base is end_seq_q + delta_qk, and the height is end_seq_q - begin_seq_q + 1. Then the trapezoidal area S1 of core1 can be expressed as:
[0073] The above analysis examines the actual computational load generated by the effective region from an algorithm implementation perspective. However, when implementing the algorithm in hardware, to ensure that the computational tasks of each processing core meet the hardware parallel instruction requirements (e.g., SIMD instructions require data lengths to be integer multiples of the vector width), it is necessary to round and align the computational tasks. For example, the processing capacity of a single arithmetic unit is... Where tile_q represents the height of the standard data block. This indicates the size of a standard data block, such as 32×32 or 64×64. When the arithmetic unit processes computational tasks, a square represents one standard data block.
[0074] As shown in Figure 8(a), the squares in the upper left and lower right corners contain both valid and invalid regions (i.e., gridded regions). Data from the invalid regions inevitably participates in the calculation. The redundant operations caused by the mismatch between the invalid regions defined by the causal mask and the hardware-defined blocks are a real issue in the hardware-implemented algorithm. For example, during alignment calculations, after cutting the attention score matrix into standard data block sizes, the boundary block spans both valid and invalid regions. The entire boundary block must be multiplied and added by the arithmetic unit in one operation. Although the invalid regions are eventually masked to zero, the multiplication-addition instructions have already been issued, consuming computational power and bandwidth. Clearly, the computational load at the boundary block is unavoidable. Therefore, after calculating qk, adding causal_mask, and finally performing the Softmax operation and subsequent calculations, the actual total computational load generated by the processing core's arithmetic unit should be... .
[0075] To better reflect the actual computational load generated during hardware execution, this disclosed embodiment proposes that when estimating the actual computational load, both the computational load corresponding to the effective part of causal_mask should be considered from the perspective of algorithm computation principle, and the redundant computational load generated by tile_q alignment should be considered from the perspective of artificial intelligence processor hardware implementation.
[0076] In some embodiments, the processing core includes a computation unit, the attention score matrix further includes an invalid region, and the actual computational cost of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the effective range of the causal mask. It also includes splitting the attention score matrix according to a standard data block to obtain a first redundant computational cost generated at the boundary between the effective and invalid regions. The standard data block is obtained based on the ability of the computation unit to process data in a single operation. The sum of the first computational cost and the first redundant computational cost is calculated as the actual computational cost of the current processing core.
[0077] In some embodiments, the attention score matrix is split according to standard data blocks to obtain the first redundant computational cost generated at the boundary between the effective and invalid regions. The number of standard data blocks located at the boundary between the effective and invalid regions can be determined based on the sequence length of the input sequence and the height of the standard data blocks. The first redundant computational cost is calculated based on the number of standard data blocks at the boundary, the height of the standard data blocks, and the remainder of the sequence length of the input sequence relative to the length of the standard data blocks. For example, returning to Figure 8(a), due to the alignment processing, the size of each square is... The area of the grid pattern represents the redundant computational cost resulting from the alignment operation. Considering both multiplication and addition operations, and multiplying the actual computational cost by 2, the area of the grid pattern can be calculated using pseudocode to determine the redundant computational cost corresponding to the data blocks on the boundary line: int q_loop = (end_seq_q - begin_seq_q) / tile_q; int remain_q = (end_seq_q - begin_seq_q) % tile_q; int mask_ops = ; Furthermore, to ensure that the starting address of data accessed by each processing core (such as fragments of the Q, K, and V matrices) in memory meets the hardware-required alignment boundaries (e.g., 64-bit alignment, 128-bit alignment, 256-bit alignment, etc.), data alignment needs to be achieved through padding or adjusting data pointers. For example, each processing core's arithmetic unit moves / computes data in a fixed vector width or standard data block size as the smallest unit. The starting address must be aligned to this width, and zeros are automatically padded when the number of rows / columns is insufficient. The tail must also be fully executed for multiplication and addition operations. Fixed vector widths include, for example, 32B, 64B, and 128B. After determining the sequence interval for each processing core, if the boundaries of the sequence interval do not meet the data alignment requirements, the interval boundaries need to be adjusted to meet the alignment requirements. This allows each processing core to use aligned memory access instructions when loading data, thereby improving bandwidth utilization.
[0078] When implementing computational tasks through hardware, the actual computational workload involved by the hardware mainly includes two parts: effective computation and redundant computation. The embodiments disclosed in this paper incorporate the redundant computational workload generated by alignment as part of the actual computational workload of the processing core, accurately reflecting the actual hardware operands. This approach can better optimize resource allocation, thereby improving overall utilization.
[0079] In some embodiments, the actual computational cost of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask. It also includes calculating a second redundant computational cost based on the number of zero-padding elements when the length of the sequence interval corresponding to the sequence interval computation task does not meet the hardware alignment requirements. The zero-padding elements refer to the zero elements at the end of the query sequence sub-block and key sequence sub-block corresponding to the sequence interval computation task allocated to the processing core. The sum of the first computational cost, the first redundant computational cost, and the second redundant computational cost is calculated as the actual computational cost of the current processing core.
[0080] When the number of rows or the main dimension length of the query sequence and key sequence matrix do not meet the hardware alignment requirements, zeros are padded at the end of the matrix. These padded elements are still fully multiplied and added by the arithmetic units. The resulting extra computational cost is called the second redundancy computational cost, or tail redundancy computational cost. In other words, the second redundancy computational cost is the number of additional computational operations caused by the participation of the padded positions in the calculation. It can include three parts: calculation of the effective query position and the padded key position; calculation of the padded query position and the effective key position; and calculation of the padded query position and the padded key position. Considering the causal masking effect, during the calculation process, one or more of the aforementioned three parts can be retained as the second redundancy computational cost, depending on the actual situation.
[0081] Because different data alignment methods may result in different numbers of padding positions, the amount of computational redundancy will also vary. When the length of the input sequence cannot be divided evenly by the height of the standard data block, the sequence needs to be padded to achieve data alignment.
[0082] Figure 9(a) illustrates a top-down alignment method in some embodiments of this disclosure. Assuming the query sequence seq_q and key sequence seq_k have lengths of 100 query positions, and the standard data block height tile_q is 64 query positions, and the query sequence is calculated in blocks from top to bottom according to the alignment method shown in Figure 9(a), there are 36 remaining query positions below the query sequence. During the alignment process, padding is required so that the length of the remaining query positions is divisible by the height of the standard data block. Therefore, 28 padding positions are added at the bottom of the query sequence seq_q, resulting in a redundant calculation amount exceeding 28 × 100 for the query sequence seq_q.
[0083] In some embodiments, the zero elements at the tail of the query sequence sub-blocks and key sequence sub-blocks corresponding to the sequence interval computation tasks allocated to the processing core are obtained by hardware alignment of the query sequence and key sequence using a bottom-first approach. The bottom-first approach can be implemented using a bottom-up data alignment method.
[0084] Figure 9(b) illustrates a bottom-up alignment method in some embodiments of this disclosure. Following the alignment method shown in Figure 9(b), the query sequence is divided into blocks for calculation in a bottom-up manner, leaving 36 query positions above the query sequence. During alignment, padding is required so that the length of the remaining query positions is divisible by the height of the standard data block. Therefore, 28 padding positions are added at the top of the query sequence seq_q, resulting in 28 × 36 redundant calculations beyond the query sequence seq_q. Clearly, adopting a bottom-first alignment strategy can reduce the redundant calculations beyond the range of seq_q.
[0085] Especially for shorter query sequences, considering the performance loss due to alignment, the actual computational cost is calculated by aligning upwards to the size of `tile_q`, where `seq_q = max(seq_q, tile_q)`. For example, if `seq_q` is only 10, the actual block size of `tile_q` is typically 64 during kernel computation, so the computational cost is calculated based on the aligned `tile_q` size.
[0086] In summary, the embodiments disclosed herein utilize the effective amount of the effective region and the redundant computational amount generated by alignment to estimate the actual computational amount of the current processing core. This can more accurately reflect the actual computational load on the hardware, thereby enabling better decisions in task allocation, resource scheduling, and performance optimization to maximize hardware utilization.
[0087] In some embodiments, the target computational cost refers to the average computational cost per processing core participating in the multi-head attention computation task. The target computational cost can be calculated based on the remaining computational cost and the number of remaining processing cores. The remaining computational cost refers to the total computational cost that has not yet been allocated to any processing core during the task allocation process, and the number of remaining processing cores refers to the total number of processing cores that have not yet been allocated computational tasks. The total computational cost corresponding to the input sequence can be calculated based on the data structure shape of the input sequence.
[0088] In some embodiments, after the initial task allocation is completed, the target computational load is dynamically updated based on the cumulative computational load already allocated or the actual computational load corresponding to a single attention head. The remaining computational load can be calculated in different ways during task allocation. For example, during core-level task splitting, the remaining computational load of the current processing core can be updated using the following pseudocode description: local_remain_attn_ops=task_avg_attn_ops-local_attn_ops Where local_remain_attn_ops represents the remaining computational cost of the current processing core, task_avg_attn_ops represents the initial average computational cost of each processing core, and local_attn_ops represents the actual computational cost of the current attention head of the current processing core.
[0089] During the task decomposition process at the attention head level, the remaining computational load of the current processing core can be updated using the following pseudocode: local_remain_attn_ops-=head_attn_ops Here, `local_remain_attn_ops` represents the remaining computational cost of the current processing core, and `head_attn_ops` represents the actual computational cost of the current attention head.
[0090] During the task decomposition process, the target computational load changes dynamically as the computational tasks are decomposed.
[0091] This disclosure proposes to dynamically adjust the target computational load of each processing core, enabling the task allocation process to adaptively cope with the unevenness of computational tasks and the heterogeneity of hardware processing cores, thereby achieving more refined load balancing and higher hardware utilization.
[0092] In some embodiments, the sequence interval computation task is dynamically split based on the comparison result between the actual computation amount of the current processing core and the target computation amount of each processing core. When the comparison result indicates that the target computation amount is greater than the actual computation amount, the sequence interval computation task can be determined based on the position of the attention head and the batch position; or, when the comparison result indicates that the target computation amount is less than or equal to the actual computation amount, the sequence interval is searched in the input sequence to obtain the sequence interval computation task.
[0093] If the comparison result indicates that the target computational cost is greater than the actual computational cost, it means that the current processing core's current attention head still has remaining computational cost that can be allocated further. If the comparison result indicates that the target computational cost is less than or equal to the actual computational cost, it means that the actual computational cost still has room for further partitioning. In this case, the function to find the optimal interval position is called to search for a sequence interval in the input sequence that matches the target computational cost.
[0094] This disclosure further proposes a hierarchical dynamic splitting strategy to achieve precise load balancing and resource scheduling during the multi-head attention computing task allocation process.
[0095] In some embodiments, the sequence interval computation task is determined based on the position of the attention head and the batch position. The starting index of the attention head can be updated to the next attention head. If it is determined that the starting index of the updated attention head is less than the total number of attention heads in the current batch, a batch switch is checked based on the starting index of the current batch. If a batch switch is determined to exist, the actual computation amount of the new attention head is recalculated. Alternatively, if it is determined that the starting index of the updated attention head is greater than or equal to the total number of attention heads in the current batch, the sequence interval of the input sequence is used as the sequence interval computation task.
[0096] If the starting index of the updated attention head is less than the total number of attention heads in the current batch, and the actual computational cost of the new attention head is less than the remaining computational cost of the current processing core, the starting index of the attention head is updated to bind multiple attention heads to the current processing core, thus obtaining a sequence interval computation task; or, if the starting index of the updated attention head is less than the total number of attention heads in the current batch, and the actual computational cost of the new attention head is greater than or equal to the remaining computational cost of the current processing core, a sequence interval is searched in the input sequence to obtain a sequence interval computation task.
[0097] Determining the sequence interval computation task based on the position of the attention head and the batch position is a first-level judgment based on the remaining computational capacity of the current processing core and a preset threshold. It further allocates tasks at the attention head level based on the comparison between the current attention head index position and the total number of attention heads in the current batch. When the actual computational capacity of a single attention head is less than the remaining computational capacity of the current processing core, the computational task of the current attention head is further adapted to improve the hardware utilization rate.
[0098] In some embodiments, when the target computational cost is less than or equal to the actual computational cost, it indicates that the current attention head of the current processing core has no remaining computational cost, or when the starting index of the updated attention head is determined to be less than the total number of attention heads in the current batch, and the actual computational cost of the new attention head is greater than or equal to the remaining computational cost of the current processing core, a sequence interval is searched in the input sequence. Searching for a sequence interval in the input sequence can be based on a binary search method to find the task segmentation position in the input sequence, so that the theoretical computational cost obtained based on the task segmentation position is closest to the target computational cost; the sequence interval is then determined from the input sequence based on the task segmentation position.
[0099] For example, the task segmentation position in the input sequence can be found by calling the function to find the optimal seq_q position. This function uses a binary search to find the optimal local_seq_end position between local_seq_begin and max_seq_q, so that the computational cost obtained from the optimal local_seq_end is closest to the target computational cost local_remain_attn_ops. The function interface for finding the optimal seq_q position can be implemented, for example, as described in the following pseudocode: int64_t findSeqQ(int64_t attn_ops, int begin_seq_q, int end_seq_q, int delta_qk, int tile_q); begin_seq_q represents the starting index of seq_q, for example, it can be represented as local_seq_begin.
[0100] end_seq_q represents the end position of seq_q, for example, it is represented as max_seq_q.
[0101] delta_qk represents the difference between seq_k and seq_q, i.e., delta_qk = seq_k - seq_q.
[0102] local_seq_end is the result of calling the int64_t findSeqQ function.
[0103] To adapt to input sequences of varying lengths and achieve efficient computation, this disclosure proposes a multi-head attention computation task allocation method. This method can allocate tasks based on the data structure shape of the input sequence, the total sequence length corresponding to the input sequence, and the total number of processing cores participating in the multi-head attention computation task, satisfying both static configuration and dynamic perception requirements, thereby improving the efficiency of task allocation. The following describes a method for allocating multi-head attention computation tasks. Figure 10 It describes in detail the task allocation method for multi-head attention computation. Figure 10 An exemplary flowchart of a multi-head attention computation task allocation method in some embodiments of this disclosure is shown.
[0104] like Figure 10 As shown, in step S1001, the data structure shape of the input sequence, the total sequence length corresponding to the input sequence, and the total number of processing cores participating in the multi-head attention computation task are determined.
[0105] In step S1002, when the data structure shape of the input sequence represents a filling pattern and the total sequence length corresponding to the input sequence is divisible by the total number of processing cores participating in the multi-head attention computation task, a static splitting strategy is adopted to split the input sequence. The static splitting strategy is to split the input sequence uniformly.
[0106] In step S1003, when the structure shape of the input sequence represents a compression mode or the data structure shape of the input sequence represents a filling mode, and the total sequence length corresponding to the input sequence cannot be divided by the total number of processing cores participating in the multi-head attention computation task, a dynamic splitting strategy is adopted to split the input sequence into sequence interval computation tasks. The dynamic splitting strategy refers to the strategy of dynamically splitting the input sequence based on the comparison result between the actual computation amount of the current processing core and the target computation amount of each processing core.
[0107] In step S1004, after determining the sequence interval calculation task, the unique identifier of each processing core is traversed in a loop. When the loop variable is equal to the unique identifier of the current processing core, the sequence interval calculation task is assigned to the current processing core until the calculation tasks of all processing cores are assigned.
[0108] In step S1005, when the loop variable is not equal to the unique identifier of the current processing core, it is determined whether all sequences in the current batch have been allocated. If yes, the remaining processing cores do not need to process computational tasks, and the computational task splitting of all processing cores ends. If no, the parameter update operation is performed until the computational tasks of all processing cores are allocated.
[0109] In the above steps, it is determined whether the input sequence is in Pad mode and whether the total sequence length corresponding to the input sequence is divisible by the total number of processing cores participating in the multi-head attention computation task. When the determination result is yes, that is, when the data structure shape of the input sequence represents the Pad mode, and it is determined that the total sequence length corresponding to the input sequence is divisible by the total number of processing cores participating in the multi-head attention computation task, the static splitting strategy is configured to split the input sequence evenly and distribute the splitting result equally to each processing core participating in the multi-head attention computation task.
[0110] The data structure shape of the input sequence can be, for example, a four-dimensional shape [batch_size, head_num, max_seq_q, head_size] or a three-dimensional shape [total_seq_q, head_num, head_size]. With four-dimensional input, each batch has the same seq_q, which is Pad mode. With three-dimensional input, it may be Pad mode or non-Pad mode; it is necessary to traverse the batches and determine if the seq_q of each batch is the same. If they are the same, it is Pad mode; otherwise, it is non-Pad mode.
[0111] When the result is negative, i.e. when the structure of the input sequence represents a compressed mode (Pack mode) or the data structure of the input sequence represents a filled mode, and the total sequence length corresponding to the input sequence cannot be divided by the total number of processing cores participating in the multi-head attention computation task, a dynamic splitting strategy is adopted to split the input sequence.
[0112] This disclosure also proposes a judgment method based on multi-level cascading to implement a dynamic splitting strategy. The following is in conjunction with... Figures 11-13 This will further describe the idea of implementing a dynamic splitting strategy based on a multi-level cascading judgment method. Specifically, Figure 11 It describes the idea of splitting the processing core level, which triggers the splitting of computational tasks at the processing core level by judging the actual workload of the processing core and the target computational workload of each processing core. Figure 12 This paper describes the idea of splitting attention heads at the hierarchical level. It determines whether all attention heads have been allocated by judging the relationship between the position of each attention head and the total number of attention heads, and determines whether intra-head splitting is necessary by considering the remaining computational cost of the current processing core and the actual computational cost of the current head to be allocated. Finally, Figure 13 The description describes how the splitting results are sequentially assigned to each processing core participating in the multi-head attention computation task. The steps of generating the task descriptor list and storing it in the specified storage location are not shown in the figure.
[0113] Figure 11 A partial exemplary flowchart of a computational task splitting method in some embodiments of this disclosure is shown. For example... Figure 11 As shown, the method includes: In step S1, the total number of queries in the query sequence seq_q, total_seq_q, and the total global computation amount, global_total_attn_ops, are calculated.
[0114] Depending on the scenario being split, the calculation methods for the total global computation and total number of queries also differ. In Pad mode, the seq of each batch is the same, assuming the shape of the input q is... [batch_size, head_num, max_seq_q, head_size].
[0115] In non-Pad mode, the sequence number (seq) of each batch is different, assuming the shape of the input q is... [total_seq_q, head_num, head_size].
[0116] In Pad mode and If the value is an integer multiple of the total number of tasks (taskDim), it means that seq_q can be evenly split. The total global computational load can be expressed by the following formula:
[0117] The total number of seq_q can be expressed by the following formula:
[0118] The number of seq_q values allocated to each processing core can be obtained using the following formula:
[0119] The actual computational cost of each head can be calculated by calling a function that tallies the actual computational cost. For example, the actual computational cost corresponding to an attention head can be implemented using the following pseudocode: head_attn_ops=computeAttnOps(0,max_seq_q,delta_qk,tile_q) Where delta_qk=max_seq_q.
[0120] In non-Pad mode, or and When the number of tasks is not an integer multiple of the total number of tasks (taskDim), the seq_q of each batch is different, and the actual computational cost (head_attn_ops) of each attention head is also different. The total global computational cost (global_total_attn_ops) can be expressed by the following formula:
[0121] The total number of seq_q can be expressed by the following formula:
[0122] When not in Pad mode, or in Pad mode and If the number of tasks is not an integer multiple of the total number of tasks (taskDim), then it needs to be dynamically split based on the actual amount of computation and the average amount of computation per processing core.
[0123] In step S2, it is determined whether the device is in non-Pad mode or Pad mode and When the number of tasks is not an integer multiple of the total number of tasks (taskDim), the scenario where splitting is not possible is entered. First, initialize the various scalar parameters used to implement the splitting strategy. These scalar parameters include, but are not limited to: the number of query blocks allocated to the current processing core (bhq_num=0); the starting index of the query blocks handled by the current processing core (bhq_begin=0); the cumulative computation (accumulate_attn_ops=0); the starting index of the attention head handled by the current processing core (head_begin=0); the starting index of the current batch handled by the current processing core (batch_begin=0); the starting index of the previous batch (last_batch_begin=0); the starting index of the query sequence handled by the current processing core (local_seq_begin=0); and the current loop variable (i=0), etc.
[0124] In step S3, it is determined whether the loop variable i is less than the total number of tasks, taskDim, which is the total number of processing kernels participating in the multi-head attention computation task. If yes, proceed to step S4; otherwise, proceed to step S5.
[0125] In step S4, the starting index of the query blocks handled by the current processing core is determined through an update operation, the number of query blocks processed by the current core is reset, and the starting index of the current batch handled by the current processing core is also reset. Then, proceed to step S6.
[0126] For example, variable updates can be implemented using the following pseudocode: bhq_begin += bhq_num; bhq_num=0; batch_begin=head_begin / head_num; After step S4, steps S6-S8 are executed to dynamically determine the target computational load of each processing core and the estimated actual computational load of the current processing core based on the remaining total task load and the remaining processing cores.
[0127] In step S6, in pack mode, the current batch is used as the processing object, and the maximum query sequence length and the maximum key sequence length are updated.
[0128] For example, by checking if a new batch has been reached (last_batch_begin != batch_begin), the following update is made: last_batch_begin==batch_begin; Update max_seq_q and max_seq_k in pack mode; In step S7, the remaining total computational load global_remain_attn_ops and the average workload per processing core (i.e., the target computational load) are calculated.
[0129] For example, the remaining total computational cost calculated according to the following pseudocode description can also be called the global remaining computational cost:
[0130] Among them, global_total_attn_ops is the total global computation, which is the computation of all attention heads on all sequences; accumulate_attn_ops is the accumulated computation, that is, the computation that has been allocated to each processing core; and global_remain_attn_ops is the remaining global computation.
[0131] During dynamic task allocation, the exact amount of computation allocated to each processing core cannot be determined. This is because the computational load for each attention head may differ, especially when the sequence length varies, such as in pack mode where different sequences have different lengths. Based on the already allocated computational load, the remaining computational load, and the number of remaining processing cores, the target computational load that should be allocated to each processing core can be dynamically adjusted.
[0132] The target computational cost per core can be obtained by dividing the remaining total computational cost by the number of remaining processing cores. Alternatively, the target computational cost per core (i.e., average computational cost) can be obtained using the following pseudocode:
[0133] Here, DIV_UP represents the floor function, where i is the value of the loop variable i in this iteration.
[0134] This disclosed embodiment calculates the global remaining computational load in real time during the core task allocation process. Based on the global remaining computational load global_remain_attn_ops and the number of remaining unallocated cores taskDim-i, it dynamically calculates the target computational load for the current core. It then allocates computational tasks close to the target computational load to the current core and tracks the cumulative computational load accumulate_attn_ops to reflect the allocation status in real time, avoiding estimation errors, until all cores are allocated.
[0135] In step S8, the actual computational load of the currently processing core and the remaining computational load after the current core has been processed are calculated. After step S8, proceed to step S9.
[0136] By precisely controlling the computational load of the current processing core (i.e., a single processing core), the computational load can be dynamically adapted to changes in computational load at different sequence positions and attention heads. This is achieved by calling predefined functions to estimate the actual computational load for the current head handled by the current processing core (local_attn_ops) and the remaining computational load of the current processing core (local_remain_attn_ops). For example, this can be implemented using the following pseudocode: local_attn_ops=computeAttnOps(local_seq_begin,max_seq_q,…); local_remain_attn_ops= task_avg_attn_ops- local_attn_ops; The predefined function interface for calculating the actual computational load can be implemented, for example, using the following pseudocode description: int64_t computeAttnOps(int begin_seq_q, int end_seq_q, int delta_qk,int tile_q); begin_seq_q represents the starting index of the sequence interval seq_q used to calculate the actual computational quantity; end_seq_q represents the end index of the sequence interval seq_q used to calculate the actual computational quantity; delta_qk represents the length difference between the key sequence interval seq_k and the query sequence seq_q used to calculate the actual computational cost, i.e., delta_qk = seq_k - seq_q.
[0137] The function used to calculate the actual computational cost can include two parts: the computational cost of the effective portion of the causal_mask and the computational cost of alignment redundancy. The specific calculation method is detailed in the description of the actual computational cost and redundancy cost of the effective region. It will not be repeated here.
[0138] In step S5, when it is determined that the current processing core is the last processing core, the remaining sequence is assigned to the last processing core.
[0139] For example, `taskid == i - 1` is used to determine if the current processing core is the last processing core participating in the multi-head attention calculation task, and the last processing core is responsible for processing all remaining `seq_q`. The starting index `bhq_begin` of the query block handled by the current processing core is moved to the ending index of the query block of the previous processing core, and all remaining sequences are assigned to the last core. This can be implemented, for example, using the following pseudocode: bhq_begin += bhq_num; bhq_num=total_seq_q- bhq_begin.
[0140] In step S9, it is determined whether the remaining computational load of the current processing core, local_remain_attn_ops, is less than or equal to 0; if yes, proceed to step S10; if no, proceed to step S11. For example, the judgment condition can be expressed as local_remain_attn_ops<=0.
[0141] In step S10, the findSeqQ function is called to obtain the optimal seq_q task segmentation position local_seq_end; and a parameter update operation is performed. For example, the parameter update operation can be implemented using the following pseudocode: accumulate_attn_ops+=computeAttnOps(local_seq_begin,max_seq_q,…); bhq_num=local_seq_end- local_seq_begin; local_seq_begin=( local_seq_end== max_seq_q)?0: local_seq_end; The size of the sequence block being processed by the current processing core is updated cumulatively to determine the length of the sequence segment that the current processing core is responsible for, and to determine whether the starting index of the sequence segment that the current processing core is responsible for has reached the end of the sequence that the current processing core is responsible for. If it is determined that the end of the sequence has been reached, the process switches to the next attention head.
[0142] In this disclosed embodiment, by comparing the actual computational load of the current processing core with the target computational load of each processing core, it is possible to effectively prevent a bottleneck problem caused by assigning too many computational tasks to a single processing core.
[0143] Figure 12 A partial exemplary flowchart of a computational task splitting method in some embodiments of this disclosure is shown. For example... Figure 12 As shown, if the determination in step S9 is negative, proceed to step S11. The method includes: In step 11, based on the starting index `local_seq_begin` and the maximum sequence length `max_seq_q` of the sequence segment handled by the current processing core, the actual computational cost of the remaining sequence for a single attention head is calculated, and this actual computational cost is accumulated into the global cumulative computational cost `accumulate_attn_ops`. The computational tasks corresponding to the sequence interval [local_seq_begin, max_seq_q) are then assigned to the current processing core. For example, this can be implemented using the following pseudocode: bhq_num=max_seq_q- local_seq_begin; accumulate_attn_ops+=computeAttnOps(local_seq_begin,max_seq_q,…); Here, max_seq_q can be understood as the maximum sequence of the current attention head, and local_seq_begin is the starting index of the sequence segment of the current attention head of the current processing kernel.
[0144] In step S12, the current attention head responsible for the current processing core is updated to the next attention head, and the actual computational cost of the new attention head is recalculated; The current attention head is updated using `head_begin++;`, and the actual computational cost from the start position of the new attention head to `max_seq_q` is calculated by calling the `computeAttnOps` function. For example, the actual computational cost of the new attention head can be recalculated using the following pseudocode: head_attn_ops= computeAttnOps(0,max_seq_q,…); In step S13, it is determined whether there are attention points to be assigned. If yes, proceed to step S14; otherwise, proceed to step S15.
[0145] This step can be achieved by determining whether the starting index of the current attention head (head_begin, i.e., the updated starting index of the current attention) handled by the current processing core is less than the total number of attention heads. To achieve this, for example, the judgment condition can be expressed as: head_begin< .
[0146] In step S14, the batch index to which the current attention head belongs is calculated. The batch index `batch_begin` to which the current attention head belongs is compared with the batch index `last_batch_begin` of the previous record. If they are different, it indicates that a batch switch has occurred, and a parameter update operation is performed. The parameter update operation includes, but is not limited to, loading the maximum query sequence length `max_seq_q` and the maximum key sequence length `max_seq_k` corresponding to the current batch index `batch_begin`; recalculating the actual computational cost `head_attn_ops` of the current attention head based on the updated `max_seq_q`; and updating the previous batch index of the record: `last_batch_begin = batch_begin`.
[0147] The batch index to which the current attention head belongs can be calculated using batch_begin=head_begin / head_num.
[0148] Update the batch index of the previous batch by setting last_batch_begin==batch_begin.
[0149] The function `head_attn_ops=computeAttnOps(0,max_seq_q, ...)` is recalculated based on the maximum sequence length of the new batch. When dealing with multi-batch, variable-length sequence scenarios, accurately calculating which batch the current attention head belongs to ensures the accuracy of dynamically updated data and guarantees that the correct sequence length is used to calculate the actual computational cost of a single attention head.
[0150] After step S14, proceed to step S16.
[0151] In step S16, it is determined whether the actual computational cost of the current attention head is less than the remaining computational cost of the current processing core. If not, proceed to step S17; otherwise, proceed to step S18.
[0152] The condition is determined by whether the actual computational cost of the current attention head, head_attn_ops, is less than the remaining computational cost of the current processing core, local_remain_attn_ops. <local_remain_attn_ops。
[0153] In step S17, a parameter update operation is performed, and then the process returns to step S13. For example, the update operation can be implemented using the following pseudocode: local_remain_attn_ops-= head_attn_ops; accumulate_attn_ops += head_attn_ops; bhq_num += max_seq_q; head_begin++; In step S18, the findSeqQ function is called to obtain the task splitting position local_seq_end of the optimal seq_q; and parameter update operations. For example, the parameter update operations can be implemented in the following way described by the pseudo code: accumulate_attn_ops += computeAttnOps(0, local_seq_end,...); bhq_num += local_seq_end; local_seq_begin = (local_seq_end == max_seq_q)? 0 : local_seq_end; In the embodiments of the present disclosure, three levels of decision-making judgments are constructed through step S9, step S13, and step S16. Among them, step S9 determines whether the current processing core can still accept more computing tasks. local_remain_attn_ops > 0 indicates that there is still remaining computing power in the current processing core and further allocation can continue. local_remain_attn_ops <= 0 indicates that the current processing core has reached or exceeded the target computing power, and the allocation to the current processing core should be stopped. This effectively avoids a single processing core being overloaded and becoming a performance bottleneck. Step S13 further determines whether there are still attention heads to be allocated. head_begin < indicates that there are still unallocated attention heads. head_begin >= indicates that all attention heads have been allocated, which can ensure that all attention heads can be allocated. Finally, step S16 determines whether the actual computing power of the current attention head can fully match the remaining computing power of the processing core. head_attn_ops < local_remain_attn_ops indicates that the current attention head can be fully allocated to the current processing core. head_attn_ops >= local_remain_attn_ops indicates that the current head needs to be further split or adjusted. Through the processing of step S16, it can be ensured that the computing power of a single processing core can be reasonably matched, thereby improving hardware utilization. Through the three-level decision cascade judgment, the scheduling efficiency can be effectively improved.
[0154] Figure 13 shows a partial exemplary flowchart of the computing task splitting method in some embodiments of the present disclosure. As Figure 12As shown, after executing step S10, or step S13, or step S18, the process proceeds to step S15. The method includes: In step S15, it is determined whether the current processing core is the designated core for this round. If so, the sequence interval calculation task is written into the task descriptor of the current processing core to complete the task acquisition; otherwise, proceed to step S19.
[0155] The determination of whether the current processing kernel is the designated processing kernel for this round is achieved by checking whether the unique identifier taskid of the processing kernel is equal to the current loop variable i.
[0156] If the unique identifier taskid of the processing kernel is not equal to the current loop variable i, the current processing kernel skips the task pick-up operation and only performs state update calculations to ensure that the same sequence interval is not assigned repeatedly.
[0157] In step S19, it is determined whether all sequences in the current batch have been allocated. If yes, the remaining processing cores do not need to process computation tasks, bhq_num is set to 0, and the computation task splitting of all processing cores ends. If not, proceed to step S20.
[0158] The determination of whether all sequences in the current batch have been allocated is achieved by checking whether bhq_begin + bhq_num >= total_seq_q.
[0159] In step S20, a parameter update operation is performed, and then the process returns to step S3 until the computational tasks for all processing cores are allocated. For example, the update operation can be implemented using the following pseudocode: If local_seq_begin == max_seq_q, update head_begin++ and i++.
[0160] In some embodiments, the data structure shape of the input sequence is a three-dimensional input or a four-dimensional input.
[0161] In some embodiments, the actual computational load and the target computational load are the number of floating-point operations.
[0162] In summary, steps S1 to S14, and steps 16 to 18 can be understood as sub-steps that implement step S1003. Steps S15 to S20 can be understood as sub-steps that implement steps S1004 and S1005. These sub-steps can be partially combined into one step, or they can be further broken down into multiple sub-steps.
[0163] The method provided in this disclosure not only effectively prevents processing core overload and ensures full task coverage, but also significantly improves the overall throughput and hardware resource utilization of the parallel computing system through precise matching of task granularity and core processing capabilities.
[0164] First, by comparing the actual computational load of the current processing core with the dynamically changing target computational load (i.e., the average computational load), it can adjust and split the computational tasks in real time and dynamically to obtain the computational tasks in the sequence interval, which ensures that the computational load of a single processing core is not too large.
[0165] Secondly, the actual computational cost is estimated based on the specific data of the input sequence and the effective range of the causal mask. Although the causal mask causes the attention computation to exhibit a non-uniform computational density of "trapezoidal" or "triangular" (the later the sequence position, the more preceding correlations need to be processed), the actual computational cost estimation method provided in this disclosure can more accurately reflect the true computational complexity of each sequence interval.
[0166] Finally, based on the dynamic partitioning mechanism described above, the computational load of each processing core is maximized to be close to the target computational load. This effectively avoids the problem that some processing cores may complete their computations prematurely and enter an idle waiting state (i.e., unbalanced load) that may occur with traditional static partitioning, thereby fully tapping the parallel computing potential of all processing cores and improving hardware utilization and overall computational throughput.
[0167] Based on the foregoing description, those skilled in the art will understand that the artificial intelligence processor of the embodiments disclosed herein can be implemented in a single chip. Therefore, this disclosure also discloses a chip that includes the artificial intelligence processor as described in any of the foregoing embodiments. Additionally, this disclosure also discloses a circuit board that includes the aforementioned chip.
[0168] Depending on the application scenario, the electronic devices or apparatus disclosed herein may include servers, cloud servers, server clusters, data processing devices, robots, computers, printers, scanners, tablets, smart terminals, PC devices, IoT terminals, mobile terminals, mobile phones, dashcams, navigators, sensors, cameras, video cameras, projectors, watches, headphones, mobile storage, wearable devices, visual terminals, autonomous driving terminals, vehicles, home appliances, and / or medical devices. The vehicles include airplanes, ships, and / or vehicles; the home appliances include televisions, air conditioners, microwave ovens, refrigerators, rice cookers, humidifiers, washing machines, lights, gas stoves, and range hoods; the medical devices include MRI scanners, ultrasound machines, and / or electrocardiographs. The electronic devices or apparatus disclosed herein can also be applied in fields such as the Internet, IoT, data centers, energy, transportation, public management, manufacturing, education, power grids, telecommunications, finance, retail, construction sites, and healthcare. Furthermore, the electronic devices or apparatus disclosed herein can also be used in application scenarios related to artificial intelligence, big data, and / or cloud computing, such as cloud computing, edge computing, and terminal applications. In one or more embodiments, the high-computing-power electronic devices or apparatuses according to the present disclosure can be applied to cloud devices (e.g., cloud servers), while the low-power electronic devices or apparatuses can be applied to terminal devices and / or edge devices (e.g., smartphones or cameras). In one or more embodiments, the hardware information of the cloud devices and the hardware information of the terminal devices and / or edge devices are compatible with each other, so that suitable hardware resources can be matched from the hardware resources of the cloud devices to simulate the hardware resources of the terminal devices and / or edge devices based on the hardware information of the terminal devices and / or edge devices, so as to complete the unified management, scheduling and collaborative work of the end-to-cloud or cloud-edge-end integration.
[0169] It should be noted that, for the sake of brevity, this disclosure describes some methods and their embodiments as a series of actions and combinations thereof. However, those skilled in the art will understand that the solutions disclosed herein are not limited by the order of the described actions. Therefore, based on the disclosure or teachings of this document, those skilled in the art will understand that some steps can be performed in a different order or simultaneously. Furthermore, those skilled in the art will understand that the embodiments described in this disclosure can be considered optional embodiments, that is, the actions or modules involved are not necessarily essential for the implementation of one or more solutions disclosed herein. In addition, depending on the solution, the description of some embodiments in this disclosure may have different emphases. In view of this, those skilled in the art will understand that parts not described in detail in a certain embodiment of this disclosure can also be referred to the relevant descriptions of other embodiments.
[0170] In terms of specific implementation, based on the disclosure and teachings of this document, those skilled in the art will understand that several embodiments disclosed herein can also be implemented in other ways not disclosed herein. For example, regarding the various units in the electronic device or apparatus embodiments described above, this document divides them based on logical functions, but in actual implementation, there may be other division methods. As another example, multiple units or components can be combined or integrated into another system, or some features or functions in a unit or component can be selectively disabled. Regarding the connection relationships between different units or components, the connections discussed above in conjunction with the accompanying drawings can be direct or indirect couplings between units or components. In some scenarios, the aforementioned direct or indirect couplings involve communication connections utilizing interfaces, where the communication interface can support electrical, optical, acoustic, magnetic, or other forms of signal transmission.
[0171] In this disclosure, the units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units. The aforementioned components or units may be located in the same location or distributed across multiple network units. Furthermore, depending on actual needs, some or all of the units can be selected to achieve the purpose of the solution described in the embodiments of this disclosure. Additionally, in some scenarios, multiple units in the embodiments of this disclosure may be integrated into one unit or each unit may exist physically independently.
[0172] In some implementation scenarios, the integrated unit described above can be implemented as a software program module. If implemented as a software program module and sold or used as an independent product, the integrated unit can be stored in a computer-readable storage device (CMSDD). Therefore, when the disclosed solution is embodied in a software product (e.g., a computer-readable storage medium), the software product can be stored in a memory, which may include several instructions to cause a computer device (e.g., a personal computer, server, or network device) to execute some or all of the steps of the method described in the embodiments of this disclosure. The aforementioned memory may include, but is not limited to, various media capable of storing program code, such as USB flash drives, flash drives, read-only memory (ROM), random access memory (RAM), portable hard drives, magnetic disks, or optical disks.
[0173] In other implementation scenarios, the integrated units described above can also be implemented in hardware, i.e., as specific hardware circuits, which may include digital circuits and / or analog circuits. The physical implementation of the circuit's hardware structure may include, but is not limited to, physical devices, which may include, but are not limited to, transistors or memristors. Therefore, the various devices described herein (e.g., artificial intelligence processors or other processing devices) can be implemented using appropriate hardware processors, such as CPUs, GPUs, FPGAs, DSPs, and ASICs. Furthermore, the aforementioned storage units or storage devices can be any suitable storage medium (including magnetic storage media or magneto-optical storage media), such as resistive random access memory (RRAM), dynamic random access memory (DRAM), static random access memory (SRAM), enhanced dynamic random access memory (EDRAM), high-bandwidth memory (HBM), hybrid memory cube (HMC), ROM, and RAM.
[0174] While numerous embodiments of this disclosure have been shown and described herein, it will be apparent to those skilled in the art that such embodiments are provided by way of example only. Many modifications, alterations, and alternatives will occur to those skilled in the art without departing from the spirit and intent of this disclosure. It should be understood that various alternatives to the embodiments of this disclosure described herein may be employed in the practice of this disclosure. The appended claims are intended to define the scope of this disclosure and therefore cover equivalents or alternatives within the scope of these claims.
Claims
1. A method for allocating multi-head attention computing tasks, characterized in that, This method is applied to an artificial intelligence processor comprising multiple processing cores, and is executed by each processing core. The method includes: Load the sequence interval computation task of the multi-head attention allocated to the processing core. The sequence interval computation task is dynamically split according to the comparison result of the actual computation amount of the current processing core and the target computation amount of each processing core. The actual computation amount of the current processing core is estimated according to the input sequence of the multi-head attention computation task and the scope of the causal mask. The target computation amount refers to the average computation amount of each processing core participating in the multi-head attention computation task. Perform the sequence interval calculation task.
2. The computational task allocation method according to claim 1, characterized in that, The input sequence includes a query sequence and a key sequence. The actual computational cost of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask, including: Determine the effective region in the attention score matrix based on the scope of the causal mask; The estimated computational cost required to generate the effective region is used as the actual computational cost of the current processing core. The attention score matrix is obtained by performing a dot product operation on the query sequence and the key sequence.
3. The computational task allocation method according to claim 2, characterized in that, The estimated computational cost required to generate this effective region is taken as the actual computational cost of the current processing core, including: The first computational quantity is calculated based on the area of the geometry of the effective region, and the first computational quantity is used as the actual computational quantity of the current processing core.
4. The computational task allocation method according to claim 3, characterized in that, When the sequence length of the query sequence is the same as the sequence length of the key sequence, the geometry of the effective region is a triangle.
5. The computational task allocation method according to claim 3, characterized in that, When the length of the query sequence is less than the length of the key sequence, the geometry of the effective region is trapezoidal.
6. The computational task allocation method according to claim 3, characterized in that, The processing kernel includes a computation unit, the attention score matrix further includes an invalid region, and the actual computational complexity of the current processing kernel is estimated based on the input sequence of the multi-head attention computation task and the effective range of the causal mask. It also includes: The attention score matrix is split according to the standard data block to obtain the first redundant computational amount generated at the boundary between the effective and ineffective regions. The standard data block is obtained based on the single data processing capability of the computing unit. The sum of the first computational cost and the first redundant computational cost is calculated as the actual computational cost of the current processing core.
7. The computational task allocation method according to claim 6, characterized in that, Based on the standard data block splitting of the attention score matrix, the first redundant computational cost generated at the boundary between the effective and ineffective regions is obtained, including: Based on the sequence length of the input sequence and the height of the standard data blocks, determine the number of standard data blocks located on the boundary line between the valid and invalid regions; The first redundancy computation is calculated based on the remainder of the input sequence length relative to the length of the standard data block, the number of standard data blocks on the boundary line, and the height of the standard data block.
8. The computational task allocation method according to claim 6, characterized in that, The actual computational cost of the current processing core is estimated based on the input sequence of the multi-head attention computation task and the scope of the causal mask, and also includes: When the length of the sequence interval corresponding to the sequence interval calculation task does not meet the hardware alignment requirements, the second redundant calculation amount is calculated based on the number of zero padding elements. The zero padding elements refer to the zero elements at the end of the query sequence sub-block and key sequence sub-block corresponding to the sequence interval calculation task allocated to the processing core. The sum of the first computational load, the first redundant computational load, and the second redundant computational load is calculated as the actual computational load of the current processing core.
9. The computational task allocation method according to claim 8, characterized in that, The zero elements at the tail of the query sequence sub-blocks and key sequence sub-blocks corresponding to the sequence interval computation tasks assigned to the processing core are obtained by performing hardware alignment operations on the query sequence and key sequence using the bottom-first principle.
10. The computational task allocation method according to claim 1, characterized in that, The target computational load is calculated based on the remaining computational load and the number of remaining processing cores. The remaining computational load refers to the total computational load that has not yet been allocated to any processing core during the task allocation process, and the number of remaining processing cores refers to the total number of processing cores that have not yet been allocated computing tasks.
11. The computational task allocation method according to claim 10, characterized in that, After the initial task allocation is completed, the target computational volume is dynamically updated based on the cumulative computational volume of the allocated tasks or the actual computational volume corresponding to a single attention head.
12. The computational task allocation method according to claim 1, characterized in that, The sequence interval computation task is dynamically divided based on a comparison between the actual computational load of the current processing core and the target computational load of each processing core, including: When the comparison result indicates that the target computational load is greater than the actual computational load, the sequence interval computation task is determined based on the position of the attention head and the batch position; or, The comparison result indicates that the target computational amount is less than or equal to the actual computational amount. The sequence interval is then searched in the input sequence to obtain the sequence interval computation task.
13. The computational task allocation method according to claim 12, characterized in that, The sequence interval calculation task is determined based on the position of the attention head and the batch position, including: Update the starting index of the attention head to the next attention head; If the starting index of the updated attention head is determined to be less than the total number of attention heads in the current batch, check for a batch switch based on the starting index of the current batch. If a batch switch is determined to exist, recalculate the actual computational cost of the new attention head; or, When the starting index of the updated attention head is determined to be greater than or equal to the total number of attention heads in the current batch, the sequence interval of the input sequence is used as the sequence interval calculation task.
14. The computational task allocation method according to claim 13, characterized in that, The sequence interval calculation task, which determines the position of the attention head and the batch position, also includes: If the starting index of the updated attention head is determined to be less than the total number of attention heads in the current batch, and the actual computational cost of the new attention head is less than the remaining computational cost of the current processing core, then the starting index of the attention head continues to be updated to bind multiple attention heads to the current processing core, thus obtaining a sequence interval computation task; or, When it is determined that the starting index of the updated attention head is less than the total number of attention heads in the current batch, and the actual computational cost of the new attention head is greater than or equal to the remaining computational cost of the current processing core, the sequence interval is searched in the input sequence to obtain the sequence interval computation task.
15. The computational task allocation method according to claim 12, characterized in that, Find the sequence interval in the input sequence, including: The task segmentation position is found in the input sequence using a binary search method, so that the actual computational cost obtained based on the task segmentation position is closest to the target computational cost. The sequence interval is determined from the input sequence based on the task segmentation position.
16. The computational task allocation method according to claim 1, characterized in that, After defining the sequence interval computation task, the method further includes: The loop iterates through the unique identifier of each processing core. When the loop variable is equal to the unique identifier of the current processing core, the sequence interval calculation task is assigned to the current processing core.
17. The computational task allocation method according to claim 16, characterized in that, The method also includes: If the loop variable is not equal to the unique identifier of the current processing core, determine whether all sequences in the current batch have been allocated.
18. The computational task allocation method according to claim 1, characterized in that, The dynamic splitting of the input sequence based on the comparison between the actual computational cost of the attention head and the target computational cost of each processing core is determined by a dynamic splitting strategy. This dynamic splitting strategy is determined when the data structure shape of the input sequence represents a compression mode or a filling mode, and the total sequence length corresponding to the input sequence cannot be divided by the total number of processing cores participating in the multi-head attention computation task.
19. The computational task allocation method according to claim 18, characterized in that, The data structure shape of the input sequence is either three-dimensional or four-dimensional.
20. The computational task allocation method according to claim 1, characterized in that, The actual computational load and the target computational load are the number of floating-point operations.
21. An artificial intelligence processor, characterized in that, The artificial intelligence processor includes multiple processing cores, each processing core including a control unit and at least two computing units, wherein... The control unit is configured to perform the step of loading the sequence interval computation task of the multi-head attention allocated to the processing core in the multi-head attention computation task allocation method as described in any one of claims 1-20; The computing unit is configured to perform the sequence interval calculation task.
22. A chip configured to include the artificial intelligence processor as described in claim 21.
23. A circuit board comprising the chip according to claim 22.
Citation Information
Cited By
Load balanced distribution system for attention calculation
CN122111690A
A load balancing assignment system for attention computation
CN122111690B