KV cache management method, scheduler, medium, equipment and product

Abstract

The invention relates to a KV cache management method, a scheduler, a medium, equipment and a product, and the method comprises the steps: determining a target slave node from N slave nodes if a master node is monitored to have a reasoning sequence meeting a migration condition in a process of executing a reasoning task; a migration instruction is issued to each master node, the migration instruction is used for instructing the master node to send a first KV cache to a target slave node, and the first KV cache is a KV cache corresponding to a second number of continuous tokens generated by the master node; and issuing a receiving instruction to the target slave node, wherein the receiving instruction is used for indicating the target slave node to receive and store the first KV cache sent by each master node. According to the technical scheme, the configuration mode that the master node and the slave nodes are heterogeneous is adopted, the total cost of the inference system can be reduced, the KV cache generated by the master node is migrated to the slave nodes, and the explosion problem of the KV cache when the master node conducts long sequence inference is avoided.

Description

Technical Field

[0001] The present disclosure relates to the field of computer technology, and in particular to a KV cache management method, scheduler, medium, device and product. Background Art

[0002] The reasoning process of the large language model refers to reasoning using a trained large language model. The user inputs question information, and the large language model can reason based on the question information and output corresponding answer information. The input question information and the output answer information are collectively referred to as context.

[0003] The Transformer model plays an important role in the Large Language Model (LLM), and its core module is the attention mechanism. The reasoning process of the large language model can be divided into two stages: prefill and decoding. In the prefill stage, the large language model processes the input text prompt and generates the first token and KV cache. In the decoding stage, the large language model generates output tokens one by one based on the last generated token and the historical KV cache. As the context grows, the amount of computation required for large language model reasoning will increase linearly. Therefore, the KV cache is currently used to save the already obtained KV matrix to avoid repeated calculations. However, the memory space occupied by the KV cache will also increase linearly with the growth of the context length. In the scenario of long sequence reasoning and multiple parallel reasoning tasks, the KV cache explosion problem is very likely to occur. Summary of the invention

[0004] The purpose of the present disclosure is to provide a KV cache management method, scheduler, medium, device and product, which can reduce the total cost of the reasoning system and avoid the explosion problem of the KV cache when the master node performs long sequence reasoning.

[0005] In order to achieve the above-mentioned object, in a first aspect, the present disclosure provides a management method of KV cache in a Transformer model, which is applied to a scheduler in an inference system based on a Transformer model, wherein the inference system comprises the scheduler, M master nodes and N slave nodes, wherein the M master nodes and the N slave nodes are used to perform inference tasks, wherein the M master nodes are used to load the Transformer model, wherein the computing power of the master nodes is higher than that of the slave nodes, and M≥1 and N≥1; wherein the method comprises: If it is monitored that there is an inference sequence that meets the migration condition in the process of the master node executing the inference task, a target slave node is determined from the N slave nodes, wherein the migration condition is that the number of tokens included in the inference sequence reaches a first number; Sending a migration instruction to each of the master nodes respectively, the migration instruction is used to instruct the master node to send a first KV cache to the target slave node, the first KV cache is a KV cache corresponding to a second number of consecutive tokens generated by the master node, the second number of consecutive tokens is the tokens included in the inference sequence, and the second number is less than or equal to the first number; A receiving instruction is sent to the target slave node, where the receiving instruction is used to instruct the target slave node to receive and store the first KV caches respectively sent by each of the master nodes.

[0006] In a second aspect, the present disclosure provides a scheduler, characterized in that it is applied to an inference system based on a Transformer model, the inference system comprising the scheduler, M master nodes and N slave nodes, the M master nodes and the N slave nodes are used to perform inference tasks, the M master nodes are used to load the Transformer model, the computing power of the master node is higher than the computing power of the slave node, M≥1, N≥1; the scheduler comprises: A determination module, configured to determine a target slave node from the N slave nodes if it is monitored that there is an inference sequence that satisfies a migration condition during the execution of the inference task by the master node, wherein the migration condition is that the number of tokens included in the inference sequence reaches a first number; A first instruction sending module, used to send migration instructions to each of the master nodes respectively, wherein the migration instructions are used to instruct the master node to send a first KV cache to the target slave node, wherein the first KV cache is a KV cache corresponding to a second number of consecutive tokens generated by the master node, wherein the second number of consecutive tokens is the tokens included in the inference sequence, and the second number is less than or equal to the first number; The second instruction sending module is used to send a receiving instruction to the target slave node, and the receiving instruction is used to instruct the target slave node to receive and store the first KV cache sent by each of the master nodes.

[0007] In a third aspect, the present disclosure provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the method provided in the first aspect of the present disclosure.

[0008] In a fourth aspect, the present disclosure provides an electronic device, including: a memory having a computer program stored thereon; A processor is used to execute the computer program in the memory to implement the steps of the method provided in the first aspect of the present disclosure.

[0009] In a fifth aspect, the present disclosure provides a computer program product, comprising a computer program, which, when executed by a processor, implements the steps of the method provided in the first aspect of the present disclosure.

[0010] Through the above technical solution, the reasoning system based on the Transformer model may include a scheduler, M master nodes and N slave nodes, the computing power of the master node is higher than that of the slave node, and the M master nodes are used to load the Transformer model. By adopting a heterogeneous configuration mode of the master node and the slave node, the total cost of the reasoning system can be reduced and the resource utilization rate can be improved. If the scheduler detects that there is a reasoning sequence that meets the migration condition in the process of the master node executing the reasoning task, the migration instruction can be issued to each master node respectively, wherein the migration condition is that the number of tokens included in the reasoning sequence reaches the first number, so that when there is a long sequence in the reasoning task executed by the master node, the KV cache generated by the master node can be migrated to the slave node, and the first KV cache migrated from the master node is the KV cache corresponding to the second number of continuous tokens included in the reasoning sequence, and the complete KV cache corresponding to the continuous token can be migrated to the target slave node, which is convenient for addressing the KV cache based on the token in the target slave node. In this way, the memory usage of the KV cache in the master node is reduced, the explosion problem of the KV cache when the master node performs long sequence reasoning is avoided, and the normal execution of the reasoning task is ensured.

