KV length sensing large model reasoning scheduling method and system and medium

Abstract

The invention discloses a KV length sensing large model reasoning scheduling method and system and a medium, and the method comprises the steps: deploying a KV pool for a two-stage separation type large model reasoning framework comprising a pre-filling instance and a decoding instance, when a reasoning request arrives, carrying out the pre-filling stage processing of the reasoning request through the pre-filling instance, and carrying out the decoding stage processing of the decoding instance; the key value cache generated in the pre-filling stage is sent to the KV pool; and then, the reasoning requests with similar key value lengths are organized into a batch through a batch generator, the batch is asynchronously prefetched into the pre-filling instance to serve as a candidate batch, and the candidate batch is delivered to the decoding instance through a communication link between the pre-filling instance and the decoding instance so that the decoding instance can execute the processing of the decoding stage. The method aims to construct a KV length sensing large model reasoning scheduling mechanism for two-stage separation type large model reasoning so as to provide high-throughput and low-delay large model reasoning service, the large model reasoning efficiency is improved, and the large model operation cost is reduced.

Description

Technical Field

[0001] This invention belongs to the field of large model inference, specifically relating to a KV length-aware large model inference scheduling method, system, and medium. Background Technology

[0002] In recent years, large language models (LLMs) have become one of the most active research topics in the field of artificial intelligence, and their application in various industries has reshaped the landscape of AI technology to some extent. In particular, large models, represented by ChatGPT, have achieved tremendous success in scenarios such as text generation and dialogue systems, demonstrating their unprecedentedly powerful natural language understanding capabilities. Following ChatGPT, many other well-known large models have emerged, such as OPT, BLOOM, and LLaMA, further reinforcing the consensus that larger models generally offer greater performance. Driven by this understanding, models with hundreds of billions of parameters are becoming increasingly common.

[0003] Currently, most large-scale models employ a Transformer decoder architecture, utilizing an autoregressive method to generate output sentences. Specifically, in each generation step, the large model takes the entire token sequence (including the original input tokens and previously generated tokens) as input, generates the next token, and appends it to the end of the input sequence. Driven by the autoregressive model, all subsequent tokens are generated sequentially until the sentence is generated. During this process, since the generation of each token depends on all preceding tokens, the length of the input sequence increases rapidly as tokens are continuously generated, leading to a rapid increase in computational overhead for generating each token. Because the generation of each token depends on all preceding tokens, in principle, calculating new attention weights requires repeatedly calculating the key-value (KV) vectors already calculated in previous steps, resulting in significant redundant computation. To reduce computational overhead, researchers have proposed a key-value caching (KVcache) technique, storing previously generated KV vectors in memory to trade storage for computation and reduce overall computational cost.

[0004] The large model inference process described above can be divided into two stages: (1) Prefill stage: The large model calculates and stores the KVcache of the original prompt and generates the first output word. (2) Decode stage: The large model uses the KVcache to generate output words one by one and updates the KVcache according to the key-value pairs of the newly generated words. These two stages exhibit different characteristics. Generally speaking, the prefill stage is a computationally intensive load, while the decode stage is a memory-intensive load. The prefill stage can make full use of the GPU's computing resources with a smaller batch size, while the decode stage requires a larger batch size to achieve higher throughput. However, due to the limited GPU memory capacity, it is difficult to set a large batch size for the decode stage, resulting in most inference requests spending most of the computation time in the decode stage. It is precisely because the prefill and decode stages exhibit different characteristics that the split inference architecture has emerged. In the split framework, the prefill and decode stages are not merged into one batch, and may even be scheduled on different GPUs, with the two stages executing independently. After the pre-filled instance has finished calculating, the calculated KVcache is sent to the decoding instance, which then continues the inference process.

[0005] Current large-scale model inference frameworks generally face two challenges: excessive memory consumption and low utilization of computing resources. Regarding memory consumption, taking the LLaMA-2-70B model as an example, it contains 70 billion parameters. Storing its weights in FP16 format requires 140 GB of memory, necessitating at least six RTX 3090Ti GPUs (each with 24 GB of VRAM) or two NVIDIA A100 GPUs (each with 80 GB of VRAM) for inference. Besides the memory overhead caused by model parameters, the memory consumption from KVcache is even more significant. Taking inference under the LLaMA-2-7B model as an example, assuming KVcache is stored in FP16 format, each term requires an average of approximately 521KB of VRAM. An A10 GPU with 24GB of VRAM, after loading 14GB of model parameters, only has 10GB of VRAM remaining, ideally accommodating a maximum of 20,480 terms in KVcache. Limited GPU memory capacity means that fewer inference requests can be executed simultaneously on a GPU, making it difficult to form large batches to increase the computational density of decoding, resulting in low GPU resource utilization. Regarding computational resource utilization, large models use the original prompt word and the currently generated word sequence as input to generate the next word. This sequential generation process, which involves dependencies, is difficult to generate a sufficiently dense computational load to fully utilize the GPU's computational performance under memory constraints. Batch processing techniques are widely used for large model training and inference to improve GPU utilization. By processing multiple inputs simultaneously, model parameters loaded from GPU memory can be reused for different inputs, making the overall workload more computationally intensive. Current mainstream inference frameworks all employ batch processing techniques to improve the inference efficiency of large models. For example, early static batching combined inference requests of different lengths into a single batch. As inference progresses, some requests gradually finish, causing the batch size to gradually decrease, and the GPU resource utilization to gradually decline. To address this, researchers proposed Continuous Batching, a technique that appends a new request to a batch as soon as a request in the batch ends, ensuring the batch size remains constant and thus significantly improving GPU resource utilization. However, while Continuous Batching maintains a constant batch size by continuously appending new inference requests, it still cannot guarantee full utilization of GPU resources. Current Continuous Batching techniques do not consider the key-value (KV) lengths of each request within a batch when grouping multiple requests into a single batch, resulting in requests with significantly different KV lengths being combined into one batch.Since the key-value (KV) length determines the computational cost of attention, a request containing a long cue word and having generated a large number of tokens during the decoding phase has a relatively long KV length, resulting in a higher cost for calculating the attention value and a longer time to generate a single token. Conversely, a request containing a short cue word and having generated fewer tokens during the decoding phase has a relatively short KV length and a shorter time to generate a single token. In each iteration of a batch, all requests in that batch generate one token. This means that the iteration time is based on the request that takes the longest time to generate a single token; that is, requests with shorter KV lengths must wait for requests with longer KV lengths. During this process, GPU computing resources are not fully utilized.

