Large language model reasoning method based on distributed KV cache pool

Abstract

The invention discloses a large language model reasoning method based on a distributed KV cache pool. According to the method, a platform receives requests which are sent by a large number of users and need a large language model service, the requests of the users and machines in a cluster are modeled correspondingly, and then corresponding strategies are used for processing. In addition, the use condition of machine resources in the cluster is also considered, and machines with more idle resources are preferentially considered. In this way, interference caused by resource competition is reduced to a certain extent. Meanwhile, by abstracting memories of numerous NPU cards into a distributed KV cache pool, elastic expansion and contraction are facilitated during request processing. Through the method, an efficient cluster based on a large language model can be constructed, and a corresponding platform can better understand the demands and intentions of users so as to provide more timely and personalized services.

Description

Technical Field

[0001] The present application relates to the field of computer technology, and in particular to a large language model reasoning method of a distributed KV cache pool. Background Art

[0002] Most popular large language models use the Transformer architecture. A model typically consists of a series of Transformer layers, each of which includes an attention layer and a feed-forward neural network (FFN) layer. The attention layer allows the token in a request to interact with other tokens, and the FFN layer intelligently processes the request token. In each iteration, the model needs to predict the next token. To avoid redundant computation, large language model serving systems cache intermediate states of tokens, such as key-value caches, and reuse them for future token generation. An optimization divides the entire generation process into two stages: a pre-filling stage and a decoding stage. The pre-filling stage processes all input tokens in a single iteration to build a key-value cache and generate the first output token. And for the newly generated output token, the decoding stage only needs to calculate the key-value cache. Therefore, the pre-filling stage is more computationally intensive than the decoding stage.

[0003] However, the large language model services provided by some existing platforms do not adequately meet user requirements. Specifically, large language model service systems need to handle dynamically changing inference traffic and support the unique characteristics introduced by the decoding stage, including the unique computing model of the decoding stage and key-value cache management. In addition, existing cluster setups have scalability and load balancing issues. Therefore, solving these existing problems is also an important research direction in related fields. Summary of the Invention

[0004] The purpose of the present invention is to address the deficiencies of the existing technology and provide a large language model reasoning method based on a distributed KV cache pool.

[0005] The object of the present invention is achieved through the following technical solutions: a large language model inference method based on a distributed KV cache pool, including a global administrator and several NPUs;

[0006] The global administrator includes an extended information repository, a distributor, and a scheduler; the memory of several NPU cards is combined to form a distributed KV cache pool for caching key-value tensors;

[0007] The extended information base contains tasks, monitoring data, and cluster topology; tasks correspond to user requests;

[0008] The extended information library sends tasks, monitoring data, and cluster topology to the distributor and the scheduler;

[0009] Based on the data sent by the extended information library, the scheduler generates an optimal flow, extracts the locations of processing nodes from the optimal flow, and sends the prompt words, corresponding models, and the extracted location information to the distributor;

[0010] Based on the data sent by the extended information library and the scheduler, the distributor distributes user requests to elastic instances; the elastic instances dynamically organize the instances corresponding to the NPUs into a series of non - overlapping elastic sequence parallel groups to process requests in parallel batches, with different degrees of parallelism configurations;

[0011] During the large - language model inference phase, corresponding calculations are performed on the NPU cards; the key - value generated during the inference is placed in the machine HBM of the node where the NPU is located.

[0012] Furthermore, the global administrator further includes an elastic controller; the elastic controller sorts the elastic instances to update the configurations of the elastic instances to form corresponding elastic sequence parallel groups and process requests in parallel; and expands and contracts the instances according to the resource usage of the tasks.

[0013] Furthermore, in the case of the elastic contraction phase, after the pre - filling phase of a batch processing is completed, the computational requirements for the decoding phase will decrease. In this case, the size of the parallel group is contracted to release resources for other batch processes; specifically, for a parallel group R with a corresponding degree of parallelism d, the elastic contraction mechanism needs to reduce the group to a new parallel group R' with a corresponding degree of parallelism d', where d' < d; the key - value tensors of the instances that are no longer used are transferred to the machines of the instances that are still in use, ensuring that the key - value tensors of the entire request in the parallel group R are efficiently transferred to the new parallel group R';

[0014] For the elastic expansion phase, if the generated key - value tensors exceed the capacity of the KV cache pool in the parallel group, new instances need to be added to expand the capacity; in addition, if the batch size is large enough, the calculations in the decoding phase become computationally intensive. In this case, deploying more NPUs can reduce the latency; deploying more NPUs includes: first, the newly added NPUs are other idle NPU cards in the same machine; second, the NPU cards in remote machines are considered only then.

[0015] Furthermore, it also includes a load balancing strategy, specifically: before request distribution, consider the resource load of each instance in the cluster, sort the resources or load of each node from small to large, and dynamically update the status of each node according to the use of the request, that is, evenly distribute the load pressure to each computer in the cluster.

[0016] Furthermore, it also includes a strategy to avoid over-submitting machine network bandwidth. Specifically, for network bandwidth requests, each task is connected to a request aggregator; the request aggregator has an arc for each task, and this arc will adapt to each machine with sufficient spare bandwidth; when the observed bandwidth usage changes, these arcs will dynamically adapt.