[0011] Other features and advantages of the present disclosure will be described in detail in the following detailed description. BRIEF DESCRIPTION OF THE DRAWINGS

[0012] The accompanying drawings are used to provide a further understanding of the present disclosure and constitute a part of the specification. Together with the following specific embodiments, they are used to explain the present disclosure but do not constitute a limitation of the present disclosure. In the accompanying drawings: Figure 1 The diagram is a schematic diagram of a heterogeneous distributed reasoning system according to an exemplary embodiment.

[0013] Figure 2 The present invention is a flowchart of a method for managing a KV cache in a Transformer model according to an exemplary embodiment.

[0014] Figure 3 The diagram is a schematic diagram of a master node according to an exemplary embodiment.

[0015] Figure 4 The figure is a schematic diagram showing a KV cache migration according to an exemplary embodiment.

[0016] Figure 5 The figure is a schematic diagram showing a KV cache migration according to an exemplary embodiment.

[0017] Figure 6 It is a block diagram of a scheduler according to an exemplary embodiment.

[0018] Figure 7 It is a block diagram of an electronic device according to an exemplary embodiment.

[0019] Figure 8 It is a block diagram of an electronic device according to an exemplary embodiment. DETAILED DESCRIPTION

[0020] The specific implementation of the present disclosure is described in detail below in conjunction with the accompanying drawings. It should be understood that the specific implementation described herein is only used to illustrate and explain the present disclosure, and is not used to limit the present disclosure.

[0021] At present, for long sequence reasoning tasks, the main method used is the sequence parallel method, which distributes the KV cache and computing load to different nodes. In the related art, the KV cache explosion problem of long sequence reasoning is solved by a distributed system. Among them, each node manages the local KV cache through a local manager, and then when the local available memory is insufficient, the local manager can send a request for memory to the global manager. In this way, the memory management method and task scheduling in the related art are highly complex, and each local manager needs to work closely with the global manager and communicate in a timely manner to ensure the timely allocation of memory. However, in a distributed scenario, the complex protocol interaction between the local and the global also means that it is easier to generate performance bottlenecks, and even operation errors, which will affect the performance of the distributed system. In addition, the distributed system in the related art adopts a homogeneous design, that is, each node in the reasoning system uses the same hardware configuration, each node loads the model to perform the reasoning task, and the node may also receive the reasoning task transmitted from the remote node when it is running. How to adjust the priority of local and remote tasks will also affect the performance of the reasoning system. At the same time, each node needs to be configured with a high-speed GPU and a large-capacity high-bandwidth memory, and the system cost is also high.

[0022] In view of this, the present disclosure provides a management method, scheduler, medium, device and product for KV cache in a Transformer model, which can reduce the total cost of the reasoning system and avoid the explosion problem of KV cache when the master node performs long sequence reasoning.

[0023] In the present disclosure, the reasoning system based on the Transformer model includes a scheduler, M master nodes and N slave nodes, the M master nodes and N slave nodes are used to perform reasoning tasks, the M master nodes are used to load the Transformer model, the computing power of the master node is higher than that of the slave node, M≥1, N≥1. Among them, the scheduler and each master node, the scheduler and each slave node, the master node and the slave node, each master node, and each slave node can communicate, and there is no restriction on the specific communication method adopted.

[0024] Figure 1 is a schematic diagram of a heterogeneous distributed reasoning system according to an exemplary embodiment. Figure 1 As shown, the reasoning system includes a scheduler, master node 0, master node 1, master node 2, master node 3, slave node 0, and slave node 1. It should be noted that Figure 1 The number of master nodes and the number of slave nodes in the reasoning system shown are only for illustration. The present disclosure does not limit the number of master nodes and the number of slave nodes, which can be set as needed.

[0025] In the present disclosure, the master node and the slave node are heterogeneous, that is, the computing power of the master node is higher than that of the slave node. The master node and the slave node jointly perform the same reasoning task. Only the master node needs to load the Transformer model, and the slave node does not need to load the Transformer model. The slave node only needs to store the KV cache and perform some attention calculations that are not related to the model parameters. The calculation amount of the slave node is much smaller than that of the master node. Through the heterogeneous configuration of the master node and the slave node, the master node can be configured with a high-performance GPU and a large-bandwidth memory. The slave node has lower requirements for computing performance and memory bandwidth than the master node. It mainly needs to configure a large-capacity memory to save the KV cache. For example, the slave node can use a low-performance GPU, or other hardware computing accelerators, such as DSA (Domain Specific Architecture, domain-specific architecture), FPGA (Field-Programmable Gate Aray, field-programmable gate array), which can reduce the total cost of the reasoning system.

[0026] Therefore, for long sequence reasoning, if the configuration of each node in the reasoning system is the same, the number of nodes will be very large, which will lead to idle resources in practical applications. Therefore, the present disclosure adopts a heterogeneous configuration method of master nodes and slave nodes, which can reduce the total cost of the reasoning system and improve resource utilization.

[0027] Figure 2is a flowchart of a method for managing a KV cache in a Transformer model according to an exemplary embodiment. The method can be applied to a scheduler in an inference system based on a Transformer model, such as Figure 2 As shown, the method may include steps 21 to 23.

[0028] In step 21, if it is monitored that there is an inference sequence that meets the migration condition during the execution of the inference task by the master node, a target slave node is determined from the N slave nodes.