[0006] Regarding the optimization of inference efficiency for large models, although dynamic batching technology achieves iterative scheduling through continuous appending of requests, it still faces many challenges. The core difficulty lies in computational efficiency. Since the computational complexity of the attention mechanism is quadratic with the sequence length, long-sequence requests take significantly longer to compute than short-sequence requests. The batching mechanism requires all requests within the same batch to synchronously generate each lexical unit, forcing short-sequence requests to wait for long-sequence requests, resulting in idle computing resources. In terms of system throughput, the response times of different requests vary significantly, and long-sequence requests can block the entire batch processing flow, reducing the overall system throughput. In short, current batching technologies struggle to effectively balance the computational load of requests with different lengths, leading to underutilization of GPU computing resources. This not only limits the inference efficiency of large models but also increases system operating costs. Summary of the Invention

[0007] The technical problem to be solved by this invention is to provide a KV length-aware large model inference scheduling method, system and medium to address the above-mentioned problems in the prior art. This invention aims to build a KV length-aware large model inference scheduling mechanism for two-stage separate large model inference to provide high-throughput, low-latency large model inference services, improve the efficiency of large model inference and reduce the operating cost of large models.

[0008] To solve the above-mentioned technical problems, the technical solution adopted by the present invention is as follows: A key-value length-aware large-model inference scheduling method includes the following steps: S101, Deploy a KV pool residing in CPU memory for a two-stage separate large model inference framework including pre-filled instances and decoding instances. The pre-filled instances consist of one GPU, multiple GPUs on the same server, or multiple GPUs on different servers. The decoding instances consist of one GPU, multiple GPUs on the same server, or multiple GPUs on different servers. S102, when an inference request arrives, the pre-filling instance processes the inference request in the pre-filling phase and stores the inference request and the key-value cache generated in the pre-filling phase into the KV pool; inference requests with similar key-value lengths in the KV pool are organized into a batch, the batch is asynchronously prefetched into the pre-filling instance as a candidate batch, and the candidate batch is delivered to the decoding instance through the inter-GPU communication link for the decoding instance to perform the decoding phase processing.

[0009] Optionally, when organizing inference requests with similar key-value lengths in the KV pool into a batch, the similar key-value lengths mean that the difference in key-value lengths of the inference requests organized into a batch is less than a preset threshold.

[0010] Optionally, it also includes: obtaining the key-value length of the batch currently being decoded by the decoding instance; filtering out candidate inference requests in the KV pool that are close to the key-value length of the batch currently being decoded; whenever the decoding instance completes an inference request and releases the corresponding GPU memory space, if the computational resource consumption of the decoding instance or the size of the batch currently being decoded is greater than a preset threshold, then selecting a candidate inference request from the candidate inference requests and adding it to the batch currently being decoded; otherwise, unloading the batch currently being decoded to the swap-out buffer of the pre-filled instance and adding the unfinished inference requests back to the KV pool, triggering scheduling to deliver the candidate batch to the decoding instance through the inter-GPU communication link.

[0011] Optionally, when the decoding instance performs the decoding phase processing, the decoding instance prioritizes executing inference requests with shorter key-value lengths in the current decoding batch.

[0012] Optionally, the KV pool uses a multi-way tree structure to manage all inference requests waiting to be executed and the key-value cache generated during the pre-filling stage. The inference requests and the key-value cache generated during the pre-filling stage are stored in the leaf nodes of the multi-way tree. In the internal nodes of the multi-way tree other than the leaf nodes, each internal node is responsible for a corresponding key-value length coverage range. Except for the root node, the key-value length coverage ranges corresponding to each internal node and its sibling internal nodes with the same parent node are combined to obtain the key-value length coverage range corresponding to its parent node.

[0013] Optionally, the multi-way tree is a quadtree, where the key value length coverage of each internal node (excluding the root node) and its three sibling internal nodes with the same parent node are combined to obtain the key value length coverage of its parent node.

[0014] Optionally, each internal node in the multi-way tree has a density attribute and a time attribute. The value of the density attribute is equal to the number of inference requests under that internal node divided by the key-value length coverage range of that internal node. The time attribute is used to record the timestamp of batch formation on the subtree rooted at that internal node. Organizing inference requests with similar key-value lengths in the KV pool into a batch includes: S201, take the level where the root node of the multi-way tree is located as the current level; S202: Starting from the current layer, select the subtree with the highest density and the highest priority. The highest density means that the density attribute value of the subtree root node is the highest, and the highest priority means that the time difference between the timestamp and the current time is the largest in the time attribute of the subtree root node. S203, determine whether the selected subtree meets the limit. If the number of inference requests in the selected subtree exceeds the preset threshold and the memory space occupied by the inference requests does not exceed the video memory limit, then it is determined that the limit is met, and the inference requests in the selected subtree are organized into a batch as inference requests of similar length, and the process ends and exits; otherwise, jump to step S204. S204, determine whether the memory space occupied by the inference request in the selected subtree exceeds the video memory limit. If it exceeds the video memory limit, take the next layer of the current layer as the new current layer to continue searching downwards and jump to step S202; otherwise, jump to step S205. S205, determine if the selected subtree has insufficient inference requests. If the number of inference requests is insufficient, backtrack to the parent node of the root node of the selected subtree, find the sibling node of the root node of the selected subtree from the parent node, and supplement the inference requests from the sibling node. If the number of inference requests in the sibling node is insufficient, continue to search upwards for other adjacent internal nodes to supplement the inference requests. Organize the supplemented inference requests and the inference requests in the selected subtree together as inference requests of similar length into a batch, end and exit.

