An MLA operator load balancing method suitable for tokens short sequence

Abstract

The present application relates to the field of heterogeneous parallel computing, and particularly relates to a MLA operator load balancing method suitable for tokens short sequence, comprising the following steps: firstly, obtaining the task parameters of the MLA operator, and judging whether the task of the MLA operator is a short sequence task; for the short sequence task, traversing each batch, calculating the alignment length of the key-value sequence, the total block number, the allocated block number of each core, the single block split length and the required core number for completing the batch calculation; storing the split parameters to the specified address space of the memory; for the non-full core batch with remaining load after full core split, using the back-and-forth core distribution strategy to dynamically distribute the remaining load; each core processes the corresponding key-value sequence block in parallel according to the read split parameters and executes the MLA operator core calculation. The present application can effectively solve the inter-core load imbalance problem in the non-full core scenario, improve the execution efficiency of the MLA operator in the short sequence scenario, and fully utilize the multi-core computing power of the general multi-core computing architecture.

Description

Technical Field

[0001] This invention relates to the field of heterogeneous parallel computing, and more specifically to a load balancing method for MLA operators suitable for short sequences of tokens. Background Technology

[0002] With the development of artificial intelligence technology, the proportion of short sequence tasks in fields such as natural language processing and computer vision, such as short text classification and small-sized image feature extraction, is gradually increasing. In such tasks, the MLA operator is the core computing unit, and its execution efficiency directly affects the overall task response speed.

[0003] General-purpose multi-core computing architectures employ multi-core architectures to improve computing power by processing computing tasks in parallel across multiple cores. However, in short-sequence scenarios, when MLA operators process multiple batches of key-value sequences (i.e., sequences of "key-value" data pairs storing feature association information), they face the following key problems: First, the lengths of key-value sequences in each batch vary significantly in short-sequence scenarios. Traditional load allocation methods only divide the key-value sequences based on a fixed block size, without considering the maximum utilization capacity of all cores. This results in the inability to achieve "full-core" partitioning in most cases, meaning that not all available cores are fully utilized, leading to a waste of hardware computing power. Second, even if core allocation is performed, due to the irregular length of the key-value sequences, some cores bear only a small amount of load or are even idle, while other cores are overloaded, resulting in a load imbalance between cores and further reducing computing efficiency. In addition, existing load allocation strategies lack a dynamic adjustment mechanism. For short-sequence MLA tasks executed multiple times, the remaining load is consistently allocated to fixed cores, causing some cores to remain in a low-load state for extended periods, failing to leverage the advantages of multi-core parallelism in general-purpose multi-core computing architectures. To improve the ecosystem support of general-purpose multi-core computing architectures for short-sequence tasks, it is necessary to ensure the high efficiency of MLA operators, and load balancing is a core aspect of improving operator efficiency. Therefore, a load balancing method for MLA operators is needed that can strive to achieve full-core partitioning for each key-value sequence across multiple batches and compensate for load imbalances in non-full-core scenarios through dynamic core allocation strategies. This will fully utilize the multi-core computing power of general-purpose multi-core architectures and improve the computational performance of short-sequence tasks. Summary of the Invention

[0004] To address the technical problems existing in the prior art, this invention provides a load balancing method for MLA operators in short sequence scenarios. This method performs full-core priority splitting on each key-value sequence of multiple batches based on the number of computing cores and block size. At the same time, it adopts a back-and-forth core splitting strategy for non-full-core scenarios to achieve load balancing between cores, improve the execution efficiency of MLA operators in short sequence scenarios, and make full use of the multi-core computing power of general multi-core computing architecture.

[0005] The first objective of this invention is to provide an MLA operator load balancing method suitable for short token sequences.

[0006] A second objective of this invention is to provide a computer device.

[0007] The first objective of this invention can be achieved by adopting the following technical solution:

[0008] A load balancing method using MLA operators for short token sequences includes the following steps:

[0009] S1. Obtain the task parameters of the MLA operator for processing multiple batches of key-value sequences, and determine whether the task of the MLA operator for processing multiple batches of key-value sequences is a short sequence task;

[0010] S2. For short sequence tasks, traverse each batch of key-value sequences and calculate the splitting parameters based on the key-value sequence length of the current batch and the fixed block size: the alignment length of the key-value sequence of the current batch, the total number of blocks, the number of blocks per core, the splitting length, and the number of cores required.

