Processing method for improving batch reasoning efficiency of large language model

Abstract

The invention discloses a processing method for improving batch inference efficiency of a large language model, which comprises the following steps of: 1) executing pre-filling processing on all inference tasks to be processed by utilizing the large language model to obtain a key vector and a value vector of each lexical element in a sequence of each inference task and a lexical element, and caching the key vector and the value vector and the lexical element into a task pool; 2) during the first round of iterative calculation, forming a tensor TB by using the latest single lexical elements T corresponding to n reasoning tasks selected from a task pool, forming a tensor KVB by using corresponding key value vectors, and forming a tensor AB by using attention mask vectors; inputting the TB, the KVB and the AB into a large language model for reasoning, calculating and updating the TB, the AB and the KVB for the next round of iterative calculation; and 3) after the iteration of a reasoning task is finished, splicing the lexical elements output by the reasoning task in each iteration to obtain a reasoning result of the reasoning task, and selecting a to-be-processed task from the task pool to execute the next iteration. The task processing efficiency and the computing power resource utilization rate are improved.

Description

Technical Field

[0001] The present invention belongs to the field of computer software technology and relates to a processing method for improving the efficiency of batch reasoning of a large language model. Background Art

[0002] The application of large models plays a vital role in the modern technology environment. Service providers complete the large model reasoning calculation according to the user's request and return the results. The computing efficiency determines the service quality. Usually, in order to make full use of the parallel processing capabilities of computing devices such as GPUs, reasoning tasks are usually performed in batches. However, the reasoning of large language models (LLMs) has an autoregressive nature, that is, each reasoning task requires an indefinite number of iterative calculations to generate the final result. In the current run-to-completion mode of batch processing, tasks with fewer iterations will still occupy computing resources after completion, and continue unnecessary iterative calculations with other unfinished tasks until all tasks in the batch are completed. This idle calculation not only causes a waste of computing resources, but also increases the task processing delay, and ultimately reduces the quality of service. Therefore, how to improve the reasoning efficiency of large models based on the autoregressive reasoning characteristics of LLM and the needs of batch processing has become a key technical problem in this field. The traditional solutions to improve large model reasoning are as follows:

[0003] 1. PD separation reasoning

[0004] Reasoning for large models involves two stages: first, context initialization of the input content and caching of intermediate states, with the output of the first word as the end mark of the stage, which is called the prefilling stage; then, output is generated step by step based on the context information, with one word generated in each iteration until the reasoning task is completed, which is called the sequence decoding stage. The computational overhead of the prefilling stage is positively correlated with the initial sequence length of the input model, and with the help of the KV cache mechanism, the computational overhead of each iteration of the decoding stage is fixed. In batch processing, the computational overhead of all tasks in the prefilling stage is consistent with the overhead of the longest sequence task in the batch, resulting in unnecessary processing delays. To solve this problem, a strategy of separating prefilling and decoding (PD) is proposed to achieve more flexible task scheduling and batch combination. The latest PD separation schemes are DistServe and Splitwise, but they still face the problem of differences in the iteration rounds of tasks in the batch in the decoding stage, that is, there is idle calculation for tasks with shorter output sequences, which will cause waste of resources.

[0005] 2. Model Inference Service

[0006] TensorFlow Serving and Triton Inference Server are widely used production-level inference service systems in the industry. As an abstraction layer above the execution engine, they lack model-specific optimization. Existing solutions such as Clipper, Clockwork, Shepherd, and Nexus focus on serving relatively small models in clusters, such as traditional deep neural network models such as ResNet that are not autoregressive. Due to the lack of perception of the inference latency of models with autoregressive properties, such methods are difficult to apply to the inference service of LLM. Orca is the latest strategy for the autoregressive generation mode of LLM. In order to eliminate the requirement of ensuring the consistency of tensor dimensions corresponding to each task in batch inference, this method dynamically sets the number of self-attention layer replicas according to the batch size. Each replica calculates each task independently and completes the calculation of the linear layer by vector concatenation. Orca breaks the limitation of batch inference running to the end and supports batch task deletion and insertion during the inference process. On this basis, FastServe supplements the task update strategy with priority. However, these solutions have the following problems: On the one hand, the self-attention layer has a large number of parameters, setting up copies will cause additional resource overhead and make it difficult to support a large batch size; on the other hand, the calculation of each attention copy is synchronous, that is, all copies need to complete the calculation before subsequent linear calculations can be performed. This results in that each time a new task is added, the calculation of the decoding stage of other tasks will be blocked due to the pre-filling processing of the task, that is, there is a problem equivalent to idling calculation. Summary of the invention

