Task scheduling method and device for attention calculation, medium, equipment and product

Abstract

The application discloses a task scheduling method and device for attention calculation, a medium, equipment and products, the method comprises the following steps: decomposing the loading task of query block into N1 loading subtasks, and decomposing the calculation task of the jth attention score block into N2 calculation subtasks; decomposing the update task of the intermediate accumulation block corresponding to the attention output block in the current iteration round into N3 parallel update subtasks; asynchronously submitting all subtasks to different task slots of N4 synchronous channels; polling the task slots in each synchronous channel to execute the subtasks stored in the task slots, and controlling the execution timing between the subtasks with a dependency relationship in each synchronous channel by configuring the number threshold of incomplete subtasks in the waiting instruction. The application can improve the masking ability and parallel computing ability of hardware, thereby improving the efficiency of attention calculation.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence, and in particular to a task scheduling method and device for attention calculation, a computer readable storage medium, an electronic device and a computer program product. BACKGROUND

[0002] With the rapid development of artificial intelligence technology, the attention mechanism has become a core operator in many fields such as natural language processing and computer vision. It realizes dynamic weighted aggregation of input data by calculating the interaction between query, key and value. With the continuous expansion of model size and the continuous growth of input sequence length, the complexity of attention calculation is also increasing significantly, resulting in a sharp increase in computing time and high hardware resource occupation. Therefore, how to deeply tap the parallel potential of attention calculation and optimize the calculation process to improve the calculation efficiency is a problem to be solved. SUMMARY

[0003] The purpose of the embodiments of the present application is to provide a task scheduling method and device for attention calculation, a computer readable storage medium, an electronic device and a computer program product. By finely decomposing the loading task of query block, the calculation task of attention score block and the update task of intermediate accumulation block in the attention calculation process, and then using multiple task slots of a synchronization channel to realize asynchronous submission and polling execution, and by dynamically configuring the threshold parameter of the waiting instruction to manage the dependency relationship between subtasks, the masking ability and parallel computing ability of the hardware can be improved, and the efficiency of attention calculation can be improved.

[0004] The first embodiment of the present application provides a task scheduling method for attention calculation, applied to a consumer thread group, wherein the consumer thread group is configured with one or more synchronization channels, and each synchronization channel is configured with multiple task slots; the method comprises:

[0005] decomposing the loading task of the query block from shared memory to register into N1 loading subtasks, and decomposing the calculation task of the jth attention score block corresponding to the query block into N2 calculation subtasks;

[0006] decomposing the update task of the intermediate accumulation block corresponding to the attention output block in the current iteration round into N3 parallel update subtasks;

[0007] asynchronously submitting the N1 loading subtasks, the N2 calculation subtasks and the N3 update subtasks to different task slots of N4 synchronization channels; wherein the subtasks with a dependency relationship are submitted to the same synchronization channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1.

[0008] Polling the task slots in each of the synchronization channels to execute the sub-tasks stored in the task slots, and controlling the execution timing between the sub-tasks with the dependency relationship in each of the synchronization channels by configuring a threshold of the number of incomplete sub-tasks in the waiting instruction.

[0009] Optionally, the asynchronous submission of the N1 loading sub-tasks, the N2 computing sub-tasks and the N3 updating sub-tasks into different task slots of the N4 synchronization channels comprises:

[0010] In the case of j = 1, the N1 loading sub-tasks and the N2 computing sub-tasks are asynchronously submitted into different task slots of the first synchronization channel group according to the dependency relationship;

[0011] In the case of j > 1, the N2 computing sub-tasks are asynchronously submitted into different task slots of the first synchronization channel group;

[0012] The N3 updating sub-tasks are asynchronously submitted into different task slots of the second synchronization channel group; wherein the first synchronization channel group and the second synchronization channel group each consist of at least one of the synchronization channels.

[0013] Optionally, each sub-task is composed of at least one execution instruction.

[0014] Optionally, the i-th loading sub-task is used to load an i-th query sub-block in the query block from the shared memory into the register; wherein 1≤i≤N1; the size of the query sub-block satisfies the first calculation size of the matrix multiplication and addition instruction.

[0015] Optionally, the calculation task of the j-th attention score block corresponding to the query block is decomposed into N2 computing sub-tasks, comprising:

[0016] Based on the second calculation size of the matrix multiplication and addition instruction, the j-th key transpose block is divided into N1×M key transpose sub-blocks; wherein M≥1; the size of the key transpose sub-block supports the expansion of the splicing of multiple second calculation sizes in the horizontal direction;

[0017] In the case of j = 1, the index of the key transpose sub-block is traversed in row-major order, and the matrix multiplication and addition calculation of the n1-th query sub-block and the key transpose sub-block in the n1-th row and the m-th column is taken as a computing sub-task to form N2 computing sub-tasks; wherein 1≤n1≤N1, 1≤m≤M, N2=N1×M.

[0018] Optionally, the calculation task of the j-th attention score block corresponding to the query block is decomposed into N2 computing sub-tasks, further comprising:

[0019] In the case of j>1, the indexes of the key transpose sub-blocks are traversed in column priority order, and the combination of N1 query sub-blocks and all key transpose sub-blocks in the mth column are calculated as one calculation sub-task to form N2 calculation sub-tasks; wherein N2=M.

[0020] Optionally, the first attention score sub-block is calculated by the following steps:

[0021] When the first query sub-block is loaded into the register monitored by the waiting instruction, the calculation sub-tasks corresponding to all key transpose sub-blocks in the first row are started in turn;

[0022] When the nth2 query sub-block is loaded into the register monitored by the waiting instruction, and the calculation sub-task corresponding to the key transpose sub-block in the mth column of the nth2-1 row is executed, the calculation sub-task corresponding to the key transpose sub-block in the mth column of the nth2 row is started; wherein 2≤n2≤N1;

[0023] When the calculation sub-tasks corresponding to all key transpose sub-blocks in the mth column are executed, the mth score sub-block of the first attention score sub-block is obtained.

[0024] Optionally, the method further comprises:

[0025] When the mth score sub-block of the jth attention score sub-block is calculated, the mth score sub-block is locally statistically calculated by the waiting instruction.

[0026] Optionally, in the current iteration round, the kth update sub-task is used to multiply the exponential score sub-block and the corresponding sub-block in the value sub-block, and the obtained sub-block product result is accumulated to the corresponding region of the scale-aligned intermediate accumulation sub-block to obtain the updated kth intermediate accumulation sub-block; wherein k≥1.

[0027] Optionally, the method further comprises:

