A processing method for improving batch inference efficiency of a large language model
By using PD separation and task tensor shaping, the problem of idle computation in batch inference of large language models is solved, achieving efficient task processing and resource utilization, and improving the inference efficiency of large models.
Patent Information
- Application Number
- CN202411401959.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-10-09
- Publication Date
- 2025-11-11
- Estimated Expiration
- 2044-10-09
AI Technical Summary
The batch inference process of Large Language Models (LLM) suffers from idle computation, which leads to wasted computing resources and increased task processing latency. Existing solutions are unable to effectively improve inference efficiency.
By employing PD separation and task tensor shaping, and caching the task key vectors and attention mask vectors, flexible task replacement and zero-idle computation are achieved, thereby improving the inference efficiency and resource utilization of large models.
Without increasing additional resource overhead, it achieves flexible task processing and near-zero idle computation during the batch inference process of large models, improving task processing efficiency and computing resource utilization.
Smart Images
Figure CN119558398B_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of computer software technology and relates to a processing method for improving the efficiency of batch reasoning of large language models. Background Technology
[0002] Large models play a crucial role in modern technological environments. Service providers perform large model inference calculations based on user requests and return the results; computational efficiency determines service quality. Typically, to fully utilize the parallel processing capabilities of computing devices such as GPUs, inference tasks are performed in batches. However, Large Language Model (LLM) inference exhibits autoregressive properties, meaning each inference task requires an indefinite number of iterative calculations to generate the final result. In current run-to-completion batch processing, tasks with fewer iterations continue to consume computational resources after completion, unnecessarily iterating with other incomplete tasks until all tasks in the batch are finished. This idle computation not only wastes computational resources but also increases task processing latency, ultimately reducing service quality. Therefore, improving the inference efficiency of large models, considering the autoregressive inference characteristics of LLMs and the needs of batch processing, has become a key technical challenge in this field. Traditional solutions for improving large model inference are as follows:
[0003] 1. PD Separation Reasoning
[0004] Large-scale model inference involves two phases: the first is the prefilling phase, which initializes the input context and caches intermediate states, ending with the first output term; the second is the decoding phase, which generates the output step-by-step based on the context information, one term per iteration until the inference task ends. The computational cost of the prefilling phase is directly proportional to the initial sequence length of the input model, while the computational cost of each iteration in the decoding phase is fixed thanks to the key-value caching mechanism. In batch processing, the computational cost of all tasks in the prefilling phase is consistent with the cost of the longest sequence task within the batch, resulting in unnecessary processing latency. To address this issue, 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 varying iteration rounds within the batch during the decoding phase, meaning that tasks with shorter output sequences may have idle computations, leading to resource waste.
[0005] 2. Model Inference Service
[0006] TensorFlow Serving and Triton Inference Server are widely used production-grade inference service systems in industry. As abstraction layers above the execution engine, they lack model-specific optimizations. Existing solutions such as Clipper, Clockwork, Shepherd, and Nexus focus on serving relatively small models in clusters, such as traditional deep neural network models without autoregressive behavior like ResNet. Due to the lack of awareness of inference latency for autoregressive models, these methods are difficult to apply to LLM inference services. Orca is a state-of-the-art strategy for autoregressive generation patterns in LLM. To eliminate the requirement of ensuring consistent tensor dimensions for each task in batch inference, this method dynamically sets the number of self-attention layer replicas based on the batch size. Each replica independently computes each task and completes the computation of linear layers through vector concatenation. Orca breaks through the limitation of batch inference running to completion, supporting the deletion and insertion of batch tasks during inference. Building on this, FastServe supplements it with a priority-based task update strategy. However, these solutions have the following problems: First, the self-attention layer has a large number of parameters, and setting up replicas will cause additional resource overhead and make it difficult to support large batch sizes; second, the computation of each attention replica is synchronous, meaning that all replicas must complete their computations before subsequent linear computations can be executed. This results in other tasks' decoding phases being blocked each time a new task is added due to the pre-filling process of that task, which is equivalent to idle computation. Summary of the Invention
[0007] To address the problem of idle computation in existing batch inference processes for large language models, this invention aims to provide a method to improve the efficiency of batch inference for large language models. This invention further improves the utilization of computing resources and task processing efficiency, and can be widely applied in the field of inference services based on large language models.
[0008] The technical solution of this invention is as follows:
[0009] A method for improving the efficiency of batch inference in large language models includes the following steps:
[0010] 1) Use a large language model to perform pre-filling processing on all the reasoning tasks to be processed, to obtain the key vector and value vector of each word in the sequence of each reasoning task i, as well as the word t. i Let the key vector and value vector of each word in the sequence of reasoning task i be denoted as the key-value vector KV. i , key-value vector KV i and word element t i Cache the task in the memory pool;
[0011] 2) During the first round of iteration calculation, the latest word units corresponding to each of the n inference tasks selected from the task pool will be used to form a tensor T with a dimension of n*1. B The key-value vectors corresponding to the selected n inference tasks are combined to form a tensor KV. B The attention mask vectors corresponding to the selected n inference tasks are used to form a tensor A of dimension n*l. B l is the tensor KV B The length of the tensor T; B KV B and A B The large language model is input for inference calculation to obtain the new lexical unit T corresponding to each inference task b in the n inference tasks. b ', and the new word 'T' b The key vector and value vector between the given words in the sequence corresponding to reasoning task b are denoted as the key-value vector kv. b ';Transfer the key-value vector kv b 'Add to tensor KV B The latest tensor KV is then formed. B In the current tensor A B Then append a vector of all 1s to update tensor A. B ; The new word tensor T B Tensor A B and tensor KV B Used for the next round of iteration calculation;
[0012] 3) After an iteration of a reasoning task is completed, the words output by the reasoning task in each iteration are concatenated to obtain the reasoning result of the reasoning task.
[0013] Furthermore, if the inference task i in the current batch ends its 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 inference task j, when the key-value vector KV corresponding to inference task j... j Length l q Less than or equal to the latest tensor KV B In the length l of the sequence dimension kv At that time, the latest lexical unit t of the reasoning task j is used. j Directly covering the new lexical tensor T B The latest lexical unit in reasoning task i; KV j Cover KV in a right-aligned manner B The key vector and value vector of reasoning task i will be converted into tensor KV. B The positions in inference task i that were not covered by task j were modified to negative infinity to ensure the correctness of the calculation; tensor A was overwritten in a right-aligned manner.B The right l of the row corresponding to reasoning task i in the middle q If one element has a value of 1, the values of all other elements in that row are overwritten with 0 to ensure that the latest term t is based on the reasoning task j. j Attention Relationship KV j Correctly calculate the latest word t j The key-value vector kv between the existing words in the task sequence of reasoning task j and the given words. j 'and the new morpheme t j '.
[0015] Furthermore, release the key-value vector KV corresponding to the newly added inference task j. j The front end has a length of O i Placeholder data and its position in tensor A B The corresponding data in the database.
[0016] Furthermore, for the newly added inference task j, when the key-value vector KV corresponding to inference task j... j Length l q Greater than the latest tensor KV B In the length l of the sequence dimension kv At that time, the latest lexical unit t of the reasoning task j is first used. j Directly covering the new lexical tensor T B The latest lexical unit of reasoning task i will be the tensor KV. B The length is expanded to l q Using KV j Overwrite the values in the row containing the original inference task i, and fill the left side of the key-value vectors of other inference tasks in the current batch with negative infinity; secondly, tensor A... B Expand to a length of l q and tensor A B In the inference task j, all positions are assigned a value of 1, indicating that task j has no filler symbols; for tensor A B The attention mask vectors corresponding to other inference tasks are padded with 0s on the left, indicating that the positions of the key and value vectors of the task are padded values, and ensuring that the attention mask vector A corresponding to inference task j is padded with 0s. j The lengths are consistent.
[0017] Furthermore, the pre-padded method for all pending inference tasks is as follows: The inference tasks are divided into multiple batches based on their sequence lengths; the sequence lengths of each batch of inference tasks are aligned using placeholder padding, and an attention mask vector and an initialized key-value vector are generated for each aligned inference task; then, based on the original sequence S of each inference task i... i Attention mask vector A iThe initialized key-value vector outputs the key-value vector KV between each word in the sequence of inference task i. i And the word element t i , key-value vector KV i and word element t i The task is cached in the memory-based task pool.
[0018] Furthermore, the task pool is located in the memory of the system where the large language model resides.
[0019] Furthermore, the large language model is a generative model based on a self-attention mechanism.
[0020] A server is characterized by comprising a memory and a processor, the memory storing a computer program configured to be executed by the processor, the computer program including instructions for performing the methods described above.
[0021] A computer-readable storage medium having a computer program stored thereon, characterized in that the computer program implements the above-described method when executed by a processor.
[0022] The advantages of this invention are as follows:
[0023] This invention supports flexible replacement of tasks within a batch being processed without incurring additional resource overhead through PD separation and task tensor shaping, enabling large-scale model batch inference with near-zero idle computation, thereby improving task processing efficiency and computing resource utilization. Attached Figure Description
[0024] Figure 1 This is a schematic diagram of tensor shaping.
[0025] Figure 2 This is a diagram illustrating resource release.
[0026] (a) KV cache release, (b) attention mask resource release. Detailed Implementation
[0027] The present invention will now be described in further detail with reference to the accompanying drawings. The examples given are only for explaining the present invention and are not intended to limit the scope of the present invention.
[0028] The main steps of this invention include:
[0029] 1. PD separation
[0030] Large model inference tasks originate from user inference requests, such as questions asked by users to ChatGPT; each question constitutes a inference task. During large model inference computation, to fully utilize the parallel computing capabilities of the GPU, several tasks are typically grouped into a batch for parallel computation. The number of tasks within a batch is called the batch size. The inference system usually has a default batch size, such as 5 tasks per batch.
[0031] Pre-padding is performed on all tasks that need to be processed. Since the computational cost of pre-processing is proportional to the initial sequence length of the tasks, tasks with similar sequence lengths can be grouped into batches. Given that the model's computational requirements dictate that the vector dimension of each task within each batch is known, the sequences can be aligned using placeholder padding, and an attention mask vector can be generated for each task. The attention mask vector is a 0-1 vector with the same length as the sequence, filled with 0s and unfilled bits with 1s. Furthermore, the batch size can be dynamically adjusted based on the sequence length and computational resources within the current batch, maximizing the batch size without causing memory overflow, thereby improving the efficiency of pre-padding.
[0032] Suppose the original sequence of task i is S i The attention mask is A i After pre-padding, the Key and Value vectors used to calculate attention weights for the terms in the sequence are stored in a KV cache, denoted as KV. i And to obtain a new output term, denoted as T. i T i and KV i They will be stored in the task pool. To conserve video memory, the task pool can be stored in memory.
[0033] 2. Tensor Shaping
[0034] Tensor shaping, such as Figure 1 As shown, the initial state involves selecting a batch size of tasks from the task pool and assembling their corresponding key-value caches and sequences into a batch processing tensor. Since the tasks have already been pre-filled, it is only necessary to assemble the latest single-word tensor T from the selected n tasks into a sequence tensor T. B Cache key-value pairs B Tensor and Attention Mask Tensor A B Input model, where A B Let n*l be a 0-1 tensor, and l be a KV tensor. B The length of the tensor is either 1 or 0 depending on whether it is padded. Through inference and calculation using a large model, a new word tensor T with dimension n*1 can be obtained.B (For each task, output a new lexical unit to form a new lexical tensor T) B ), and the new lexical tensor T B Keys and values kv between existing terms in an existing task sequence B Among them, kv B It will be appended to the current KV cache tensor to form the latest KV. B The computation is used for the next iteration. Furthermore, in the current attention mask tensor A... B A vector consisting entirely of 1s is appended to represent the new lexical tensor T. B This is a non-padding character. At this point, KV... B and A B The lengths of all become l+1. Then, T B The corresponding attention mask tensor A B and KV B This will serve as the model input for a new iteration, repeatedly performing inference calculations. When a task's output terminology within a batch is a predefined inference termination identifier (…),… <eos>When the result is complete, it indicates that the reasoning task has ended and the accumulated output words can be concatenated into a complete reasoning result and returned.
[0035] Assuming that task i in a batch completes and returns its result, a new task j can be added to the batch. The core of this operation is that the value corresponding to the returned task in the key-value cache tensor can be overwritten. We combine the key-value pairs of task j... j The cache reshapes the batch of key-value cache tensors to make the key-value pairs... j It can embed and override the cached KV content of task i that has already been returned. i Large models are unaware of tensor reshaping and can continuously perform inference during the decoding phase. Specific reshaping can be divided into two cases.
[0036] Case 1: When KV j Length l q Less than or equal to the current batch KV B The length l of the tensor in the sequence dimension kv At this time, the latest word for task j is t. j It can directly cover T B Medium task t i The value of T B Updated to T B ';KV j Cover KV in a right-aligned manner B The values at positions not covered in the original key-value vector of task i need to be modified to negative infinity (-inf) to ensure the correctness of the calculation. B Updated to KV B The attention mask tensor has a dimension of 2, and the attention mask for each task is represented by a single row. For A... B The update overwrites the right l of the line containing task i in a right-aligned manner. q One element is 1 (indicating that the position corresponding to task j is an existing word in the sequence), and the part that was not covered (left l) kv -l q The element is overwritten with 0 (indicating that the position is a fill sign), A B Updated to A B This ensures that the large model can adapt to the word t input to task j. j KV j t was calculated correctly. j The key values of existing words in the sequence and the new word t j From a batch perspective, the updated three tensor variables will be input into the large model to complete a new round of inference calculations.
[0037] Scenario 2: When KV j Length l q Greater than the current batch KV B The length l of the tensor in the sequence dimension kv Lexical tensor T B The processing remains 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 Therefore, a placeholder-like filling method is used for the KV (Key-Value) database. B and A B Perform plastic surgery. First, from KV... B The left side extends its length to l q To maintain KV B The key-value caches for each task can maintain consistency in the sequence length dimension, and the key-value pairs of task j can be stored in the cache. j Overwrite the line containing task i, and fill the newly expanded positions on the left side of the KV cache for other incomplete tasks with negative infinity (-inf). B Updated to KV B '. Secondly, A B It was also expanded to a length of l. q Set all rows corresponding to task j (the original row containing task i) to 1, indicating that there are no padding symbols in the task j sequence. For the attention masks corresponding to other incomplete tasks, fill the newly expanded positions on the left with 0, indicating that the positions corresponding to the KV cache of these tasks are padding symbols. B Updated to A B '。 In A B The padding with zero values corresponds to padding with negative infinity in the KV, indicating that these positions are placeholders introduced by the padding operation and have no computational significance. After shaping and embedding are completed, the three tensors required for model computation are input to complete the next round of iterations.
[0038] Until a task in the batch completes inference, the same tensor shaping operation is performed to support the insertion of new tasks; or if there are no tasks to process in the task pool, no operation is performed according to the most traditional inference method, and all resources are released after all tasks in the batch have completed inference.
[0039] 3. Resource Release
[0040] For scenario 1, the newly inserted task's key-value vector will have placeholder data of variable length at the front. Therefore, after all positions within the batch have been updated, each task's cached front within the key-value tensor will contain placeholder data of a certain length. These placeholders do not contribute to the computation, so some of the space occupied by the placeholders can be released while maintaining dimension alignment. Assume the length of each task's front-end placeholder is O. i Then the KV cache tensor has the first Min(O) in the sequence dimension. i ) elements can be released, such as Figure 2 As shown in (a). Correspondingly, the positions corresponding to the attention mask tensor need to be released, such as... Figure 2 As shown in (b).
[0041] Although specific embodiments of the invention have been disclosed for illustrative purposes to aid in understanding and implementing the invention, those skilled in the art will understand that various substitutions, variations, and modifications are possible without departing from the spirit and scope of the invention and the appended claims. Therefore, the invention should not be limited to the content disclosed in the preferred embodiments, and the scope of protection claimed by the invention is defined by the claims.< / eos>
Claims
1. A method for improving the efficiency of batch inference in large language models, comprising the following steps: 1) Use a large language model to perform pre-filling processing on all the reasoning tasks to be processed, to obtain the key vector and value vector of each word in the sequence of each reasoning task i, as well as the word t. i Let the key vector and value vector of each word in the sequence of reasoning task i be denoted as the key-value vector KV. i , key-value vector KV i and word element t i Cache the task in the memory pool; 2) During the first round of iteration calculation, the latest word units corresponding to each of the n inference tasks selected from the task pool will be used to form a tensor T with a dimension of n*1. B The key-value vectors corresponding to the selected n inference tasks are combined to form a tensor KV. B The attention mask vectors corresponding to the selected n inference tasks are used to form a tensor A of dimension n*l. B l is the tensor KV B The length of the tensor T; B KV B and A B The large language model is input for inference calculation to obtain the new lexical unit T corresponding to each inference task b in the n inference tasks. b ', and the new word 'T' b The key vector and value vector between the given words in the sequence corresponding to reasoning task b are denoted as the key-value vector kv. b ';Transfer the key-value vector kv b 'Add to tensor KV B The latest tensor KV is then formed. B In the current tensor A B Append a vector consisting entirely of 1s; update the new lexical tensor T. B Tensor A B and tensor KV B Used for the next round of iteration calculation; 3) After an iteration of a reasoning task is completed, the word units output by the reasoning task in each iteration are concatenated to obtain the reasoning result of the reasoning task; 4) If the inference task i in the current batch ends its iteration, remove it from the current batch and add a new inference task j to the current batch; 5) For the newly added reasoning task j, when the key-value vector KV corresponding to reasoning task j... j Length l q Less than or equal to the latest tensor KV B In the length l of the sequence dimension kv At that time, the latest lexical unit t of the reasoning task j is used. j Directly covering the new lexical tensor T B The latest lexical unit in reasoning task i; KV j Cover KV in a right-aligned manner B The key vector and value vector of reasoning task i will be converted into tensor KV. B The positions in inference task i that were not covered by task j were modified to negative infinity to ensure the correctness of the calculation; tensor A was overwritten in a right-aligned manner. B The right l of the row corresponding to reasoning task i in the middle q If one element has a value of 1, the values of all other elements in that row are overwritten with 0 to ensure that the latest term t is based on the reasoning task j. j Attention Relationship KV j Correctly calculate the latest word t j The key-value vector kv between the existing words in the task sequence of reasoning task j and the given words. j 'and the new morpheme t j Then release the key-value vector KV corresponding to the newly added inference task j. j The front end has a length of O i Placeholder data and its position in tensor A B The corresponding data in; 6) For the newly added reasoning task j, when the key-value vector KV corresponding to reasoning task j... j Length l q Greater than the latest tensor KV B In the length l of the sequence dimension kv At that time, the latest lexical unit t of the reasoning task j is first used. j Directly covering the new lexical tensor T B The latest lexical unit of reasoning task i will be the tensor KV. B The length is expanded to l q Using KV j Overwrite the values in the row containing the original inference task i, and fill the left side of the key-value vectors of other inference tasks in the current batch with negative infinity; secondly, tensor A... B Expand to a length of l q and tensor A B In the inference task j, all positions are assigned a value of 1, indicating that task j has no filler symbols; for tensor A B The attention mask vectors corresponding to other inference tasks are padded with 0s on the left, indicating that the positions of the key and value vectors of the task are padded values, and ensuring that the attention mask vector A corresponding to inference task j is padded with 0s. j The lengths are consistent.
2. The method according to claim 1, characterized in that, The pre-padded method for all pending inference tasks is as follows: Divide the inference tasks into multiple batches based on their sequence length; align the sequence lengths of each batch of inference tasks using placeholder padding, and generate an attention mask vector and an initialized key-value vector for each aligned inference task; then, based on the original sequence S of each inference task i... i Attention mask vector A i The initialized key-value vector outputs the key-value vector KV between each word in the sequence of inference task i. i And the word element t i , key-value vector KV i and word element t i The task is cached in the memory-based task pool.
3. The method according to claim 1, characterized in that, The task pool is located in the memory of the system containing the large language model.
4. The method according to claim 1, characterized in that, The large language model is a generative model based on a self-attention mechanism.
5. A server, characterized in that, It includes a memory and a processor, the memory storing a computer program configured to be executed by the processor, the computer program including instructions for performing the method of any one of claims 1 to 4.
6. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the method of any one of claims 1 to 4.