[0007] In view of the idle computing problem existing in the existing large language model batch reasoning process, the purpose of the present invention is to provide a processing method for improving the efficiency of large language model batch reasoning. The present invention further improves the computing resource utilization and task processing efficiency, and can be widely used in the field of reasoning services based on large language models.

[0008] The technical solution of the present invention is:

[0009] A processing method for improving the efficiency of batch inference of a large language model, the steps of which include:

[0010] 1) Use the large language model to perform pre-filling processing on all pending inference tasks, and obtain the key vector and value vector of each word in the sequence of each inference task i and a word t i , the key vector and value vector of each word in the sequence of reasoning task i are recorded as key-value vector KV i , the key-value vector KV i and word t i Cache into the task pool in memory;

[0011] 2) In the first round of iterative calculation, the latest word element corresponding to each reasoning task selected from the n reasoning tasks in the task pool is composed into a tensor T with a dimension of n*1 B , the key-value vectors corresponding to the selected n inference tasks are combined into a tensor KV B , the attention mask vectors corresponding to the selected n inference tasks are combined into a tensor A with a dimension of n*l B , l is the tensor KV B The length of the tensor T B , KV B and A B Input the large language model to perform reasoning calculation to obtain the new word T corresponding to each reasoning task b in the n reasoning tasks b ', and the new word T b 'The key vector and value vector between the existing word in the sequence of the corresponding reasoning task b are recorded as the key-value vector kv b '; The key value vector kv b 'Append to tensor KV B Then form the latest tensor KV B , in the current tensor A B Then append a column of vectors with all values ​​1 to update tensor A B ; The new word tensor T B , tensor A B and tensor KV B Used for the next round of iterative calculation;

[0012] 3) When an iteration of a reasoning task is completed, the word units output in each iteration of the reasoning task are concatenated to obtain the reasoning result of the reasoning task.

[0013] Furthermore, if the inference task i in the currently processed batch finishes iteration, it is removed from the current batch and a new inference task j is added to the current batch.

[0014] Furthermore, for the newly added reasoning task j, when the key-value vector KV corresponding to the reasoning task j j The length l q Less than or equal to the latest tensor KV B In the sequence dimension, length l kv When the latest word t of the inference task j is used j Directly overwrite the new word tensor T B The latest word element of reasoning task i in KV j Overwrite KV in right-aligned way B The key vector and value vector of inference task i, tensor KV B The positions of inference task i that are not covered by task j are modified to negative infinity to ensure the correctness of the calculation; tensor A is overwritten in a right-aligned mannerB The right side of the row corresponding to reasoning task i q The value of the element is 1, and the values ​​of other elements in the row are overwritten to 0 to ensure that the latest word element t of the inference task j is j , Attention Relation KV j Correctly calculate the current latest word t j The key-value vector kv between the existing word units in the task sequence of reasoning task j j ' and the new word t j '.

[0015] Furthermore, the key-value vector KV corresponding to the newly added reasoning task j is released j The front end has a length of O i The placeholder data and its position in tensor A B The corresponding data in .

[0016] Furthermore, for the newly added reasoning task j, when the key-value vector KV corresponding to the reasoning task j j The length l q Greater than the latest tensor KV B In the sequence dimension, length l kv When , we first use the latest word t of the inference task j j Directly overwrite the new word tensor T B The latest word element of the inference task i in the tensor KV B The length of l q , use KV j Overwrite the value of the row where the original inference task i is located, and fill the key value vectors of other inference tasks in the current batch with negative infinity on the left; secondly, tensor A B Expand to length l q , and transform the tensor A B The positions of the inference task j in the matrix are all assigned 1, indicating that the task j has no fill symbol; for the tensor A B The left side of the attention mask vector corresponding to other reasoning tasks in is filled with 0, indicating that the position of the key vector and value vector of the task is filled with a value, and ensuring that the attention mask vector A corresponding to the reasoning task j j The length is consistent.