[0011] S3. Store the splitting parameters into memory, traverse each query sequence in the current batch, determine the offset position in memory used to store the splitting parameters of the current batch, store the key-value sequence splitting length and the required number of cores of the current batch into the offset position, and update the number of processed tasks at the same time.

[0012] S4. Determine whether the key-value sequence splitting of the current batch is a non-full-core batch key-value sequence splitting based on the number of cores required for the current batch key-value sequence splitting. For non-full-core batch key-value sequence splitting, use a back-and-forth core splitting strategy to allocate the remaining load so that the load of each core is balanced in the long term.

[0013] S5. Each core reads the key-value sequence block of the corresponding batch according to the segmentation parameters stored in memory and performs multiply-add calculations of the MLA operator until all batches of key-value sequences have been processed.

[0014] Specifically, the task parameters include: the total number of batches of key-value sequences, the query sequence length of each batch, the fixed block size, and the maximum number of key-value splitting cores;

[0015] The determination of whether the task of the MLA operator processing multiple batches of key-value sequences is a short sequence task includes:

[0016] If the length of the key-value sequence in a single batch is less than or equal to the product of the maximum number of key-value splitting cores and the fixed block size, then this batch is considered a short-sequence batch; if more than 80% of the batches in a task are short-sequence batches, then this task is determined to be a short-sequence task.

[0017] Specifically, step S2 includes:

[0018] S21. For each batch of key-value sequences, the RoundUp function is used to calculate the key-value sequence alignment length based on the current batch's key-value sequence length and the fixed block size.

[0019] S22. Calculate the total number of blocks required to length the current key-value sequence based on the alignment length and fixed block size;

[0020] S23. Based on the maximum number of key-value splitting cores and the total number of blocks, the number of blocks per core is calculated using the CeilDiv function;

[0021] S24. Based on the number of blocks per core and the length of the key-value sequence, use the CeilDiv function to calculate the number of cores required for the key-value sequence of the current batch.

[0022] Specifically, step S3 includes: calculating the offset position D based on a preset base offset and the number of processed tasks, wherein the offset position D is represented as:

[0023] D = T + P × L;

[0024] Where T is the preset base offset, P is the number of processed tasks, and L is the parameter storage length for each task.

[0025] Specifically, step S4 includes:

[0026] Determine whether the number of cores required for the key-value sequence splitting of the current batch is less than the maximum number of cores required for key-value splitting. If the number of cores required is less than the maximum number of cores required for key-value splitting, then the key-value sequence splitting of the current batch is considered to be a non-full-core batch key-value sequence splitting.

[0027] For key-value sequence splitting of non-full-core batches, the initial core splitting direction is set to forward, and the core numbers are sorted from smallest to largest. When processing key-value sequence splitting of non-full-core batches, the load is distributed in ascending order of core number. When processing the next key-value sequence splitting of non-full-core batches, the core splitting direction is switched to reverse, the core numbers are sorted from largest to smallest, and the load is distributed in descending order of core number.

[0028] Specifically, step S5 includes:

[0029] Each core reads the segmentation length and the required number of cores from the corresponding field in memory based on its own sequence number, and obtains the starting address of the key-value sequence block based on the segmentation length;

[0030] Through the efficient memory read interface of the general multi-core computing architecture, key-value sequence blocks are read from global memory into L1 cache using continuous memory copy instructions;

[0031] The matrix multiplication and addition operation of the MLA operator is executed, and the calculation result is written to the preset output memory area. After the current core finishes processing the allocated load, if there are still unprocessed key-value sequence blocks in the current batch, the next part of the data is read according to the core direction mark until all key-value sequence blocks in the current batch are processed.

[0032] Specifically, the starting address is represented as:

[0033] Q = DZ + CH×CD;

[0034] Where DZ represents the base address of the key-value sequence of the batch, CH represents the core number, and CD represents the segment length.

[0035] The second objective of this invention can be achieved by adopting the following technical solution:

[0036] A computer device includes a processor and a memory for storing a processor-executable program, wherein when the processor executes the program stored in the memory, it implements the above-described MLA operator load balancing method applicable to short token sequences.

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