[0017] Furthermore, the distributed KV cache pool is used to flexibly store requested key-value tensors at the granularity of a token.

[0018] Furthermore, the scheduler includes a scheduling policy module, a network flow module, and a minimum-cost maximum-flow parser. Incoming data first passes through the scheduling policy module, where it is modified into a corresponding network flow. It is then submitted to the minimum-cost maximum-flow parser to generate the optimal flow. The optimal flow result is then transmitted to the network flow module, which extracts the node location of each NPU used from the optimal flow and returns the extracted location to the scheduling policy module.

[0019] The vertices in the network flow represent requests, machines, sinks, unscheduled aggregation nodes, aggregator nodes, racks, and request aggregators, and the directed edges connecting the vertices represent costs.

[0020] Furthermore, the scheduling policy is any one of a group scheduling policy, a fair share scheduling policy, a queue scheduling policy, a preemptive scheduling policy, a topology-based scheduling policy, a recycling policy, a backfill scheduling policy, and a resource reservation scheduling policy.

[0021] The present invention also provides a large language model inference device based on a distributed KV cache pool, comprising one or more processors for implementing the above-mentioned large language model inference method.

[0022] The present invention also provides a computer-readable storage medium having a program stored thereon, which, when executed by a processor, is used to implement the above-mentioned large language model inference method.

[0023] The beneficial effects of the present invention are as follows: to address the problem that the corresponding clusters in some platforms cannot respond to user requests in a timely manner or requests are not responded to, the present invention modifies tasks, monitoring data, and cluster topology into corresponding network flows. By designing a corresponding minimum cost maximum flow algorithm for this network flow structure, a load balancing strategy based on resources is implemented, and the memory of numerous NPU cards is abstracted into a distributed KV cache pool, thereby enabling user requests to be responded to in a timely manner and reducing the number of unprocessed requests. The method of the embodiments of the present disclosure can build a corresponding model, thereby providing users with a better experience. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] To more clearly illustrate the technical solutions of the embodiments of the present disclosure, the following briefly introduces the drawings required for describing the embodiments. Obviously, the drawings described below are only some exemplary embodiments of the present disclosure, and those skilled in the art can derive other drawings based on these drawings without inventive effort.

[0025] Figure 1 2 is a schematic diagram illustrating computing using a distributed KV cache pool according to an embodiment of the present disclosure.

[0026] Figure 2 FIG. 1 is a schematic diagram illustrating corresponding processing of an incoming request according to an embodiment of the present disclosure.

[0027] Figure 3 is a schematic diagram illustrating the interaction between an extension information repository and a scheduler according to an embodiment of the present disclosure.

[0028] Figure 4 FIG. 4 is a schematic diagram showing a network flow for five requests on a cluster consisting of four machines according to an embodiment of the present disclosure.

[0029] Figure 5 FIG. 4 is a network flow diagram illustrating a load balancing strategy that considers evenly allocating requests to each machine according to an embodiment of the present disclosure.

[0030] Figure 6 FIG. 4 is a network flow diagram illustrating a Quincy strategy according to an embodiment of the present disclosure.

[0031] Figure 7 is a network flow diagram illustrating a network awareness policy according to an embodiment of the present disclosure.

[0032] Figure 8 FIG. 4 is a schematic diagram illustrating load balancing of machine resource usage in a cluster according to an embodiment of the present disclosure. DETAILED DESCRIPTION

[0033] In order to make the purpose, technical solutions and advantages of the present disclosure more apparent, the exemplary embodiments of the present disclosure will be described in detail below with reference to the accompanying drawings. Obviously, the described embodiments are only part of the embodiments of the present disclosure, rather than all the embodiments of the present disclosure, and it should be understood that the present disclosure is not limited to the exemplary embodiments described herein.

[0034] In this specification and the accompanying drawings, substantially the same or similar steps and elements are denoted by the same or similar reference numerals, and repeated descriptions of these steps and elements will be omitted. At the same time, in the description of the present disclosure, the terms "first", "second", etc. are only used to distinguish the description and are not to be understood as indicating or implying relative importance or ranking.

[0035] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as those commonly understood by those skilled in the art to which the present disclosure pertains. The terms used herein are for the purpose of describing embodiments of the present invention only and are not intended to limit the present invention.

[0036] To facilitate description of the present disclosure, concepts related to the present disclosure are introduced below.

[0037] The method based on distributed NPU pooling disclosed in the present invention can be based on a large language model (LLM). A large language model (LLM) is a general model for natural language processing based on deep learning technology. It achieves understanding and generation of natural language by learning the patterns and rules of large-scale text data. LLM is highly versatile and can perform well in a wide range of tasks, rather than being limited to specific tasks. It is large in scale, and the number of parameters usually reaches billions or even more. These models include but are not limited to foreign models such as GPT-3, GPT-4, PaLM, Galactica, LLaMA, as well as domestic models such as ChatGLM and Wenxin Yiyan. With the increase in computing power, the scale and capabilities of LLM are also constantly expanding. For example, GPT-3 has 175 billion parameters, and PaLM has reached 540 billion parameters. These large models exhibit capabilities that are completely different from small models (such as BERT's 330 million parameters and GPT-2's 1.5 billion parameters), especially when processing complex tasks. For example, GPT-3 is able to perform well in few-shot learning scenarios, which is difficult for previous models to achieve. An important application of LLM is ChatGPT, a conversational application based on the GPT series that demonstrates smooth and natural human-computer interaction capabilities.