[0017] Furthermore, the method for performing pre-filling processing on all pending inference tasks is as follows: dividing the inference tasks into multiple batches according to the sequence length of the inference tasks; aligning the sequence length of each batch of inference tasks by filling placeholders, and generating an attention mask vector and an initialized key value vector for each aligned inference task; and then performing pre-filling processing on the original sequence S of each inference task i. i , attention mask vector A iAnd the initialized key-value vector outputs the key-value vector KV between each word in the sequence of the inference task i i and the term t i , the key-value vector KV i and word t i Cache into the task pool in memory.

[0018] Furthermore, the task pool is located in the memory of the system where the large language model is located.

[0019] Furthermore, the large language model is a generative model based on a self-attention mechanism.

[0020] A server, characterized in that it comprises a memory and a processor, wherein the memory stores a computer program, the computer program is configured to be executed by the processor, and the computer program comprises instructions for executing the above method.

[0021] A computer-readable storage medium stores a computer program, wherein the computer program implements the above method when executed by a processor.

[0022] The advantages of the present invention are as follows:

[0023] Through PD separation and task tensor shaping, the present invention supports flexible replacement of tasks in a processing batch without causing additional resource overhead, realizes large-model batch reasoning with close to zero idle computing, and improves task processing efficiency and computing resource utilization. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 Schematic diagram of tensor shaping.

[0025] Figure 2 This is a schematic diagram of resource release.

[0026] (a) KV cache release, (b) attention mask resource release. DETAILED DESCRIPTION

[0027] The present invention is further described in detail below in conjunction with the accompanying drawings. The examples given are only used to explain the present invention but not to limit the scope of the present invention.

[0028] The main steps of the present invention include:

[0029] 1. PD separation

[0030] The task of large model inference is the inference request from the user, such as the question asked by the user to ChatGPT. Each question is an inference task. When calculating large model inference, in order to fully utilize the parallel computing capability of GPU, several tasks are usually grouped into a batch for parallel calculation. The number of tasks in a batch is called the batch size. The batch size usually has a default value in the inference system, such as 5 tasks forming a batch.

[0031] Pre-filling is performed on all tasks that need to be processed. Given that the computational overhead of pre-processing is proportional to the initial sequence length of the task, tasks with similar sequence lengths can be prioritized to be grouped into a batch. Given that the model's calculation requires that the vector dimension of each task in each batch is known, the sequence can be aligned through a padding operation, and an attention mask vector is generated for each task. The attention mask vector is a 0-1 vector with the same length as the sequence length, filled with corresponding 0s and non-filled bits with 1s. At the same time, the batch size can be dynamically adjusted according to the sequence length and computing power resources in the current batch, and the batch size can be increased as much as possible without causing video memory overflow, thereby improving the processing efficiency of pre-filling.

[0032] Assume that the original sequence of task i is S i , the attention mask is A i After being pre-filled, the Key vector and Value vector of the word in the sequence used to calculate the attention weight are stored in the KV cache, denoted as KV i , and get a new output word, denoted by T i . T i and KV i Will be stored in the task pool. In order to save video memory, the task pool can be stored in the main memory.

[0033] 2. Tensor Reshaping

[0034] The tensor is reshaped as Figure 1 As shown in the figure, the initial state is to select the batch size of tasks from the task pool and compose the corresponding KV cache and sequence into the corresponding batch tensor. Since the tasks have been pre-filled, only the latest single word T of the selected n tasks needs to be composed into the sequence tensor T B , cache KV B tensor and attention mask tensor A B Input model, where A B It is a 0-1 tensor of n*l, where l is KV B The length of , the specific value is 1 or 0 depending on whether it is filled. After the inference calculation of the large model, a new word unit tensor T with a dimension of n*1 can be obtainedB (Output a new word for each task to form a new word tensor T B ), and the new word tensor T B The key and value kv between the existing tokens in the existing task sequence B Among them, kv B It will be appended to the current KV cache tensor to form the latest KV B For the calculation of the next round of iteration. In addition, in the current attention mask tensor A B Then append a column of vectors with all values ​​1 to represent the new word tensor T B Is a non-filling character. At this time, KV B and A B The length of all becomes l+1. Then, T B , the corresponding attention mask tensor A B and KV B It will be used as the model input for the next round of iterations, and the reasoning calculation will be repeated. When a word element output by a task in the batch is a pre-defined reasoning end identifier ( <eos>), it means that the task has been inferred and the accumulated output tokens can be concatenated into a complete inference result and returned.