[0038] This invention provides a load balancing method for MLA operators in short sequence scenarios. By calculating splitting parameters based on a fixed block size and the maximum number of splitting cores, it achieves full-core priority splitting of each batch of key-value sequences, ensuring that the split key-value sequence blocks are fully compatible with cache storage to avoid cache waste. It fully utilizes the multi-core computing power of a general multi-core computing architecture, reducing computational waste. For non-full-core scenarios, a back-and-forth core allocation strategy is designed, alternating the load distribution direction to avoid long-term idleness or concentrated load on fixed cores, solving the problem of load imbalance between cores and ensuring long-term load balance. Simultaneously, the splitting parameters are stored in a designated memory location for easy configuration retrieval by each core, improving parameter acquisition efficiency. This method effectively optimizes the execution efficiency of MLA operators in short sequence scenarios, solving the problems of low computational utilization and uneven load between cores in existing methods. It fully leverages the multi-core parallel advantages of a general multi-core computing architecture, improving the overall computational performance of short sequence tasks. Attached Figure Description

[0039] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on the structures shown in these drawings without creative effort.

[0040] Figure 1 This is a flowchart of an MLA operator load balancing method applicable to short token sequences in an embodiment of the present invention;

[0041] Figure 2 This is a schematic diagram of the key-value sequence segmentation method for each batch in an embodiment of the present invention;

[0042] Figure 3 This is a schematic diagram showing the storage location of the segmentation parameters in memory in an embodiment of the present invention;

[0043] Figure 4 This is a schematic diagram of a load distribution example in a non-full-core scenario according to an embodiment of the present invention. Detailed Implementation

[0044] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments, and the implementation of the present invention is not limited thereto. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0045] Example 1:

[0046] This embodiment provides a load balancing method for MLA (Multi-Head LatentAttention) operators in short-sequence scenarios, which is applicable to general multi-core computing architectures. It can solve the problems of uneven multi-core load and low computing power utilization during the execution of MLA operators in short-sequence scenarios, and improve the computing efficiency of general multi-core computing architectures in short-sequence tasks.

[0047] like Figure 1 The flowchart shown illustrates a load balancing method for MLA operators in short-sequence scenarios. The load balancing method for MLA operators in short-sequence scenarios described in this invention specifically includes:

[0048] S1. Obtain the task parameters for the MLA operator to process multiple batches of key-value sequences, calculate the total number of tokens for the MLA operator based on the task parameters, and determine whether the task of processing multiple batches of key-value sequences is a short sequence task.

[0049] Before executing short-sequence MLA operators to process multiple batches of key-value sequences, the general-purpose multi-core computing architecture needs to read the complete task parameters of the MLA operator from the task configuration module to ensure the accuracy of subsequent load balancing. The total number of batches of the key-value sequence is the number of batches the MLA operator needs to process. The original sequence number of each batch is obtained through an index array, which records the actual processing order of each batch, avoiding errors in the original data correspondence due to batch sorting. The length of each batch's key-value sequence (the number of tokens in the key-value sequence) is read from the key-value array; this length is the core basis for load balancing. The fixed block size, denoted as BLOCK_SIZE, is a fixed block size preset during the MLA operator compilation. This size is preset during the operator compilation stage and is adapted to the L1 cache line size of the general-purpose multi-core computing architecture to ensure cache hit rate during data reading. The maximum number of key-value splitting cores is denoted as MAX_KEY-VALUE_SPLIT_NUM. This number is determined by the hardware specifications of a general multi-core computing architecture, which is the total number of AICores currently available for parallel processing of MLA operators.

[0050] Obtain the task parameters for the MLA operator when processing multiple batches of key-value sequences. Calculate the total number of tokens for the MLA operator based on these parameters. The total number of tokens for the MLA operator reflects the overall computational scale of this task. The calculation formula is as follows:

[0051] ;

[0052] in, The total number of tokens for the MLA operator. Let B be the number of tokens in the i-th batch of key-value sequences, and B be the total number of batches.

[0053] Specifically, the task of determining whether a multi-batch key-value sequence is a short-sequence task is classified into short-sequence tasks and long-sequence tasks. The classification of short and long-sequence tasks is based on two criteria: the length of a single batch of key-value sequences and the maximum processing capacity of the hardware. The specific determination rules are as follows:

[0054] A batch is considered a short-sequence batch if the length of its key-value sequence is less than or equal to the product of the maximum number of key-value splitting cores and the fixed block size. Specifically, a batch is considered short-sequence if the length of its key-value sequence is ≤ MAX_KEY_VALUE_SPLIT_NUM × BLOCK_SIZE. If more than 80% of the batches in a task are short-sequence batches, the task is considered a short-sequence task. Tasks that do not meet the above conditions are considered long-sequence tasks and are handled using existing conventional load balancing methods.