[0028] In the next iteration round, when the kth update sub-task is executed, the kth intermediate accumulation sub-block is scale-aligned by the waiting instruction.

[0029] The second aspect embodiment of the application provides a task scheduling device for attention calculation, which is applied to a consumer thread group, the consumer thread group is configured with one or more synchronization channels, each synchronization channel is configured with a plurality of task slots; the device comprises:

[0030] A first decomposition module is configured to decompose the loading task of the query sub-block from the shared memory to the register into N1 loading sub-tasks, and decompose the calculation task of the jth attention score sub-block corresponding to the query sub-block into N2 calculation sub-tasks;

[0031] A second decomposition module is configured to decompose an update task of the intermediate accumulation block corresponding to the attention output block in the current iteration round into N3 parallel update subtasks.

[0032] A task submission module is configured to asynchronously submit the N1 loading subtasks, the N2 computing subtasks and the N3 update subtasks into different task slots of the N4 synchronous channels; wherein, the subtasks with the dependency relationship are submitted into the same synchronous channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1.

[0033] A task execution module is configured to perform a polling operation on the task slots in each synchronous channel to execute the subtasks stored in the task slots, and control the execution timing sequence between the subtasks with the dependency relationship in each synchronous channel by configuring a number threshold of the unfinished subtasks in the waiting instruction.

[0034] The third aspect embodiment of the present application provides a computer readable storage medium, which comprises a stored computer program; wherein, the computer program controls the device where the computer readable storage medium is located to execute the task scheduling method for attention calculation of any one of the first aspect embodiments when running.

[0035] The fourth aspect embodiment of the present application provides a computer program product comprising computer instructions, which realize the task scheduling method for attention calculation of any one of the first aspect embodiments when executed by a processor.

[0036] The fifth aspect embodiment of the present application provides an electronic device comprising a processor, a memory and a computer program stored in the memory and configured to be executed by the processor, wherein the processor realizes the task scheduling method for attention calculation of any one of the first aspect embodiments when executing the computer program.

[0037] Compared with the prior art, the embodiments of the present application provide a task scheduling method, device, computer readable storage medium, electronic device and computer program product for attention calculation, wherein the three stages of query block loading, attention score block calculation and intermediate accumulation block update are finely decomposed in the attention calculation process, and the dependency relationship is managed based on the waiting instruction through the asynchronous submission of multiple task slots, so as to enhance the masking ability and parallel computing ability of the hardware, make the pipeline arrangement of the attention calculation more compact, and finally significantly improve the computing efficiency of the attention. In addition, the computing task of the attention score block and the update task of the intermediate accumulation block are the tasks with the highest computing complexity and the most intensive hardware resource occupation in the attention calculation process, and the embodiments of the present application finely schedule and optimize them, so as to obtain significant performance improvement through small scheduling overhead. BRIEF DESCRIPTION OF DRAWINGS

[0038] Figure 1 is a flow diagram of one embodiment of a task scheduling method for attention computation provided by the present application;

[0039] Figure 2 is a diagram of one embodiment of a correlation size of each block in attention computation provided by the present application;

[0040] Figure 3 is a flow diagram of one embodiment of a block for computing attention score provided by the present application;

[0041] Figure 4 is a diagram of another embodiment of a block for computing attention score provided by the present application.

[0042] Figure 5 is a diagram of yet another embodiment of a block for computing attention score provided by the present application;

[0043] Figure 6 is a diagram of one embodiment of an intermediate accumulated block update provided by the present application;

[0044] Figure 7 is a diagram of one embodiment of an intermediate accumulated block size alignment provided by the present application;

[0045] Figure 8 is a structural diagram of one embodiment of a task scheduling apparatus for attention computation provided by the present application;

[0046] Figure 9 is a structural diagram of one embodiment of an electronic device provided by the present application. DETAILED DESCRIPTION

[0047] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only a part of the embodiments of the present application, rather than all the embodiments of the present application. Based on the embodiments in the present application, all other embodiments obtained by a person of ordinary skill in the art without creative work fall within the scope of protection of the present application.

[0048] The artificial intelligence processor involved in the present application can be any one of a CPU (Central Processing Unit), a GPU (Graphics Processing Unit), a TPU (Tensor Processing Unit), an NPU (Neural network Processing Unit), a DPU (Deeplearning Processing Unit), an APU (Accelerated Processing Unit), and a GPGPU (General-Purpose computing on Graphics Processing Unit), as determined when the embodiments of the present application are applied to specific products or technologies.

[0049] In the embodiments of the present application, the calculation size of the matrix multiply accumulate instruction mma (matrix multiply accumulate) is a×b×K B ; wherein K B is the number of bytes, ranging from 32B to 256B; when the data type of the matrix is FP16, a×b×K B =64×16×16c (c=1, 2,…, 8), that is, the first calculation size of the A matrix is a×K B =64×16c, and the second calculation size of the B matrix is K B ×b=16c×16. Based on this, in the “QK T ” multiplication, the query sub-block in the query block belongs to the A matrix, and the maximum size thereof is 64×128; the key transpose sub-block in the key transpose block belongs to the B matrix, and the maximum size of the key transpose sub-block is 128×16; in the “PV” multiplication, the sub-block in the exponentized fraction block belongs to the A matrix, and the sub-block in the value block belongs to the B matrix; wherein i≥1; j≥1. In the mma calculation of the embodiments of the present application, the A matrix is in the thread level register (TLR), and the B matrix is in the shared memory (GSM).

[0050] Furthermore, this invention employs the Flash Attention method for attention calculation, decomposing a large matrix into many smaller blocks for computation through a "blocking" approach. In this embodiment, after the producer thread group loads the query block (Q), key block (K), and value block (V) from Global Memory (GLM) to GSM in batches, the consumer thread group performs attention calculation through the following main operations / processes:

[0051] (1) In the outer loop of Q, the query is divided into blocks. Loaded into the TLR for iterative computation within the KV inner loop to generate the i-th attention output block. );

[0052] (2) In the iterative calculation of the KV inner loop, the query blocks in the TLR are processed by the specified matrix multiplication and addition instruction (mma). Key block in GSM Conduct "QK" T Multiplication (i.e.) ), to obtain the corresponding attention score blocks ;

[0053] (3) Obtain attention score blocks The current maximum value tensor And combined with the historical cumulative maximum tensor Calculate the current cumulative maximum value tensor ;

[0054] (4) Obtain the update correction factor Modifier= Used for scale alignment;

[0055] (5) Segmenting attention scores Perform exponential operations (i.e.) ), to obtain indexed fractional blocks Where sf is the scaling factor;