[0029] In step 22, a migration instruction is issued to each master node respectively.

[0030] In step 23, a receiving instruction is sent to the target slave node.

[0031] Among them, the migration condition is that the number of tokens included in the reasoning sequence reaches a first number. The first number is, for example, 1024. The number of tokens included in the reasoning sequence reaches the first number, which can characterize the reasoning sequence as a long sequence. For the reasoning of long sequence tasks, a parallel strategy can be adopted, such as a pipeline parallel strategy. The pipeline parallel strategy distributes different layers in the Transformer model to different computing devices to achieve parallel execution of reasoning tasks. For example, master node 0 and master node 1 are used to execute reasoning tasks of layer 0 and layer 1 in the Transformer model, and master node 2 and master node 3 are used to execute reasoning tasks of layer 2 and layer 3 in the Transformer model.

[0032] It should be noted that since the master nodes jointly execute reasoning tasks and perform reasoning on the same sequence, the number of tokens generated by each master node and the number of KV caches are the same, and the number of tokens corresponding to each layer in the Transformer model and the number of KV caches are also the same. In addition, token is a structure, and this disclosure mentions that the number of tokens reaches the first number, which can be considered to mean that the number of tokens generated by the master node at each layer of the model reaches the first number.

[0033] If the scheduler detects that there is an inference sequence that meets the migration conditions during the execution of the inference task by the master node, and the inference sequence is a long sequence, the KV cache explosion problem is likely to occur when inferring the long sequence.

[0034] In the present disclosure, the slave node does not need to load the Transformer model, and the slave node can play the role of storing the KV cache. Therefore, the target slave node can be determined from the N slave nodes, and the migration instruction is issued to each master node respectively. The migration instruction is used to instruct the master node to send the first KV cache to the target slave node. The first KV cache is the KV cache corresponding to the second number of consecutive tokens generated by the master node and included in the reasoning sequence, and the second number is less than or equal to the first number. For example, the second number is 512. In this way, the KV cache generated by the master node can be migrated to the slave node, which can avoid the problem of KV cache explosion when the master node performs long sequence reasoning. The scheduler can issue a receiving instruction to the target slave node, which is used to instruct the target slave node to receive and store the first KV cache sent by each master node. In addition, the scheduler can also record the storage location of the KV cache sent by each master node in the target slave node. When the master node needs to call this part of the KV cache, the scheduler can know the distribution of the KV cache, thereby ensuring that the master node can correctly address this part of the KV cache.

[0035] In addition, if there is no reasoning sequence that meets the migration condition in the reasoning task, that is, there is no long sequence, the slave node can be in an idle waiting state.

[0036] Through the above technical solution, the reasoning system based on the Transformer model may include a scheduler, M master nodes and N slave nodes, the computing power of the master node is higher than that of the slave node, and the M master nodes are used to load the Transformer model. By adopting a heterogeneous configuration mode of the master node and the slave node, the total cost of the reasoning system can be reduced and the resource utilization rate can be improved. If the scheduler detects that there is a reasoning sequence that meets the migration condition in the process of the master node executing the reasoning task, the migration instruction can be issued to each master node respectively, wherein the migration condition is that the number of tokens included in the reasoning sequence reaches the first number, so that when there is a long sequence in the reasoning task executed by the master node, the KV cache generated by the master node can be migrated to the slave node, and the first KV cache migrated from the master node is the KV cache corresponding to the second number of continuous tokens included in the reasoning sequence, and the complete KV cache corresponding to the continuous token can be migrated to the target slave node, which is convenient for addressing the KV cache based on the token in the target slave node. In this way, the memory usage of the KV cache in the master node is reduced, the explosion problem of the KV cache when the master node performs long sequence reasoning is avoided, and the normal execution of the reasoning task is ensured.

[0037] In the present disclosure, the parallel strategies adopted by M master nodes to execute reasoning tasks include tensor parallelism and pipeline parallelism, where M is the tensor parallelism n TP The parallelism of the pipeline is n PP The product of .

[0038] The pipeline parallel strategy has been explained above, that is, the pipeline parallel strategy distributes different layers in the Transformer model to different computing devices to achieve parallel execution of reasoning tasks. Pipeline parallelism (Pipemne Parallelsm, PP) is used to determine the degree of splitting of the Transformer model in the pipeline parallel strategy. For example, if the Transformer model has 32 layers, the pipeline parallelism n PP =4, the inference task can be assigned to 4 groups of computing devices, each group of computing devices is responsible for 8 layers of calculation. Tensor parallelism means that for each layer in the Transformer model, its parameters can be divided into multiple parts vertically, and the multiple parts are set on different computing devices. Tensor parallelism can refer to the division of each layer in the Transformer model into several parts.

[0039] Figure 3 is a schematic diagram of a master node according to an exemplary embodiment. Figure 3 As shown, the Transformer model has four layers, namely layer 0, layer 1, layer 2, and layer 3. The reasoning system has four main nodes and the tensor parallelism is n. TPP =2, pipeline parallelism n PPP =2. Among them, master nodes 0 and 1 are responsible for the reasoning tasks of the 0th and 1st layers of the model. Master node 0 is responsible for the reasoning tasks of the first part of the 0th and 1st layers of the model, and master node 1 is responsible for the reasoning tasks of the second part of the 0th and 1st layers of the model. Master nodes 2 and 3 are responsible for the reasoning tasks of the 2nd and 3rd layers of the model. Master node 2 is responsible for the reasoning tasks of the first part of the 2nd and 3rd layers of the model, and master node 3 is responsible for the reasoning tasks of the second part of the 2nd and 3rd layers of the model.

[0040] The following is an introduction to the implementation of the storage and migration of the KV cache in the present disclosure.