[0015] Furthermore, the present invention also provides a KV length-aware large model inference scheduling system, including interconnected microprocessors and memory, wherein the microprocessors are programmed or configured to execute the KV length-aware large model inference scheduling method.

[0016] Furthermore, the present invention also provides a computer-readable storage medium storing a computer program or instructions that are programmed or configured to execute the KV length-aware large model inference scheduling method by a processor.

[0017] In addition, the present invention also provides a computer program product, including a computer program or instructions that are programmed or configured to execute the KV length-aware large model inference scheduling method via a processor.

[0018] Compared with the prior art, the present invention can mainly achieve the following beneficial effects: (1) This invention can significantly improve inference efficiency. Generally, the length of the input request sequence varies greatly in various application scenarios, and the length of the output sequence cannot be determined in advance. If sequences of different lengths are scheduled into the same batch, the difference in sequence length causes a difference in the amount of computation required to generate a word for each request in the batch. Specifically, whether in the pre-filling or decoding stage, the computation of short sequences needs to wait for the computation of long sequences to be completed, resulting in some computing resources being idle during the waiting process. Based on dynamic batch processing, this embodiment selects requests with similar KV lengths (and therefore similar computational amounts) to schedule into the currently running batch each time. This scheduling method can effectively avoid the problem of short sequences waiting for long sequences, and improve the throughput of inference while improving resource utilization.

[0019] (2) This invention can support online services with a large number of requests. This embodiment simultaneously employs a two-stage separation architecture, pipelined parallel processing, dynamic batch processing, and memory offloading technologies to fully utilize the large memory on the host side. The application of these technologies ensures that the inference framework can serve a large number of user requests, while selecting requests with similar KV lengths from numerous candidate requests and combining them into a batch. Through this innovative design, not only is the resource utilization of the GPU significantly improved, but the memory bottleneck problem is also alleviated, which is of great significance for improving the throughput of large model inference and thus reducing the cost of expensive inference services.

[0020] (3) This invention can significantly reduce the waiting time of inference scheduling. When traditional inference frameworks use memory offloading technology to alleviate the video memory bottleneck, they generally face the problems of low PCIe bandwidth and slow speed of key-value cache transfer between CPU and GPU. In this embodiment, if the key-value cache in CPU memory needs to be transferred to the video memory of the decoding instance, it is not directly transferred from CPU memory to the video memory of the decoding instance through the PCIe interface. Instead, the key-value cache to be transferred is first asynchronously prefetched into the video memory of the prefilled instance, and then low-latency transmission is achieved through the communication link (e.g., NVLink) between the prefilled instance and the decoding instance, which minimizes the latency on the critical path of task scheduling and significantly improves the scheduling efficiency. Attached Figure Description

[0021] Figure 1 This is a schematic diagram of the basic process of the method in an embodiment of the present invention.

[0022] Figure 2 This is a schematic diagram of the two-stage split large model inference framework in an embodiment of the present invention.

[0023] Figure 3 This is a schematic diagram of the structure of the pre-filling instance and the decoding instance in the embodiments of the present invention.

[0024] Figure 4 This is a schematic diagram illustrating the scheduling principle of candidate batches and candidate requests in an embodiment of the present invention.

[0025] Figure 5 This is a schematic diagram illustrating the prefetching and eviction principle of the key-value cache in the decoding instance of this invention.

[0026] Figure 6 This is a schematic diagram illustrating the process of organizing reasoning requests into a batch in an embodiment of the present invention. Detailed Implementation

[0027] To enable those skilled in the art to better understand the technical solutions of the present invention, the technical solutions of the present invention will be further described in detail below with reference to the accompanying drawings in the embodiments of the present invention.

[0028] like Figure 1 As shown, the KV length-aware large model inference scheduling method in this embodiment includes the following steps: S101 deploys a KV pool residing in CPU memory for a two-stage, separate large-model inference framework that includes pre-filled instances and decoded instances. S102, when an inference request arrives, the pre-filling instance processes the inference request in the pre-filling phase and stores the inference request and the key-value cache generated in the pre-filling phase into the KV pool; inference requests with similar key-value lengths in the KV pool are organized into a batch, the batch is asynchronously prefetched into the pre-filling instance as a candidate batch, and the candidate batch is delivered to the decoding instance through the inter-GPU communication link (e.g., NVLink) for the decoding instance to perform the decoding phase processing.