[0056] (6) Obtain the indexed score blocks The current exponent and tensor And combining historical accumulation and tensors Calculate the current accumulated sum tensor ;

[0057] (7) Segmenting attention output Corresponding intermediate accumulation blocks Perform scale alignment operation, i.e. = ;

[0058] (8) Blocking of indexed fractions and value blocks performing “PV” multiplication (i.e. ), and accumulating the product result to the scale-aligned intermediate accumulation block (i.e. ), to obtain the iteratively updated intermediate accumulation block , i.e. the scale-aligned intermediate accumulation block = ;

[0059] (9) after the end of the inner loop iteration calculation, according to the current accumulation and tensor obtained by the last iteration, performing a rescaling operation on the intermediate accumulation block to obtain the attention output block .

[0060] The above (2) to (8) are located in the iteration calculation of the KV inner loop, and (1) and (9) are located in the Q outer loop.

[0061] The embodiment of the application performs fine-grained scheduling optimization on the above attention calculation process, and proposes a task scheduling method, device, computer readable storage medium, electronic equipment and computer program product for attention calculation.

[0062] Referring to Figure 1 is a flowchart of an embodiment of a task scheduling method for attention calculation provided by the application.

[0063] The first aspect embodiment of the application provides a task scheduling method for attention calculation, applied to a consumer thread group, the consumer thread group being configured with one or more synchronization channels, each of the synchronization channels being configured with a plurality of task slots; the method comprises steps S1 to S4, and specifically as follows:

[0064] Step S1: decomposing a loading task of a query block from shared memory to a register into N1 loading sub-tasks, and decomposing a calculation task of a jth attention score block corresponding to the query block into N2 calculation sub-tasks;

[0065] Step S2: decomposing an updating task of an intermediate accumulation block corresponding to an attention output block in a current iteration round into N3 parallel updating sub-tasks;

[0066] Step S3: asynchronously submitting the N1 loading sub-tasks, the N2 calculation sub-tasks and the N3 updating sub-tasks to different task slots of N4 synchronization channels; wherein sub-tasks with a dependency relationship are submitted to the same synchronization channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1;

[0067] Step S4: Poll the task slots in each synchronization channel to execute the subtasks stored in the task slots, and control the execution sequence of dependent subtasks in each synchronization channel by configuring a threshold for the number of incomplete subtasks in the waiting instructions.

[0068] For example, for a certain artificial intelligence chip, the consumer thread group is configured with 3 synchronization channels (SyncChannel), namely sc0, sc1 and sc2; each synchronization channel is configured with 8 task slots (slot0~slot7).

[0069] The embodiments of the present invention, through the aforementioned hardware conditions, perform query segmentation. The loading phase (corresponding to operation (1) above), attention score block The calculation phase (corresponding to operation (2) above) involves "QK T Multiplication), intermediate accumulation blocks The iterative update phase (corresponding to the operation (8) above, involving “PV” multiplication) is optimized and executed with fine granularity.

[0070] During the query chunk loading phase, the query chunks are... The loading task from GSM to TLR is broken down into N1 loading subtasks. For example, a query block of size 64×256 is divided into subtasks. (Data type is FP16) Logically divided into 2 query sub-blocks (size is 64×128), then for the query sub-blocks The loading task is then broken down into N1=2 loading subtasks, each responsible for loading the corresponding query subblock into the register, so as to further implement "QK" through the mma instruction. T "multiplication.