[0041] The i-th master node is used to perform the inference task of the j-th part of the m-th layer to the n-th layer of the Transformer model, i∈[1,M], j∈[1,n TP]. The i-th master node includes a first storage unit group corresponding to each layer from the m-th layer to the n-th layer, and the first storage unit group is composed of multiple first storage units. The first storage unit corresponding to the k-th layer is used to store the KV cache corresponding to the third number of tokens generated by the i-th master node performing a local reasoning task. The local reasoning task is the reasoning task of the j-th part in the k-th layer of the Transformer model, k∈[m,n]. Among them, the third number of tokens are continuous tokens included in the reasoning sequence, the third number is less than the second number, and the second number is a first preset multiple of the third number, the first number is a second preset multiple of the third number, the first preset multiple and the second preset multiple are both greater than 2, and the first preset multiple is less than or equal to the second preset multiple. As an example, the first number is 1024, the second number is 512, and the third number is 16.

[0042] like Figure 3 As shown, taking the master node 0 as an example, the master node 0 includes the first storage unit groups corresponding to the 0th layer and the 1st layer respectively. Taking the first storage unit group corresponding to the 0th layer as an example, the first storage unit group is composed of the first storage unit 31, the first storage unit 32, the first storage unit 33, and the first storage unit 34. Taking the first storage unit 31 corresponding to the 0th layer as an example, the first storage unit 31 is used to store the KV cache corresponding to the third number of tokens generated by the master node 0 performing the local reasoning task, which is the reasoning task of the first part in the 0th layer of the Transformer model.

[0043] In the present disclosure, in order to improve the management efficiency of the KV cache, in the master node, the KV cache corresponding to the third number, i.e., b tokens, can be stored in a first storage unit, and b tokens can constitute a block (sub-block), that is, a first storage unit can be regarded as a block, and the KV cache is stored and allocated at the granularity of block.

[0044] For example, Figure 3 As shown, the first storage unit filled with gray indicates that the storage unit currently stores KV cache, and the first storage unit not filled with gray indicates that the storage unit currently does not store KV cache, wherein the first storage unit 31 can store KV cache corresponding to b tokens, and the first storage unit 34 can store KV cache corresponding to another b tokens. The KV cache method in other master nodes is similar.

[0045] Reference Figure 4 , introduces the first implementation method of KV cache migration in the present disclosure.

[0046] In this embodiment, the target slave nodes include n TPA second storage unit group includes a second storage unit corresponding to each layer in the Transformer model.

[0047] like Figure 4 As shown, slave node 0 is used as the target slave node, and slave node 0 includes two second storage unit groups, wherein the left side of the dotted line shown in slave node 0 is the first second storage unit group, which is used to store the first part of the KV cache split according to the tensor parallel strategy, and the right side of the dotted line is the second second storage unit group, which is used to store the second part of the KV cache split according to the tensor parallel strategy. The second storage unit group includes the second storage units corresponding to each layer in the Transformer model, such as Figure 4 As shown, the second storage unit group includes second storage units corresponding to each layer from the 0th layer to the 3rd layer, and the number of second storage units corresponding to each layer can be multiple and greater than or equal to the fourth number. For example, the 0th layer in the first second storage unit group includes second storage unit 41, second storage unit 42 and other second storage units.

[0048] Among them, in the master node, the first storage unit is used to store the KV cache corresponding to b tokens, and the KV cache corresponding to b tokens constitutes a block (sub-block), and a first storage unit can be regarded as a block. The KV cache can be stored in a similar manner in the slave node. In order to reduce the complexity of memory management in the slave node and improve the memory access efficiency, the KV cache can be stored in the slave node with a coarser granularity, for example, it can be stored according to the granularity of super block, wherein a super block corresponds to the fourth number of blocks. When migrating the KV cache, it can be migrated according to the granularity of super block, that is, the KV cache in the fourth number of first storage units in the master node can be migrated to the target slave node, and the fourth number is equal to the first preset multiple. In this way, the KV cache in the fourth number of first storage units is the KV cache corresponding to the second number of continuous tokens. Based on this, the migration instruction issued by the scheduler to the i-th master node is used to indicate: for each layer from the m-th layer to the n-th layer, the first KV cache corresponding to the layer is sent to the target slave node. Among them, the first KV cache corresponding to the kth layer includes the KV caches in the fourth number of first storage units corresponding to the kth layer in the ith master node, and the first KV cache is also the KV cache of a super block mentioned above.

[0049] The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored, wherein the first KV cache corresponding to the k-th layer is continuously stored in the second storage unit corresponding to the k-th layer in the j-th second storage unit group.

[0050] like Figure 4 As shown, taking the master node 0 as an example, the migration instruction sent by the scheduler to the master node 0 is used to instruct the master node 0 to send the first KV cache corresponding to the 0th layer to the slave node 0, and send the first KV cache corresponding to the 1st layer to the slave node 0. The slave node 0 can store the KV cache corresponding to each layer sent by the master node 0.

[0051] Taking the first KV cache corresponding to the 0th layer in the master node 0 as an example, the first KV cache may include the KV cache in the first storage unit 31 and the first storage unit 34. The slave node 0 may store the KV cache in the first storage unit 31 and the KV cache in the first storage unit 34 continuously in the second storage unit corresponding to the 0th layer in the first second storage unit group, for example, continuously in the second storage unit 41 and the second storage unit 42. The storage granularity of the second storage unit in the slave node is the same as the storage granularity of the first storage unit in the master node, that is, one second storage unit corresponds to one block, so for example, the KV cache in the first storage unit 31 can be directly stored in the second storage unit 41 to improve storage efficiency, and the KV cache constituting a block can be stored as a whole.