[0038] The distributed NPU pooling method disclosed in the present invention can be based on a neural processing unit (NPU). A neural processing unit is an artificial intelligence chip designed to perform artificial intelligence tasks faster than a GPU (graphics processing unit) and a CPU (computer processing unit). It reduces some of the load on the GPU and CPU by taking on small-scale repetitive processing, allowing the computer to work more efficiently when meeting artificial intelligence-driven requests. For example, an NPU can blur the background in a video call or perform object detection in video or photo editing, thereby helping the computer's GPU and CPU to keep running efficiently, so that the CPU and GPU can free up their hands to handle other tasks.

[0039] The distributed NPU pooling method disclosed herein can be based on KV Cache. KV Cache is a very common technology in LLM reasoning acceleration. This technology can improve reasoning performance by trading space for time without compromising any computational accuracy. The reasoning process of a generative model is very unique. When we give it an input text, the model will output an answer (length N), which actually executes N reasoning processes in this process. That is, GPT-type models only output a single token per reasoning. The output token is concatenated with the input tokens and then used as the input for the next reasoning, and this process is repeated until a terminator is encountered. This poses a problem - the input tokens for each reasoning process become longer, resulting in a corresponding increase in the reasoning computing power. How to solve this problem? By opening the process, the output tokens of the current round are concatenated with the input tokens and used as the input tokens for the next round, and this is repeated multiple times. It can be seen that the input data for round i+1 only has one additional token compared to the input data for round i, while everything else is the same. Therefore, the reasoning for round i+1 necessarily includes some of the calculations for round i. This is the starting point of KV Cache, which caches the reusable calculation results of the current round and directly reads the cached results in the next round of calculation. This concept is not new and is the Save and Restore concept often used in AI accelerator design.

[0040] The distributed NPU pooling method disclosed in the present invention can be based on sequence parallelism. The basic idea of ​​sequence parallelism is as follows. Before processing, the input sequence is rearranged and divided into multiple segments. And then distributed to different instances. Except for the attention layer, the instances execute without communication. In the attention layer, each instance simultaneously (1) uses the local query (q i ) and key-value (kv j ) tensor, and (2) send key-value (kv j) tensor to its neighbor instance (i+1)％n ). After multiple rounds of this process, the query tensor in each instance will interact with all key-value tensors, and then the final output will be given to the next layer. Sequence parallelism is compatible with popular attention mechanisms, such as multi-head attention, multi-query attention, and grouped query attention. It also has the same computational complexity as tensor parallelism and consumes less NPU memory for cache activation. In addition, it can also be used in combination with other parallelisms. For example, tensor parallelism can further accelerate LLM training.

[0041] In summary, the solutions provided by the embodiments of the present disclosure involve technologies such as large language models, NPU, KV Cache, and sequence parallelism. The embodiments of the present disclosure will be further described below in conjunction with the accompanying drawings.

[0042] A large language model inference method based on a distributed KV cache pool of the present invention includes a global administrator and several NPUs;

[0043] The global administrator includes an extended information repository, a distributor, and a scheduler; the memory of several NPU cards is combined to form a distributed KV cache pool for caching key-value tensors;

[0044] The extended information base contains tasks, monitoring data, and cluster topology; tasks correspond to user requests;

[0045] The extended repository sends tasks, monitoring data, and cluster topology to the dispatcher and scheduler;

[0046] The scheduler generates the optimal flow based on the data sent by the extended information library, extracts the location of the processing node from the optimal flow, and sends the prompt word, corresponding model and extracted location information to the distributor;

[0047] The distributor distributes user requests to elastic instances based on the data sent by the extension information repository and the scheduler. Elastic instances dynamically organize the corresponding NPU instances into a series of disjoint elastic sequence parallel groups to process batch requests in parallel, with different configurations of parallelism. They can also efficiently support elastic expansion and contraction without additional overhead to adapt to the changing resource requirements of requests. The NPU card memory of the elastic instance can also form a unified distributed key-value cache pool, which can be flexibly used to store the requested key-value tensors at the granularity of a token. This method reduces NPU memory fragmentation through elastic instances.

[0048] During the large language model inference phase, corresponding calculations are performed on the NPU card; the key-value pairs generated during inference are placed in the HBM of the machine at the node where the NPU is located.

[0049] In each iteration, based on the analysis results from the extended information repository, the global administrator dynamically adjusts the parallelism level of existing batches, the grouping strategy for elastic instances (short prompts are typically grouped together for processing, and long prompts are grouped together for processing), the batching and forwarding strategy for newly arrived requests (generally forwarding requests together after a certain number of incoming requests), and the key-value placement strategy (usually considered within the HBM) to improve the throughput and latency of real-time requests. Here is a brief explanation of the specific process for analyzing the extended information repository results mentioned above. First, the corresponding tasks are processed through the scheduling strategy to obtain the corresponding network flow. Then, the optimal flow is obtained through the minimum cost maximum flow algorithm. The corresponding node positions in the optimal flow are then extracted from the corresponding network flow to obtain the corresponding analysis results. During this process, the global administrator's dispatcher distributes the newly arrived requests to a series of specific elastic instances according to the global administrator's requirements. Simultaneously, based on the parallelism level specified by the global administrator and the generated scaling plan, the elastic controller updates the configuration of the elastic instances to form the corresponding elastic sequence parallel groups and process requests in parallel. The key-value pairs generated from large language model inference are placed in the HBM of the machine on the node where the NPU resides; this location efficiently supports fast data access to improve inference performance. A global administrator monitors the progress of requests, resource usage of elastic instances, and the key-value cache pool for update decisions. The scaling plan determines the number of NPUs based on the resources required by the request. If fewer resources are needed to process a request, the number of NPUs used is reduced; if more resources are required to process a request, the number of NPUs used is increased.