[0029] like Figure 2As shown, the two-stage split large model inference framework is a scalable large model inference scheduling framework that supports high-throughput large model inference. The two-stage split large model inference framework mainly consists of three parts: a key-value pool, pre-filled instances, and decoding instances. Pre-filled instances consist of a single GPU, multiple GPUs on the same server, or multiple GPUs on different servers. Decoding instances also consist of a single GPU, multiple GPUs on the same server, or multiple GPUs on different servers. Specifically: The key-value pool, residing in CPU memory, stores the key-value cache (KVCache) unloaded from the GPU. The massive capacity of CPU memory ensures sufficient inference requests are available for processing. Compared to the limited video memory capacity of the GPU, the main memory managed by the CPU can be as large as several terabytes, capable of accommodating key-value caches generated from millions of tokens (thousands of inference requests). Decoding instances, like traditional inference service systems, are responsible for accepting packaged batches and generating new tokens. It should be noted that decoding instances may consist of multiple GPUs and can employ any existing parallelization strategy, such as data parallelism, tensor parallelism, and pipeline parallelism. However, how to perform parallel decoding is beyond the scope of this invention. The pipelined parallelism utilizes the memory of multiple GPUs, ensuring that the decoding instance can still support large batches even when inference requests generate large key-value caches. Pre-filled instances: Traditionally, pre-filled instances are only responsible for handling input prompts and initializing the key-value cache for the decoding instance. However, in the framework proposed in this invention, the pre-filled instance must further act as an intermediate buffer for key-value cache exchange between the decoding instance and host memory. Figure 2 As shown, the KVCache for executing requests is mostly stored in host memory. When a batch of requests is inferred, the traditional strategy is to directly migrate the batch's key-value cache from host memory to the decoding instance's video memory via the PCIe interface between the CPU and GPU. However, it is well known that PCIe bandwidth is limited. Instead, in this embodiment, the prepared batch of key-value cache is prefetched from host memory to the pre-filled instance's video memory, as shown... Figure 3 As shown, during batch scheduling execution, the key-value cache corresponding to the batch is transmitted to the decoding instance via the NVLink inter-GPU communication link. Since the bandwidth of the NVLink inter-GPU communication link is significantly higher than that of PCIe, the above innovative measures help accelerate batch scheduling execution and reduce GPU idle time.

[0030] In traditional LLM inference frameworks, the pre-filling and decoding stages are typically executed by the same GPU. The inference engine's scheduler switches between the pre-filling and decoding stages based on memory usage and request queue status to complete the entire inference process. In a split architecture, these two stages are separated and run independently on different GPU instances. All GPUs are divided into two categories: pre-filling instances and decoding instances. Due to the limited memory space of GPUs, excessive key-value caching cannot be implemented in memory. Therefore, the cluster is equipped with large memory to store a large number of key-value caches for inference requests. Both pre-filling and decoding instances use NVLink for connectivity. NVLink is a high-speed interconnect technology specifically designed for connecting NVIDIA GPUs. It allows GPUs to communicate point-to-point, bypassing the traditional PCIe bus, achieving higher bandwidth and lower latency. NVLink can be used to connect two or more GPUs for high-speed data transfer and sharing, providing higher performance and efficiency for multi-GPU systems. Compared to traditional PCIe system solutions, it offers a faster alternative for multi-GPU systems. NVLink technology enables memory and performance expansion by connecting two NVIDIA graphics cards, thereby meeting the needs of larger computing workloads.

[0031] like Figure 2As shown, the three main components (KV pool, decoding instance, and pre-filled instance) are coordinated and managed by batch generation and scheduling strategies. Specifically, when an inference request arrives, it is first processed by the pre-filled instance (step ①), and then the generated key-value cache is sent to the KV pool (step ②). The KV pool contains key-value caches for a large number of inference requests, where each inference request depends on a different length of KV (key-value pair). Generally speaking, if an inference request depends on a longer KV, the overhead of calculating attention during inference is greater, and the computation time is longer; conversely, if an inference request depends on a shorter KV, the overhead of calculating attention during inference is smaller, and the computation time is shorter. Traditional batch generators do not consider the length of the KV that each inference request depends on, and generally organize batches according to the order in which requests arrive, which leads to low GPU resource utilization. This invention organizes inference requests with similar KV lengths into one batch (step ③), ensuring that the computational overhead of each inference request within the same batch is similar, and there is no situation where one request waits for other requests, thus helping to improve GPU resource utilization. The batch to be scheduled for execution is asynchronously prefetched into the pre-filled instance (step ④), waiting to be quickly forwarded to the decoding instance via NVLink (step ⑤). The scheduler monitors the decoding instance during runtime and triggers scheduling in two scenarios. The first scenario is when the running batch generates a large number of tokens, exhausting the decoding instance's memory capacity. In this case, considering the limited size of the key-value cache for a single request and the sufficient bandwidth of PCIe, the scheduler will offload the request's key-value cache to the KV pool via the PCIe interface (step ⑥). The second scenario is when, as some requests complete inference, fewer inference requests are running in the decoding instance, and the running batch cannot fully utilize the decoding instance's computational power. In this case, the scheduler will offload the entire batch processing to the swap-out buffer (step ⑦) and deliver the candidate batch to the decoding instance (step ⑤).

[0032] Computing components, such as GPUs, possess powerful concurrency capabilities. A single inference request may not fully utilize their computational efficiency, making batching multiple inference requests into a common optimization method. Current large-model inference frameworks primarily employ sequential batching techniques to group multiple inference requests into batches based on their arrival order. As decoding progresses, some inference requests are completed successively. This method removes completed requests from the batch and dynamically adds new requests, maintaining a consistent batch size to ensure GPU computational efficiency. However, the above methods do not consider the key-value (KV) lengths required for each inference request, leading to situations where inference requests within the same batch wait for each other, reducing GPU computational efficiency. Therefore, this embodiment proposes a KV length-aware batch generation method, effectively avoiding the problem of short sequences waiting for long sequences, improving both system resource utilization and inference throughput. This embodiment organizes inference requests with similar key-value lengths in the KV pool into a single batch, thereby limiting the difference in request lengths within the same batch to alleviate the problem of short KV requests waiting for long KV requests in the same iteration. Specifically, in this embodiment, when inference requests with similar key-value lengths in the KV pool are organized into a batch, the similar key-value lengths mean that the difference in key-value lengths of the inference requests organized into a batch is less than a preset threshold. The preset threshold can be set according to actual needs, or other means can be used as long as the above constraints can be met.