[0035] Assume that after task i in a batch is completed and the result is returned, a new task j can be added to the batch. The core of the operation is that the value corresponding to the returned task in the KV cache tensor can be overwritten. We combine the KV of task j j Cache, reshape the batch KV cache tensor so that KV j Can embed and overwrite the cache content KV of the returned task i i The large model is unaware of the tensor reshaping and can perform reasoning in the decoding phase without interruption. The specific reshaping can be divided into two cases.

[0036] Case 1: When KV j The length l q Less than or equal to the current batch KV B The length of the tensor in the sequence dimension is l kv At this time, the latest word t of task j j Can directly cover T B Medium Task i The value of T B Update to T B ';KV j Overwrite KV in right-aligned way B The key-value vector of the original task i, the value of the position that is not covered needs to be modified to negative infinity (-inf) to ensure the correctness of the calculation, KV B Update to KV B '; The dimension of the attention mask tensor is 2, and the attention mask of each task is represented by one row. B The update overwrites the right l of the row where task i is located in a right-aligned manner q The first element is 1 (indicating that the corresponding position of task j is the original word in the sequence), and the part that cannot be covered (left l kv -l q elements) is overwritten to 0 (indicating that the position is a fill symbol), A B Update to A B ', this ensures that the large model can be based on the word t input by task j j , KV j , correctly calculate t j and the key values ​​of the existing tokens in the sequence and the new token t j '. From a batch perspective, the three updated tensor variables are input into the large model to complete a new round of inference calculations.

[0037] Case 2: When KV j The length l q Greater than the current batch KV B The length of the tensor in the sequence dimension is l kv . Term tensor T B The processing is consistent with case 1 and is updated to T B '. However, due to l q Greater than l kv , the key-value vector KV of task j j Cannot be directly embedded into KV B Then, use a similar placeholder filling method to KV B and A B First, from KV B The left side extends its length to l q , to maintain KV B The KV cache of each task in can be kept consistent in the sequence length dimension, and the KV of task j j Overwrite the original row where task i is located, and fill the newly expanded position on the left side of the KV cache of other unfinished tasks with negative infinity (-inf). B Update to KV B '. Secondly, A B It is also expanded to a length of l q , set all rows corresponding to task j (the rows where the original task i is located) to 1, indicating that there is no fill symbol in the task j sequence, and for the corresponding attention masks of other unfinished tasks, fill the newly expanded positions on the left with 0, indicating that the positions corresponding to the KV cache of this part of the task are fill symbols, A B Update to A B '. In A B Filling 0 values ​​in corresponds to filling negative infinity in KV, indicating that this part of the position is a placeholder symbol introduced by the filling operation that has no meaning for the calculation. After the shaping and embedding are completed, the three tensors required for the model calculation are input to complete the next round of iterative calculation.

[0038] Until a task in the batch completes reasoning, the same tensor reshaping operation is performed to support the insertion of new tasks; or if there are no tasks that can be processed in the task pool, no operation is performed according to the most traditional reasoning method, and all resources are released after all tasks in the batch have completed reasoning.

[0039] 3. Resource release

[0040] For case 1, there will be placeholder data of indefinite length at the front end of the KV vector corresponding to the newly inserted task. Then, after all positions in the batch have been updated once, there will be placeholder data of a certain length at the front end of the cache corresponding to each task in the KV tensor. This part of the placeholder does not contribute to the calculation, so the space occupied by some placeholders can be released under the premise of ensuring dimension alignment. Assume that the length of the placeholder at the front end of each task is O i , then the KV cache tensor is in the first Min(O i ) elements can be released, such as Figure 2 (a). Correspondingly, the position corresponding to the attention mask tensor needs to be released, such as Figure 2 (b) as shown.

[0041] Although the specific embodiments of the present invention are disclosed for the purpose of illustration, the purpose is to help understand the content of the present invention and implement it accordingly, those skilled in the art will understand that various substitutions, changes and modifications are possible without departing from the spirit and scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the content disclosed in the best embodiment, and the scope of the present invention is subject to the scope defined in the claims.< / eos>

Claims

