KV cache management system and method for importance perception

Abstract

The invention discloses an importance perception KV cache management system and method, and aims to solve the problems of long reasoning time and low throughput rate caused by low efficiency of a traditional K and V cache management strategy in a large model reasoning process. According to the method, the importance is introduced into the management strategy of the K and V caches, the cache hit rate on a relatively fast medium in storage is improved by tracking the importance of the K and V, combining the read frequency of the K and V and using the importance score as the basis of cache management, the time for reading the K and V is shortened, the overall time delay of an inference system is reduced, and the throughput is improved. The method is suitable for importance-based large model reasoning tasks with shared prefixes, the model reasoning precision can be maintained, meanwhile, the output time delay is shortened, and the throughput capacity is improved by about two times.

Description

Technical Field

[0001] The present invention belongs to the field of computer science artificial intelligence systems, and in particular relates to an importance-aware KV cache management system and method. Background Art

[0002] Large Language Model (LLM) is a large-scale model based on neural networks, which is specially used to process natural language data. By learning a large amount of text data, LLM can capture the complex grammar, semantics and contextual relationships in the language and generate coherent and logical language output. Large language model reasoning refers to the process of using input data to make predictions on a trained model. The reasoning process of LLM usually includes the following stages: ① Data preparation: First, the input text data needs to be preprocessed. This includes word segmentation of the text, converting vocabulary into corresponding word vectors or encodings (such as through word embedding models or dictionary mapping) to adapt to the input requirements of the model. ② Forward propagation: Pass the processed input data to the large language model. The model processes the data layer by layer through its multi-layer network structure, and generates prediction results after a large number of linear transformations, activation functions and attention mechanisms. ③ Post-processing: Perform necessary post-processing on the model output (such as deduplication, denoising or format adjustment).

[0003] The reasoning of a large language model is an autoregressive process, that is, the generation of the current prediction result depends on the previously generated prediction result (a prediction result is called a token). The key intermediate results (K, V) in these generation processes are reusable, so they can be managed in a cached manner (KV cache).

[0004] Existing studies have pointed out that not all tokens’ K and V contribute equally to maintaining the accuracy of reasoning. If the K and V of some tokens are missing, the reasoning accuracy will drop significantly. Such tokens are called “important” tokens. However, if the K and V of some tokens are missing, it will not have a significant impact on the reasoning accuracy. Such tokens are called “unimportant” tokens.

[0005] In large model generation tasks, multiple input sequences share the same starting part (prefix). The K and V of these prefixes can be calculated only once and cached so that other input sequences can directly skip the calculation of K and V to perform subsequent calculations, without having to repeatedly calculate K and V on each input sequence. Combined with importance, the model can calculate K and V only once on multiple input sequences with the same prefix and only load important K and V for subsequent calculations. However, we cannot only store important K and V, because the important K and V of different input sequences with the same prefix are often different. Therefore, we need to store K and V for all tokens. Due to the large size of K and V cache data, we use GPU memory and main memory (storage capacity is small, large, and reading speed is fast, slow) to store KV cache, and package K and V cache into chunks for storage. A chunk contains K or V of multiple consecutive tokens, so as to make full use of PCIe bandwidth. This forms a two-layer storage structure, and GPU memory can be regarded as a cache of main memory. How to make K and V read requests hit the GPU memory as much as possible has become the main problem to be solved by this multi-level K and V cache management system. The existing cache management strategy is based on the traditional cache replacement method: LFU (least used), which does not consider importance, resulting in suboptimal cache hit rate, further leading to longer K and V read time, increasing the delay in processing input sequences, and reducing system throughput. Summary of the invention

[0006] In view of the problem that the cache management system of the KV cache is inefficient when reasoning with a large language model, the purpose of the present invention is to provide an importance-aware KV cache management system and method.

[0007] The object of the present invention is to achieve the following technical solution: an importance-aware KV cache management system, comprising:

[0008] Storage module: saves K and V of all prefixes in the form of blocks; saves the required model parameters in the GPU memory, and reserves a part of the space for caching K and V; in the main memory, maintains a small top heap to manage the metadata of the blocks in the GPU memory, the top of the heap is the block with the smallest score in the GPU memory, and the remaining space is used to cache K and V; wherein the metadata includes the location of each block, which K and V are contained in the block, and the access frequency and score of the block. The K and V stored in the main memory and GPU memory are not redundant;

[0009] Importance identification module: used to identify which important K and V are needed in the current generation stage of a large model generation task, obtain the blocks corresponding to the important K and V through the metadata of the blocks, and inform the storage module to read the blocks corresponding to the important K and V;

[0010] Cache management module: swaps K and V in and out of GPU memory and main memory according to the block score to ensure that read requests for K and V hit the GPU memory, reduce the latency of reading K and V, further reduce the latency of inference, and increase the throughput of the system.

[0011] Furthermore, the working method of the module is specifically as follows: when the importance identification module tells the storage module which blocks to read, the storage module performs the reading, and for the blocks that hit in the GPU memory, the cache management module updates the access frequency and score information of the block and retains it in the GPU memory; for the blocks that do not hit in the GPU memory but are in the main memory, the cache management module updates the access frequency and score information of this block, and checks whether the score of the top block of the mini-heap is lower than the current block. If so, it exchanges its position with the current block, otherwise it remains in the main memory.

[0012] The present invention also provides an importance-aware KV cache management method, comprising the following steps:

[0013] Data preparation stage: The system receives an inference task input by the user, encodes the input into a form that can be calculated by the large language model, and puts it into a first-in-first-out task queue to wait for execution;

[0014] Importance identification phase: If the current iteration is the first iteration, take out the input of a new task from the head of the task queue, and determine the important K and V required for the current iteration of the input; if the current iteration is not the first iteration, determine the important K and V required for the current iteration based on the accumulated results of all previous iterations of the current task;

[0015] Important K, V reading stage: by searching metadata, find the blocks corresponding to the important K, V required for the current iteration, read the important K, V from these blocks for calculation, when each block is read, the frequency in the metadata corresponding to the block is increased by 1, and the score of the block is recalculated. For the block hit in the GPU memory, the block is retained in the GPU memory. For the block hit in the main memory, check whether the score of the top block of the mini-top heap is lower than that of the block. If so, swap it with the block.

[0016] Calculation phase: Use the important K and V read to perform the calculation of the current iteration, obtain the prediction result of the current iteration, and save the K and V generated by the current iteration for use in subsequent iterations;

[0017] The steps from the importance identification phase to the calculation phase are repeated until all iterations are completed.

[0018] The beneficial effect of the present invention is that in a large model reasoning system based on importance, for read requests of K and V, the present invention makes their hit rate on GPU memory higher, reduces the read latency, further reduces the waiting time of calculation, reduces the delay of the system, and improves GPU utilization and throughput. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 Schematic diagram of cache management strategy based on frequency and cache management strategy based on importance score;

[0020] Figure 2 It is a schematic diagram of the system of the present invention;

[0021] Figure 3 The throughput comparison diagrams of the traditional cache management strategy LFU (white column) and the present invention (black column), wherein (a) is a throughput comparison diagram of the traditional cache management strategy LFU and the present invention on the PIQA dataset, (b) is a throughput comparison diagram of the traditional cache management strategy LFU and the present invention on the RTE dataset, (c) is a throughput comparison diagram of the traditional cache management strategy LFU and the present invention on the COPA dataset, and (d) is a throughput comparison diagram of the traditional cache management strategy LFU and the present invention on the OpenBookQA dataset;

[0022] Figure 4 1 and 1 , which are I / O time comparison diagrams of the traditional cache management strategy LFU (white column) and the present invention (black column), wherein (a) is an I / O time comparison diagram of the traditional cache management strategy LFU and the present invention on the RTE dataset, and (b) is an I / O time comparison diagram of the traditional cache management strategy LFU and the present invention on the COPA dataset. DETAILED DESCRIPTION

[0023] The technical solution provided by the present invention will be described in detail below in conjunction with specific embodiments. It should be understood that the following specific implementation methods are only used to illustrate the present invention and are not used to limit the scope of the present invention.