[0033] Computing components, such as GPUs, possess powerful concurrency capabilities. A single inference request may not fully utilize their computational efficiency, making batch processing of multiple inference requests a common optimization method. Current large-model inference frameworks primarily employ sequential batch processing techniques to package multiple inference requests into batches according to their arrival order. As decoding progresses, some inference requests complete sequentially. This method removes completed requests from the batch and dynamically adds new requests, maintaining a consistent batch size to ensure GPU computational efficiency. However, these methods do not consider the key-value (KV) lengths required for each inference request, leading to situations where inference requests within the same batch wait for each other, thus reducing GPU computational efficiency. To address this, this embodiment proposes a KV length-aware batch generation method, effectively avoiding the problem of short sequences waiting for long sequences, thereby improving system resource utilization and inference throughput.

[0034] This embodiment also includes: obtaining the key-value length of the batch currently being decoded by the decoding instance; filtering candidate inference requests in the KV pool that are close to the key-value length of the batch currently being decoded; whenever the decoding instance completes an inference request and releases the corresponding video memory space, if the computational resource consumption of the decoding instance or the size of the batch currently being decoded is greater than a preset threshold, then selecting a candidate inference request from the candidate inference requests and adding it to the batch currently being decoded; otherwise, unloading the batch currently being decoded to the swap-out buffer of the pre-filled instance and adding the unfinished inference requests back to the KV pool, triggering scheduling to deliver the candidate batch to the decoding instance through the inter-GPU communication link.

[0035] like Figure 4 As shown, this embodiment limits the difference in request length within the same batch to a certain range to alleviate the problem of short key-value requests waiting for long key-value requests in the same iteration. This approach includes two aspects: First, in this embodiment, when generating batches, inference requests with similar key-value lengths are aggregated into one batch. Such batches will not experience requests waiting for each other during decoding. Second, this embodiment always prepares a candidate request with a KV length similar to the currently running batch in the pre-filled instance, for example... Figure 4In batch processing, inference requests are Request 1 to Request y, with a key-value length of Len. Therefore, a candidate request with a key-value length close to Len (with a difference less than a preset value) is always prepared in the pre-filled instance based on Request 1 to Request x in the KV pool. When a request in the decoding instance completes and exits, the decoding instance releases the corresponding GPU memory. The free GPU memory can then be filled by waiting candidate requests in the pre-filled instance. This "filling" mechanism ensures that the batch size in the decoding instance is large while maintaining close KV lengths for each request in the current batch. However, if a large number of requests in a batch exit due to inference completion, but no candidate request with a close KV length can be added to the batch, the GPU's computational efficiency continuously decreases as the batch shrinks. In this case, the invention prepares a new batch in the pre-filled instance to replace the small batch currently running in the decoding instance. The switching between the two batches is entirely handled by the NVLink between the pre-filled instance and the decoding instance, ensuring low-latency batch switching. Note that the new batches prepared in the pre-filling process also follow the principle of consistent KV length proposed in this invention, that is: select the batch of requests with the most similar KV length in the current system to generate the batch. After the decoding instance receives a large number of key-value caches transmitted by the pre-filling instance, the decoding instance performs the following in each iteration: 1) Check the batch size in the currently running batch. If it does not exceed the set threshold and a large number of requests end in the same iteration, then switch to a new batch; 2) There are new batches prepared in advance in the decoding instance. The principle of batching is based on the similarity of KV length (i.e., similar computational load); 3) If it is found that only a small number of requests end in this iteration, then the scheduler counts the key-value caches of the requests in the currently running batch and finds a request in the maintained KV pool that is closest to the batch in the current iteration and schedules it to the currently running batch; 4) If no such request is found in the GPU memory, considering the overhead of migrating requests from the host memory, a new batch is rescheduled. The currently running batch is put back into the GPU, and at the same time, it is searched on the host side for requests with similar computational load to be prefetched into the GPU and batched for scheduling. Based on the basic scheduling process described above, it can be seen that during each iteration of the decoding instance, a new batch is not triggered every time. Instead, the scheduler only switches to a new batch already prepared in the current GPU and restarts execution when a request that cannot find a similar key-value cache in the current video memory is encountered. This scheduling mechanism effectively reduces the situation where short sequences of computation need to wait for long sequences of computation to complete, resulting in some computing resources being idle during the waiting process, thereby greatly improving the throughput of the inference system.

[0036] Furthermore, in this embodiment, when the decoding instance performs the decoding phase, it prioritizes the inference requests with shorter key-value lengths in the current decoding batch. By prioritizing the processing of requests with shorter KV values, a larger batch size can be achieved, thereby improving the utilization of the system's computing resources.

[0037] As an optional implementation method, such as Figure 5 As shown, the KV pool in this embodiment is a distributed KV pool, and its cached key-value cache can be distributed in SSD, RAM and CPU (on-chip cache) as needed, so that the key-value cache of the distributed KV pool is obtained according to the priority of CPU (on-chip cache), RAM and SSD.

[0038] To reduce inference latency caused by loading key-value cache from host memory to GPU memory, this embodiment reserves a buffer in GPU memory, thereby overlapping the loading process with the inference computation process, allowing the computation process of the current batch and the key-value cache loading process of the next batch to occur simultaneously. To reduce latency caused by unloading key-value cache from GPU memory to host memory, this embodiment maintains a write stream during inference to continuously unload newly generated key-value cache, achieving simultaneous key-value cache unloading and inference computation through asynchronous unloading. To expand the KV length-aware inference scheduling space, the available storage space of the cached key-value cache needs to be increased. Since requests arrive randomly, their corresponding key-value caches are more likely to reside in a larger capacity of host memory, resulting in poor access performance. To reduce the impact of migrating key-value caches from host memory on inference performance, this embodiment proposes a request-aware acquisition scheme. This scheme prefetches requests that need to access key-value caches from host memory to GPU memory based on the current free space in GPU memory and the batch status of requests. Specifically, the scheduler maintains a request queue to monitor requests waiting to be executed, thereby fully understanding the requests that are waiting. If the key-value cache length of waiting requests in host memory meets the batching requirements of the current GPU memory, the system prefetches the key-value cache of waiting requests from host memory to GPU memory before executing these waiting requests. A buffer is included on the GPU to allow seamless fetching of key-value cache from host memory to GPU memory, thus preventing latency when GPU memory is full. When the capacity of free GPU memory reaches a defined threshold, the key-value cache is evicted from GPU memory to ensure the continuous availability of the GPU memory buffer. Specifically, to effectively utilize the limited GPU memory space, requests in GPU memory are monitored to fully understand the current request status on the GPU, identify the key-value caches that are least likely to be batched, and evict them from host memory. With this design, the system can handle a large number of inference requests simultaneously while expanding the scheduling selection space. Combining the large capacity of host memory and the high bandwidth of GPU memory, the storage capacity of the key-value cache is expanded, solving the problem of limited GPU memory.