1. A method for improving the efficiency of batch reasoning of a large language model, comprising: 1) Use the large language model to perform pre-filling processing on all pending inference tasks, and obtain the key vector and value vector of each word in the sequence of each inference task i and a word t i , the key vector and value vector of each word in the sequence of reasoning task i are recorded as key-value vector KV i , the key-value vector KV i and word t i Cache into the task pool in memory; 2) In the first round of iterative calculation, the latest word element corresponding to each reasoning task selected from the n reasoning tasks in the task pool is composed into a tensor T with a dimension of n*1 B , the key-value vectors corresponding to the selected n inference tasks are combined into a tensor KV B , the attention mask vectors corresponding to the selected n inference tasks are combined into a tensor A with a dimension of n*l B , l is the tensor KV B The length of the tensor T B , KV B and A B Input the large language model to perform reasoning calculation to obtain the new word T corresponding to each reasoning task b in the n reasoning tasks b ', and the new word T b 'The key vector and value vector between the existing word in the sequence of the corresponding reasoning task b are recorded as the key-value vector kv b '; The key value vector kv b 'Append to tensor KV B Then form the latest tensor KV B , in the current tensor A B Then append a column of vectors with all values ​​1; update the new word element tensor T B , tensor A B and tensor KV B Used for the next round of iterative calculation; 3) When an iteration of a reasoning task is completed, the word units output in each iteration of the reasoning task are concatenated to obtain the reasoning result of the reasoning task.

2. The method according to claim 1, characterized in that If the inference task i in the currently processed batch finishes its iteration, it is removed from the current batch and a new inference task j is added to the current batch.

3. The method according to claim 2, characterized in that For the newly added reasoning task j, when the key-value vector KV corresponding to the reasoning task j j The length l q Less than or equal to the latest tensor KV B In the sequence dimension, length l kv When the latest word t of the inference task j is used j Directly overwrite the new word tensor T B The latest word element of reasoning task i in KV j Overwrite KV in right-aligned way B The key vector and value vector of inference task i, tensor KV B The positions of inference task i that are not covered by task j are modified to negative infinity to ensure the correctness of the calculation; tensor A is overwritten in a right-aligned manner B The right side of the row corresponding to reasoning task i q The value of the element is 1, and the values ​​of other elements in the row are overwritten to 0 to ensure that the latest word element t of the inference task j is j , Attention Relation KV j Correctly calculate the current latest word t j The key-value vector kv between the existing word units in the task sequence of reasoning task j j ' and the new word t j '.

4. The method according to claim 3, characterized in that Release the key-value vector KV corresponding to the newly added reasoning task j j The front end has a length of O i The placeholder data and its position in tensor A B The corresponding data in .

5. The method according to claim 2, characterized in that: For the newly added reasoning task j, when the key-value vector KV corresponding to the reasoning task j j The length l q Greater than the latest tensor KV B In the sequence dimension, length l kv When , we first use the latest word t of the inference task j j Directly overwrite the new word tensor T B The latest word element of the inference task i in the tensor KV B The length of l q , use KV j Overwrite the value of the row where the original inference task i is located, and fill the key value vectors of other inference tasks in the current batch with negative infinity on the left; secondly, tensor A B Expand to length l q , and transform the tensor A B The positions of the inference task j in the matrix are all assigned 1, indicating that the task j has no fill symbol; for the tensor A B The left side of the attention mask vector corresponding to other reasoning tasks in is filled with 0, indicating that the position of the key vector and value vector of the task is filled with a value, and ensuring that the attention mask vector A corresponding to the reasoning task j j The length is consistent.

6. The method according to claim 1, characterized in that The method for performing pre-filling processing on all pending inference tasks is as follows: divide the inference tasks into multiple batches according to the sequence length of the inference tasks; align the sequence length of each batch of inference tasks by filling placeholders, and generate an attention mask vector and an initialized key value vector for each aligned inference task; then, according to the original sequence S of each inference task i i , attention mask vector A i And the initialized key-value vector outputs the key-value vector KV between each word in the sequence of the inference task i i and the term t i , the key-value vector KV i and word t i Cache into the task pool in memory.

7. The method according to claim 1, characterized in that The task pool is located in the memory of the system where the large language model is located.

8. The method according to claim 1, characterized in that The large language model is a generative model based on a self-attention mechanism.

9. A server, characterized in that: The invention comprises a memory and a processor, wherein the memory stores a computer program, the computer program is configured to be executed by the processor, and the computer program comprises instructions for executing the method according to any one of claims 1 to 8.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the method according to any one of claims 1 to 8 is implemented.

Images