[0055] S2. For short sequence tasks, traverse each batch of key-value sequences and calculate the splitting parameters based on the key-value sequence length of the current batch and the fixed block size: the alignment length of the key-value sequence of the current batch, the total number of blocks, the number of blocks per core, the splitting length, and the required number of cores.

[0056] S21. For each batch of key-value sequences, the RoundUp function is used to calculate the alignment length of the key-value sequence based on the length of the current batch of key-value sequences and the fixed block size. The RoundUp function outputs the minimum length that is not less than the length of the key-value sequence and is an integer multiple of the fixed block size.

[0057] The RoundUp function is a mathematical function that rounds up to a specified multiple. For example, if the key-value sequence length of the current batch is 256 and the fixed block size is 128, then the alignment length is 256 (256 is an integer multiple of 128); if the key-value sequence length is 280 and the fixed block size is 128, then the alignment length is 384 (384 is the minimum length that is not less than 280 and is an integer multiple of 128).

[0058] S22. Calculate the total number of blocks required to lengthen the current key-value sequence based on the alignment length and fixed block size.

[0059] Specifically, the total number of blocks = alignment length / fixed block size. For example, if the alignment length is 384 and the fixed block size is 128, then the corresponding total number of blocks is 3.

[0060] S23. Based on the maximum key-value splitting core number and the total number of blocks, the CeilDiv function is used to calculate the number of blocks per core.

[0061] The CeilDiv function is a common floor function that calculates the division of two integers and rounds the result up to the nearest integer. If the total number of blocks is m, and the maximum number of key-value splitting cores MAX_KEY-VALUE_SPLIT_NUM in current general-purpose multi-core computing architectures is n, then the number of blocks per core = CeilDiv(m, n), where CeilDiv(m, n) is the quotient of m divided by n and rounded up.

[0062] S24. Based on the number of blocks per core and the length of the key-value sequence, use the CeilDiv function to calculate the number of cores required for the key-value sequence of the current batch.

[0063] For each core, the split length of the current batch can be calculated from the number of blocks per core, using the formula: Split length = Number of blocks per core × Fixed block size. The CeilDiv function is used: if the key-value sequence length is L and the split length is l, then the required number of cores = CeilDiv(L, l), ensuring that all key-value sequence data can be allocated to the corresponding cores.

[0064] like Figure 2 The diagram illustrates the key-value sequence segmentation method for each batch in this embodiment of the invention, showing the three core steps of full-core priority segmentation: the first step is to align the total key-value sequence length upwards according to a fixed block size; the second step is to segment the aligned total key-value sequence length into multiple equal-length blocks according to a fixed block size; and the third step is to divide the total number of blocks upwards by the total number of hardware cores to obtain the minimum number of blocks required to be processed by each core.

[0065] S3. Store the splitting parameters into memory, traverse each query sequence in the current batch, calculate the offset position used to store the splitting parameters of the current batch, store the key-value sequence splitting length and the required number of cores of the current batch into the offset position, and update the number of processed tasks at the same time.

[0066] Specifically, the offset position D is calculated based on the preset base offset and the number of tasks processed. The offset position D can be expressed as:

[0067] D = T + P × L;

[0068] Where T is the preset base offset, which is the starting position; P is the number of processed tasks, starting from 0; and L is the parameter storage length for each task.

[0069] The current batch's segmentation length and required number of cores are stored in the specified field at this offset. The number of processed tasks, P, is then incremented by 1 to prepare for storing parameters for the next query sequence. This area of ​​memory is shared by all NPU cores, allowing each core to quickly read its own required segmentation configuration.

[0070] like Figure 3 The diagram shown illustrates the storage location of the partitioning parameters in memory in an embodiment of the present invention, demonstrating the layout of the partitioning parameters in shared memory: the memory starts with a basic information storage area of ​​size TILING_HEAD_SIZE; followed by consecutive storage units of size PARA_TILING_ELENUM, each unit corresponding to the partitioning parameters of a task.

[0071] S4. Determine whether the key-value sequence splitting of the current batch is a non-full-core batch key-value sequence splitting based on the number of cores required for the current batch key-value sequence splitting. For non-full-core batch key-value sequence splitting, adopt a back-and-forth core splitting strategy to allocate the remaining load, so that the load of each core is balanced in the long term.