[0024] Under the existing management strategy based on block granularity, there are both important and unimportant K or V in a block. The present invention combines the importance of a block with the access frequency of the block and proposes a score strategy. The score of each block is equal to the proportion of important K or V in the block multiplied by the access frequency of the block. The present invention uses the score as the basis for cache management and places blocks with higher scores on faster storage media. Figure 1As shown in the figure, the difference in read latency between the traditional cache management strategy and the cache management strategy that takes importance into consideration is shown. Assuming that both the GPU memory and the main memory can accommodate only one block, and there are existing blocks 1 and 2, both storing K, with access frequencies of 10 and 15 respectively, the traditional cache management strategy (a) will place block 1 in the main memory and block 2 in the GPU memory based on the frequency, while the score-based strategy of the present invention (b) is the opposite. Based on the historical access frequency, it can be assumed that there will be 10 requests to access block 1 and 15 requests to access block 2 in the future. Under the traditional method, in order to obtain all the important Ks, the traditional method needs to transfer 30 Ks from the main memory to the GPU memory via PCIe, while the score-based method only needs to transfer 15 Ks, which greatly reduces the burden of PCIe bandwidth, thereby reducing the read latency of K and V and improving the system throughput.

[0025] like Figure 2 As shown, an importance-aware KV cache management system proposed in an embodiment of the present invention includes:

[0026] Storage module: K and V of all prefixes are saved in the form of blocks. The required model parameters are saved in the GPU memory, and a part of the space is reserved for caching K and V. In the main memory, a small top heap is maintained to manage the metadata of the blocks, which includes the location of each block and the access frequency and score of the block. The top of the heap is the block with the smallest score, and the remaining space is used to cache K and V. Among them, K and V stored in the main memory and GPU memory are not redundant.

[0027] Importance identification module: It is used to identify the important K and V of the current iteration of a task, obtain the blocks corresponding to the important K and V through the metadata of the blocks, and inform the storage module which blocks need to be read.

[0028] Cache management module: According to the block score dynamics, K and V are swapped in and out from the GPU memory and the main memory to ensure that the read requests for K and V hit the GPU memory as much as possible, reduce the delay of reading K and V, further reduce the delay of inference, and increase the throughput of the system. Specifically, after the importance recognition module tells the storage module which blocks to read, the storage module performs the read. For the blocks that hit the GPU memory, the cache management module updates its access frequency and score information, and still keeps it in the GPU memory; for the blocks that do not hit the GPU memory, that is, when the required block is in the main memory, the cache management module updates the access frequency and score information of this block, and checks whether the score of the top block of the mini-heap is lower than the current block. If so, it is replaced with it, otherwise it is still kept in the main memory.

[0029] An importance-aware KV cache management method proposed in an embodiment of the present invention includes the following steps:

[0030] (1) Data preparation stage. The system receives an inference task input by the user. The data preparation module encodes the input into a form that can be calculated by the model and puts it into a first-in-first-out task queue, waiting for the system to execute.

[0031] (2) Importance identification phase. If the current iteration is the first iteration, the importance identification module takes a new task input from the head of the task queue and determines the important K and V required for the current iteration. If the current iteration is not the first iteration, the importance identification module obtains the important K and V required for the current iteration based on the accumulated results of all previous iterations of the current task.

[0032] (3) Important K and V reading stage: The cache management module ensures that the system reads important K and V with minimal latency. At this time, the score-based cache replacement strategy takes effect.

[0033] (4) Calculation phase: Execute inference calculations to obtain the prediction results of the current iteration.

[0034] (5) Repeat steps (2) to (4) until all iterations are completed.

[0035] The design of the present invention is implemented based on pytorch, and performance tests are performed on a computing node with two AMD EPYC 7763 CPUs, 128GB DRAM, an A100 GPU with 80GB memory, a 2TB solid-state drive with 5GB / S read and write bandwidth, and GPU and CPU connected via PCIe 4.0×16. Among them, the SSD stores all K and V, 32GB in DRAM, and 10GB in GPU are used to cache K and V. Each block can accommodate 64 K or V. The data sets used to test throughput are PIQA, RTE, COPA, and OpenBookQA, and the models are OPT-6.7B, OPT-13B, and OPT-30B. The data sets used to test I / O time are RTE and COPA, and the models are OPT-13BH and OPT-30B.