[0052] In this way, although the KV cache is stored in a dispersed manner in the master node (i.e., the first storage unit 31 and the first storage unit 34 are not continuous storage units), the KV cache migrated to the slave node can be stored continuously (i.e., the second storage unit 41 and the second storage unit 42 are continuous storage units), which can improve the addressing efficiency of the KV cache. Moreover, the KV cache in the first storage unit 31 and the KV cache in the first storage unit 34 are only stored in a dispersed manner, and from the perspective of the token, they are KV caches corresponding to continuous tokens.

[0053] like Figure 4 As shown, taking the master node 1 as an example, the master node 1 is used to execute the reasoning tasks of the second part in the 0th layer and the 1st layer of the model. For example, for the 0th layer, the first KV cache corresponding to the 0th layer migrated from the master node 1 includes the KV cache in the first storage unit 37 and the KV cache in the first storage unit 38. When storing from node 0, it can be continuously stored in the second storage unit 43 and the second storage unit 44 in the second second storage unit group.

[0054] It should be noted that the fourth number of 2 is used as an example in the implementation of the present disclosure for explanation only and does not constitute a limitation on the implementation of the present disclosure. The present disclosure does not limit the value of the fourth number. In addition, the KV cache migrated to the slave node can be the KV cache corresponding to all or part of the tokens in the inference sequence generated by the master node, for example Figure 4 The KV cache in the first storage unit 31 and the first storage unit 34 migrated from the master node 0 may be all or part of the KV cache generated by the master node 0 executing the reasoning task of the first part in the 0th layer of the Transformer model.

[0055] Through the above technical solution, the KV caches corresponding to different layers in the model generated by the master node can be migrated to the target slave node separately. The target slave node can store the KV caches corresponding to different layers in the model separately, and can be divided into multiple second storage unit groups for storage according to the segmentation method of the tensor parallel strategy. That is, the KV caches of different layers and parts in the model are stored separately in the slave node. In this way, when the scheduler needs to schedule the KV cache in the target slave node, it can search according to the model layer and the segmentation method of the tensor parallel strategy, and can quickly locate the KV cache to be found from the target slave node.

[0056] Reference Figure 5 , introduces the second implementation method of KV cache migration in the present disclosure.

[0057] In this embodiment, the target slave node includes a third storage unit group corresponding to each master node, and the third storage unit group corresponding to the i-th master node includes a third storage unit corresponding to each layer from the m-th layer to the n-th layer.

[0058] like Figure 5 As shown, the third storage unit group corresponding to the master node 0 includes the third storage unit group 51 and the third storage unit group 52, and the third storage unit group corresponding to the master node 1 includes the third storage unit group 53. Taking the third storage unit group 51 as an example, it includes the third storage units corresponding to the 0th layer and the 1st layer respectively, for example, the third storage unit 511 and the third storage unit 512 corresponding to the 1st layer.

[0059] The migration instruction issued by the scheduler to the i-th master node is used to instruct: to send the first KV cache corresponding to each layer from the m-th layer to the n-th layer to the target slave node together, wherein the first KV cache corresponding to the k-th layer includes the KV caches in the fourth number of first storage units corresponding to the k-th layer in the i-th master node.

[0060] The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored in the third storage unit group corresponding to the i-th master node, wherein the first KV cache corresponding to the k-th layer is continuously stored in the third storage unit corresponding to the k-th layer in the third storage unit group. The number of the third storage units corresponding to each layer can be multiple and greater than or equal to the fourth number.

[0061] In this embodiment, since the target slave node includes the third storage unit group corresponding to each master node, the master node can send the first KV cache corresponding to each layer to the target slave node together, and the master node only needs to perform one data transmission to complete the migration of the KV cache.

[0062] like Figure 5 As shown, for example, for master node 0, the first KV cache corresponding to layer 0, that is, the KV cache in the first storage unit 31 and the first storage unit 34, and the first KV cache corresponding to layer 1, that is, the KV cache in the first storage unit 35 and the first storage unit 36, can be transmitted together to slave node 0. Slave node 0 can store the KV cache sent by master node 0 in the third storage unit group 51. For example, for layer 1, the KV cache in the first storage unit 35 can be stored in the third storage unit 511, and the KV cache in the first storage unit 36 ​​can be stored in the third storage unit 512.

[0063] Among them, the storage granularity of the third storage unit in the slave node is the same as the storage granularity of the first storage unit in the master node, that is, one third storage unit corresponds to one block. Therefore, for example, the KV cache in the first storage unit 35 can be directly stored in the third storage unit 511, which improves the storage efficiency and stores the KV cache that constitutes a block as a whole.

[0064] like Figure 5 As shown, for example, for master node 1, the first KV cache corresponding to layer 0, i.e., the KV cache in first storage unit 37 and first storage unit 38, and the first KV cache corresponding to layer 1, i.e., the KV cache in first storage unit 39 and first storage unit 30, can be transmitted to slave node 0 together. Slave node 0 can store the KV cache sent by master node 1 in the third storage unit group 53.

[0065] In this implementation, although the KV cache is stored in a dispersed manner in the master node, the KV cache migrated to the slave node can be stored continuously. In addition, the master node only needs to perform a transfer operation once to migrate out the first KV cache corresponding to each layer, thereby improving the migration efficiency of the KV cache.

[0066] Preferably, in the above two embodiments, considering that the KV cache at the head may be shared by more sequences, retaining the KV cache at the head in the master node can improve the KV cache hit rate. Therefore, the first KV cache migrated from the master node may be the KV cache corresponding to the second number of consecutive tokens at the tail of the inference sequence.

[0067] For example, an implementation of determining a target slave node from N slave nodes may be: According to the polling scheduling algorithm, the currently polled slave node is used as the target slave node.