[0050] Figure 1 2 is a schematic diagram illustrating computing using a distributed NPU memory pool according to an embodiment of the present disclosure. The distributed KV cache pool includes a distributed NPU memory pool.

[0051] like Figure 1As shown, the input token x undergoes a QKV linear transformation. This transformation is then distributed and cached across different machines in the cluster (i.e., KV Cache blocks). For the corresponding query, the corresponding node is scheduled to perform computation, and the intermediate results are aggregated to the local node. This transformation then passes through the Feed-Forward Network. Linear transformation is a common data transformation method in machine learning. It can be used to reduce dimensionality, decouple, filter, and refine data. The linear transformation in the Transformer is unique and crucial, serving as the foundation for the successful operation of the Multi-Head Attention mechanism. Regarding the QKV linear transformation, QKV simply performs three linear transformations on X (through three different fully connected layers), resulting in the three transformed outputs of X. During the calculation, any one of these three can be designated as the QKV (of course, once designated, it cannot be changed). Once the QKV is obtained, the following formula represents the actual self-attention calculation process. QKV simply introduces trainable parameters and simultaneously transforms the feature space of X.

[0052]

[0053] The specific process of Attention calculation is as follows:

[0054] 1. Similarity calculation. Here Q is a size (n q ,d q ) matrix, K is a matrix of size (n k ,d k ) matrix, where d q =d k =d, refers to the dimension of Key or Query. The dot product result is the attention score matrix AS, which is (n q ,n k ). Therefore AS i,j It is understood as the attention score of the i-th Query on the j-th Key. The larger the value, the greater the association between the i-th Query and the j-th Key.

[0055] AttentionScore(AS)=QK T =Q×K T

[0056] 2. Scaling. In order to prevent the attention value from being too large, the Attention mechanism usually divides the result by to zoom in or out.

[0057] This scaling step is done to stabilize the training process.

[0058]

[0059] 3. Mask. Mask operations are used to block out certain elements that do not need attention. For example, in NLP problems, masking the attention paid to a token relative to the token that follows it. Mask is usually a matrix of the same dimension as the Scores matrix, containing values ​​of 0 or -∞ (0: indicates elements that are allowed to be paid attention to, -∞: indicates elements that need to be blocked, so that these elements have a weight of 0 after Softmax). In this way, when Softmax is performed, the blocked parts will be assigned a weight close to 0.

[0060] MaskedScores=ScaledAttentionScores+Mask

[0061] 4. Normalization. Input ScaledAttentionScores into Softmax to normalize the scaled attention scores into a probability distribution.

[0062]

[0063] 5. Weighted sum. In 3, we get an attention weight matrix AttentionWeight that obeys the probability distribution, with a size of (n q ,n k ). The size of the Value matrix V is (n v ,d v ), where n v =n k , it can be understood that 1 key corresponds to 1 value. Therefore, the size of the attention matrix is ​​(n q ,d v ), which can be understood as a query to find a Value. (i) Represents the Value found by the i-th Query.

[0064] Attention=AttentionWeight×V

[0065] Figure 2 is a schematic diagram illustrating corresponding processing of an incoming request according to an embodiment of the present disclosure.

[0066] like Figure 2As shown, the global administrator module includes an extended information repository, a dispatcher, a scheduler, and an elastic controller. The extended information repository contains tasks, monitoring data, and cluster topology. For a large number of requests from users, we regard it as the corresponding task in the extended information repository module of the global administrator. During the processing, corresponding monitoring data and cluster topology will be generated. During this period, the extended information repository will send the data (tasks, monitoring data, cluster topology) to the dispatcher and the scheduler. The scheduler will send the data (prompt, corresponding model, optimal flow) to the dispatcher and the elastic controller. Specifically, the elastic controller will receive these data of prompt, corresponding model, and optimal flow. The prompt corresponds to the user's requirements, the model used gives corresponding answers according to the user's requirements, and the optimal flow contains the positions of each NPU used in the node. The elastic controller will expand and contract instances accordingly according to the resource usage of the tasks. During the large language model inference stage, corresponding calculations will be performed on the NPU cards. In addition, the memories of many instance NPU cards together constitute a distributed KV cache pool.

[0067] For the elastic contraction stage, after a batch completes the pre-filling stage, the computational requirements for the decoding stage will be greatly reduced. In this case, it is usually beneficial to shrink the size of the parallel group to release resources for other batches. Specifically, for a parallel group R with a corresponding parallelism of d, the elastic contraction mechanism needs to shrink the group to a new parallel group R' with a corresponding parallelism of d', where d' < d. The main challenge is to ensure that the key-value tensors of the entire request in the parallel group R are efficiently transmitted to the new parallel group R'. To address this challenge, we will no longer transfer the key-value tensors of the instances to the machines that are still using the instances. In addition, a multi-stage migration method is adopted to overlap the decoding calculation and the KV cache transmission.