[0072] The total bandwidth of a general multi-core computing architecture is fixed, and the bandwidth of each core is shared. However, the computing power of each core is not shared. Therefore, when dealing with non-full-core scenarios, the computing power is not fully utilized. In this case, adopting a back-and-forth alternating core distribution strategy can make the load of each core relatively balanced.

[0073] S41. Determine whether the number of cores required for the key-value sequence splitting of the current batch is less than the maximum number of cores required for key-value splitting. If the number of cores required is less than the maximum number of cores required for key-value splitting, then the key-value sequence splitting of the current batch is considered to be a non-full-core batch key-value sequence splitting.

[0074] In this embodiment, it is first determined whether the number of cores required for the key-value sequence splitting of the current batch is less than the maximum number of key-value splitting cores MAX_KEY-VALUE_SPLIT_NUM. If the number of cores required is equal to the maximum number of key-value splitting cores, it is considered a full-core batch key-value sequence splitting; if the number of cores required is less than the maximum number of key-value splitting cores, it is considered that the key-value sequence splitting of the current batch is a non-full-core batch key-value sequence splitting.

[0075] S42. For key-value sequence splitting of non-full-core batches, a back-and-forth core allocation strategy is adopted to distribute the remaining load, including: setting a core allocation direction marker, setting the initial core allocation direction to positive, sorting the core numbers from smallest to largest, and allocating the load in ascending order of core numbers when processing key-value sequence splitting of non-full-core batches; when processing the next key-value sequence splitting of non-full-core batches, the core allocation direction is switched to reverse, the core numbers are sorted from largest to smallest, and the load is allocated in descending order of core numbers.

[0076] In this embodiment, a core allocation direction marker is set, initially in the forward direction: For the first batch of non-full-core tasks, tasks are allocated in the order of core number 1→2→…→n (where n is the required number of cores), allowing each core to bear the load sequentially. Then, when the next non-full-core batch arrives, the core allocation direction marker switches to "reverse," allocating the load in the order of core number n→n-1→…→1; the next non-full-core batch is allocated in the "forward" direction, and so on. This alternating allocation method ensures that the total load borne by each NPU core tends to be consistent across multiple task executions, avoiding long-term under- or over-loaded loads on any single core.

[0077] Figure 4 This is a schematic diagram of a load distribution example in a non-full-core scenario according to an embodiment of the present invention, demonstrating the execution effect of the back-and-forth core distribution strategy, which achieves long-term load balancing of each core through alternating distribution.

[0078] S5. Each core reads the key-value sequence block of the corresponding batch according to the segmentation parameters stored in memory and performs multiply-add calculations of the MLA operator until all batches of key-value sequences have been processed.

[0079] S51. Each core reads the segmentation length and the required number of cores from the corresponding field in memory according to its own sequence number, and determines the starting address of the key-value sequence block. The starting address Q can be represented as:

[0080] Q = DZ + CH×CD.

[0081] Where DZ represents the base address of the batch's key-value sequence, CH represents the core sequence number, and CD represents the segmentation length. The base address DZ of the batch's key-value sequence refers to the starting virtual address of all key-value data in global memory, allocated by the task scheduling module during the data preprocessing stage. The required number of cores for reading has two core functions: 1. Participation in computation determination: Each core compares its own sequence number CH with the required number of cores. If CH ≥ the required number of cores, the core does not participate in the computation of the current batch and directly enters an idle state to wait for the next batch of tasks. 2. Core allocation strategy execution: For batches with less than full cores, each core determines whether it needs to handle the remaining load and the reading location of the remaining load based on the required number of cores and the core allocation direction marker.

[0082] S52: Through the efficient memory read interface of the general multi-core computing architecture, the key-value sequence block is read from global memory into the L1 cache using the continuous memory copy instruction, avoiding bandwidth loss caused by non-aligned access.

[0083] S53. Perform matrix multiplication and addition operations of the MLA operator and write the calculation results to the preset output memory area. After the current core finishes processing its allocated load, if there are still unprocessed key-value sequence blocks in the current batch, continue reading the next part of the data according to the core distribution direction mark until all key-value sequence blocks in the current batch have been processed. Repeat the above steps until all batches of MLA operator tasks are completed.