[0068] Among them, a round robin scheduling algorithm can be used to schedule the N slave nodes in turn, that is, the N slave nodes take turns to serve as target slave nodes in a certain order.

[0069] For example, another implementation of determining the target slave node from N slave nodes may be: The slave node with the largest remaining memory capacity among the N slave nodes is selected as the target slave node.

[0070] Among them, the KV cache in the slave node is controlled by the scheduler, so the scheduler can know the current remaining capacity of each slave node and use the slave node with the largest current memory remaining capacity as the target slave node.

[0071] It should be noted that the present disclosure does not limit the implementation method of determining the target slave node from N slave nodes, and the above two implementation methods are only used as examples. For example, any other scheduling algorithm can also be used to determine the target slave node, and the present disclosure does not limit it.

[0072] Through the above solution, only one of the slave nodes can be selected as the target slave node, avoiding the need for a large number of slave nodes in the inference system and reducing the network overhead during the KV cache migration process.

[0073] In the present disclosure, the method further comprises: In response to the completion of the reasoning task, a first release instruction is sent to each master node, and a second release instruction is sent to the target slave node, wherein the first release instruction is used to instruct the master node to release the currently stored KV cache generated during the execution of the reasoning task, and the second release instruction is used to instruct the target slave node to release the KV cache received from each master node.

[0074] In this way, after the inference task is completed, both the master node and the slave node can release the KV cache related to this inference task.

[0075] By adopting the above technical solution, the total cost of the reasoning system can be reduced and resource utilization can be improved by adopting a heterogeneous configuration mode of the master node and the slave node. In addition, for long sequence reasoning, the KV cache generated by the master node can be migrated to the slave node, reducing the memory usage of the KV cache in the master node, avoiding the explosion problem of the KV cache when the master node performs long sequence reasoning, and ensuring the normal execution of the reasoning task.

[0076] Based on the same inventive concept, the present disclosure also provides a scheduler, which is applied to an inference system based on a Transformer model. The inference system includes the scheduler, M master nodes and N slave nodes. The M master nodes and the N slave nodes are used to perform inference tasks. The M master nodes are used to load the Transformer model. The computing power of the master nodes is higher than that of the slave nodes, M≥1, N≥1. Figure 6 is a block diagram of a scheduler according to an exemplary embodiment. Figure 6 As shown, the scheduler 60 includes: A determination module 61 is configured to determine a target slave node from the N slave nodes if it is detected that, during the process of the master node executing the reasoning task, there is a reasoning sequence that satisfies a migration condition, wherein the migration condition is that the number of tokens included in the reasoning sequence reaches a first number; A first instruction sending module 62 is used to send migration instructions to each of the master nodes respectively, wherein the migration instructions are used to instruct the master node to send a first KV cache to the target slave node, wherein the first KV cache is a KV cache corresponding to a second number of consecutive tokens generated by the master node, wherein the second number of consecutive tokens is the tokens included in the inference sequence, and the second number is less than or equal to the first number; The second instruction sending module 63 is used to send a receiving instruction to the target slave node, and the receiving instruction is used to instruct the target slave node to receive and store the first KV cache sent by each of the master nodes.

[0077] Optionally, the parallel strategies adopted by the M master nodes to execute the reasoning task include tensor parallelism and pipeline parallelism, where M is the tensor parallelism n TP The parallelism of the pipeline is n PP The product of Among them, the i-th master node is used to execute the reasoning task of the j-th part in the m-th layer to the n-th layer of the Transformer model, and the i-th master node includes a first storage unit group corresponding to each layer in the m-th layer to the n-th layer, and the first storage unit group is composed of multiple first storage units. The first storage unit corresponding to the k-th layer is used to store the KV cache corresponding to the third number of tokens generated by the i-th master node performing the local reasoning task, and the local reasoning task is the reasoning task of the j-th part in the k-th layer of the Transformer model. The third number of tokens is the continuous tokens included in the reasoning sequence, the third number is less than the second number, and the second number is a first preset multiple of the third number, i∈[1, M], j∈[1, n TP ], k∈[m,n].

[0078] Optionally, the target slave node includes n TP a second storage unit group, wherein the second storage unit group includes a second storage unit corresponding to each layer in the Transformer model; The migration instruction issued by the scheduler to the i-th master node is used to instruct: for each layer from the m-th layer to the n-th layer, the first KV cache corresponding to the layer is sent to the target slave node, wherein the first KV cache corresponding to the k-th layer includes the KV caches in the fourth number of the first storage units corresponding to the k-th layer in the i-th master node, and the fourth number is equal to the first preset multiple; The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored, wherein the first KV cache corresponding to the k-th layer is continuously stored in the second storage unit corresponding to the k-th layer in the j-th second storage unit group.

[0079] Optionally, the target slave node includes a third storage unit group corresponding to each of the master nodes, and the third storage unit group corresponding to the i-th master node includes a third storage unit corresponding to each layer from the m-th layer to the n-th layer; The migration instruction issued by the scheduler to the i-th master node is used to instruct: the first KV cache corresponding to each layer from the m-th layer to the n-th layer is sent to the target slave node together, wherein the first KV cache corresponding to the k-th layer includes the KV caches in the fourth number of the first storage units corresponding to the k-th layer in the i-th master node, and the fourth number is equal to the first preset multiple; The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored in the third storage unit group corresponding to the i-th master node, wherein the first KV cache corresponding to the k-th layer is continuously stored in the third storage unit corresponding to the k-th layer in the third storage unit group.

[0080] Optionally, the determining module 61 is used to: According to the polling scheduling algorithm, the slave node currently polled is used as the target slave node; or, The slave node with the largest current remaining memory capacity among the N slave nodes is used as the target slave node.