[0068] During elastic scaling, due to the lightweight computational model of the decoding phase, decoding batch execution often uses a low degree of parallelism to optimize overall efficiency. However, as the decoding process progresses, the generated key-value tensors may exceed the capacity of the KV cache pool in the parallel group, necessitating the addition of new instances to expand capacity. Furthermore, if the batch size is sufficiently large, the computation during the decoding phase may become computationally intensive. In this case, deploying more NPUs can reduce latency, but this also requires efficient scaling. When scaling a parallel group, the primary challenge is ensuring that the newly added instances can efficiently participate in the ongoing computation without incurring additional overhead. To address this challenge, we ensure that the newly added NPUs are idle NPUs in the same machine. Only then are NPUs in remote machines considered. Regardless of whether the NPUs are in the same machine or remotely, a multi-stage migration approach is used to overlap the decoding computation and KV cache transfer.

[0069] Figure 3 is a schematic diagram illustrating the interaction between an extension information repository and a scheduler according to an embodiment of the present disclosure.

[0070] like Figure 3 As shown in the figure, the extended information base includes task, monitoring data, and cluster topology modules. The scheduler includes scheduling policies, network flows, and a minimum-cost maximum-flow parser. For the scheduler, incoming data (tasks, monitoring data, and cluster topology) first passes through the scheduling policy module and is modified into the corresponding network flow. In this network flow, the vertices in the graph can be understood as requests, machines, sinks, unscheduled aggregation nodes, aggregator nodes, racks, and request aggregators. The values ​​on the directed edges connecting the vertices in the graph represent the cost (or the required resources). The dotted lines in the figure represent the optimal flow. It is then submitted to the minimum-cost maximum-flow parser to generate the optimal flow. The optimal flow result is in turn transmitted to the network flow module. After corresponding processing (i.e., extracting the node location of each NPU used from the optimal flow), the extracted location is returned to the scheduling policy module.

[0071] There are many commonly used scheduling strategies, such as: Gang-scheduling, Fair-sharescheduling, Queue scheduling, Preemptionscheduling, Topology-based scheduling, Reclaims, Backfill, and Resource Reservation.

[0072] Below, I will briefly explain the principles and scenarios of some of the listed scheduling strategies. Gang Scheduling: Gang scheduling (also known as coscheduling) is a scheduling strategy used to ensure that all processes in a job receive resources simultaneously. If a job cannot be scheduled as a whole, the entire job is postponed until all resources are available. This strategy is commonly used in parallel computing, where a group of processes must run simultaneously to complete a task. Gang Scheduling: It is suitable for scenarios where multiple tasks need to work together, such as high-performance computing (HPC) and certain distributed computing frameworks, where a group of related tasks must run simultaneously, such as machine learning training jobs or large-scale data processing jobs. Fair-share Scheduling: Fair-share scheduling is a scheduling strategy for multi-user environments that aims to fairly allocate computing resources to all users. It dynamically adjusts resource allocation to ensure that each user or user group receives a reasonable share of resources. Fair-share scheduling: It is suitable for multi-user shared computing environments, such as cloud computing platforms and data centers, where it is necessary to ensure fair distribution of resources among different users or user groups. Queue Scheduling: Queue scheduling is a queue-based scheduling mechanism where jobs are placed into a queue according to arrival order or priority, and then removed from the queue and executed according to certain rules. Queue scheduling works in systems that need to process a large number of jobs, such as batch processing systems and task scheduling systems, where jobs can wait for execution until system resources are available. The principle of preemption scheduling: Preemptive scheduling is a scheduling strategy that allows high-priority tasks to preempt the CPU time of currently executing tasks. This strategy ensures that more important tasks can be processed faster. Preemption scheduling works in systems that have high real-time requirements, such as real-time operating systems and network services, where high-priority tasks need to respond quickly. The principle of topology-based scheduling: Topology-based scheduling is a scheduling strategy that takes into account the data transmission path and network topology. This strategy aims to optimize the efficiency of data transmission and reduce communication delays. Topology-based scheduling works in distributed systems that need to optimize data transmission and process network traffic, such as data center networks and stream processing systems.

[0073] The minimum-cost maximum-flow parser module uses the minimum-cost maximum-flow algorithm. I will briefly explain this algorithm below. If a node is the starting point of some links (or directed arcs) and not the end point of any link, it is a source node. If a node is the end point of some links and not the starting point of any link, it is a sink node. Each link can have a certain flow rate, which must be in the same direction as the link. The maximum flow rate allowed on a link is called the capacity of the link. Based on the current flow rate, links can be divided into three categories: 1. Full (the flow rate has reached the maximum value, i.e., equal to the capacity); 2. Empty (the flow rate is zero); 3. Partial (the flow rate is greater than zero but less than the capacity). The maximum flow problem: For a network with source and sink nodes, select a source node s and a sink node t. Require that the outgoing and incoming flows on all nodes except these two are equal. Arrange the flow rates on each link so that the outgoing flow from s and the incoming flow from t (obviously, they must be equal) reach their maximum values. Minimum cost maximum flow problem: Assume that each unit of traffic will incur a certain cost when passing through a link. Given the cost of each link, give a solution that minimizes the total cost while achieving the above maximum flow.