[0039] In this embodiment, the KV pool uses a multi-way tree structure to manage all inference requests waiting to be executed and the key-value cache generated during the pre-filling stage. This allows for the search within the tree structure to obtain a candidate batch consisting of requests with similar KV lengths. The inference requests and the key-value cache generated during the pre-filling stage are stored in the leaf nodes of the multi-way tree. In the internal nodes of the multi-way tree, excluding the leaf nodes, each internal node is responsible for a corresponding key-value length coverage range. The key-value length coverage ranges corresponding to each internal node and its sibling internal nodes with the same parent node, excluding the root node, are combined to obtain the key-value length coverage range corresponding to its parent node.

[0040] The number of branches in a multi-branch tree can be selected as needed. For example, as an optional implementation, the multi-branch tree in this embodiment is a quadtree. Except for the root node, the key-value length coverage of each internal node and its three sibling internal nodes with the same parent node are combined to obtain the key-value length coverage of its parent node. Each internal node is responsible for a specific KV length range. When a node generates a child node, the current node's interval is also divided into four smaller intervals, and these child nodes inherit and are responsible for these subdivided intervals. Based on the quadtree structure constructed above, all inference requests to be executed are stored in the leaf nodes of the tree in order of KV length. Inference requests with similar KV lengths can be obtained simply by performing a nearest neighbor search on the tree.

[0041] During the search process, this embodiment first employs a greedy strategy to find enough inference requests to form a batch in the subtree with the highest density, thus ensuring that the key-value lengths of all requests in the generated batch have minimal differences. To achieve this, a counter `count1` needs to be maintained on each internal node to record the total number of inference requests contained in the subtree rooted at that internal node. Based on this design, this embodiment further proposes a density-first search method to traverse the entire tree. That is, when encountering an internal node during traversal, the child node with the largest counter is always selected as the next level traversal node, because the child node with the largest counter often indicates that it is the subtree with the highest request density, and the requests represented by the leaf nodes under the densest subtree are the batch of similar requests found in this embodiment. However, the above greedy strategy may lead to some inference requests under certain subtrees failing to form a batch, resulting in a "starvation" situation. To solve this problem, this embodiment sets another counter `count2` for each internal node to record the timestamp of batch formation on the subtree rooted at that node. If inference requests under a subtree have not formed a batch for a long time, they are given higher priority during the search process. In other words, it searches subtrees that have not been visited for a long time with a higher probability, thus preventing some requests from "starving" due to prolonged lack of service. This design not only efficiently locates and extracts the most densely populated batch of requests but also ensures that all requests in the system receive reasonable attention and processing, significantly improving overall resource management and scheduling efficiency. Previously, no research addressed the problem of short sequences waiting for long sequences due to different KV lengths within the same iteration step, meaning that existing inference frameworks generally suffer from wasted computational resources. The KV length-aware large-model inference scheduling method proposed in this embodiment effectively avoids the problem of short sequences waiting for long sequences, improving both system resource utilization and inference throughput. This method avoids grouping requests with large KV length differences together, reducing the waiting time between different requests in the same batch, thereby improving GPU resource utilization. In this embodiment, each internal node in the multi-way tree has a density attribute and a time attribute. The value of the density attribute is equal to the number of inference requests under that internal node divided by the key-value length coverage range of that internal node. The time attribute is used to record the timestamps of batches formed on the subtree rooted at that internal node. Therefore, if inference requests under a subtree have not formed batches for a long time, they are given higher priority during the search process. Figure 6 As shown, organizing inference requests with similar key-value lengths in the KV pool into a batch includes: S201, take the level where the root node of the multi-way tree is located as the current level; S202: Starting from the current layer, select the subtree with the highest density and the highest priority. The highest density means that the density attribute value of the subtree root node is the highest, and the highest priority means that the time difference between the timestamp and the current time is the largest in the time attribute of the subtree root node. S203, determine whether the selected subtree meets the limit. If the number of inference requests in the selected subtree exceeds the preset threshold and the memory space occupied by the inference requests does not exceed the video memory limit, then it is determined that the limit is met, and the inference requests in the selected subtree are organized into a batch as inference requests of similar length, and the process ends and exits; otherwise, jump to step S204. S204, determine whether the memory space occupied by the inference request in the selected subtree exceeds the video memory limit. If it exceeds the video memory limit, take the next layer of the current layer as the new current layer to continue searching downwards and jump to step S202; otherwise, jump to step S205. S205, determine if the selected subtree has insufficient inference requests. If the number of inference requests is insufficient, backtrack to the parent node of the root node of the selected subtree, find the sibling node of the root node of the selected subtree from the parent node, and supplement the inference requests from the sibling node. If the number of inference requests in the sibling node is insufficient, continue to search upwards for other adjacent internal nodes to supplement the inference requests. Organize the supplemented inference requests and the inference requests in the selected subtree together as inference requests of similar length into a batch, end and exit.