[0036] Figure 3 The system throughput of the present invention and the system throughput of the traditional method are shown. The throughput is quantified by the number of requests completed per second by the system (requests / second). (a) is a throughput comparison chart of the traditional cache management strategy LFU and the present invention on the PIQA dataset, (b) is a throughput comparison chart of the traditional cache management strategy LFU and the present invention on the RTE dataset, (c) is a throughput comparison chart of the traditional cache management strategy LFU and the present invention on the COPA dataset, and (d) is a throughput comparison chart of the traditional cache management strategy LFU and the present invention on the OpenBookQA dataset. The results show that on all datasets and models, the throughput of the present invention is better than that of the traditional method, and the average throughput is 2.1 times that of the traditional method.

[0037] Figure 4 The average I / O time of reading K and V by the system and the traditional method is shown. (a) is a comparison chart of the I / O time of the traditional cache management strategy LFU and the present invention on the RTE dataset, and (b) is a comparison chart of the I / O time of the traditional cache management strategy LFU and the present invention on the COPA dataset. The results show that the I / O time of the system is half of that of the traditional method on average.

Claims

1. An importance-aware KV cache management system, characterized in that: include: Storage module: saves K and V of all prefixes in the form of blocks; The required model parameters are saved in the GPU memory, and a portion of space is reserved for caching K and V. In the main memory, a small top heap is maintained to manage the metadata of the blocks in the GPU memory. The top of the heap is the block with the smallest score in the GPU memory, and the remaining space is used to cache K and V. The metadata includes the location of each block, which K and V are contained in the block, and the access frequency and score of the block. The K and V stored in the main memory and GPU memory are not redundant. Importance identification module: used to identify which important K and V are needed in the current generation stage of a large model generation task, obtain the blocks corresponding to the important K and V through the metadata of the blocks, and inform the storage module to read the blocks corresponding to the important K and V; Cache management module: swaps K and V in and out of GPU memory and main memory according to the block score to ensure that read requests for K and V hit the GPU memory, reduce the latency of reading K and V, further reduce the latency of inference, and increase the throughput of the system.

2. The importance-aware KV cache management system according to claim 1, characterized in that: The working method of the module is specifically as follows: when the importance identification module tells the storage module which blocks to read, the storage module performs the reading. For blocks that hit in the GPU memory, the cache management module updates the access frequency and score information of the block and retains it in the GPU memory; for blocks that do not hit in the GPU memory but are in the main memory, the cache management module updates the access frequency and score information of this block, and checks whether the score of the top block of the mini-heap is lower than that of the current block. If so, it exchanges its position with the current block, otherwise it remains in the main memory.

3. An importance-aware KV cache management method, characterized in that: The method is implemented by an importance-aware KV cache management system as described in any one of claims 1 to 2, comprising the following steps: Data preparation stage: The system receives a reasoning task input by the user, encodes the input into a form that can be calculated by the large language model, and puts it into a first-in-first-out task queue, waiting for the system to execute; Importance identification phase: If the current iteration is the first iteration, take the input of a new task from the head of the task queue and determine whether the input is important K and V required for the current iteration; If the current iteration is not the first iteration, determine the important K and V required for the current iteration based on the accumulated results of all previous iterations of the current task; Important K, V reading stage: by searching metadata, find the blocks corresponding to the important K, V required for the current iteration, read the important K, V from these blocks for calculation, when each block is read, the frequency in the metadata corresponding to the block is increased by 1, and the score of the block is recalculated. For the block hit in the GPU memory, the block is retained in the GPU memory. For the block hit in the main memory, check whether the score of the top block of the mini-top heap is lower than that of the block. If so, swap it with the block. Calculation phase: Use the important K and V read to perform the calculation of the current iteration, obtain the prediction result of the current iteration, and save the K and V generated by the current iteration for use in subsequent iterations; The steps from the importance identification phase to the calculation phase are repeated until all iterations are completed.

Images