[0074] Figure 4 FIG. 4 is a schematic diagram showing a network flow for five requests (tasks) on a cluster consisting of four machines according to an embodiment of the present disclosure.

[0075] A network flow is a directed graph whose edges carry flows from source nodes to sink nodes. Each edge is associated with a cost and capacity that constrains the flow and specifies a route for the network flow. Figure 4 As shown in Figure 1, it shows an example of a network flow, which expresses a simple cluster scheduling problem. Each task node T i On the left side (representing the i-th task), it represents the i-th task, which is the source of a unit flow. For the optimal problem of feasible solution, all these flows must flow into the sink node (S). In order to reach the sink node (S), the flow from T i The flow can proceed through a machine node (M m ), which will schedule the task to machine m (e.g., T2 is processed on M1). Alternatively, the flow can continue to the sink node through an unscheduled aggregation node (U j For task j), this means that the task was not scheduled (T1) or was preempted if it was running.

[0076] In this example, a task's location preference is represented as the cost of the directed arc to the machine. The cost of having the task unscheduled or preempted while running is the cost of the arc to the unscheduled aggregator node (e.g., 7 for T4). Given a network flow, an MCMF solver finds a globally optimal (e.g., minimum cost) flow ( Figure 4 The optimal flow expresses the best trade-off between the unscheduled cost of tasks and their preferred locations. Task locations are extracted from the flow traced back to the task by the machine.

[0077] Figure 5 FIG. 4 is a network flow diagram illustrating a load balancing strategy that considers evenly allocating requests to each machine according to an embodiment of the present disclosure.

[0078] like Figure 5 The figure shows a trivial use of an aggregator: all tasks have arcs pointing to the cluster-wide aggregator (X). The cost of the outgoing arcs from X to each machine node is distributed proportionally to the number of tasks already running on that machine (one task on M3). The effect is that the number of tasks on a machine only increases once all other machines have at least that many tasks (e.g., in Docker SwarmKit). This strategy neither requires nor uses the full complexity of stream-based scheduling. In the MCMF algorithm, we use it to highlight specific edge cases.

[0079] Figure 6 FIG. 4 is a network flow diagram illustrating a Quincy strategy according to an embodiment of the present disclosure.

[0080] like Figure 6 As shown, Quincy's original locality-oriented strategy is described, which uses the rack aggregator (R r ) and a cluster aggregator (X) to represent data locality for batch tasks. Tasks have low-cost preference arcs pointing to machines and racks on which they have data locality. However, if their preference arc is inaccessible (e.g., T2), scheduling falls back to the cluster aggregator. This strategy is suitable for batch tasks and optimizes the trade-off between data locality, task waiting time, and preemption cost. We use it to characterize the performance of the MCMF algorithm and compare it to head-to-head with Quincy.

[0081] Figure 7 is a network flow diagram illustrating a network awareness policy according to an embodiment of the present disclosure.