[0081] Optionally, the device 60 further includes: The third instruction sending module is used to send a first release instruction to each of the master nodes and a second release instruction to the target slave node in response to the completion of the reasoning task, wherein the first release instruction is used to instruct the master node to release the KV cache currently stored and generated during the execution of the reasoning task, and the second release instruction is used to instruct the target slave node to release the KV cache received from each of the master nodes.

[0082] Regarding the device in the above embodiment, the specific manner in which each module performs operations has been described in detail in the embodiment of the method, and will not be elaborated here.

[0083] Figure 7 FIG. 7 is a block diagram of an electronic device 700 according to an exemplary embodiment. Figure 7 As shown, the electronic device 700 may include: a processor 701 , a memory 702 . The electronic device 700 may also include one or more of a multimedia component 703 , an input / output (I / O) interface 704 , and a communication component 705 .

[0084] The processor 701 is used to control the overall operation of the electronic device 700 to complete all or part of the steps in the management method of the KV cache in the above-mentioned Transformer model. The memory 702 is used to store various types of data to support the operation of the electronic device 700, and these data may include, for example, instructions for any application or method used to operate on the electronic device 700, and application-related data, such as contact data, sent and received messages, pictures, audio, video, etc. The memory 702 can be implemented by any type of volatile or non-volatile storage device or a combination thereof, such as static random access memory (Static Random Access Memory, referred to as SRAM), electrically erasable programmable read-only memory (Electrically Erasable Programmable Read-Only Memory, referred to as EEPROM), erasable programmable read-only memory (Erasable Programmable Read-Only Memory, referred to as EPROM), programmable read-only memory (Programmable Read-Only Memory, referred to as PROM), read-only memory (Read-Only Memory, referred to as ROM), magnetic memory, flash memory, magnetic disk or optical disk. The multimedia component 703 may include a screen and an audio component. The screen may be, for example, a touch screen, and the audio component is used to output and / or input audio signals. For example, the audio component may include a microphone, which is used to receive external audio signals. The received audio signal may be further stored in the memory 702 or sent through the communication component 705. The audio component also includes at least one speaker for outputting audio signals. The I / O interface 704 provides an interface between the processor 701 and other interface modules, and the above-mentioned other interface modules may be a keyboard, a mouse, buttons, etc. These buttons may be virtual buttons or physical buttons. The communication component 705 is used for wired or wireless communication between the electronic device 700 and other devices. Wireless communication, such as Wi-Fi, Bluetooth, Near Field Communication (NFC), 2G, 3G or 4G, or a combination of one or more of them, so the corresponding communication component 705 may include: Wi-Fi module, Bluetooth module, NFC module.

[0085] In an exemplary embodiment, the electronic device 700 can be implemented by one or more application specific integrated circuits (ASIC), digital signal processors (DSP), digital signal processing devices (DSPD), programmable logic devices (PLD), field programmable gate arrays (FPGA), controllers, microcontrollers, microprocessors or other electronic components to execute the above-mentioned KV cache management method in the Transformer model.

[0086] In another exemplary embodiment, a computer-readable storage medium including program instructions is also provided, and when the program instructions are executed by a processor, the steps of the management method of the KV cache in the Transformer model are implemented. For example, the computer-readable storage medium can be the memory 702 including the program instructions, and the program instructions can be executed by the processor 701 of the electronic device 700 to complete the management method of the KV cache in the Transformer model.

[0087] In another exemplary embodiment, a computer program product is also provided. The computer program product includes a computer program that can be executed by a processor. When the computer program is executed by the processor, the steps of the method for managing the KV cache in the Transformer model are implemented.

[0088] Figure 8 1 is a block diagram of an electronic device 1900 according to an exemplary embodiment. For example, the electronic device 1900 may be provided as a server. Figure 8 The electronic device 1900 includes a processor 1922, which may be one or more, and a memory 1932 for storing a computer program executable by the processor 1922. The computer program stored in the memory 1932 may include one or more modules each corresponding to a set of instructions. In addition, the processor 1922 may be configured to execute the computer program to perform the above-mentioned method for managing the KV cache in the Transformer model.

[0089] In addition, the electronic device 1900 may further include a power supply component 1926 and a communication component 1950. The power supply component 1926 may be configured to perform power management of the electronic device 1900, and the communication component 1950 may be configured to implement communication of the electronic device 1900, for example, wired or wireless communication. In addition, the electronic device 1900 may further include an input / output (I / O) interface 1958. The electronic device 1900 may operate based on an operating system stored in the memory 1932, such as Windows Server 2000. TM , Mac OS X TM , Unix TM , Linux TM etc.

[0090] In another exemplary embodiment, a computer-readable storage medium including program instructions is also provided, and when the program instructions are executed by a processor, the steps of the management method of the KV cache in the Transformer model are implemented. For example, the computer-readable storage medium can be the memory 1932 including the program instructions, and the program instructions can be executed by the processor 1922 of the electronic device 1900 to complete the management method of the KV cache in the Transformer model.

[0091] In another exemplary embodiment, a computer program product is also provided. The computer program product includes a computer program that can be executed by a processor. When the computer program is executed by the processor, the steps of the method for managing the KV cache in the Transformer model are implemented.

[0092] The preferred embodiments of the present disclosure are described in detail above in conjunction with the accompanying drawings; however, the present disclosure is not limited to the specific details in the above embodiments. Within the technical concept of the present disclosure, a variety of simple modifications can be made to the technical solution of the present disclosure, and these simple modifications all fall within the protection scope of the present disclosure.

[0093] It should also be noted that the various specific technical features described in the above specific embodiments can be combined in any suitable manner without contradiction. In order to avoid unnecessary repetition, the present disclosure will not further describe various possible combinations.

[0094] In addition, various embodiments of the present disclosure may be arbitrarily combined, and as long as they do not violate the concept of the present disclosure, they should also be regarded as the contents disclosed by the present disclosure.