[0042] When determining whether the memory space occupied by inference requests in the selected subtree exceeds the GPU memory limit, it can be determined whether the memory space occupied by inference requests in the selected subtree exceeds a preset threshold. This preset threshold can be a parameter value that is either the GPU memory size or lower than the GPU memory size. All requests under the optimal subtree are collected, and the memory space required by these requests is calculated. There are three cases: Case ①: If the memory space occupied by all requests is lower than a specific threshold (e.g., the GPU memory size), and the number of requests is greater than the specific threshold, then this batch is directly returned as the best batch found. Case ②: If the memory space occupied by all requests exceeds a specific threshold (e.g., exceeding the GPU memory size), then return to step S202 to continue searching to the next lower level node until the condition is met. Case ③: If the number of requests is insufficient, and the memory space occupied by these requests is significantly lower than the given threshold (e.g., the GPU memory size), then enter the supplementation stage, as detailed in step S205.

[0043] The pre-filling and decoding phases together constitute the core process of large-scale model inference. The pre-filling phase is responsible for generating key-value caches from the input sequence, while the decoding phase is responsible for efficiently generating the output sequence. They work together to achieve efficient and accurate large-scale model inference tasks. In this process, efficient task scheduling is crucial, directly impacting key performance indicators such as resource utilization, inference latency, and system throughput. Specifically, traditional inference frameworks employ the following process: For the pre-filling phase, the input text is first segmented and encoded into vectors. Then, the large model calculates the key-value caches for all input tokens through forward propagation. This embodiment prioritizes scheduling the first arriving requests to the pre-filling instance, calculating the initial key-value caches for the input tokens, and storing the generated key-value caches in a KV pool for use in the subsequent decoding phase. For the decoding phase, based on the results of the pre-filling phase, the large model gradually generates output tokens, feeding them back to the model after each generation, iterating until completion. Due to the order constraint of autoregressive generation during large-scale model inference, the decoding phase generates only one token per iteration. Before each iteration, the scheduler checks the status of all requests in the decoding instance. If a request in the batch has been inferred, the key-value cache of the completed request is released, and the newly arrived request is added to the batch. The current approach does not consider the length of the request when adding a new request to the batch. However, this approach leads to the problem of short requests waiting for long requests, resulting in a waste of system resources. To address the above problem, the KV length-aware large model inference scheduling method in this embodiment dynamically schedules the key-value cache in the KV pool by analyzing the batch processing situation in the decoding stage, thereby achieving adaptive optimization of task scheduling. The scheduling strategy proposed in this embodiment is triggered in two situations: (1) Assuming that there is a batch in the decoding instance that is performing inference, some requests exit due to completion in a certain iteration step, releasing a certain amount of GPU memory space. At this time, the scheduler schedules several inference requests from the GPU memory of the pre-filled instance to the decoding instance, ensuring that the batch being executed does not become smaller, thereby ensuring that the computational efficiency of the GPU does not decrease. In this embodiment, the key-value cache of the request most likely to be added to the current batch is prefetched into the GPU memory of the prefilled node based on the monitoring of the batch running in the decoding instance. When the request on the decoding instance finishes and exits, releasing the space, the candidate KV length similar requests are quickly migrated to the decoding instance through NVLink. (2) For the first scenario above, if a request exits in the batch being executed, and no request of similar length can be found in the KV pool to fill the GPU memory space released in the decoding instance, the batch being executed continues to shrink, and the computing efficiency of the GPU continues to decline. At this time, the batch running in the decoding instance cannot fully utilize the computing power of the decoding instance. The scheduler will release the space of the decoding instance by reverse migration of the current batch to the swap-out buffer on the prefilled instance (step ⑦), and then migrate the candidate batch on the prefilled instance to the decoding instance (step ⑤).Candidate batches are generated by the batch generator before scheduling. Note that whether it is a small number of requests scheduled in the first case or a complete batch scheduled in the second case, both are generated by the scheduler triggering the batch generator to package them from a large number of requests in the host memory. Figure 1 Steps ③ and ④ in the process, the selection of candidate requests and batches, all follow the principle of similar KV lengths. In summary, this embodiment introduces a KV length-aware large-model inference scheduling method for the pre-filling and decoding stages, achieving efficient management of inference requests. During the pre-filling stage, the key-value cache of requests is continuously calculated and stored in the KV pool, providing selection space for the subsequent decoding stage. During the decoding stage, the model gradually generates output tokens through autoregression, while the scheduler monitors the batch processing status in real time, prioritizing requests with similar KV lengths for batch processing, effectively avoiding resource waste caused by short requests waiting for long requests. Furthermore, this embodiment asynchronously prefetches upcoming inference requests into the GPU memory of the pre-filling instance, and then quickly transmits them to the decoding instance for inference via NVLink, significantly reducing task scheduling latency and achieving overlap between computation and communication, thereby further improving system throughput.

[0044] Furthermore, this embodiment also provides a KV length-aware large model inference scheduling system, including an interconnected microprocessor and a memory, wherein the microprocessor is programmed or configured to execute the KV length-aware large model inference scheduling method.

[0045] Furthermore, this embodiment also provides a computer-readable storage medium storing a computer program or instructions that are programmed or configured to execute the KV length-aware large model inference scheduling method via a processor.

[0046] In addition, this embodiment also provides a computer program product, including a computer program or instructions that are programmed or configured to execute the KV length-aware large model inference scheduling method via a processor.