[0082] like Figure 7As shown, we describe a strategy that avoids overcommitting machine network bandwidth (which degrades task response time). Each task connects to a request aggregator (RA) for its network bandwidth requests. The RA has an arc for each task that adapts to each machine with sufficient spare bandwidth (e.g., 650MB / s out of 1.25GB / s on an M2's 10G link). These arcs adapt dynamically as observed bandwidth usage changes. The cost of an arc for a machine is the sum of the request and the currently used bandwidth, which incentivizes balanced utilization. We use this strategy to illustrate its potential for making high-quality decisions and extend it to include notions of priority and additional resource dimensions (e.g., CPU / RAM).

[0083] A flow-based scheduler can use any minimum-cost maximum-flow algorithm. However, some algorithms are better suited to this scheduling problem than to other problems. Here, we explain the MCMF algorithms and empirically compare them, explaining their sometimes surprising performance. A minimum-cost maximum-flow algorithm takes as input a directed network flow G = (N, A). Each arc(i, j) ∈ A has a cost C i,j And the maximum capacity U i,j Each node i∈N also has an associated supply b(i). Informally, the MCMF algorithm must optimize the routing of traffic from all sources (task nodes T i ) to a sink node (e.g., sink node S), without exceeding the capacity limit on any edge. To understand the differences between MCMF algorithms, we need a slightly more formal definition: the goal is to find a flow f that minimizes its cost. The formulas related to mass balance and flow feasibility limits for capacity are as follows. (1)

[0085] (2)

[0087] (3)

[0089]

[0090] Some algorithms use an equivalent definition of network flow, called a residual network. In this residual network, each line with a cost C ij The arc (i, j)∈A and the maximum capacity U ij is replaced by two edges: (i,j) and (j,i). The cost of edge (i,j) is C ij , and the residual capacity is r ij =U ij -f ij。And the cost of arc (j, i) is -C ij and the corresponding residual capacity is r ji = f ij 。Feasibility constraints also apply to the residual network.

[0091] Embodiments of the present disclosure provide a performance analysis of the minimum cost maximum flow algorithm. Where N represents the number of nodes, M represents the number of arcs, C is the maximum cost of an edge, and U is the maximum capacity of an edge. Where M > N > C > U. The worst-case time complexity of the corresponding algorithms is analyzed as follows: Relaxation O(M 3 CU 2 ), Cycle canceling O(NM 2 CU), Costscaling O(N 2 M * log(NC)), Successive shortest path O(N 2 U * log(N)). The complexity indicates that the Successive shortest path should perform the best as long as U * log(N) < M * log(NC) and U << M and C ≥ 1. However, since the MCMF algorithm has varying running times, which depend on the input graph.

[0092] The embodiments of the present disclosure provide relevant algorithm performance analysis and test performance analysis. We take a small sample of Google traces and run them on simulated clusters of different sizes. For batch processing tasks, we use the Quincy scheduling policy and prioritize tasks (batch size greater than 1). Considering the average runtime for each minimum cost maximum flow algorithm, despite having the best and worst case time complexity, the successive shortest path algorithm is only slightly better than cycle canceling. Even on a cluster of 1250 machines, the successful shortest path algorithm runs for over 100 seconds. Furthermore, the relaxation algorithm generally performs best in practice, but it has the highest worst case time complexity. The relaxation algorithm outperforms the cost scaling algorithm (using the Quincy policy) by two orders of magnitude: on average, even on a cluster of 12,500 machines, the relaxation algorithm completes within 200 milliseconds. The key reason for this impressive performance is that although most scheduling choices are straightforward, the relaxation algorithm performs minimal work. This occurs if the destination of the task flow is uncontested. For example, several new tasks may have edges to the same location and attempt to be scheduled there. In this case, the relaxation algorithm routes most of the traffic through the graph through a single node. Furthermore, for large numbers of incoming requests, the scheduler analyzes the resource usage of each machine in the cluster and forwards them to the appropriate machines for processing. When distributing requests, the load and resource availability of the machines in the cluster are taken into consideration, prioritizing incoming requests to machines with low loads and sufficient resources.

[0093] Figure 8 FIG. 4 is a schematic diagram illustrating load balancing of machine resource usage in a cluster according to an embodiment of the present disclosure.

[0094] like Figure 8As shown in the figure, the global administrator distributes incoming requests to machines in the cluster for processing. However, before distributing requests, it considers the resource load of each instance in the cluster. Instances with more idle resources are prioritized. Specifically, the resources or load of each node are sorted from smallest to largest, and the status of each node is dynamically updated based on the request usage. This evenly distributes the load across all machines in the cluster, alleviating the pressure on some servers and reducing their hardware and software requirements. Furthermore, during the pre-population and decoding phases, additional instances may be needed to handle the processing. The pre-population phase is computationally intensive, while the decoding phase is memory intensive. The pre-population phase is from the user entering the prompt to the generation of the first token, while the decoding phase is from the generation of the first token to the termination of inference. During the pre-population phase, the large model calculates the QKV for all tokens in the prompt simultaneously. Since the calculation of different tokens is independent, this process can be parallelized. In the attention phase, the calculated QKV is further used to calculate the output matrix, which is then passed through the subsequent FFN layer and decoding to obtain the initial token. During the decoding phase, the computational principles are identical to those of the pre-filling phase, but the calculation method differs. This is due to two reasons: First, as seqlen increases, the computational complexity of attention increases quadratically. Direct computation is very expensive, resulting in extremely slow or even infeasible inference times for long sequences. Second, a brief analysis of the decoding phase reveals that it can reuse the KV results from the pre-filling phase as well as the KV results already generated during the decoding phase. This allows the generated KV to be stored, eliminating the need for recalculation. This is known as the KV cache. For the Q matrix, only the last row q of Q needs to be calculated each time. This involves calculating the attention associated with the qKV, rather than the attention associated with the QKV itself. This reduces complexity by an order of magnitude, achieving computation-by-storage.

[0095] The method provided by the embodiment of the present disclosure is for a platform that receives a large number of requests for large model services. Utilizing the method of the embodiment of the present disclosure, the processed requests can respond to users in a timely manner, reducing the number of requests that cannot be processed. In addition, the number of requests received by the platform at the same time is also increased. The method of the embodiment of the present disclosure can build a platform based on large language model services. By embedding knowledge in various fields into LLM, the platform can have a better understanding and performance capability for tasks in these fields, thereby better understanding the needs and intentions of users to provide more accurate and personalized services.

[0096] The present invention also provides a large language model inference device based on a distributed KV cache pool, comprising one or more processors for implementing the above-mentioned large language model inference method.

[0097] The present invention also provides a computer-readable storage medium having a program stored thereon, which, when executed by a processor, is used to implement the above-mentioned large language model inference method.

[0098] It should be noted that the diagram illustrates the possible implementation architecture, functions and operations of the apparatus, method and computer program product according to various embodiments of the present disclosure. In this regard, each box in the flow chart or block diagram can represent a module, program segment, or a part of code, and the module, program segment, or a part of code contains at least one executable instruction for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions marked in the box can also occur in an order different from that marked in the accompanying drawings. For example, two boxes represented in succession can actually be executed substantially in parallel, and they can sometimes be executed in the opposite order, depending on the functions involved. It should also be noted that each box in the block diagram can be implemented with a dedicated hardware-based system that performs a specified function or operation, or can be implemented with a combination of dedicated hardware and computer instructions.

[0099] In general, various example embodiments of the present disclosure may be implemented in hardware or dedicated circuitry, software, firmware, logic, or any combination thereof. Certain aspects may be implemented in hardware, while other aspects may be implemented in firmware or software that may be executed by a controller, microprocessor, or other computing device. When various aspects of the embodiments of the present disclosure are illustrated or described as block diagrams, flow charts, or using some other graphical representation, it will be understood that the blocks, devices, systems, techniques, or methods described herein may be implemented, as non-limiting examples, in hardware, software, firmware, dedicated circuitry or logic, general-purpose hardware or a controller or other computing device, or some combination thereof.

[0100] The exemplary embodiments of the present disclosure described in detail above are merely illustrative and not restrictive. It should be understood by those skilled in the art that various modifications and combinations may be made to these embodiments or their features without departing from the principles and spirit of the present disclosure, and such modifications should fall within the scope of the present disclosure.

Claims

1. A large language model inference method based on a distributed KV cache pool, characterized in that: It includes a global administrator and several NPUs; The global administrator includes an extended information library, a dispatcher, and a scheduler; The memories of several NPU cards are together constituted into a distributed KV cache pool for caching key-value tensors; The extended information library contains tasks, monitoring data, and cluster topology; Tasks correspond to user requests; The extended information library sends tasks, monitoring data, and cluster topology to the dispatcher and the scheduler; The scheduler generates an optimal flow based on the data sent by the extended information library, extracts the positions of processing nodes from the optimal flow, and sends the prompt words, corresponding models, and the extracted position information to the dispatcher; The dispatcher distributes user requests to elastic instances based on the data sent by the extended information library and the scheduler; The elastic instances dynamically organize the instances corresponding to the NPUs into a series of disjoint elastic sequence parallel groups to process requests in parallel batches, with different configured parallel degrees; During the large language model inference stage, corresponding calculations are performed on the NPU cards; The key-values generated during the inference are placed in the machine HBM of the node where the NPU is located.

2. The method according to claim 1, characterized in that The global administrator further includes an elastic controller; The elastic controller sorts the elastic instances to update the configurations of the elastic instances to form corresponding elastic sequence parallel groups and process requests in parallel; and expands and contracts the instances accordingly according to the resource usage of the tasks.

3. The method according to claim 2, characterized in that For the elastic contraction stage, after a batch processing completes the pre-filling stage, the computational requirements for the decoding stage will decrease. In this case, the size of the parallel group is contracted to release resources for other batch processes; Specifically, for a parallel group R with a corresponding parallel degree of d, the elastic contraction mechanism needs to reduce the group to a new parallel group R' with a corresponding parallel degree of d', where d' < d; The key-value tensors of the instances that are no longer used are transmitted to the machines of the instances that are still in use, ensuring that the key-value tensors of the entire request in the parallel group R are efficiently transmitted to the new parallel group R'; For the elastic expansion stage, if the generated key-value tensors exceed the capacity of the KV cache pool in the parallel group, new instances need to be added to expand the capacity; In addition, if the batch size is greater than the preset threshold, the calculations in the decoding stage become computationally intensive. In this case, more NPUs are deployed to reduce latency; Deploying more NPUs includes: First, the newly added NPUs are other idle NPU cards in the same machine; Second, the NPU cards in remote machines are only considered.

4. The method according to claim 1, wherein It also includes a load balancing strategy, specifically: Before request distribution, considering the resource loads of each instance in the cluster, sorting the resources or loads of each node from smallest to largest, and dynamically updating the status of each node according to the usage of the requests, that is, evenly distributing the load pressure to each computer in the cluster.

5. The method according to claim 1, characterized in that It also includes a strategy to avoid over-submitting machine network bandwidth. Specifically, each task connects to a request aggregator for network bandwidth requests; the request aggregator has an arc for each task, which adapts to each machine with sufficient spare bandwidth; these arcs adapt dynamically when the observed bandwidth usage changes.

6. The method according to claim 1, characterized in that The distributed KV cache pool is used to flexibly store requested key-value tensors at the granularity of a token.

7. The method according to claim 1, characterized in that The scheduler includes a scheduling policy module, a network flow module, and a minimum cost maximum flow parser; The incoming data first passes through the scheduling strategy module, is modified into the corresponding network flow, and then submitted to the minimum cost maximum flow parser to generate the optimal flow. The optimal flow result is in turn transmitted to the network flow module, which extracts the node location of each NPU used from the optimal flow and returns the extracted location to the scheduling strategy module. The vertices in the network flow represent requests, machines, sinks, unscheduled aggregation nodes, aggregator nodes, racks, and request aggregators, and the directed edges connecting the vertices represent costs.

8. The method according to claim 7, characterized in that The scheduling policy is any one of a group scheduling policy, a fair share scheduling policy, a queue scheduling policy, a preemptive scheduling policy, a topology-based scheduling policy, a recycling policy, a backfill scheduling policy, and a resource reservation scheduling policy.

9. A large language model inference device based on a distributed KV cache pool, characterized in that: The method comprises one or more processors for implementing the large language model reasoning method according to any one of claims 1 to 8.

10. A computer-readable storage medium having a program stored thereon, characterized in that: When the program is executed by a processor, it is used to implement the large language model inference method described in any one of claims 1-8.

Images