[0071] Attention score segmentation During the computation phase, key transpose-based block partitioning is performed. Granularity of division (e.g., dividing) The logic is divided into N1×M key transpose sub-blocks, and the attention score is divided into blocks. computational tasks ( The matrix is ​​decomposed into one or more query sub-blocks and corresponding key transpose sub-blocks.

[0072] For example, transpose a 256×32 key into blocks. The logic is divided into N1×M=2×2 key transpose sub-blocks (size 128×16), corresponding to attention score blocks. The two calculation tasks are decomposed in the following two ways: ① decomposed into four calculation sub-tasks, each of which is a matrix multiplication and addition calculation of one query sub-block and one key transpose sub-block (here, one calculation sub-task corresponds to one mma instruction), which can fully mask the gap between data loading and calculation when processing the first key transpose sub-block; ② decomposed into two calculation sub-tasks, each of which is a combined calculation of two query sub-blocks and all key transpose sub-blocks in a single column (here, one calculation sub-task corresponds to two mma instructions), which can directly calculate the score sub-block corresponding to the column in the attention score block, and can reduce the overhead of task slots. The above two decomposition methods of calculation tasks are suitable for all key transpose sub-blocks, but for the case of j = 1, the first decomposition method is preferred.

[0073] In the iteration update phase of the intermediate accumulation block , the update task is the intermediate accumulation block after scale alignment , the intermediate accumulation block is logically divided into N3 intermediate accumulation sub-blocks, so that the update task of the intermediate accumulation block is decomposed into N3 parallel update sub-tasks, each of which is used to update the corresponding intermediate accumulation sub-block.

[0074] In step S4, N1 loading sub-tasks, N2 calculation sub-tasks and N3 update sub-tasks are asynchronously submitted to different task slots of N4 synchronous channels, so that each sub-task is stored in a task slot slot, and sub-tasks with a dependency relationship are submitted to the same synchronous channel to avoid calculation delay caused by cross-channel synchronization. In an embodiment of the present application, there is a dependency relationship between the loading sub-task and the calculation sub-task corresponding to the first attention score block (calculation needs to wait for the corresponding query sub-block to be loaded); there is no dependency relationship between the loading sub-tasks; similarly, there is also no dependency relationship between the update sub-tasks.

[0075] In step S5, a round-robin operation is performed on the task slots in each synchronous channel, and the sub-tasks in each task slot are started in turn (such as the cycle execution from slot0 to slot7 in sc0), to ensure that the sub-tasks of all task slots in the synchronous channel are scheduled.

[0076] ​The execution timing between sub-tasks with dependencies is arranged through a wait instruction wait_group(scX, n_pending) (i.e., a waitle instruction); wherein scX is a synchronization channel numbered X; n_pending is a threshold value of the number of incomplete sub-tasks, ranging from 0 to the total number of slots of the synchronization channel-1, for example, sc0 has 8 slots, and the range of n_pending is 0~7.

[0077] The waitle instruction wait_group(scX, n_pending) is used for the current thread to wait until the number of pending tasks in the synchronization channel scX equals the preset n_pending, and then unblock to continue executing subsequent instructions. For example, in the case of j=1, by configuring the waitle instruction, it is ensured that the calculation sub-task needs to wait for the completion of the slot where the previous loading sub-task is located before execution, avoiding calculation errors caused by unready data.

[0078] As can be seen from the above, in the attention calculation process, the three stages of query block loading, attention score block calculation, and intermediate accumulation block updating are finely divided, and the dependency relationship is managed based on the wait instruction, which can fully utilize the hardware characteristics of the multi-task slot in the synchronization channel, enhance the masking ability and parallel computing ability of the hardware, make the pipeline arrangement of attention calculation more compact, and finally significantly improve the calculation efficiency of attention. In addition, the calculation task of the attention score block (involving "QK T " multiplication) and the update task of the intermediate accumulation block (involving "PV" multiplication) are the most computationally complex and hardware resource-intensive tasks in the attention calculation process. The fine-grained scheduling optimization of the embodiment can significantly improve the performance with relatively small scheduling overhead.

[0079] In an optional embodiment, the asynchronous submission of the N1 loading sub-tasks, the N2 calculation sub-tasks, and the N3 update sub-tasks to different task slots of N4 synchronization channels includes:

[0080] In the case of j=1, the N1 loading sub-tasks and the N2 calculation sub-tasks are asynchronously submitted to different task slots of the first synchronization channel group according to the dependency relationship;

[0081] In the case of j>1, the N2 calculation sub-tasks are asynchronously submitted to different task slots of the first synchronization channel group;

[0082] The N3 update sub-tasks are asynchronously submitted to different task slots of the second synchronization channel group; wherein the first synchronization channel group and the second synchronization channel group each consist of at least one synchronization channel.

[0083] It should be noted that the first synchronization channel group is used to carry the loading subtask (only when j=1) and calculation subtask directly related to attention score calculation; the second synchronization channel group is used to carry the update subtask related to intermediate accumulation block iterative update; the number of channels contained in the two synchronization channel groups can be dynamically configured according to the task size.

[0084] In the processing stage of the first key transpose block (i.e., j=1), since the computation subtask has a strong dependency on the loading subtask (it needs to wait for the query subblock to be loaded into the register before the corresponding mma computation can be started), the loading subtask needs to occupy the preceding task slot of the synchronization channel to ensure that the computation subtask can perceive its execution status through the waittle instruction.

[0085] In the processing stage of the second and subsequent key transpose blocks (i.e., when j > 1), all query sub-blocks in the query block have been loaded into the registers, eliminating the need to repeatedly execute the loading subtasks. Therefore, only N2 computation subtasks need to be asynchronously submitted to different task slots in the first synchronization channel group. In other words, when j > 1, there is no need to reserve slots for loading subtasks, and computation subtasks can be submitted continuously starting from the initial slot in the first synchronization channel group.

[0086] Regardless of whether j=1 or j>1, all N3 update subtasks must be submitted asynchronously to different task slots in the second synchronization channel group. When the number of update subtasks is large (e.g., N3=16), and the 8 slots of a single synchronization channel are insufficient, the second synchronization channel group can be expanded into 2 synchronization channels (e.g., sc1 and sc2), with each synchronization channel carrying 8 update subtasks.

[0087] In an optional embodiment, the i-th loading subtask is used to load the i-th query sub-block in the query block from shared memory into a register; where 1≤i≤N1; the size of the query sub-block satisfies the first computation size of the matrix multiply-add instruction.

[0088] For example, such as Figure 2 The diagram shown is a schematic representation of an embodiment of the block association dimensions in attention calculation provided by the present invention. Figure 2 In this context, Q chunk, K chunk, and V chunk represent the query chunk, key chunk, and value chunk, respectively; S represents the attention score chunk calculated from the query chunk and key chunk; P represents the exponential score chunk of S after softmax; and O chunk is the attention output chunk corresponding to the query chunk.

[0089] As mentioned earlier, when the data type is FP16, the first computation size of the matrix multiply-add instruction (mma) is defined as 64×16c (c=1, 2,…, 8). Figure 2The size of the Q chunk is 64×256, therefore the Q chunk can be logically divided into N1=2 query sub-blocks, each with a size of 64×128 (where c=8); for example Figure 3 The diagram shown is a schematic representation of an embodiment of the attention score block calculation provided by the present invention. Figure 3 In this context, "Half Q" represents a query subblock (64×128 in size), and "Half K" represents a key transpose subblock (128×16 in size).

[0090] In an optional embodiment, the task of decomposing the calculation task of the j-th attention score block corresponding to the query block into N2 calculation subtasks includes:

[0091] The second computational size based on matrix multiplication and addition instructions divides the j-th key transpose block into N1×M key transpose sub-blocks; where M≥1; the size of the key transpose sub-blocks supports the horizontal splicing and expansion of multiple second computational sizes;

[0092] When j=1, the indexes of the key transpose subblocks are traversed in row-major order. The matrix multiplication and addition of the n1th query subblock and the key transpose subblock in the n1th row and mth column is taken as a computational subtask to form N2 computational subtasks; where 1≤n1≤N1, 1≤m≤M, and N2=N1×M.

[0093] like Figure 3 As shown, when the data type is FP16, the second computation size of the mma instruction is defined as 16c×16. Combined with the query sub-block size of 64×128, and the number of task slots in a single synchronization channel meeting the computation scheduling requirements of the current attention score block, the key transpose block (size 256×32) is logically divided into N1×M=2×2 key transpose sub-blocks. Each key transpose sub-block has a size of 128×16 (the second computation size when c=8), and no horizontal expansion of the second computation size is required.

[0094] In the processing stage of the first key-transposed block (j=1), the indices of the key-transposed sub-blocks are traversed in row-major order. The matrix multiplication and addition calculation of the n1-th query sub-block and the key-transposed sub-block in the n1-th row and m-th column is treated as an independent computational subtask, forming N2=N1×M computational subtasks. At this time, each computational subtask corresponds to one mma instruction, which can ensure that, in the case of j=1, the gap between data loading and computation is fully masked by fine-grained task decomposition.

[0095] by Figure 3As an example of the scenario where N1=2 and M=2, the first query sub-block (front half Q) is multiplied and added with the matrix of the transpose key sub-block in the first row and the first column as the first calculation sub-task, and multiplied and added with the matrix of the transpose key sub-block in the first row and the second column as the second calculation sub-task; the second query sub-block (back half Q) is multiplied and added with the matrix of the transpose key sub-block in the second row and the first column as the third calculation sub-task, and multiplied and added with the matrix of the transpose key sub-block in the second row and the second column as the fourth calculation sub-task.

[0096] The consumer thread group implements the first attention score block by means of instruction example 1. Figure 3 The first attention score block is calculated; of course, the decomposition manner of the calculation task in instruction example 1 can also be used for the calculation of the second and subsequent attention score blocks. The specific instruction example 1 is as follows:

[0097] Loop{ / KV inner loop /

[0098] ldmatrix front halfQ (2 instr) / / the first loading sub-task, which loads the first query sub-block (front half Q) from the GSM into the TLR, and the loading sub-task includes 2 ldmatrix instructions;

[0099] commit sc0 / / the first loading sub-task is committed to slot0 of sc0;

[0100] ldmatrix back halfQ (2 instr) / / the second loading sub-task, which loads the second query sub-block (back half Q) from the GSM into the TLR, and the loading sub-task also includes 2 ldmatrix instructions;

[0101] commit sc0 / / the second loading sub-task is committed to slot1 of sc0;

[0102] waitle sc0, 1 / / when the number of sub-tasks not completed in sc0 is equal to 1, the blocking is released to continue execution, at this time, the first query sub-block has been loaded into the TLR;

[0103] regS1 = front mma (halfQ, halfK) / / the first calculation sub-task, which performs mma calculation on the first query sub-block and the transpose key sub-block in the first row and the first column (half K), and stores the result in the register regS1;

[0104] commit sc0 / / commit the 1st compute subtask to slot2 of sc0;

[0105] regS2 = front mma(halfQ, halfK) / / the 2nd compute subtask, perform mma computation between the 1st query subblock and the key transpose subblock in the 1st row and 2nd column, and store the result into regS2;

[0106] commit sc0 / / commit the 2nd compute subtask to slot3 of sc0;

[0107] waitl sc0, 1 / / wait until the number of unfinished subtasks in sc0 is equal to 1, then execute the 3rd compute subtask, to ensure that the result data provided by the 1st compute subtask is used for subsequent accumulation operation;

[0108] regS1 = mma(halfQ, halfK, regS1) / / the 3rd compute subtask, perform mma computation between the 2nd query subblock and the key transpose subblock in the 2nd row and 1st column, and accumulate the result into regS1 to obtain the 1st score subblock of the 1st attention score tile;

[0109] commit sc0 / / commit the 3rd compute subtask to slot4 of sc0;

[0110] waitl sc0, 1 / / wait until the number of unfinished subtasks in sc0 is equal to 1, then execute the 4th compute subtask, to ensure that the result data provided by the 2nd compute subtask is used for subsequent accumulation operation;

[0111] regS2 = mma(halfQ, halfK, regS2) / / the 4th compute subtask, perform mma computation between the 2nd query subblock and the key transpose subblock in the 2nd row and 2nd column, and accumulate the result into regS2 to obtain the 2nd score subblock of the 1st attention score tile;

[0112] commit sc0 / / commit the 4th compute subtask to slot5 of sc0;

[0113] waitl sc0, 1 / / wait until the 1st score subblock is computed, and then perform local statistics computation on it, such as row_max instruction to obtain the maximum value in row direction;

[0114] row_max(regS1)

[0115] waitle sc0, 0 / / Wait until the number of unfinished subtasks in sc0 equals 0, then the second fractional sub-block is completed and local statistical calculations can be performed on it;

[0116] row_max(regS2)

[0117] …

[0118] }

[0119] Based on the above embodiments, the first attention score block is calculated through the following steps:

[0120] When the waiting instruction monitor detects that the first query sub-block has been loaded into the register, the computation subtasks corresponding to all key transpose sub-blocks in the first row are started sequentially.

[0121] When the waiting instruction monitors that the n2nd query subblock is loaded into the register and the computation subtask corresponding to the key transpose subblock in row n2-1 and column m is completed, the computation subtask corresponding to the key transpose subblock in row n2 and column m is started; where 2≤n2≤N1;

[0122] When all the computational subtasks corresponding to the key transpose subblocks in column m are completed, the mth fraction subblock of the first attention fraction block is obtained.

[0123] Specifically, with Figure 3 For example, in the corresponding instruction example 1, when the waiter instruction detects that the first query subblock (fronthalfQ) is loaded from GSM to TLR, it immediately starts the computation subtasks corresponding to all key transpose subblocks in the first row (i.e., the first computation subtask and the second computation subtask).

[0124] When the waittle instruction detects that the second query sub-block (back halfQ) has been loaded and the first computation subtask has been completed, it immediately starts the third computation subtask. Similarly, when the waittle instruction detects that the second query sub-block (back halfQ) has been loaded and the second computation subtask has been completed, it immediately starts the fourth computation subtask.

[0125] Figure 3 After the computational subtasks corresponding to the key transpose subblocks in the first row and first column and the second row and first column (i.e., the first and third computational subtasks) are completed, the accumulated result stored in register regS1 is the first fractional subblock. Similarly, after the computational subtasks corresponding to the key transpose subblocks in the first row and second column and the second row and second column (i.e., the second and fourth computational subtasks) are completed, the accumulated result stored in register regS2 is the second fractional subblock.

[0126] In an alternative embodiment, in the case of j>1, the indexes of the key transpose sub-blocks are traversed in column priority order, and the combination of N1 query sub-blocks with all key transpose sub-blocks of the m-th column is calculated as one calculation sub-task to form N2 calculation sub-tasks; where N2=M.

[0127] As shown in FIG. 2, is a schematic diagram of another embodiment of the calculation of the attention score sub-blocks provided by the present application. In the processing stage of the 2nd and subsequent key transpose sub-blocks (j>1), the query sub-blocks have been stored in the TLR through the loading sub-task of the j=1 stage and do not need to be loaded again, so there is no need for masking of the loading and calculation, and the mma instructions can be integrated at the column level to obtain the corresponding calculation sub-tasks. Figure 4

[0128] Taking the scenario of N1=2 and M=2 in FIG. 3 as an example, the matrix multiply-add calculation of the 1st query sub-block (the former "Half Q") with the key transpose sub-block of the 1st row and 1st column, and the matrix multiply-add calculation of the 2nd query sub-block (the latter "Half Q") with the key transpose sub-block of the 2nd row and 1st column, are jointly taken as the 1st calculation sub-task; and the matrix multiply-add calculation of the 1st query sub-block with the key transpose sub-block of the 1st row and 2nd column, and the matrix multiply-add calculation of the 2nd query sub-block with the key transpose sub-block of the 2nd row and 2nd column, are jointly taken as the 2nd calculation sub-task. At this time, each calculation sub-task corresponds to 2 mma instructions. Figure 4 The consumer thread group implements the calculation of the 2nd and subsequent attention score sub-blocks in FIG. 4 through instruction example 2. The specific instruction example 2 is as follows:

[0129] Figure 4 Loop{ / KV inner loop

[0130] / waitl sc0, 1 / / When the number of unfinished sub-tasks in sc0 is equal to 1, the blocking is released;

[0131] regS1 = front mma(halfQ, halfK) / / One mma instruction of the 1st calculation sub-task, which performs the mma calculation of the 1st query sub-block with the key transpose sub-block of the 1st row and 1st column, and stores the result in regS1;

[0132] regS1 = mma(halfQ, halfK, regS1) / / The other mma instruction of the 1st calculation sub-task, which performs the mma calculation of the 2nd query sub-block with the key transpose sub-block of the 2nd row and 1st column, and accumulates the result in regS1 to obtain the 1st score sub-block of the attention score sub-blocks;

[0133]

[0134] ​​​commit sc0 / / commit the 1st compute subtask to slot0 of sc0;

[0135] regS2 = mma(halfQ, halfK) / / the 2nd compute subtask, a mma instruction, performs mma computation between the 1st query subblock and the key transpose subblock at row 2, column 1, and stores the result to regS2;

[0136] regS2 = mma(halfQ, halfK, regS2) / / the 2nd compute subtask, another mma instruction, performs mma computation between the 2nd query subblock and the key transpose subblock at row 2, column 2, and accumulates the result to regS2 to obtain the 2nd score subblock of the attention score tile;

[0137] commit sc0 / / commit the 2nd compute subtask to slot1 of sc0;

[0138] wait sc0, 1 / / wait until the number of unfinished subtasks in sc0 is equal to 1, the 1st score subblock computation is completed, and local statistical computation, such as row-wise maximum (row_max instruction), can be performed thereon;

[0139] row_max(regS1)

[0140] wait sc0, 0 / / wait until the number of unfinished subtasks in sc0 is equal to 0, the 2nd score subblock computation is completed, and local statistical computation can be performed thereon;

[0141] row_max(regS2)

[0142] …

[0143] }

[0144] As can be seen from instruction example 1 and instruction example 2, each subtask is composed of at least one execution instruction.

[0145] In an optional embodiment, the method further comprises:

[0146] When the jth attention score tile is monitored by the waiting instruction to be completed, local statistical computation is performed on the mth score subblock of the jth attention score tile.

[0147] As shown in Instruction Example 1 and Instruction Example 2, once it is confirmed that the mth fractional sub-block calculation is completed through the waitle instruction, the register storing the mth fractional sub-block (such as regS1 corresponding to the 1st fractional sub-block, regS2 corresponding to the 2nd fractional sub-block) is directly accessed, and the row_max(regSm) instruction is called to obtain the local maximum value by calculating the maximum value in the row direction of the fractional sub-block in the register regSm. After the local maximum values of all fractional sub-blocks are calculated, the global maximum value of the current attention score block is obtained.

[0148] It is worth noting that when the number of task slots in a single synchronization channel cannot meet the calculation scheduling requirements of the current attention score block, the size of the key transpose sub-block in the embodiment of the present application supports the expansion of multiple second calculation sizes in the horizontal direction, thereby reducing the consumption of task slots.

[0149] As shown in Instruction Example 1 and Instruction Example 2, once it is confirmed that the mth fractional sub-block calculation is completed through the waitle instruction, the register storing the mth fractional sub-block (such as regS1 corresponding to the 1st fractional sub-block, regS2 corresponding to the 2nd fractional sub-block) is directly accessed, and the row_max(regSm) instruction is called to obtain the local maximum value by calculating the maximum value in the row direction of the fractional sub-block in the register regSm. After the local maximum values of all fractional sub-blocks are calculated, the global maximum value of the current attention score block is obtained. Figure 5 As shown in Instruction Example 1 and Instruction Example 2, once it is confirmed that the mth fractional sub-block calculation is completed through the waitle instruction, the register storing the mth fractional sub-block (such as regS1 corresponding to the 1st fractional sub-block, regS2 corresponding to the 2nd fractional sub-block) is directly accessed, and the row_max(regSm) instruction is called to obtain the local maximum value by calculating the maximum value in the row direction of the fractional sub-block in the register regSm. After the local maximum values of all fractional sub-blocks are calculated, the global maximum value of the current attention score block is obtained.

[0150] According to Figure 3 As shown in Instruction Example 1 and Instruction Example 2, once it is confirmed that the mth fractional sub-block calculation is completed through the waitle instruction, the register storing the mth fractional sub-block (such as regS1 corresponding to the 1st fractional sub-block, regS2 corresponding to the 2nd fractional sub-block) is directly accessed, and the row_max(regSm) instruction is called to obtain the local maximum value by calculating the maximum value in the row direction of the fractional sub-block in the register regSm. After the local maximum values of all fractional sub-blocks are calculated, the global maximum value of the current attention score block is obtained. Figure 5 In the configuration of N1×M=2×4 key transpose sub-blocks in

[0151] Of course, in the calculation task decomposition stage of the 1st attention score block, the key transpose block in Figure 5 can be logically divided into N1×M=2×2 key transpose sub-blocks, so that the task slot occupancy (6) does not exceed the total slot number; the size of each key transpose sub-block is 96×32, which is obtained by horizontally splicing 2 second calculation sizes, so two mma instructions are needed to complete the matrix multiplication and addition calculation of one query sub-block and one key transpose sub-block. In other words, when calculating Figure 5 the corresponding 1st attention score block, 4 calculation tasks are needed, and each calculation task includes 2 mma instructions.

[0152] In an alternative embodiment, in the current iteration round, the kth update subtask is configured to multiply the exponentialized fractional tile with the corresponding subtile in the value tile, and accumulate the resulting subtile product result to the corresponding region of the scale-aligned intermediate accumulation tile to obtain an updated kth intermediate accumulation subtile; where k≥1.

[0153] Further, the method further comprises:

[0154] In the next iteration round, when the kth update subtask is monitored to be executed by the waiting instruction, the kth intermediate accumulation subtile is scale-aligned.

[0155] As Figure 6 shown in FIG. 1, is a schematic diagram of an embodiment of the intermediate accumulation tile update provided by the present application. The size of the exponentialized fractional tile (“P”) is 64x32, which directly matches the first calculation size of the mma instruction and does not need to be divided; the size of the value tile is 32x256, which needs to be logically divided into 16 sub-tiles according to the second calculation size of the mma, and each sub-tile (size 32x16) corresponds to an update subtask, so the update task of the intermediate accumulation tile is decomposed into 16 parallel executable update subtasks.

[0156] Since the number of slots in a single synchronization channel is 8, in the current iteration round (denoted as the jth round), the update operation of the intermediate accumulation tile (corresponding to the operation (8) in the foregoing, involving “PV” multiplication) requires 2 synchronization channels (sc1 and sc2); the 8 slots of sc1 are used to carry 8 update subtasks corresponding to the first half HalfV (containing 8 sub-tiles) of the value tile; the 8 slots of sc2 are used to carry 8 update subtasks corresponding to the second half HalfV (containing 8 sub-tiles) of the value tile; each slot only stores one update subtask. After all the update subtasks of sc1 are executed, the first half of the updated intermediate accumulation tile can be obtained in the register regO1; similarly, after all the update subtasks of sc2 are executed, the second half of the updated intermediate accumulation tile can be obtained in the register regO2.

[0157] As Figure 7 shown in FIG. 2, is a schematic diagram of an embodiment of the scale alignment of the intermediate accumulation tile provided by the present application. Before the iteration update of the intermediate accumulation tile in the next iteration round (denoted as the (j+1)th round), the scale alignment of the intermediate accumulation tile in the jth round is needed. In Figure 7 the case of sc1, each yellow block represents an update subtask (occupying 1 slot) of sc1 in the jth round; each blue block represents the scale alignment operation (corresponding to the operation (7) in the foregoing) of the kth intermediate accumulation subtile obtained in the jth round in the (j+1)th round.

[0158] Specifically, in the j+1th round, when the kth update subtask of sc1 is monitored to be completed by the waitle instruction, it indicates that the kth intermediate accumulation subblock update in the jth round is completed, and the scale alignment is immediately performed on it by the update modifier in the j+1th round to further execute the kth update subtask in the j+1th round.

[0159] In summary, the embodiment of the application forms the time overlap of the update subtask and the alignment operation through the asynchronous submission of the update subtask and the wait trigger of the scale alignment. Figure 7 As shown, the update task submitted to the synchronization channel sc1 in the jth round can be partially overlapped in time with the scale alignment task of sc1 in the j+1th round, without waiting for all the update subtasks in the jth round to be completed before performing the subsequent operation.

[0160] It is worth noting that when the hardware synchronization channel resource is limited, such as the hardware channel resource is occupied by other tasks, or the hardware itself has a small number of synchronization channels, which cannot meet the independent grouping demand, the first synchronization channel group and the second synchronization channel group allow to share the same synchronization channel. Because the update subtask is started after the loading / computing subtask is completed (they are not executed in the same period), the scale alignment operation of the intermediate accumulation subblock in the j+1th round cannot be refined, and needs to wait for all the computing subtasks to be completed. In addition, if the intermediate accumulation matrix in the jth round is updated in batches, the scale alignment operation of the intermediate accumulation subblock in the j+1th round also needs to wait for all the computing subtasks to be completed.

[0161] Referring to Figure 8 is a structural schematic diagram of one embodiment of the task scheduling device for attention calculation provided by the application.

[0162] The second aspect embodiment of the application provides a task scheduling device for attention calculation, which is applied to a consumer thread group, the consumer thread group is configured with one or more synchronization channels, each synchronization channel is configured with a plurality of task slots; the device comprises:

[0163] The first decomposition module 11 is configured to decompose the loading task of the query subblock from the shared memory to the register into N1 loading subtasks, and decompose the computing task of the jth attention score subblock corresponding to the query subblock into N2 computing subtasks.

[0164] The second decomposition module 12 is configured to decompose the update task of the intermediate accumulation subblock corresponding to the attention output subblock in the current iteration round into N3 parallel update subtasks.

[0165] The task submitting module 13 is configured to submit N1 loading sub-tasks, N2 computing sub-tasks and N3 updating sub-tasks into different task slots of N4 synchronization channels asynchronously, wherein sub-tasks with a dependency relationship are submitted into the same synchronization channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1.

[0166] The task executing module 14 is configured to perform a polling operation on the task slots in each synchronization channel to execute the sub-tasks stored in the task slots, and control the execution timing between the sub-tasks with a dependency relationship in each synchronization channel by configuring a threshold of the number of incomplete sub-tasks in the waiting instruction.

[0167] It should be noted that the task scheduling apparatus for attention calculation provided by the second aspect of the present application can implement all processes of the task scheduling method for attention calculation described in any one of the first aspect embodiments, and the functions and technical effects of each module and unit in the apparatus are the same as those of the task scheduling method for attention calculation described in any one of the first aspect embodiments, which will not be repeated here.

[0168] The third aspect of the present application provides a computer readable storage medium, which includes a stored computer program; wherein the computer program controls the device where the computer readable storage medium is located to execute the task scheduling method for attention calculation of any one of the first aspect embodiments when running.

[0169] The fourth aspect of the present application provides a computer program product, which includes computer instructions, and the computer instructions implement the task scheduling method for attention calculation of any one of the first aspect embodiments when executed by a processor.

[0170] Referring to Figure 9 is a structural schematic diagram of an embodiment of the electronic device provided by the present application.

[0171] The fifth aspect of the present application provides an electronic device, which includes a processor 21, a memory 22, and a computer program stored in the memory 22 and configured to be executed by the processor 21, wherein the processor 21 implements the task scheduling method for attention calculation of any one of the first aspect embodiments when executing the computer program.

[0172] Preferably, the computer program can be divided into one or more modules / units (such as computer program one, computer program two, …), which are stored in the memory 22 and executed by the processor 21 to complete the present application. The one or more modules / units can be a series of computer program instruction segments capable of completing a specific function, which are used to describe the execution process of the computer program in the electronic device.

[0173] The processor 21 can be any one of a CPU (Central Processing Unit), a GPU (Graphics Processing Unit), a TPU (Tensor Processing Unit), a NPU (Neural network Processing Unit), a DPU (Deep learning Processing Unit), an APU (Accelerated Processing Unit), and a GPGPU (General-Purpose computing on Graphics Processing Unit). The processor 21 is the control center of the electronic device, which connects various parts of the electronic device through various interfaces and lines.

[0174] The memory 22 mainly includes a program storage area and a data storage area, wherein the program storage area can store an operating system, at least one application required by a function, etc., and the data storage area can store related data, etc. In addition, the memory 22 can be a high-speed random access memory, and can also be a non-volatile memory, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, etc., or the memory 22 can also be other volatile solid-state storage devices.

[0175] It should be noted that the above-mentioned electronic device can include, but is not limited to, a processor, a memory, and the like, and those skilled in the art can understand that Figure 9 The structural block diagram shown is only an example of the structure of the above-mentioned electronic device, and does not constitute a limitation on the structure of the above-mentioned electronic device. The above-mentioned electronic device can include more or fewer components than those shown, or combine certain components, or different components.

[0176] The above merely describes the preferred embodiments of the present application, and it should be pointed out that, for those skilled in the art, several improvements and modifications can be made without departing from the technical principles of the present application, and these improvements and modifications should also be considered as the protection scope of the present application.

Claims

1. A task scheduling method for attention computation, characterized in that, Applied to a consumer thread group, the consumer thread group being configured with one or more synchronization channels, each synchronization channel being configured with multiple task slots; the method includes: The task of loading the query block from shared memory to registers is decomposed into N1 loading subtasks, and the task of computing the j-th attention score block corresponding to the query block is decomposed into N2 computing subtasks. The update task of the intermediate accumulation block corresponding to the attention output block in the current iteration round is decomposed into N3 parallel update subtasks; N1 loading subtasks, N2 computation subtasks, and N3 update subtasks are asynchronously submitted to different task slots of N4 synchronization channels; wherein, subtasks with dependencies are submitted to the same synchronization channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1; Within each synchronization channel, a polling operation is performed on the task slots to execute the subtasks stored in the task slots. By configuring a threshold for the number of incomplete subtasks in the waiting instructions, the execution sequence of dependent subtasks within each synchronization channel is controlled.

2. The task scheduling method for attention calculation as described in claim 1, characterized in that, The step of asynchronously submitting N1 loading subtasks, N2 computation subtasks, and N3 update subtasks to different task slots of the N4 synchronization channels includes: When j=1, N1 loading subtasks and N2 computation subtasks are asynchronously submitted to different task slots in the first synchronization channel group according to their dependencies. When j > 1, N2 computational subtasks are asynchronously submitted to different task slots in the first synchronization channel group; N3 update subtasks are asynchronously submitted to different task slots in the second synchronization channel group; wherein both the first synchronization channel group and the second synchronization channel group consist of at least one of the synchronization channels.

3. The task scheduling method for attention calculation as described in claim 1, characterized in that, Each subtask consists of at least one execution instruction.

4. The task scheduling method for attention calculation as described in claim 1, characterized in that, The i-th loading subtask is used to load the i-th query sub-block from shared memory into a register; where 1≤i≤N1; the size of the query sub-block satisfies the first computation size of the matrix multiply-add instruction.

5. The task scheduling method for attention calculation as described in claim 4, characterized in that, The task of decomposing the calculation task of the j-th attention score block corresponding to the query block into N2 calculation subtasks includes: The second computational size based on matrix multiplication and addition instructions divides the j-th key transpose block into N1×M key transpose sub-blocks; where M≥1; the size of the key transpose sub-blocks supports the horizontal splicing and expansion of multiple second computational sizes; When j=1, the indexes of the key transpose subblocks are traversed in row-major order. The matrix multiplication and addition of the n1th query subblock and the key transpose subblock in the n1th row and mth column is taken as a computational subtask to form N2 computational subtasks; where 1≤n1≤N1, 1≤m≤M, and N2=N1×M.

6. The task scheduling method for attention computation as described in claim 5, characterized in that, The step of decomposing the computation task of the j-th attention score block corresponding to the query block into N2 computation subtasks further includes: When j > 1, the indexes of the key transpose subblocks are traversed in column priority order. The combination of N1 query subblocks and all key transpose subblocks in column m is calculated as a computational subtask to form N2 computational subtasks; where N2 = M.

7. The task scheduling method for attention computation as described in claim 5, characterized in that, The first attention score block is calculated using the following steps: When the waiting instruction monitor detects that the first query sub-block has been loaded into the register, the computation subtasks corresponding to all key transpose sub-blocks in the first row are started sequentially. When the waiting instruction monitors that the n2nd query subblock is loaded into the register and the computation subtask corresponding to the key transpose subblock in row n2-1 and column m is completed, the computation subtask corresponding to the key transpose subblock in row n2 and column m is started; where 2≤n2≤N1; When all the computational subtasks corresponding to the key transpose subblocks in column m are completed, the mth fraction subblock of the first attention fraction block is obtained.

8. The task scheduling method for attention computation as described in claim 5 or 6, characterized in that, The method further includes: When the m-th fractional sub-block of the j-th attention score block is completed by waiting for the instruction to monitor, local statistical calculations are performed on the m-th fractional sub-block.

9. The task scheduling method for attention calculation as described in claim 1, characterized in that, In the current iteration round, the k-th update subtask is used to multiply the indexed fraction block with the corresponding sub-block in the value block, and accumulate the resulting sub-block product to the corresponding region of the scale-aligned intermediate accumulation block to obtain the updated k-th intermediate accumulation sub-block; where k≥1.

10. The task scheduling method for attention computation as described in claim 9, characterized in that, The method further includes: In the next iteration, when the kth update subtask is completed by waiting for the instruction to be monitored, the kth intermediate accumulation sub-block is scaled.

11. A task scheduling device for attention calculation, characterized in that, An apparatus applied to a consumer thread group, the consumer thread group being configured with one or more synchronization channels, each of the synchronization channels being configured with multiple task slots; the apparatus includes: The first decomposition module is used to decompose the loading task of the query block from shared memory to register into N1 loading subtasks, and to decompose the calculation task of the j-th attention score block corresponding to the query block into N2 calculation subtasks. The second decomposition module is used to decompose the update task of the intermediate accumulation block corresponding to the attention output block in the current iteration into N3 parallel update subtasks. The task submission module is used to asynchronously submit N1 loading subtasks, N2 calculation subtasks, and N3 update subtasks to different task slots of the N4 synchronization channels; wherein, subtasks with dependencies are submitted to the same synchronization channel; j≥1, N1≥1, N2≥1, N3≥1, N4≥1; The task execution module is used to poll the task slots in each synchronization channel to execute the subtasks stored in the task slots, and to control the execution sequence of dependent subtasks in each synchronization channel by configuring a threshold for the number of incomplete subtasks in the waiting instructions.

12. A computer-readable storage medium, characterized in that, The computer-readable storage medium includes a stored computer program; wherein, when the computer program is executed, it controls the device on which the computer-readable storage medium resides to perform a task scheduling method for attention computation as described in any one of claims 1 to 10.

13. A computer program product, characterized in that, It includes computer instructions that, when executed by a processor, implement the task scheduling method for attention computation as described in any one of claims 1 to 10.

14. An electronic device, characterized in that, The system includes a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, wherein the processor, when executing the computer program, implements the task scheduling method for attention computation as described in any one of claims 1 to 10.

Images