[0047] Those skilled in the art will understand that the technical solutions provided by this invention may take the form of a method, system, or computer program product. Therefore, this invention may take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, this invention may take the form of a computer program product embodied on one or more computer-readable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code. This invention is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the invention. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, produce an implementation of the flowchart... Figure 1 One or more processes and / or boxes Figure 1 The computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to operate in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The functions specified in one or more boxes. These computer program instructions may also be loaded onto a computer or other programmable data processing apparatus to cause a series of operational steps to be performed on the computer or other programmable apparatus to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable apparatus for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0048] The above description is merely a preferred embodiment of the present invention. The scope of protection of the present invention is not limited to the above embodiments. All technical solutions falling within the scope of the present invention's concept are within the scope of protection of the present invention. It should be noted that for those skilled in the art, any improvements and modifications made without departing from the principles of the present invention should also be considered within the scope of protection of the present invention.

Claims

1. A KV length-aware large-model inference scheduling method, characterized in that, Includes the following steps: S101, Deploy a KV pool residing in CPU memory for a two-stage separate large model inference framework including pre-filled instances and decoding instances. The pre-filled instances consist of one GPU, multiple GPUs on the same server, or multiple GPUs on different servers. The decoding instances consist of one GPU, multiple GPUs on the same server, or multiple GPUs on different servers. S102, when an inference request arrives, the pre-filling instance processes the inference request in the pre-filling phase and stores the inference request and the key-value cache generated in the pre-filling phase into the KV pool; inference requests with similar key-value lengths in the KV pool are organized into a batch, the batch is asynchronously prefetched into the pre-filling instance as a candidate batch, and the candidate batch is delivered to the decoding instance through the inter-GPU communication link for the decoding instance to perform the decoding phase processing.

2. The KV length-aware large model inference scheduling method according to claim 1, characterized in that, When organizing inference requests with similar key-value lengths in the KV pool into a batch, "similar key-value lengths" means that the difference in key-value lengths among the inference requests organized into a batch is less than a preset threshold.

3. The KV length-aware large model inference scheduling method according to claim 1, characterized in that, Also includes: Get the key-value length of the batch currently being decoded by the decoding instance; filter out candidate inference requests from the KV pool that are close to the key-value length of the batch currently being decoded; Whenever a decoding instance completes an inference request and releases the corresponding GPU memory, if the computational resource consumption of the decoding instance or the size of the current decoding batch is greater than a preset threshold, a candidate inference request is selected from the candidate inference requests and added to the current decoding batch; otherwise, the current decoding batch is unloaded to the swap-out buffer of the pre-filled instance and the unfinished inference requests are added back to the KV pool, triggering scheduling to deliver the candidate batch to the decoding instance through the inter-GPU communication link.

4. The KV length-aware large model inference scheduling method according to claim 3, characterized in that, When the decoding instance performs the decoding phase processing, it prioritizes executing inference requests with shorter key-value lengths in the current decoding batch.

5. The KV length-aware large model inference scheduling method according to claim 1, characterized in that, The KV pool uses a multi-way tree structure to manage all inference requests waiting to be executed and the key-value cache generated during the pre-filling phase. The inference requests and the key-value cache generated during the pre-filling phase are stored in the leaf nodes of the multi-way tree. In the internal nodes of the multi-way tree other than the leaf nodes, each internal node is responsible for a corresponding key-value length coverage range. The key-value length coverage ranges corresponding to each internal node and its sibling internal nodes with the same parent node are combined to obtain the key-value length coverage range corresponding to its parent node.

6. The KV length-aware large model inference scheduling method according to claim 5, characterized in that, The multi-way tree is a quadtree. Except for the root node, the key value length coverage of each internal node and its three sibling internal nodes with the same parent node are combined to obtain the key value length coverage of its parent node.

7. The KV length-aware large model inference scheduling method according to claim 5, characterized in that, Each internal node in the multi-way tree has a density attribute and a time attribute. The value of the density attribute is equal to the number of inference requests under that internal node divided by the key value length coverage range of that internal node. The time attribute is used to record the timestamps of batches formed on the subtree rooted at that internal node. The step of organizing inference requests with similar key-value lengths in the KV pool into a batch includes: S201, take the level where the root node of the multi-way tree is located as the current level; S202: Starting from the current layer, select the subtree with the highest density and the highest priority. The highest density means that the density attribute value of the subtree root node is the highest, and the highest priority means that the time difference between the timestamp and the current time is the largest in the time attribute of the subtree root node. S203, determine whether the selected subtree meets the limit. If the number of inference requests in the selected subtree exceeds the preset threshold and the memory space occupied by the inference requests does not exceed the video memory limit, then it is determined that the limit is met, and the inference requests in the selected subtree are organized into a batch as inference requests of similar length, and the process ends and exits; otherwise, jump to step S204. S204, determine whether the memory space occupied by the inference request in the selected subtree exceeds the video memory limit. If it exceeds the video memory limit, take the next layer of the current layer as the new current layer to continue searching downwards and jump to step S202; otherwise, jump to step S205. S205, determine if the selected subtree has insufficient inference requests. If the number of inference requests is insufficient, backtrack to the parent node of the root node of the selected subtree, find the sibling node of the root node of the selected subtree from the parent node, and supplement the inference requests from the sibling node. If the number of inference requests in the sibling node is insufficient, continue to search upwards for other adjacent internal nodes to supplement the inference requests. Organize the supplemented inference requests and the inference requests in the selected subtree together as inference requests of similar length into a batch, end and exit.

8. A key-value length-aware large-model inference scheduling system, comprising interconnected microprocessors and memory, characterized in that, The microprocessor is programmed or configured to execute the KV length-aware large model inference scheduling method according to any one of claims 1 to 7.

9. A computer-readable storage medium storing a computer program or instructions, characterized in that, The computer program or instructions are programmed or configured to execute the KV length-aware large model inference scheduling method of any one of claims 1 to 7 via a processor.

10. A computer program product, comprising a computer program or instructions, characterized in that, The computer program or instructions are programmed or configured to execute the KV length-aware large model inference scheduling method of any one of claims 1 to 7 via a processor.

Images