[0084] The MLA operator load balancing method provided in this embodiment for short sequence scenarios is based on the multi-core computing characteristics and memory access optimization logic of a general multi-core computing architecture. It fully utilizes the multi-core computing power of the general multi-core computing architecture by prioritizing full-core partitioning. Specifically, it determines the fixed block size by combining the maximum number of available computing cores (i.e., the maximum number of partitioned cores) and the cache line size. Each batch of key-value sequences is partitioned according to the aforementioned computing logic, ensuring that the partitioned key-value sequence blocks can fully adapt to cache storage to reduce bandwidth loss caused by non-aligned memory accesses. Simultaneously, it ensures that the initial load difference between cores is relatively small, avoiding idle computing power. A back-and-forth alternating core partitioning strategy addresses the load imbalance problem in non-full-core scenarios. A core partitioning direction status flag is set. When the number of cores required for a batch is less than the maximum number of partitioned cores, the load is initially allocated according to the core sequence number from smallest to largest. The next non-full-core batch is allocated according to the sequence number from largest to smallest. The flag is switched after each processing, avoiding long-term idleness of fixed cores or concentrated load. A standardized parameter storage method is used. This method improves the efficiency of inter-core configuration synchronization by referencing the parameter storage mechanism of general multi-core computing architectures. It stores the segmentation parameters (segmentation length, required number of cores, and original batch sequence number) in a preset area of ​​memory. The storage location is calculated using a "basic offset (starting address of the memory operator configuration area + index of processed tasks × length of single task parameters)". Each core can quickly locate and read the corresponding configuration using its own sequence number, eliminating the need for traversal searches and shortening parameter acquisition time. In short-sequence AI tasks (such as short text classification and small-size image feature extraction), this method can improve the computational utilization of MLA operators and shorten the processing time per batch. It effectively improves the execution efficiency of MLA operators in short-sequence scenarios, solving the problems of wasted computational power and uneven load distribution among cores in existing technologies. It is applicable to various short-sequence AI tasks under general multi-core computing architectures, providing key technical support for the improvement of the short-sequence task operator ecosystem for general multi-core computing architectures. Furthermore, it synergizes with existing matrix memory read optimization strategies for general multi-core computing architectures, further enhancing the overall performance of heterogeneous parallel computing systems.

[0085] Example 2:

[0086] This embodiment provides a computer device, which may be a server, computer, etc., including a processor, memory, input device, display, and network interface connected via a system bus. The processor provides computing and control capabilities. The memory includes a non-volatile storage medium and internal memory. The non-volatile storage medium stores an operating system, computer programs, and a database. The internal memory provides an environment for the operation of the operating system and computer programs in the non-volatile storage medium. When the processor executes the computer programs stored in the memory, it implements the MLA operator load balancing method for short token sequences described in Embodiment 1 above, as follows:

[0087] S1. Obtain the task parameters for the MLA operator to process multiple batches of key-value sequences, calculate the total number of tokens for the MLA operator based on the task parameters, and determine whether the task for the MLA operator to process multiple batches of key-value sequences is a short sequence task.

[0088] S2. For short sequence tasks, traverse each batch of key-value sequences, and calculate the alignment length, total number of blocks, number of blocks per core, split length, and required number of cores for the current batch of key-value sequences based on the length of the key-value sequences and the fixed block size.

[0089] S3. Traverse each query sequence in the current batch, determine the offset position in memory used to store the splitting parameters of the current batch, store the key-value sequence splitting length and the required number of cores of the current batch at the offset position, and update the number of tasks processed at the same time.

[0090] S4. Determine whether the key-value sequence splitting of the current batch is a non-full-core batch key-value sequence splitting based on the number of cores required for the current batch key-value sequence splitting. For non-full-core batch key-value sequence splitting, use a back-and-forth core splitting strategy to allocate the remaining load so that the load of each core is balanced in the long term.

[0091] S5. Each core reads the key-value sequence block of the corresponding batch according to the segmentation parameters stored in memory and performs multiply-add calculations of the MLA operator until all batches of key-value sequences have been processed.

[0092] The above embodiments are preferred embodiments of the present invention, but the embodiments of the present invention are not limited to the above embodiments. Any changes, modifications, substitutions, combinations, or simplifications made without departing from the spirit and principle of the present invention shall be considered equivalent substitutions and shall be included within the protection scope of the present invention.

Claims