Claims

1. A method for managing KV cache in a Transformer model, characterized in that: A scheduler applied to an inference system based on a Transformer model, the inference system comprising the scheduler, M master nodes and N slave nodes, the M master nodes and the N slave nodes being used to perform inference tasks, the M master nodes being used to load the Transformer model, the computing power of the master nodes being higher than that of the slave nodes, M≥1, N≥1; the method comprising: If it is monitored that there is an inference sequence that meets the migration condition in the process of the master node executing the inference task, a target slave node is determined from the N slave nodes, wherein the migration condition is that the number of tokens included in the inference sequence reaches a first number; Sending a migration instruction to each of the master nodes respectively, the migration instruction is used to instruct the master node to send a first KV cache to the target slave node, the first KV cache is a KV cache corresponding to a second number of consecutive tokens generated by the master node, the second number of consecutive tokens is the tokens included in the inference sequence, and the second number is less than or equal to the first number; A receiving instruction is sent to the target slave node, where the receiving instruction is used to instruct the target slave node to receive and store the first KV caches respectively sent by each of the master nodes.

2. The method according to claim 1, characterized in that The parallel strategies adopted by the M master nodes to execute the inference task include tensor parallelism and pipeline parallelism, where M is the tensor parallelism n TP The parallelism of the pipeline is n PP The product of Among them, the i-th master node is used to execute the reasoning task of the j-th part in the m-th layer to the n-th layer of the Transformer model, and the i-th master node includes a first storage unit group corresponding to each layer in the m-th layer to the n-th layer, and the first storage unit group is composed of multiple first storage units. The first storage unit corresponding to the k-th layer is used to store the KV cache corresponding to the third number of tokens generated by the i-th master node performing the local reasoning task, and the local reasoning task is the reasoning task of the j-th part in the k-th layer of the Transformer model. The third number of tokens is the continuous tokens included in the reasoning sequence, the third number is less than the second number, and the second number is a first preset multiple of the third number, i∈[1, M], j∈[1, n TP ], k∈[m,n].

3. The method according to claim 2, characterized in that The target slave node includes n TP a second storage unit group, wherein the second storage unit group includes a second storage unit corresponding to each layer in the Transformer model; The migration instruction issued by the scheduler to the i-th master node is used to instruct: for each layer from the m-th layer to the n-th layer, the first KV cache corresponding to the layer is sent to the target slave node, wherein the first KV cache corresponding to the k-th layer includes the KV caches in the fourth number of the first storage units corresponding to the k-th layer in the i-th master node, and the fourth number is equal to the first preset multiple; The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored, wherein the first KV cache corresponding to the k-th layer is continuously stored in the second storage unit corresponding to the k-th layer in the j-th second storage unit group.

4. The method according to claim 2, characterized in that: The target slave node includes a third storage unit group corresponding to each of the master nodes, and the third storage unit group corresponding to the i-th master node includes a third storage unit corresponding to each layer from the m-th layer to the n-th layer; The migration instruction issued by the scheduler to the i-th master node is used to instruct: the first KV cache corresponding to each layer from the m-th layer to the n-th layer is sent to the target slave node together, wherein the first KV cache corresponding to the k-th layer includes the KV caches in the fourth number of the first storage units corresponding to the k-th layer in the i-th master node, and the fourth number is equal to the first preset multiple; The receiving instruction sent by the scheduler to the target slave node is used to indicate: for the i-th master node, the first KV cache corresponding to each layer from the m-th layer to the n-th layer sent by the master node is stored in the third storage unit group corresponding to the i-th master node, wherein the first KV cache corresponding to the k-th layer is continuously stored in the third storage unit corresponding to the k-th layer in the third storage unit group.

5. The method according to claim 1, characterized in that: The determining of the target slave node from the N slave nodes comprises: According to the polling scheduling algorithm, the slave node currently polled is used as the target slave node; or, The slave node with the largest current remaining memory capacity among the N slave nodes is used as the target slave node.

6. The method according to claim 1, characterized in that The method further comprises: In response to the completion of the reasoning task, a first release instruction is sent to each of the master nodes, and a second release instruction is sent to the target slave node, wherein the first release instruction is used to instruct the master node to release the currently stored KV cache generated during the execution of the reasoning task, and the second release instruction is used to instruct the target slave node to release the KV cache received from each of the master nodes.

7. A scheduler, characterized in that: Applied to a reasoning system based on a Transformer model, the reasoning system includes the scheduler, M master nodes and N slave nodes, the M master nodes and the N slave nodes are used to perform reasoning tasks, the M master nodes are used to load the Transformer model, the computing power of the master nodes is higher than that of the slave nodes, M≥1, N≥1; The scheduler comprises: A determination module, configured to determine a target slave node from the N slave nodes if it is monitored that there is an inference sequence that satisfies a migration condition during the execution of the inference task by the master node, wherein the migration condition is that the number of tokens included in the inference sequence reaches a first number; A first instruction sending module, used to send migration instructions to each of the master nodes respectively, wherein the migration instructions are used to instruct the master node to send a first KV cache to the target slave node, wherein the first KV cache is a KV cache corresponding to a second number of consecutive tokens generated by the master node, wherein the second number of consecutive tokens is the tokens included in the inference sequence, and the second number is less than or equal to the first number; The second instruction sending module is used to send a receiving instruction to the target slave node, and the receiving instruction is used to instruct the target slave node to receive and store the first KV cache sent by each of the master nodes.

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

9. An electronic device, characterized in that: include: a memory having a computer program stored thereon; A processor, configured to execute the computer program in the memory to implement the steps of the method according to any one of claims 1 to 6.

10. A computer program product, comprising a computer program, characterized in that When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

Images