Large language model inference optimization method, system, device and storage medium
By employing KV header selective unloading and dynamic approximation caching techniques, the problems of excessive GPU memory usage and high read overhead in large language model inference are solved, achieving GPU memory optimization and improved inference performance.
Patent Information
- Application Number
- CN202511511942.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-10-22
- Publication Date
- 2026-02-17
- Estimated Expiration
- 2045-10-22
AI Technical Summary
During Large Language Model (LLM) inference, the excessive GPU memory usage of key-value (KV) cache data leads to insufficient GPU memory capacity, and the high overhead of reading KV data offloaded to CPU memory affects inference efficiency and hardware resource utilization.
It employs KV header selective offloading and dynamic approximation caching technology, where part of the KV data is cached in GPU memory and the other part is offloaded to CPU memory. It also reduces the amount of data read from CPU to GPU through top-k attention and manages KV data by combining dynamic approximation caching algorithm, prioritizing the reuse of cached data to reduce read overhead.
It effectively reduces the amount of GPU memory occupied by KV data, reduces the data reading overhead from CPU to GPU, and ensures inference efficiency and hardware resource utilization.
Smart Images

Figure CN120996208B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to the technical field of computer GPU (Graphics Processing Unit) memory and computing efficiency optimization, and particularly relates to a large language model inference optimization method, system, device and storage medium. BACKGROUND
[0002] Large language model (LLM) is the focus of the field of artificial intelligence in recent years. LLM is a generative artificial intelligence model, and the process of generating new text according to the text sequence input by the user is called inference, which is an important link of LLM application. The text sequence input by LLM is composed of tokens, which will be mapped into word embedding vectors to participate in the calculation of the attention module. Specifically, the word embedding vectors are transformed by three different matrices to generate query (Query), key (Key) and value (Value) data (Query is referred to as Q data, and the latter two are collectively referred to as KV data), and then new tokens are generated through the steps of nonlinear softmax function (normalization exponential function) calculation and activation function. The aforementioned QKV data will generate h groups in total, each of which uses a different mapping matrix, and these groups are called attention heads (Heads). According to different data types, the attention heads are divided into two types: Q heads and KV heads. In early LLM, Q heads and KV heads are one-to-one corresponding and equal in number; more advanced LLM usually uses grouped-query attention technology (Grouped-Query Attention, GQA), which reduces the number of KV heads so that m Q heads share one KV head. In addition, the structure of LLM is multi-layered, each layer has the same module structure and needs to complete the above-mentioned attention calculation process, the input of the first layer is the aforementioned word embedding vector, and the input of other layers is the output of the previous layer. These input data can be collectively referred to as hidden state vectors (referred to as hidden vectors).
[0003] The inference process of the large language model is autoregressive, and each new token generated needs to process all existing tokens according to the above calculation process, which repeatedly calculates the KV data. To solve this problem, KV cache technology is proposed, which caches the calculated KV data in the GPU (Graphics Processing Unit) memory, eliminates repeated calculation, improves inference efficiency, and has become a standard technology for LLM inference. After using the KV cache technology, the inference process of LLM is divided into two stages: the pre-filling stage and the decoding stage. In the pre-filling stage, LLM calculates the Q data and KV data of all tokens in parallel, caches the KV data, and generates the first new token. The decoding stage only calculates the Q data and KV data of the latest token, adds the KV data to the KV cache, and generates the next new token. The decoding stage will be repeated until the end of generation. As can be seen, the size of the KV cache data increases linearly with the number of tokens processed in the inference process.
[0004] In recent years, the input size of LLMs has been growing. On the task level, there are more and more tasks that need to process long text sequences. On the inference efficiency level, batch processing technology is widely used, and multiple input sequences are packaged into a batch for joint inference to improve inference throughput. The growth of sequence length and batch size both leads to more tokens that need to be processed during inference, and the size of KV data also grows rapidly. Since KV data is cached in GPU memory, the memory occupied by KV data will increase significantly, which may exceed the memory occupied by model parameters, and even exceed the memory capacity of the GPU. For example, when using the Llama3.1-8B model to infer a text sequence with a length of 256KB (kilobyte), the memory occupied by KV cache data reaches 31.25GB (gigabyte), while the memory occupied by model parameters is about 16GB, and the sum of the two exceeds the memory capacity of most mid-end GPU devices. Llama is an open-source LLM released by Meta AI, 3.1 is the model version number, and 8B (8 billion) is the model parameter amount. In summary, the high memory occupation of KV cache data has become a major bottleneck for the expansion of inference task size.
[0005] A mainstream method to solve this problem is to offload KV cache data from GPU memory to the memory of the computer central processing unit (CPU), i.e., host memory; and, by using the sparsity of attention computation involving KV data, the top-k attention technology is used to reduce the reading amount of KV data. Specifically, top-k attention generates retrieval metadata based on KV data, which has low memory occupation and can be stored in GPU memory. During inference, Q data (or its variants) and retrieval metadata are used to query key KV data, and only these key data are read from CPU memory to participate in computation. However, the high-speed serial computer expansion bus (PCIe) bandwidth of CPU-GPU memory interconnection is limited, and it will take a lot of time for GPU to read KV cache data from CPU memory, which will significantly exceed the GPU computation time, resulting in a significant reduction in inference speed and a decrease in GPU hardware resource utilization.
[0006] Therefore, how to reduce the CPU-to-GPU KV cache data reading overhead in the scenario of offloading KV cache data, optimize GPU memory occupation while ensuring inference efficiency and hardware resource utilization, is a key technical challenge in large-scale task inference of LLMs, and is one of the problems that need to be solved at the present stage.
[0007] In view of this, the present application is proposed. SUMMARY
[0008] The application aims to provide a large language model inference optimization method, system, device and storage medium, which can significantly reduce the GPU memory occupied by KV data, minimize the KV data reading overhead from CPU to GPU, and guarantee the inference efficiency and hardware resource utilization.
[0009] The application aims to achieve the above technical scheme.
[0010] A large language model inference optimization method comprises the following steps:
[0011] In the pre-filling stage, in combination with the multiplexing difficulty of the KV head, the KV data of part of the KV head is cached in the GPU memory, and the KV data of the other part of the KV head is unloaded to the CPU memory; wherein the KV head is a type of attention head in the large language model, KV is key and value, GPU is a graphics processing unit, and CPU is a central processing unit;
[0012] In the decoding stage, autoregressive decoding is performed in combination with the latest word element; the latest word element is the first new word element generated in the pre-filling stage or the new word element generated in the previous decoding step; the working process of the current layer of the large language model in a single decoding step is as follows: Q data and KV data are generated by using the hidden vector of the current layer, and the KV data is added to the GPU memory or the CPU memory; the hidden vector of the current layer is the word embedding vector of the first new word element or the calculation result of the previous layer, and Q is a query; based on the hidden vector of the current layer, it is judged whether the next layer needs to prefetch KV data in combination with the KV head as a unit; if yes, the corresponding KV data is searched in the CPU memory; after the KV data of the current layer is prefetched, the prefetching of the KV data of the next layer is started; the KV data is searched in the GPU memory by using the Q data generated by the current layer, and the calculation of the current layer is completed in combination with the prefetched KV data; the above steps are repeatedly performed until the last layer, and the calculation result of the last layer is used to obtain a new word element.
[0013] A large language model inference optimization system is used to implement the above method, and comprises the following steps:
[0014] The pre-filling unit based on selective unloading of the KV head is used to combine the multiplexing difficulty of the KV head, cache the KV data of part of the KV head in the GPU memory, and unload the KV data of the other part of the KV head to the CPU memory; wherein the KV head is a type of attention head in the large language model, KV is key and value, GPU is a graphics processing unit, and CPU is a central processing unit;
[0015] A decoding unit based on dynamic approximate caching and data prefetching is used for autoregressive decoding combined with the latest word piece, which is the first new word piece generated in the prefilling stage or the new word piece generated in the last decoding step. The working process of the current layer of the large language model in a single decoding step is as follows: Q data and KV data are generated by using the hidden vector of the current layer, and the KV data is added to the GPU memory or CPU memory. The hidden vector of the current layer is the word embedding vector of the first new word piece or the calculation result of the last layer, and Q is a query. Based on the hidden vector of the current layer, it is judged whether the next layer needs to prefetch KV data by taking the KV head as a unit. If yes, the corresponding KV data is searched in the CPU memory, and after the KV data of the current layer is prefetched, the prefetching of the KV data of the next layer is started. The KV data is searched in the GPU memory by using the Q data generated by the current layer, and the calculation of the current layer is completed by combining the prefetched KV data. The process is repeated until the last layer, and the calculation result of the last layer is used to obtain a new word piece.
[0016] A processing device, comprising: one or more processors; a memory for storing one or more programs;
[0017] Wherein, when the one or more programs are executed by the one or more processors, the one or more processors implement the foregoing method.
[0018] A readable storage medium, storing a computer program, when the computer program is executed by a processor, implementing the foregoing method.
[0019] As can be seen from the technical solutions provided by the above-mentioned application, in the inference process of the large language model, most of the KV data is offloaded to the CPU memory. For reading of the KV data from the CPU memory to the GPU memory, top-k attention is used to reduce the reading amount. Moreover, the KV data that has been read to the GPU memory is cached, and an approximate caching algorithm is used to manage the KV data by taking the KV head as a basic unit. When the KV data needs to be read in the inference process, it is firstly judged whether the data in the cache can be reused. If the cache hits, the data is directly reused, so as to reduce the data reading overhead from the CPU memory to the GPU memory. If the cache misses and the data cannot be reused, data prefetching is performed. Thanks to the above improvements, the application can effectively reduce the occupied memory of the KV data, minimize the KV data reading overhead from the CPU to the GPU, and make the inference performance reach an ideal level. BRIEF DESCRIPTION OF DRAWINGS
[0020] In order to more clearly illustrate the technical solutions of the embodiments of the present application, the following will briefly introduce the drawings needed in the embodiment description. Obviously, the drawings in the following description are only some embodiments of the present application, and other drawings can also be obtained by those skilled in the art without creative labor.
[0021] Figure 1 A flowchart of a large language model inference optimization method provided for an embodiment of the present application.
[0022] Figure 2 A flowchart of a prefill stage provided for an embodiment of the present application.
[0023] Figure 3 A flowchart of a decoding stage provided for an embodiment of the present application.
[0024] Figure 4 A schematic diagram of selective offloading of a KV head provided for an embodiment of the present application.
[0025] Figure 5 A schematic diagram of a dynamic approximate caching algorithm provided for an embodiment of the present application.
[0026] Figure 6 A schematic diagram of a GPU-centric data transfer synchronization method provided for an embodiment of the present application.
[0027] Figure 7 A schematic diagram of the overall architecture of a large language model inference optimization method provided for an embodiment of the present application.
[0028] Figure 8 A schematic diagram of a large language model inference optimization system provided for an embodiment of the present application.
[0029] Figure 9 A schematic diagram of a processing device provided for an embodiment of the present application. DETAILED DESCRIPTION
[0030] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only some of the embodiments of the present application, but not all the embodiments of the present application. Based on the embodiments of the present application, all other embodiments obtained by those of ordinary skill in the art without creative work fall within the protection scope of the present application.
[0031] First, the terms that can be used in this text are explained as follows:
[0032] The term “and / or” means either of the two or both at the same time, for example, X and / or Y means three cases including “X” or “Y” or “X and Y”.
[0033] The terms "comprising", "containing", "including", "having" or other similar semantic descriptions should be interpreted to be inclusive rather than exclusive. For example: the inclusion of an element of technical feature (e.g. a raw material, component, ingredient, carrier, dosage form, material, dimension, part, component, mechanism, device, step, process, method, reaction condition, processing condition, parameter, algorithm, signal, data, product or article, etc.) should be interpreted as including not only the explicitly recited element of technical feature but also other elements of technical feature that can or do not appear explicitly.
[0034] The term "consisting of" means excluding any element of technical feature not specifically recited. If this term is used in the context of a claim, the term will make the claim closed, meaning that the claim will not include technical features other than those specifically recited, with the exception of impurities ordinarily resulting from the manufacture of the claimed technical features. If this term is used in only some of the sub-clauses of a claim, it will limit only the sub-clauses so specified, the other sub-clauses being written in the open-ended way.
[0035] A large language model reasoning optimization method, system, device and storage medium provided by the present application are described in detail below. The content not described in detail in the embodiments of the present application belongs to the prior art known to those skilled in the art. If no specific conditions are specified in the embodiments of the present application, the conditions are performed according to the conventional conditions in the art or the conditions recommended by the manufacturer. If no manufacturer of the instruments used in the embodiments of the present application is specified, the instruments are conventional products that can be obtained by market purchase.
[0036] Embodiment one
[0037] The embodiments of the present application provide a large language model reasoning optimization method, as shown in Figure 1 The method mainly includes the following steps:
[0038] Step 1, pre-filling stage based on KV header selective offloading.
[0039] In the embodiments of the present application, in the pre-filling stage, combined with the multiplexing difficulty of the KV header, the KV data of a part of the KV header is cached in the GPU memory, and the KV data of another part of the KV header is offloaded to the CPU memory.
[0040] Specifically, in this step: the Q data and the KV data corresponding to the input token sequence are calculated, and the first new token is generated. There are multiple attention heads in each layer of the large language model, which are divided into two categories: Q heads and KV heads. Each KV head corresponds to a set of KV data, and each Q head corresponds to a set of Q data. The reuse difficulty of each KV head is quantified, and the number of KV heads that can be prefetched in each layer is calculated. In combination with the number of KV heads that can be prefetched in each layer and the reuse difficulty of the corresponding KV head, the KV heads residing in the GPU memory are selected, and the corresponding KV data is stored in the GPU memory. The remaining KV data corresponding to the KV head is unloaded to the CPU memory. The areas storing KV data in the GPU memory and the CPU memory are collectively referred to as KV cache.
[0041] The main process involved in this step is shown in Figure 2 The main process involved in this step is shown in
[0042] Step 11, generating Q data and KV data.
[0043] In the embodiment of the application, Q data and KV data are generated for the input text sequence. The principle and generation process have been introduced in the foregoing background technology, which will not be repeated here.
[0044] Step 12, generating retrieval metadata based on KV data.
[0045] In the embodiment of the application, the meaning of retrieval metadata has been introduced in the foregoing background technology, and its generation method is provided by the specific top-k attention calculation. The optimization method of the embodiment of the application focuses on the management and reading method optimization of KV data, and the top-k attention technology is not within the optimization range, which can be realized using the effective algorithm proposed in the art.
[0046] Step 13, selectively unloading KV data to CPU memory.
[0047] In the embodiment of the application, the unloading of KV data is selective. The contribution of KV data of different KV heads to the inference generation quality and the data distribution characteristics are different. By analyzing these indicators, and accordingly taking customized unloading and data reading strategies for different KV heads, specifically: using the average cosine similarity of offline analysis and the preset reuse threshold of the KV head, the reuse difficulty of the KV head is calculated; by pre-running the large language model inference, the number of KV heads that can be prefetched in each layer is calculated; for each layer, the number of KV heads with a reuse difficulty greater than 0 is counted, and the number of KV heads that can be prefetched is subtracted to determine the number of KV heads residing in the GPU memory. Then, the KV head with the highest reuse difficulty is selected to reside in the GPU memory. If the GPU memory is not sufficient to store, then within the range that the memory can accommodate, the batch of KV heads with the highest reuse difficulty is selected to reside in the GPU memory.
[0048] Step 14, the remaining calculation steps of the pre-filling stage are completed.
[0049] In the embodiment of the present application, the remaining steps of the pre-filling stage involve attention calculation, residual link, feedforward network calculation, etc., which are consistent with the existing process, so no further description is given.
[0050] Step 2, decoding stage based on dynamic approximate cache and data prefetching.
[0051] In the embodiment of the present application, in the decoding stage, autoregressive decoding is performed in combination with the latest token, which is the first new token generated in the pre-filling stage or the new token generated in the previous decoding step; the working process of the current layer of the large language model in a single decoding step is as follows: Q data and KV data are generated using the hidden vector of the current layer, and the KV data is added to the GPU memory or CPU memory; the hidden vector of the current layer is the word embedding vector of the first new token or the calculation result of the previous layer, and Q is the query; based on the hidden vector of the current layer, it is judged whether the next layer needs to prefetch KV data as a unit with the KV head, if so, the corresponding KV data is searched in the CPU memory, and after the KV data of the current layer is prefetched, the prefetching of the KV data of the next layer is started; the KV data is searched in the GPU memory using the Q data generated by the current layer, and the calculation of the current layer is completed in combination with the prefetched KV data; the process is repeated until the last layer, and the calculation result of the last layer is used to obtain a new token.
[0052] For ease of description, the current layer is denoted as the layer, the previous layer is denoted as the layer, and the next layer is denoted as the layer, and the decoding stage can be described as follows: in this step: autoregressive decoding is performed using the latest generated token and the KV cache, and the next new token is continuously generated until the generation is completed; if it is the first decoding step, the latest generated token is the first new token, which is generated by the pre-filling stage; if it is not the first decoding step, the latest generated token is the new token generated in the previous decoding step; autoregressive decoding includes the cooperation between different layers of the large language model, and for the layer, the working process is as follows: first, the corresponding Q data and KV data are generated using the input hidden vector of the layer, and the KV data is added to the KV cache; when , the hidden vector is the word embedding vector corresponding to the first new token, and when , the hidden vector is the output of the layer; then, the KV data of the layer is prefetched: based on the similarity of the hidden vectors between layers, the hidden vector of the layer is used to calculate the Approximate Q data of the layer, for KV data offloaded to CPU memory, in units of KV header, combined with the first Approximate Q data of the layer, for judging whether the first layer needs to prefetch KV data, if yes, through top-k attention technology, corresponding KV data is searched as the first layer KV data that needs to be prefetched, when the KV data prefetching of the first layer is completed, the prefetching of the first layer KV data is started, wherein, through top-k attention technology, corresponding KV data refers to searching out k KV data with the highest attention score; then, in the first layer, KV data retrieval is performed in the KV header residing in the GPU memory, when , the generated Q data and the KV data searched out from the KV header residing in the GPU memory are used to complete the calculation of the first layer, when , the generated Q data, the KV data searched out from the KV header residing in the GPU memory and the prefetched KV data of the first layer are used to complete the calculation of the first layer; the process is repeatedly performed until the last layer, and the output of the last layer is used to obtain the next new word element.
[0053] As shown in Figure 3 , the main workflow of the decoding stage is shown, the decoding stage involves the cooperation between different layers of the large language model, here, only the decoding stage workflow of the current layer (denoted as the first layer, the same below) is introduced, similarly, the process also involves the previous layer (denoted as the first layer, the same below) and the next layer (denoted as the first layer, the same below), the workflow of other layers is the same, and thus is not described herein. The main workflow is as follows:
[0054] Step 21, generating Q data and KV data of the first layer new word element.
[0055] In the embodiment of the application, the Q data and KV data generation process of the decoding stage is consistent with the existing process, which has been introduced in the foregoing background art, and thus is not described herein. The newly generated KV data is also added to the KV cache, specifically: for the offloaded KV header, the newly generated KV data is added to the data buffer thereof in the CPU memory; for the non-offloaded KV header, the newly generated KV data is added to the data buffer thereof in the GPU memory.
[0056] Step 22, generating approximate Q data of the first layer.
[0057] In the embodiment of the present application, the decoding stage will perform the reading of the KV cache data of the first layer at the same time of the calculation of the first layer, so as to reduce the KV data reading overhead exposed on the execution flow critical path, which is called KV data prefetching. Since the top-k attention algorithm is used to retrieve the key KV in the embodiment of the present application, the Q data of the first layer should be calculated in advance for the top-k retrieval and prefetching. For this purpose, the embodiment of the present application uses the similarity of the inter-layer hidden vectors of the large language model, uses the hidden vector of the first layer and the Q data mapping matrix of the first layer to calculate the approximate Q data of the first layer, and performs the top-k retrieval in advance to support the prefetching.
[0058] Step 23, using the approximate Q data to perform the cache query of the first layer.
[0059] In the embodiment of the present application, for the KV data that has been offloaded to the CPU memory, in addition to the KV prefetching technology, a dynamic approximate cache algorithm is also used, which saves the top-k KV data in the GPU memory in units of KV headers, to further reduce the data reading overhead. The dynamic cache algorithm needs to use the Q data to query whether the cached KV header in the GPU memory (specifically, the approximate cache data buffer of the first layer, which will be introduced later) hits, if the KV header hits, the KV data of the KV header can be directly reused, if it does not hit, the first layer KV data needs to be prefetched. To be combined with the prefetching, the Q data used for the query here is for the first layer, and the approximate Q data calculated in step 22 is used, specifically: in units of KV headers, combined with the approximate Q data of the first layer, the approximate cache data buffer of the first layer is queried to determine whether the cosine similarity of the Q data corresponding to the KV header and the approximate Q data of the first layer is greater than a pre-set reuse threshold, if yes, it means that the KV header hits, otherwise, it means that it does not hit; the reuse threshold is calculated by using the KV header importance algorithm, and the specific calculation method will be introduced later.
[0060] Step 24, using the approximate Q data to retrieve the first layer KV data that does not hit the top-k.
[0061] In the embodiment of the present application, top-k retrieval and prefetching are only needed when the KV data is not cached or missed (corresponding to cache initialization and cache update in turn, which will be introduced later). The retrieval here is for the first layer, and the index of the required read KV data will be generated; this part of retrieval occurs in the GPU and is based on the retrieval metadata stored in the GPU memory, which can be referred to in the foregoing background for details, and will not be repeated here.
[0062] Step 25, waiting for the KV data prefetching of the first layer to complete.
[0063] In the embodiment of the present application, except for the case of , the KV data of which may need to be prefetched, if the KV data is hit, it can be directly reused without prefetching, as described in the foregoing step 23; if it is missed, KV data prefetching is needed. Specifically, the KV data of the first layer is prefetched by the first layer, and the KV data of the first layer is prefetched by the first layer. To avoid conflicts and competition of PCIe bandwidth of CPU-GPU interconnection caused by different layers of prefetching, the KV data prefetching of the first layer needs to start after the KV data prefetching of the first layer is completed, so the first layer needs to use CPU-GPU data transmission synchronization operation to wait for the data prefetching to complete.
[0064] Step 26, triggering the data prefetching of the first layer.
[0065] In the embodiment of the present application, when the KV data prefetching of the first layer is completed, the KV data prefetching of the first layer will start immediately using the index of the KV data obtained in step 24. The data will be prefetched into the approximate cache data buffer area (located in the GPU memory) of the first layer.
[0066] Step 27, performing top-k retrieval in the KV header data of the first layer that is not offloaded.
[0067] In the embodiment of the present application, part of the KV header data will not be offloaded, which has been introduced in step 13 of the foregoing pre-population stage. Since these data are located in the GPU memory, prefetching is not needed. Therefore, in the embodiment of the present application, the real Q data (i.e. the Q data generated in step 21) of the first layer is used to retrieve the index of the top-k KV data of these KV headers.
[0068] Step 28, complete the first layer remaining calculation step.
[0069] In the embodiment of the present application, the remaining steps of the decoding stage involve attention calculation, residual link, feedforward network calculation, etc. Except for attention calculation, they are consistent with the original process before optimization. Attention calculation needs to use data from two cache areas, which are: The data buffer area of the KV head of the layer that is not unloaded and the approximate cache data buffer area of the first For the former, only the top-k KV data needs to be read, and their positions in the buffer area are identified by the index of the KV data given in step 27. The latter contains the top-k KV data required for calculation, without further screening.
[0070] The above scheme provided by the embodiment of the present application offloads most of the KV data to the CPU memory in the large language model inference process. For the reading of KV data from the CPU memory to the GPU, the system reduces the reading amount by using top-k attention. And the KV data that has been read to the GPU memory is cached, using a dynamic approximate cache algorithm, which manages KV heads as the basic unit. When KV data needs to be read during the inference process, it is first determined whether the data in the cache can be reused. If the cache hits, the data in the cache is directly reused, reducing the data reading overhead from the CPU memory to the GPU memory. If the cache misses and the data cannot be reused, data prefetching is performed. Through the three optimization methods of selective offloading, dynamic approximate caching, and data prefetching, the present application can effectively reduce the memory occupied by KV data and minimize the KV data reading overhead from the CPU to the GPU, so that the inference performance reaches an ideal level.
[0071] In order to more clearly show the technical solutions provided by the present application and the technical effects produced, the following mainly introduces the selective offloading of KV head, the dynamic approximate caching algorithm, and the prefetching and synchronization of KV data in detail. The rest of the processes not described in detail can refer to the conventional technology.
[0072] I. Selective offloading of KV head
[0073] In the embodiment of the present application, the hit rate of the dynamic approximate caching algorithm is mainly determined by the importance of the KV head and the Q data distribution characteristics corresponding to the KV head (which determines the similarity level of the Q head in the continuous decoding step). The decoding step refers to the generation process of a word element, covering the working process of all layers of the model, and defines that for the same layer, the working process of the same layer is continuous when the adjacent word element is obtained by decoding; for example, assuming that the model has 3 layers, i.e. the 0th layer, the 1st layer and the 2nd layer, in the process of decoding to obtain two word elements a and b, six processes a0, a1, a2, b0, b1 and b2 are required, wherein a0, a1 and a2 refer to the working processes of the 0th layer, the 1st layer and the 2nd layer when the word element a is decoded, which is a decoding step, and b0, b1 and b2 refer to the working processes of the 0th layer, the 1st layer and the 2nd layer when the word element b is decoded, which is also a decoding step. In this example, a0 and b0 are defined as continuous processes, and a1 and b1, a2 and b2 are also defined as continuous processes.
[0074] Considering that there may be KV heads with high importance and low similarity level in the large language model, the cache hit rate of these KV heads is difficult to reach a high level (referred to as difficult-to-reuse KV heads), and the KV cache data of these KV heads still needs to be read from the CPU memory. Even if the prefetching technology is adopted, the reading delay of the KV data of these KV heads may still be difficult to be completely masked by calculation, which will slow down the inference speed. The embodiment of the present application further uses a selective offloading method to cope with this problem, and the difficult-to-reuse KV heads are resident in the GPU video memory, and only other KV heads are offloaded, and the overall process is as shown in Figure 4 The attention calculation unit is responsible for attention calculation in the decoding stage.
[0075] In the embodiment of the present application, the reuse difficulty of each KV head is quantified in advance: for each KV head of each layer, the average cosine similarity is obtained by offline analysis, wherein in the offline analysis process, a plurality of sequences are randomly sampled from the selected evaluation data set and input into the large language model, and for each KV head of each layer, the average cosine similarity is calculated using all the outputs; for a single KV head i, the average cosine similarity is denoted as , the reuse threshold and the error tolerance difference , and the reuse difficulty is:
[0076] ;
[0077] wherein, is the reuse difficulty of a single KV head i. The physical meaning of the above formula is that the KV head with an actual cosine similarity lower than the reuse threshold is considered to be difficult to reuse, and the greater the numerical difference between the two, the higher the difficulty of cache reuse. The reuse threshold The KV head importance algorithm can be used for calculation, which will be described later.
[0078] For example, offline analysis can use 30 sequences randomly sampled from the LongBench benchmark dataset (which is a bilingual and multi-task dataset with various sequences of different lengths, distributions, patterns, languages and fields for comprehensive evaluation of long context understanding ability), The value of can be 0.05.
[0079] For each layer, after quantifying the multiplexing difficulty of each KV head, the KV heads are ranked in descending order according to the multiplexing difficulty; the number of KV heads with multiplexing difficulty greater than 0 (i.e. ) is counted .
[0080] In addition, the number of KV heads that can be prefetched (the read latency can be completely covered by calculation) needs to be analyzed in advance. Specifically, for the user-set running configuration, the large language model is pre-run once for inference, and the decoding time of each layer is counted , combined with the size of the KV data corresponding to each KV head in each layer and the PCle bandwidth, the number of KV heads that can be prefetched is calculated, represented as:
[0081] ;
[0082] Wherein, is the number of KV heads that can be prefetched in each layer, is the size of the KV data corresponding to each KV head, is the PCle bandwidth, and PCle is a high-speed serial computer expansion bus standard.
[0083] The number and can be calculated to make a selective offloading decision. For a certain layer of a large language model, the corresponding number is subtracted from the number to determine the number of KV heads resident in the GPU memory, represented as:
[0084] ;
[0085] Wherein, the max function outputs the maximum value in the parentheses.
[0086] After determining the number , the corresponding number of KV heads from the front end of the sorted results of the KV heads is selected, and the selected KV heads are used as the KV heads resident in the GPU memory, and the corresponding KV data is stored in the GPU memory. If the GPU memory is insufficient and cannot accommodate If the KV data of the KV head is selected, the KV data of the batch of KV heads with the highest multiplexing difficulty is selected from the sorting result of the KV head and is resident in the range of the display memory to maximize the performance benefit. The KV data of the remaining KV heads is all unloaded to the CPU memory to complete the selective unloading of the KV head.
[0087] II. Dynamic approximate caching algorithm.
[0088] The basis of the dynamic approximate caching strategy is that, for the same Q head of the same layer of a large language model, in the continuous decoding steps, the Q data has high directional similarity (quantified using cosine similarity, which is a floating-point number between -1 and 1, and the closer the cosine similarity of two vectors is to 1, the more similar their directions are). For example, on the Llama3-8B-1048K model, 94% of the Q heads have a Q vector cosine similarity of more than 0.85 in the continuous inference steps. The directional similarity of the Q data makes the KV data indexes generated by the top-k retrieval also have high similarity. Specifically, the top-k retrieval uses the Q data and the K data to calculate the score of each KV data by the following formula:
[0089]
[0090] where q is the Q data of the latest token in the decoding stage, which generally represents the K data of any token in the KV cache, corresponding to q, the modulus of q, is the cosine value of the angle between q and . Obviously, given a set of KV data of tokens, the relative size order of the scores is only affected by and because they share the same in the calculation process. Then, if two Q data are directionally similar, the angles between them and the same K data will also be relatively close, that is, the results of calculated by them will be highly similar. Therefore, the relative size order of the final calculated scores will also be highly similar, making the results of the top-k retrieval close. From the above analysis, an important conclusion can be drawn: in the inference process of a large language model, due to the directional similarity of the Q data, the top-k KV data required in the continuous decoding steps have high overlap. In this case, the top-k KV data read to the GPU display memory in the historical decoding steps can be reused without being read again.
[0091] Based on the above conclusion, an embodiment of the present application proposes a dynamic approximate caching algorithm based on Q similarity, as shown in Figure 5 The principle of the dynamic approximate cache algorithm is shown. It mainly includes the following steps: cache initialization; cache query; and cache update.
[0092] 1. Cache initialization.
[0093] The inference process starts with an empty cache. In the embodiment of the present application, in the first decoding step, each layer except the 0th layer is respectively initialized, that is, the top-k KV data is retrieved and read from the KV header unloaded to the CPU memory using the generated Q data, and saved in the approximate cache data buffer of the corresponding layer, to complete the cache initialization. The first decoding step refers to the first decoding step after the end of the pre-padding stage, which generates the first new token using the first new token generated in the pre-padding stage.
[0094] 2. Cache query.
[0095] After the first decoding step, for a new decoding step, each KV header of each layer except the 0th layer in the large language model needs to be queried. For the current layer (for example, the 1st layer), the approximate Q data of the next layer (for example, the 2nd layer) is calculated, and the cosine similarity is calculated using the Q data corresponding to the KV data cached in the approximate cache data buffer of the next layer. If the cosine similarity is higher than the pre-set reuse threshold, it is considered as a cache hit, indicating that the KV data of the corresponding KV header can be directly reused; otherwise, it is considered as a cache miss, and the cache update step is entered. The approximate Q data of the next layer (for example, the 2nd layer) is calculated, and the cosine similarity is calculated using the Q data corresponding to the KV data cached in the approximate cache data buffer of the next layer. If the cosine similarity is higher than the pre-set reuse threshold, it is considered as a cache hit, indicating that the KV data of the corresponding KV header can be directly reused; otherwise, it is considered as a cache miss, and the cache update step is entered.
[0096] Those skilled in the art can understand that the KV data cached in the approximate cache data buffer is partitioned by KV header, and each KV header has a Q data as its tag (saved in the initialization or update process). In a single decoding step, a Q header corresponds to only one Q data; in cache query, for each KV header, the cosine similarity between the approximate Q data and the Q data tag in the cache is calculated to determine whether it is hit; for the GQA model, a KV header corresponds to multiple Q headers, so it has multiple Q data tags, and the cosine similarity calculation will produce multiple results, which need to be aggregated. This will be introduced in the following supporting group query attention.
[0097] 3. Cache update.
[0098] For the KV header determined as a miss in the cache query, the required next layer (for example, the 2nd layer) is retrieved using the approximate Q data, and the KV data of the next layer is read and saved in the approximate cache data buffer of the next layer, to complete the cache update. After the index of the KV data of the next layer (see the description of step 24 above for details) is obtained, the corresponding KV data is read from the CPU memory by KV data prefetching (see the description of step 26 above for details) and is overwritten to the approximate cache data buffer of the next layer (see the description of step 26 above for details) to replace the KV data of the missed KV header, thereby completing the cache update.
[0099] Figure 5 The input Q data (specifically, approximate Q data) of the algorithm is referred to as Q. The Q data in the cache is referred to as Q. The K data (key data) is referred to as K. The V data (value data) is referred to as V. The Q data corresponding to the KV header is referred to as Q. The cosine similarity of Q and Q is referred to as cos. The cosine similarity of Q and Q is compared with the size of the reuse threshold to determine whether a hit occurs. If a hit occurs, it means that the corresponding KV data can be reused. Otherwise, the corresponding KV data needs to be retrieved from the KV data of the CPU (K and V in the bottom block refer to the retrieved KV data), and the approximate cache data buffer is updated.
[0100] The above three steps are collectively referred to as cache management. Compared with the traditional LRU (least recently used algorithm) / LFU (least frequently used algorithm) cache algorithm, the dynamic approximate cache algorithm provided by the present application greatly simplifies the query and update process of the cache. The query process only needs to calculate the cosine similarity based on the Q header, which is lower in cost and more suitable for the hardware characteristics of the GPU than the table lookup (such as the hash table) of the traditional cache algorithm. In the cache update process, the KV data is directly overwritten to the data buffer of the cache, which reduces one data copy within the GPU memory compared with the traditional algorithm. At the same time, the cache metadata update only needs to overwrite the Q data, while the traditional algorithm needs to perform a low concurrency and time-consuming linked list update operation by the CPU. The cache metadata here, which is the data that assists the cache in querying and updating, is the Q data possessed by each KV header as mentioned above. In the cache query, the cosine similarity is calculated between the approximate Q data of the next layer and the Q data in the approximate cache data buffer. The former is new data input, and the latter is old data existing in the cache. When a miss occurs, not only the KV data needs to be read and updated, but also the Q data in the cache is considered to be invalid, and therefore needs to be updated. In summary, the dynamic approximate cache algorithm provided by the present application significantly reduces the overhead of cache management. In the single-layer single-step reasoning of a large language model, the total time consumption is only 5-10 microseconds, which accounts for less than 1% of the total time consumption and can be ignored. Therefore, compared with the traditional algorithm, the dynamic approximate cache algorithm is more suitable for the inference of a large language model with strict time requirements.
[0101] Preferably, the application further provides a calculation scheme of multiplexing threshold for the aforementioned cache query step. Specifically, a multiplexing threshold is set for each KV head in advance to determine whether the cache hits during the cache query process and is applied to the selective offloading of the KV head. The embodiment of the application adopts an algorithm (referred to as a KV head importance algorithm) capable of quantitatively analyzing the contribution of the KV head to the inference accuracy to assist the multiplexing threshold setting. Specifically, before the inference starts, the embodiment of the application analyzes the importance of all KV heads in a given large language model using the KV head importance algorithm. For a single KV head i, the importance is denoted as , , the closer to 1, the more important the KV head is, and then is used as the upper bound of the multiplexing threshold (exemplarily, the value of may be set to 0.9), and the multiplexing threshold of each KV head is set using the following formula :
[0102] ;
[0103] ;
[0104] ;
[0105] wherein the parameter , p is an attenuation speed reinforcement index, exemplarily, the value of p is 2 or 3, which helps to highlight the difference between the heads; and are intermediate variables generated in the calculation process.
[0106] The principle of the above formula is to convert the multiplexing threshold into an angle, and then attenuate it according to the importance of the KV head. The more important the KV head is, the closer its final multiplexing threshold is to the upper bound ; the less important it is, the closer it is to the lower bound -1.
[0107] Exemplarily, the KV head importance algorithm can select the Duo-Attention algorithm, which is an algorithm capable of quantitatively analyzing the contribution of the KV head to the inference accuracy.
[0108] Preferably, the dynamic approximate cache algorithm provided by the application also supports group query attention (which has been introduced in the background section). The cosine similarity is calculated based on the Q head, but the multiplexing threshold is set based on the KV head, and the Q head and the KV head are not one-to-one corresponding in the model using the GQA technology. Therefore, it is necessary to aggregate the calculated cosine similarity within the Q head group. Based on the KV head importance algorithm, the embodiment of the application also analyzes the importance of all Q heads. In the cache query step, the importance is preferably used as the weight, and the harmonic mean method is used to aggregate the Q heads within the group, and the specific formula is:
[0109] ;
[0110] wherein, is the importance of Q head j, is the cosine similarity of Q data corresponding to Q head j and approximate Q data, m is the number of Q heads corresponding to one KV head, s is the aggregated cosine similarity, and comparison with a multiplexing threshold is performed to determine whether the m Q data corresponding to the KV head hit.
[0111] III. KV data prefetching and synchronization.
[0112] In the embodiment of the application, a KV data prefetching technology is used to read the KV head that fails to hit the cache from the CPU memory. The algorithm principle of this technology is derived from InfiniGen (it is a KV cache offload reasoning system that uses prefetching and top-k attention technology to improve reasoning performance). Compared with InfiniGen, the embodiment of the application has two innovations, one is to use a zero-copy data transmission method for data prefetching, and the other is to use a GPU-centered data transmission synchronization method.
[0113] In the embodiment of the application, the zero-copy data transmission method refers to a zero-copy method based on GPU Direct (direct communication technology) to perform KV data prefetching between CPU and GPU. Zero-copy means that data is directly copied from the source location to the target location without any temporary buffer for transfer. In the embodiment of the application, the GdrCopy library based on GPU Direct (it is an open source library mainly used to realize high-speed data transmission between GPU memory and CPU memory) is used to directly copy the required KV data from the CPU memory to the specified GPU memory buffer without first collecting the discrete distributed KV data into a continuous data block in the CPU memory for copying. This zero-copy data transmission method bypasses the high-cost CPU operation and provides a CPU-to-GPU transmission bandwidth of up to 21 GB / sec (gigabyte per second).
[0114] In the embodiment of the application, the KV data transmission process is also optimized by a GPU-centered data transmission synchronization method, such as Figure 6The data transmission between the CPU and the GPU needs to be synchronized, and the progress of the data transmission is informed to the GPU, so that the GPU is in a continuous waiting state before the data transmission is completed. The traditional data transmission synchronization method takes the CPU as the core, blocks the start of the GPU kernel function responsible by the CPU, and blocks the operation of the subsequent kernel function of the GPU, so as to make the GPU fall into a continuous waiting state. However, this will cause the GPU to be idle for a period of time after the synchronization is completed: the GPU kernel function can start to execute only after the CPU starts, and the CPU cannot start the subsequent kernel function until the synchronization is completed. After the synchronization is completed, this start process will be exposed to the critical path of the execution flow, causing the GPU to be idle. The embodiment of the application changes the data transmission synchronization to be centered on the GPU, solving the problem. Specifically, using the NVIDIA (NVIDIA) Unified Virtual Addressing (UVA) technology, a shared memory accessible directly by the CPU and the GPU is established in the CPU memory, and a signal variable is set in the shared memory to describe the data transmission state; for a KV header, when the transmission is not completed, the value of the signal variable is 0; when the transmission is completed, the value of the signal variable is 1; the GPU polls the signal variable through a kernel function, if the value of the signal variable is 0, the GPU will be in a busy waiting state, otherwise, the GPU exits the polling and executes the operation of the subsequent kernel function. This synchronization strategy does not block the start of the kernel function after the synchronization of the CPU, completely eliminating the GPU idle caused by the data transmission synchronization.
[0115] Based on the above introduction, the present application provides an overall scheme architecture as shown in the figure. Figure 7 Figure 7 In the embodiment, the KV cache manager is mainly responsible for managing the storage, reading, caching of all KV data, including selective offloading, configuration and execution of dynamic approximate caching algorithm, and retrieval of top-k KV cache data; the data transmission engine is mainly responsible for pre-fetching KV cache data from the CPU memory to the GPU memory, and performing data transmission in a zero-copy manner; it contains a plurality of data transmission threads; the data transmission synchronization controller is mainly responsible for synchronizing the CPU-GPU data transmission, and adopts a synchronization method centered on the GPU, without affecting the start of the GPU kernel function after the synchronization; the kernel function start thread is located on the CPU and is responsible by an independent thread, without being disturbed by the data transmission synchronization; the inference execution module (the weights of the large language model are included in the inference execution unit) is mainly responsible for executing the inference operation of the large language model.
[0116] Through the above description of the embodiments, those skilled in the art can clearly understand that the above embodiments can be implemented by software, or by using software plus necessary general-purpose hardware platforms. Based on this understanding, the technical solutions of the above embodiments can be embodied in the form of a software product, which can be stored in a non-volatile storage medium (such as a CD-ROM, USB flash drive, mobile hard drive, etc.), including several instructions to cause a computer device (such as a personal computer, server, or network device, etc.) to execute the methods described in the various embodiments of the present invention.
[0117] Example 2
[0118] This invention also provides a large language model inference optimization system, which is mainly used to implement the methods provided in the foregoing embodiments, such as... Figure 8 As shown, the system mainly includes:
[0119] The pre-filling unit based on selective offloading of KV headers (applied to the pre-filling stage) is used to combine the difficulty of KV header reuse, cache part of the KV data of the KV header in GPU memory, and offload the KV data of the other part of the KV header to CPU memory; where KV header is a type of attention header in large language models, KV stands for key and value, GPU stands for graphics processor, and CPU stands for central processing unit.
[0120] The decoding unit based on dynamic approximation caching and data prefetching (applied to the decoding stage) is used to perform autoregressive decoding by combining the latest lexical units. The latest lexical unit is the first new lexical unit generated in the pre-filling stage or the new lexical unit generated in the previous decoding step. The working process of the current layer of the large language model in a single decoding step is as follows: Q data and KV data are generated using the latent vector of the current layer. The KV data is added to the GPU memory or CPU memory. The latent vector of the current layer is the word embedding vector of the first new lexical unit or the calculation result of the previous layer. Q is the query. Based on the latent vector of the current layer, the KV head is used as the unit to determine whether the next layer needs to prefetch KV data. If so, the corresponding KV data is retrieved in the CPU memory. After the KV data prefetching of the current layer is completed, the prefetching of the KV data of the next layer is started. The Q data generated by the current layer is used to retrieve KV data in the GPU memory. The calculation of the current layer is completed by combining the prefetched KV data. This process is repeated until the last layer. The calculation result of the last layer is used to obtain the new lexical unit.
[0121] Since the main technical details of the above system have been described in detail in the previous embodiments, they will not be repeated here.
[0122] Those skilled in the art can clearly understand that, for the convenience and brevity of description, only the division of the above functional modules is taken as an example, and in actual application, the above functions can be completed by different functional modules according to needs, that is, the internal structure of the system is divided into different functional modules to complete all or part of the functions described above.
[0123] Embodiment three
[0124] The application further provides a processing device, as shown in the drawings, mainly comprising: one or more processors; a memory for storing one or more programs; wherein when the one or more programs are executed by the one or more processors, the one or more processors realize the method provided by the foregoing embodiments. Figure 9
[0125] Further, the processing device further comprises at least one input device and at least one output device; in the processing device, the processor, the memory, the input device and the output device are connected through a bus.
[0126] In the embodiments of the application, the specific types of the memory, the input device and the output device are not limited; for example:
[0127] The input device can be a touch screen, an image acquisition device, a physical key or a mouse, etc.
[0128] The output device can be a display terminal.
[0129] The memory can be a random access memory (RAM), and can also be a non-volatile memory, for example, a disk memory.
[0130] Embodiment four
[0131] The application further provides a readable storage medium, which stores a computer program, and when the computer program is executed by a processor, the method provided by the foregoing embodiments is realized.
[0132] In the embodiments of the application, the readable storage medium as the computer readable storage medium can be arranged in the foregoing processing device, for example, as the memory in the processing device. In addition, the readable storage medium can also be a U disk, a mobile hard disk, a read-only memory (ROM), a magnetic disk or an optical disk and various program codes can be stored in the medium.
[0133] The above description is merely preferred embodiments of the present application, but the protection scope of the present application is not limited thereto, and any changes or substitutions easily conceived by those skilled in the art within the technical scope disclosed by the present application should be covered within the protection scope of the present application. Therefore, the protection scope of the present application should be subject to the protection scope of the claims. The information disclosed in the background section of the present application is merely intended to deepen the understanding of the general background of the present application, and should not be regarded as acknowledging or implying in any form that the information constitutes the prior art known to those skilled in the art.
Claims
1. A large language model inference optimization method, characterized in that, include: In the pre-filling stage, considering the reuse difficulty of the KV header, a portion of the KV data in the KV header is cached in GPU memory, while the other portion is unloaded to CPU memory. The reuse difficulty of the KV header is calculated using the mean cosine similarity from offline analysis and the preset reuse threshold of the KV header. The KV header is a type of attention header in large language models, where KV stands for key and value, GPU stands for graphics processing unit, and CPU stands for central processing unit. During the decoding phase, autoregressive decoding is performed by combining the latest lexical units. The latest lexical unit is the first new lexical unit generated in the pre-filling phase or the new lexical unit generated in the previous decoding step. The working process of the current layer of the large language model in a single decoding step is as follows: Q-data and KV-data are generated using the latent vectors of the current layer. The KV-data is added to GPU memory or CPU memory. The latent vector of the current layer is the word embedding vector of the first new lexical unit or the calculation result of the previous layer, and Q represents the query. Based on the latent vectors of the current layer, using the KV-head as the unit, it is determined whether the next layer needs to prefetch KV-data. If so, the corresponding KV-data is retrieved from CPU memory. After the KV-data prefetching of the current layer is completed, the prefetching of the KV-data of the next layer is initiated. The KV-data generated by the current layer is used to retrieve KV-data from GPU memory, and the calculation of the current layer is completed by combining the prefetched KV-data. This process is repeated until the last layer, and the calculation result of the last layer is used to obtain a new lexical unit.
2. The large language model inference optimization method of claim 1, wherein, The method of considering the difficulty of reusing KV headers, caching part of the KV data in GPU memory and offloading the other part of the KV data to CPU memory includes: The reuse difficulty of each key-value (KV) header is quantified, and the number of KV headers that can be prefetched in each layer is calculated. Combining the number of KV headers that can be prefetched in each layer with the reuse difficulty of the corresponding KV header, the KV headers residing in GPU memory are selected, and their corresponding KV data is cached in GPU memory. The KV data of the remaining KV headers are offloaded to CPU memory. In particular, each layer of the large language model has multiple attention heads, which are divided into two categories: Q-heads and KV headers. The process of selecting key-value (KV) headers residing in GPU memory includes: for each layer, quantifying the reuse difficulty of each KV header and sorting them in descending order of reuse difficulty; counting the number of KV headers with reuse difficulty greater than 0, subtracting the number of KV headers that can be prefetched for the corresponding layer, and determining the number of KV headers residing in GPU memory; and selecting the corresponding number of KV headers from the front of the sort as the KV headers residing in GPU memory.
3. The large language model inference optimization method of claim 2, wherein, The process of quantifying the reuse difficulty of each KV head and calculating the number of KV heads that can be prefetched at each layer includes: The quantification of the reuse difficulty of each KV head includes: for each KV head of each layer, the corresponding average cosine similarity is obtained through offline analysis. In the offline analysis process, several sequences are randomly sampled from the selected evaluation dataset and input into the large language model. For each KV head of each layer, the average cosine similarity is calculated using all the outputs. For a single KV head i, let its average cosine similarity be denoted as The corresponding multiplexing threshold is calculated in advance using the KV head importance algorithm and an error tolerance difference is set, then the multiplexing difficulty is: ; wherein, multiplexing difficulty for a single KV head i; The calculating the number of KV heads that each layer can prefetch comprises: performing one inference pre-run on the large language model, and counting the decoding time of each layer In combination with the size of the KV data corresponding to a single KV head in each layer and the PCIe bandwidth, the number of KV heads that can be prefetched is calculated and represented as: ; wherein, is the number of KV heads that can be prefetched for each layer, is the size of KV data corresponding to a single KV head, is the PCle bandwidth, PCle is a high-speed serial computer expansion bus standard.
4. The large language model inference optimization method of claim 1, wherein, The current layer-based implicit vector is used to determine whether the next layer needs to prefetch KV data in units of KV headers. If so, the corresponding KV data is retrieved in the CPU memory, including: Let the current layer be denoted as the th. The next layer is denoted as the first layer. Layer, based on the similarity of latent vectors between layers, utilizes the 1st layer... The latent vectors of the layer are calculated to obtain the first layer. The approximate Q data of the layer is used with a dynamic approximation caching algorithm, taking the KV header as the unit, to determine the first... The system checks if the key-value (KV) data in the KV header of the approximate cached data buffer of the layer is matched. If the KV header is matched, it means that the corresponding KV data in the KV header can be reused; if it is not matched, it cannot be reused. This is combined with the first... The approximate Q-data of the layer, using the key-value header as the unit, is retrieved from CPU memory using top-k attention technology, and used as the first layer to be prefetched. Layer KV data, wherein retrieving the corresponding KV data through top-k attention technology means retrieving the k KV data with the highest attention scores, the approximate cache data buffer is located in GPU video memory and is used to cache prefetched KV data, wherein the KV data is partitioned according to KV header, and the approximate cache data buffer is initialized in the first decoding step; The algorithm uses a dynamic approximation caching mechanism, with the key-value header as the unit, to determine the first... Whether the KV data in the KV header of the approximate cached data buffer of the layer is hit includes: taking the KV header as the unit, combined with the first Approximate Q-data of the layer, in the first layer The query is performed in the approximate cached data buffer of the layer to determine whether the Q data corresponding to the KV header is the same as that of the first layer. If the cosine similarity of the approximate Q data of the layer is greater than a preset reuse threshold, it indicates that the KV header has been hit; otherwise, it indicates that it has not been hit. The reuse threshold is calculated using the KV header importance algorithm.
5. The large language model inference optimization method of claim 4, wherein, The dynamic approximate caching algorithm includes caching initialization, caching query, and caching update. Caching initialization: in the first decoding step, each layer except the 0th layer is initialized, that is, the top-k KV data is retrieved and read from the KV header unloaded to the CPU memory using the generated Q data, and saved in the approximate caching data buffer of the corresponding layer, completing the caching initialization; wherein the first decoding step refers to the first decoding step after the pre-filling phase ends; Caching query: for a new decoding step, each KV header of each layer in the large language model except the 0th layer performs caching query. The decoding step refers to the generation process of a word element, covering the working process of all layers of the model; for the current layer, in units of KV headers, the cosine similarity is calculated using the approximate Q data of the next layer and the Q data corresponding to the KV header in the approximate caching data buffer of the next layer to determine whether the KV data hits. If not, go to the caching update step; Caching update: for the KV header determined not to hit in the caching query, retrieve the index of the next layer KV data that needs to be prefetched, read the corresponding KV data from the CPU memory through KV data prefetching, and overwrite it to the approximate caching data buffer of the next layer, completing the caching update.
6. The large language model inference optimization method according to any one of claims 3-5, characterized in that, The calculation method of the multiplexing threshold includes: The importance of all KV heads in the large language model is analyzed using the KV head importance algorithm. For a single KV head i, the importance is recorded as , , and then, taking as the upper limit of the reuse threshold, the reuse threshold of the KV head is set as follows : ; ; ; where the parameters p is an attenuation rate enhancement exponent, and are intermediate variables generated during the calculation process.
7. The large language model inference optimization method of claim 1, wherein, In the KV data prefetching process, a zero-copy method based on GPU Direct is used to perform KV data prefetching between CPU and GPU, and a GPU-centered data transmission synchronization method is used; The zero-copy method based on GPU Direct directly copies the required KV data from the CPU memory to the specified GPU memory buffer; the GPU Direct is a GPU direct communication technology; The GPU-centered data transmission synchronization method includes: setting the data transmission synchronization to be GPU-centered, using the NVIDIA unified virtual addressing technology to establish a shared memory in the CPU memory that can be directly accessed by both CPU and GPU, and setting a signal variable in the shared memory to describe the data transmission state; for a KV header, when the transmission is not completed, the signal variable value is 0; when the transmission is completed, the signal variable value is 1; the GPU polls the signal variable through a kernel function. If the signal variable value is 0, the GPU will be in busy waiting state, otherwise, the GPU exits the polling.
8. A large language model inference optimization system, characterized in that, The method for implementing any one of claims 1-7 includes: The pre-filling unit based on KV head selective offloading is used to combine the multiplexing difficulty of KV heads, buffer a part of KV data of KV heads in GPU memory, and offload another part of KV data of KV heads to CPU memory; wherein, the average cosine similarity of offline analysis is used to calculate the multiplexing difficulty of KV heads with a preset multiplexing threshold of KV heads, the KV head is a type of attention head in a large language model, the KV is a key and value, the GPU is a graphics processing unit, and the CPU is a central processing unit; The decoding unit based on dynamic approximate caching and data prefetching is used to combine the latest word to perform autoregressive decoding, and the latest word is the first new word generated in the pre-filling stage or the new word generated in the last decoding step; the working process of the current layer of the large language model in a single decoding step is as follows: Q data and KV data are generated by using the hidden vector of the current layer, and the KV data is added to the GPU memory or the CPU memory, the hidden vector of the current layer is the word embedding vector of the first new word or the calculation result of the last layer, and Q is a query; based on the hidden vector of the current layer, it is judged whether the next layer needs to prefetch KV data by taking the KV head as a unit, if yes, the corresponding KV data is searched in the CPU memory, after the KV data of the current layer is prefetched, the prefetching of the KV data of the next layer is started; the KV data is searched in the GPU memory by using the Q data generated by the current layer, and the calculation of the current layer is completed by combining the prefetched KV data; the process is repeated until the last layer, and the calculation result of the last layer is used to obtain a new word.
9. A processing device, characterized by Comprise: One or more processors; Memory for storing one or more programs; When the one or more programs are executed by the one or more processors, the one or more processors implement the method of any one of claims 1-7.
10. A readable storage medium, storing a computer program, characterized in that, When the computer program is executed by the processor, the method of any one of claims 1-7 is implemented.
Citation Information
Patent Citations
Language model reasoning optimization method, electronic equipment, storage medium and program product
CN118396128A
Model reasoning method and device and electronic equipment
CN119129746A