1. A load balancing method using MLA operators suitable for short token sequences, characterized in that, Includes the following steps: S1. Obtain the task parameters of the MLA operator for processing multiple batches of key-value sequences, and determine whether the task of the MLA operator for processing multiple batches of key-value sequences is a short sequence task; S2. For short sequence tasks, traverse each batch of key-value sequences and calculate the splitting parameters based on the key-value sequence length of the current batch and the fixed block size: the alignment length of the key-value sequence of the current batch, the total number of blocks, the number of blocks per core, the splitting length, and the number of cores required. S3. Traverse each query sequence in the current batch, determine the offset position in memory used to store the splitting parameters of the current batch, store the key-value sequence splitting length and the required number of cores of the current batch at the offset position, and update the number of tasks processed at the same time. S4. Determine whether the key-value sequence splitting of the current batch is a non-full-core batch key-value sequence splitting based on the number of cores required for the current batch key-value sequence splitting. For non-full-core batch key-value sequence splitting, use a back-and-forth core splitting strategy to allocate the remaining load so that the load of each core is balanced in the long term. S5. Each core reads the key-value sequence block of the corresponding batch according to the segmentation parameters stored in memory and performs multiply-add calculations of the MLA operator until all batches of key-value sequences have been processed.

2. The MLA operator load balancing method for short token sequences according to claim 1, characterized in that, The specific task parameters include: the total number of batches of key-value sequences, the query sequence length of each batch, the fixed block size, and the maximum number of key-value splitting cores; The determination of whether the task of the MLA operator processing multiple batches of key-value sequences is a short sequence task includes: If the length of the key-value sequence in a single batch is less than or equal to the product of the maximum number of key-value splitting cores and the fixed block size, then this batch is considered a short-sequence batch; if more than 80% of the batches in a task are short-sequence batches, then this task is determined to be a short-sequence task.

3. The MLA operator load balancing method for short token sequences according to claim 1, characterized in that, Step S2 includes: S21. For each batch of key-value sequences, the RoundUp function is used to calculate the key-value sequence alignment length based on the current batch's key-value sequence length and the fixed block size. S22. Calculate the total number of blocks required to length the current key-value sequence based on the alignment length and fixed block size; S23. Based on the maximum number of key-value splitting cores and the total number of blocks, the number of blocks per core is calculated using the CeilDiv function; S24. Based on the number of blocks per core and the length of the key-value sequence, use the CeilDiv function to calculate the number of cores required for the key-value sequence of the current batch.

4. The MLA operator load balancing method for short token sequences according to claim 1, characterized in that, Step S3 includes: calculating the offset position D based on a preset base offset and the number of processed tasks, wherein the offset position D is represented as: D = T + P × L; Where T is the preset base offset, P is the number of processed tasks, and L is the parameter storage length for each task.

5. A load balancing method for MLA operators applicable to short token sequences according to claim 1, characterized in that, Step S4 includes: Determine whether the number of cores required for the key-value sequence splitting of the current batch is less than the maximum number of cores required for key-value splitting. If the number of cores required is less than the maximum number of cores required for key-value splitting, then the key-value sequence splitting of the current batch is considered to be a non-full-core batch key-value sequence splitting. For key-value sequence splitting of non-full-core batches, the initial core splitting direction is set to forward, and the core numbers are sorted from smallest to largest. When processing key-value sequence splitting of non-full-core batches, the load is distributed in ascending order of core number. When processing the next key-value sequence splitting of non-full-core batches, the core splitting direction is switched to reverse, the core numbers are sorted from largest to smallest, and the load is distributed in descending order of core number.

6. The MLA operator load balancing method for short token sequences according to claim 1, characterized in that, Step S5 includes: Each core reads the segmentation length and the required number of cores from the corresponding field in memory based on its own sequence number, and obtains the starting address of the key-value sequence block based on the segmentation length; Through the efficient memory read interface of the general multi-core computing architecture, key-value sequence blocks are read from global memory into L1 cache using continuous memory copy instructions; The matrix multiplication and addition operation of the MLA operator is executed, and the calculation result is written to the preset output memory area. After the current core finishes processing the allocated load, if there are still unprocessed key-value sequence blocks in the current batch, the next part of the data is read according to the core direction mark until all key-value sequence blocks in the current batch are processed.

7. A load balancing method for MLA operators applicable to short token sequences according to claim 6, characterized in that, The starting address is represented as: Q = DZ + CH×CD; Where DZ represents the base address of the key-value sequence of the batch, CH represents the core number, and CD represents the segment length.

8. A computer device comprising a processor and a memory for storing a processor-executable program, characterized in that, When the processor executes the program stored in the memory, it implements the MLA operator load balancing method for short token sequences as described in any one of claims 1-7.

Images