Key-value cache management method, controller, device, apparatus, medium and product

Abstract

The application discloses a key-value cache management method, a controller, equipment, a device, a medium and a product, which are applied to a target device and include the following steps: caching key-value data generated by each Transformer layer stacked in an inference process of a target large language model into a first storage area of a display memory of the target device, wherein the key-value data includes key-value data obtained by performing forward calculation on input target symbols of each layer by each layer; if it is monitored that the remaining capacity of the first storage area is lower than a first preset capacity threshold, discarding the key-value data of the target Transformer layer cached in the first storage area according to a preset discard strategy; and if it is monitored that the key-value data of the target Transformer layer needs to be used in the current inference process, re-executing the forward calculation of the target symbols of the target Transformer layer by using the computing resources of the display memory, so as to restore the key-value data of the target Transformer layer and write the key-value data into the first storage area.

Description

Technical Field

[0001] This application relates to the field of data storage technology, and in particular to a key-value cache management method, controller, device, apparatus, medium and product. Background Technology

[0002] During inference, the self-attention mechanism of Large Language Models (LLMs) requires caching the key and value matrices for each historical symbol, known as the KV cache. As the number of generated symbols increases, the memory usage of the KV cache grows linearly. In long text generation scenarios, the KV cache can account for 70% to 80% of the total GPU memory used for inference, becoming a core bottleneck restricting inference performance and concurrency capabilities.

[0003] Existing key-value (KV) cache management schemes divide KV data into two parts based on the symbol time dimension: new symbols (most recently generated symbols) and old symbols (earlier generated symbols). Specifically, new symbols correspond to KV data across all levels, and old symbols correspond to KV data across all levels. This relies on multi-level storage media, for example, from top to bottom: VRAM → RAM → SSD, moving KV data of old symbols to the next level of storage. When KV data moved to the next level of storage is needed, it is then moved back to the top level of storage. This KV data swapping operation generates significant input / output (I / O) latency, severely impacting the inference speed of large language models, and requires additional storage devices and bandwidth resources, resulting in low KV cache management efficiency. Summary of the Invention

[0004] The purpose of this application is to provide a key-value cache management method, controller, device, apparatus, medium, and product to solve the problem of low efficiency in existing key-value cache management.

[0005] To solve the above-mentioned technical problems, this specification is implemented as follows: Firstly, a key-value cache management method is provided, applied to a target device, the method comprising: The key-value data generated by each Transformer layer stacked during the inference process of the target large language model is cached in the first storage area of ​​the target device's video memory. The key-value data includes the key-value data obtained by each Transformer layer performing forward computation on the target symbol input to each Transformer layer. If the remaining capacity of the first storage area is found to be lower than the first preset capacity threshold, the key-value data of the first target Transformer layer cached in the first storage area will be discarded according to the preset discarding policy. If the monitoring detects that the current inference process requires the use of the key-value data of the first target Transformer layer, then the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to restore the key-value data of the first target Transformer layer and write it to the first storage area.

[0006] Secondly, a storage controller is provided, which is connected to a target device, for caching key-value data generated by the Transformer layer of a target large language model on the target device during inference. The target device, when executed, implements the steps of the method described in the first aspect.

[0007] Thirdly, a storage device is provided, including a storage controller as described in the second aspect.

[0008] Fourthly, a key-value cache management device is provided, including a processor and a memory, the memory storing a program or instructions executable on the processor, the program or instructions, when executed by the processor, implementing the steps of the method described in the first aspect.

[0009] Fifthly, a non-volatile computer-readable storage medium is provided, on which a program or instructions are stored, which, when executed by a processor, implement the steps of the method described in the first aspect.

[0010] In a sixth aspect, a computer program product is provided, comprising a non-volatile computer-readable storage medium storing a computer program operable to cause a computer to perform the steps of the method described in the first aspect.

[0011] In this embodiment, the key-value data generated by each Transformer layer stacked during the inference process of the target large language model is cached in the first storage area of ​​the target device's video memory. The key-value data includes the key-value data obtained by each Transformer layer performing forward computation on the target symbol input to each Transformer layer. If the remaining capacity of the first storage area is monitored to be lower than a first preset capacity threshold, the key-value data of the first target Transformer layer cached in the first storage area is discarded according to a preset discarding strategy. If the current inference process needs to use the key-value data of the first target Transformer layer, the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to restore the key-value data of the first target Transformer layer and write it into the first storage area. Therefore, a strategy of discarding key-value data and recompiling is adopted. The key-value data cache of the target Transformer layer is directly deleted from the GPU memory. During inference, it is recomputed and restored using the GPU memory computing resources of the target device. Since the recompiling is completed entirely within the target device and does not involve I / O transmission with the storage device, the recompiling latency is lower than the storage device swap-in latency when the GPU memory computing power of the target device is sufficient. This improves the overall inference speed of large language models and does not require additional storage devices and bandwidth resources, thereby improving the efficiency of KV cache management. Attached Figure Description

[0012] The accompanying drawings, which are included to provide a further understanding of this application and form part of this application, illustrate exemplary embodiments and are used to explain this application, but do not constitute an undue limitation of this application. In the drawings: Figure 1 This is a flowchart illustrating the key-value cache management method according to an embodiment of this application.

[0013] Figure 2 This is a schematic diagram of the Transformer layer architecture of the large language model in an embodiment of this application.

[0014] Figure 3 This is a schematic diagram illustrating the computational principle of the Transformer layer in an embodiment of this application.

[0015] Figure 4 This is a schematic diagram of the storage area structure for caching key-value data in an embodiment of this application.

[0016] Figure 5 This is a schematic diagram of the overall process of the key-value cache management method according to an embodiment of this application.

[0017] Figure 6 This is a structural block diagram of the key-value cache management device according to an embodiment of this application. Detailed Implementation

[0018] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application. The drawing numbers in this application are only used to distinguish the various steps in the solution and are not used to limit the execution order of the various steps. The specific execution order is subject to the description in the specification.

[0019] To address the problems existing in the prior art, this application provides a key-value cache management method applied to a target device. The target device is any device performing LLM inference, including but not limited to graphics processing unit (GPU), neural network processing unit (NPU), central processing unit (CPU) inference, or high-performance large model inference framework layers such as vLLM / Llama.

[0020] like Figure 1 As shown, the method includes steps 102 to 106.

[0021] Step 102: Cache the key-value data output by each Transformer layer stacked during the inference process of the target large language model into the first storage area of ​​the target device's video memory. The key-value data includes the key-value data obtained by each Transformer layer performing forward computation on the target symbol input to each Transformer layer.

[0022] The target large language model can be installed or run on the target device to receive text input from the user and output the corresponding inference results.

[0023] like Figure 2 As shown, the target large language model includes N stacked Transformer layers 100. The Transformer layer 100 is the basic computational unit in the target large language model, which consists of a self-attention sub-layer and a feedforward neural network sub-layer. The Transformer layer 100 closest to the input of the target large language model is the first layer L1 (i.e., the lowest layer), and the Transformer layer 100 closest to the output of the target large language model is the Nth layer LN (i.e., the highest layer). The N Transformer layers 100 from L1 to LN are stacked to form a complete large language model.

[0024] like Figure 3 As shown, the input to each Transformer layer 100 is a token. A token is the basic unit processed by the large language model and can be a word, subword, or character, which is represented as a vector after encoding. Transformer layer 100 performs the first stage of the large language model inference process on the input tokens: the pre-filling stage. In this stage, forward computation is performed on the input tokens to generate a KV cache corresponding to the input tokens. Then, Transformer layer 100 performs the second stage of the large language model inference process: the decoding stage. In this stage, Transformer layer 100 reads the existing KV cache (i.e., the cached KV data) and generates the corresponding new tokens. The input to Transformer layer 100 is a token, and the output is a new token.

[0025] In this embodiment, all key-value data generated by each Transformer layer 100 during the inference process of the target large language model is cached in the first storage area of ​​the target device's video memory, namely the graphics processing unit (GPU) video memory (VRAM), which is the fastest storage area in the target device. After the pre-filling stage is completed, all key-value data generated by all Transformer layers 100 are written to the first storage area.

[0026] Optionally, the first storage area is a contiguous storage area comprising multiple fixed-size memory blocks, each memory block being used to store key-value data obtained by performing forward computation on one or more symbols corresponding to the input of the same Transformer layer; wherein, when the key-value data of the first target Transformer layer cached in the first storage area is discarded, the key-value data stored in the memory block corresponding to the first target Transformer layer is released in its entirety.

[0027] In this embodiment, a contiguous storage region specifically allocated in the target device, such as the GPU's video memory (VRAM), is used to cache key-value data. The VRAM is divided into multiple fixed-size memory blocks, each of which can store one or more KV data of a certain symbol in a certain Transformer layer 100. The size of each memory block can be an integer multiple of the data size of the key-value data of a single symbol in a single Transformer layer.

[0028] Large memory block allocation in contiguous storage regions avoids memory fragmentation caused by fine-grained allocation. Key-value data of the same Transformer layer is stored in the same memory block. When discarding, the entire block is released, thus releasing all key-value data stored in the corresponding memory block of the same Transformer layer, creating a contiguous free storage region. This reduces memory fragmentation and avoids non-contiguous free regions caused by memory blocks of different sizes during release, thereby improving the utilization of VRAM storage region and supporting longer generation sequences.

[0029] Step 104: If the remaining capacity of the first storage area is lower than the first preset capacity threshold, the key-value data of the first target Transformer layer cached in the first storage area is discarded according to the preset discarding strategy.

[0030] As the inference process of the target large language model progresses, the amount of key-value data generated by each Transformer layer increases, leading to an increase in the capacity usage of the first storage area. During this step, the remaining capacity of the first storage area is continuously monitored. When new symbol generation requires appending key-value data and the remaining capacity of the first storage area is insufficient, a preset discarding strategy is adopted to select the first target Transformer layer to be discarded, and the key-value data of that first target Transformer layer cached in the first storage area is discarded. The storage space released by the first target Transformer layer is used to store the key-value data of the new symbol.

[0031] In one embodiment, the first storage area includes a low-level storage area and a high-level storage area. The low-level storage area is used to store key-value data of low-level Transformer layers with high position priority, and the high-level storage area is used to store key-value data of high-level Transformer layers with low position priority. The Transformer layers closer to the input of the target large language model have higher position priority, and the Transformer layers closer to the output of the target large language model have lower position priority.

[0032] The key-value data of the first target Transformer layer cached in the first storage area is discarded according to a preset discarding strategy, including: discarding the key-value data of the first Transformer layer in the higher storage area in order of position priority from low to high; moving the key-value data of the second Transformer layer in the lower storage area into the higher storage area; and caching the newly generated key-value data of the third Transformer layer in the lower storage area, wherein the position priority of the third Transformer layer is higher than that of the second Transformer layer.

[0033] In this embodiment, the discarding strategy is based on the positional priority of each Transformer layer. Key-value data from Transformer layers with higher positional priority are preferentially retained in the first storage area of ​​GPU memory, while key-value data from Transformer layers with lower positional priority are discarded or moved out of the first storage area of ​​GPU memory. The positional priority is determined by the distance between each Transformer layer stacked in the target large language model and the input / output end of the target large language model. The closer a Transformer layer is to the input end of the target large language model, the higher its positional priority; conversely, the closer it is to the output end of the target large language model, the lower its positional priority.

[0034] In this embodiment, such as Figure 4 As shown, the first storage area of ​​the video memory is further divided into a low-level storage area 10 and a high-level storage area 20. The low-level storage area 10 is used to store the key-value data of the lower-level Transformer layers with high position priority, and the high-level storage area 20 is used to store the key-value data of the higher-level Transformer layers with low position priority. Both the low-level and high-level storage areas can include at least one Transformer layer.

[0035] During the initial inference process, the key-value data generated by each Transformer layer is cached in the lower storage area 10. When the lower storage area 10 is saturated, the key-value data is moved to the higher storage area 20 in order of position priority from low to high, starting from the layer with the lowest position priority. The newly generated key-value data of the lower Transformer layer is also cached in the lower storage area 10.

[0036] When both the lower-level storage area 10 and the higher-level storage area 20 are saturated, i.e., when the first storage area of ​​the video memory is saturated, the key-value data of the corresponding layer in the higher-level storage area 20 is discarded, starting from the layer with the lowest positional priority, in ascending order of positional priority. Then, the process continues, moving the corresponding key-value data from the layer with the lowest positional priority into the higher-level storage area 20, and caching the newly generated key-value data of the lower-level Transformer layer in the lower-level storage area 10.

[0037] Different Transformer layers have varying importance and contribute significantly to the final output. Transformer layers closer to the input of the target large language model are responsible for basic semantic feature extraction, while Transformer layers farther from the input of the target large language model are responsible for task adaptation and abstract reasoning.

[0038] The positional priority reflects the importance of the key-value data generated by the corresponding Transformer layer. The closer the Transformer layer is to the input of the target large language model, the more valuable the corresponding key-value data is. The key-value data is preserved intact, the basic semantic features are not lost, and the effective information density of the cached data can be improved.

[0039] Combination Figure 2 The first layer (L1) is the Transformer layer closest to the input of the target large language model, and the Nth layer (LN) is the Transformer layer closest to the output of the target large language model. That is, the first layer (L1) has the highest positional priority, and the Nth layer (LN) has the lowest positional priority.

[0040] If the remaining capacity of the first storage area is found to be below a first preset capacity threshold, reaching saturation, then the key-value data of the Transformer layers with lower location priority in the higher-level storage area 20 will be discarded. The number of discarded Transformer layers can be determined based on requirements. For example... Figure 2 In this context, the k-th layer Lk is the dividing point when the capacity of the fast area (i.e., the first storage area) is saturated. Therefore, when the remaining capacity of the first storage area is lower than the first preset capacity threshold, the key-value data of the Transformer layer 100 to the N-th layer LN after the k-th layer Lk in the higher storage area 20 will be discarded first, and the key-value data of the remaining layers will be retained first.

[0041] Specifically, for example, the first layer L1 is responsible for basic semantic feature extraction, and key-value data is always kept in the fast region; the positions of the second layer L2, the third layer L3, the fourth layer L4 to the kth layer Lk have relatively high priority, and key-value data is retained first (among which, the fourth layer L4 can be discarded in even drop mode, which will be explained later); the key-value data of the (N-1)th layer LN-1 is discarded first, and the Nth layer LN is responsible for abstract reasoning and task adaptation, and key-value data is discarded first (Drop).

[0042] The storage space released in the higher-level storage area 20 after the discard operation is used to move key-value data from the Transformer layer in the lower-level storage area 10, where the data has a relatively low priority. The storage space released in the lower-level storage area 10 after the move-in operation is used to cache newly generated key-value data in the lower level.

[0043] The key-value data discarded in each layer is discarded as a whole, that is, all key-value data generated by the symbols in that layer are completely discarded.

[0044] In one embodiment, discarding the key-value data of the first target Transformer layer cached in the first storage area according to a preset discarding strategy includes: using an even-numbered-layer discarding strategy to discard the key-value data of the even-numbered Transformer layers in each Transformer layer; or using an odd-numbered-layer discarding strategy to discard the key-value data of the odd-numbered Transformer layers in each Transformer layer, wherein the odd-numbered layers do not include the first Transformer layer at the input end of the target large language model; or when the inference process includes multiple consecutive inference steps, each inference step alternately uses the even-numbered-layer discarding strategy and the odd-numbered-layer discarding strategy, correspondingly and alternately discarding the key-value data of the even-numbered Transformer layers in each Transformer layer and discarding the key-value data of the odd-numbered Transformer layers in each Transformer layer.

[0045] In this embodiment, the discarding strategy is determined based on the parity of the number of each Transformer layer. For example, key-value data of even-numbered Transformer layers is retained, while key-value data of odd-numbered Transformer layers is discarded, but the first Transformer layer at the input of the target large language model is excluded.

[0046] Alternatively, key-value data in odd-numbered Transformer layers can be retained, while key-value data in even-numbered Transformer layers can be discarded.

[0047] In this embodiment, by adopting a selective discarding strategy for even / odd layers, the target large language model only needs to recompile about half of the discarded Transformer layers in each inference step, which can avoid latency spikes caused by subsequent recompiling of discarded key-value data, resulting in small fluctuations in inference latency and a smooth user experience.

[0048] If the remaining capacity of the first storage area is found to be lower than a first preset capacity threshold, the key-value data of the corresponding even / odd Transformer layer will be discarded. Combined with... Figure 2 Layer 2 (L2) and Layer 4 (L4) are even-numbered Transformer layers, which can be discarded in even-drop mode. Similarly, Layer 3 (L3) is an odd-numbered Transformer layer, which can be discarded in odd-drop mode, but Layer 1 (L1), which is also an odd-numbered layer, must always remain in the fast zone.

[0049] When the inference process includes multiple consecutive inference steps, each inference step can alternate between even-level discarding and odd-level discarding strategies. For example, in inference steps 1-3 executed sequentially, an even-level discarding strategy can be used in inference step 1 to discard the key-value data of even-level Transformer layers; in the next inference step 2, an odd-level discarding strategy can be used to discard the key-value data of odd-level Transformer layers; and in the next inference step 3, an even-level discarding strategy can be used again to discard the key-value data of even-level Transformer layers.

[0050] Alternatively, in inference step 1, an odd-numbered layer discarding strategy can be used to discard the key-value data of the odd-numbered Transformer layers in each Transformer layer; in the next inference step 2, an even-numbered layer discarding strategy can be used to discard the key-value data of the odd-numbered Transformer layers in each Transformer layer; and in the next inference step 3, an odd-numbered layer discarding strategy can be used again to discard the key-value data of the odd-numbered Transformer layers in each Transformer layer.

[0051] Therefore, by alternately employing even-layer and odd-layer discarding strategies in multiple consecutive inference steps, and executing the discarding of key-value data in the Transformer layer of the corresponding odd-even layer in the inference step, the recomputation load corresponding to each step can be evenly distributed, the inference latency fluctuation is small, and the user experience is smooth.

[0052] In one embodiment, discarding the key-value data of the first target Transformer layer cached in the first storage area according to a preset discarding strategy includes: determining a number of Transformer layers with low position priority, wherein the position priority of the Transformer layer closer to the input end of the target large language model is higher than the position priority of the Transformer layer closer to the output end of the target large language model; and discarding the key-value data of the even-numbered / odd-numbered Transformer layers among the multiple Transformer layers with low position priority.

[0053] In this embodiment, the discarding strategy is formulated based on the position priority and the parity of the number of each Transformer layer. Key-value data of Transformer layers with high position priority are retained, while key-value data of Transformer layers with low position priority are discarded according to the odd / even number of layers.

[0054] In one embodiment, discarding the key-value data of the first target Transformer layer cached in the first storage area according to a preset discarding strategy includes: determining the importance score of the key-value data in the first target Transformer layer based on the variance of the output of the first target Transformer layer or its contribution to the final output of the target large language model; if the importance score is lower than a preset score threshold, then the key-value data of the first target Transformer layer is discarded.

[0055] In this implementation, instead of discarding data based on a fixed layer position priority, the importance score of the key-value data cached in each layer is dynamically evaluated during the inference process. For example, based on the variance of the attention output of that layer and its contribution to the final output logit, the layer with the lowest importance score is discarded first.

[0056] In other embodiments, in addition to discarding key-value data of low-priority Transformer layers based on positional priority, unimportant key-value caches of symbols in the retained high-priority Transformer layers can be eliminated according to symbol attention scores. This is a hybrid strategy of "layer-level discarding + symbol-level compression" to further save memory storage space.

[0057] Step 106: If it is detected that the current inference process requires the use of the key-value data of the first target Transformer layer, then the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to restore the key-value data of the first target Transformer layer and write it to the first storage area.

[0058] This step involves recompiling key-value data during inference. When the decoding phase reaches the inference of a discarded first-target Transformer layer, the hidden input state of that first-target Transformer layer is retrieved from the retained Transformer layers. The forward computation of the corresponding symbols for the discarded key-value data of that first-target Transformer layer is re-executed, including self-attention computation and feedforward network computation, generating new key-value data including the K matrix and V matrix. The recomputed key-value data is cached in the first storage area, or released or discarded from the first storage area after the current inference for the first-target Transformer layer.

[0059] In one embodiment, the key-value data recalculation in this application is a complete forward computation, without quantization or approximation. The recalculated key-value data is completely consistent with the original computation result, with no loss of precision. The derivation process of recalculation is as follows: KV data is discarded and recalculated → complete forward computation → no quantization error → no cumulative decrease in precision over multiple rounds of dialogue.

[0060] In another embodiment, optionally, the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to recover the key-value data of the first target Transformer layer, including: re-executing the forward computation of the target symbol of the first target Transformer layer on the value in the key-value data; obtaining the key in the key-value data of the previous Transformer layer adjacent to the first target Transformer layer; and recovering the key-value data of the first target Transformer layer based on the recalculated value and the obtained key.

[0061] In this embodiment, the discarded key-value data of the target Transformer layer is not recalculated in its entirety, but only partially. For example, only the key matrix is ​​recalculated, and the value matrix of the target Transformer layer is approximated by the value matrix of the Transformer layer above the target Transformer layer to reduce the recalculation overhead, thereby trading a small loss of precision for lower recalculation latency.

[0062] The recalculation process relies entirely on the computing resources of the target device's video memory; therefore, when the video memory's computing power is sufficient, the recalculation latency is low. Compared to key-value data swapping involving I / O operations on storage devices or RAM, the key-value cache management of this application can significantly reduce inference latency.

[0063] To reduce hardware dependence and deployment costs, this application adopts a key-value data discarding + recomputation approach. The cache of key-value data in the first target Transformer layer is directly deleted from the device's video memory storage area, and restored during inference using the target device's video memory computing resources. Since the recomputation is completed entirely within the target device and does not involve I / O transfers to RAM and / or storage devices, the recomputation latency is lower than the swap-in and swap-out latency of RAM and / or storage devices in scenarios with sufficient video memory computing power.

[0064] The latency derivation process for discarding and recompiling key-value data is as follows: discard target layer KV data → release space in the first storage area → write new symbol KV data directly to the first storage area → no need for SSD swapping in and out → reduce inference latency.

[0065] This application requires only video memory (fast zone) under default configuration, and does not mandate high-speed storage devices or dedicated data transfer interfaces. It remains effective for deployment environments with limited video memory but no high-speed storage devices (such as consumer-grade GPUs and edge devices). By requiring only device video memory, hardware requirements are reduced, expanding deployment scenarios.

[0066] In addition, when the total amount of key-value data in the cache is extremely large and the first storage area of ​​the video memory is close to saturation, memory-coordinated storage of the target device can be introduced.

[0067] Optionally, in one embodiment, the method further includes: after discarding the key-value data of the first target Transformer layer cached in the first storage area, if it is detected that the remaining capacity of the first storage area is lower than the first preset capacity threshold, then moving the key-value data of the second target Transformer layer in the first storage area into the second storage area of ​​the target device memory; if it is detected that the current inference process needs to use the key-value data of the second target Transformer layer, then moving the key-value data of the second target Transformer layer from the second storage area into the first storage area; wherein the position priority of the second target Transformer layer is higher than the position priority of the first target Transformer layer.

[0068] In this embodiment, if the key-value data of the target Transformer layer is discarded from the first storage area of ​​the video memory according to a preset discarding strategy, to address the issue of the first storage area approaching saturation due to the increased amount of key-value data generated by new symbols, the key-value data of one or more Transformer layers (i.e., the second target Transformer layer) currently cached in the first storage area can be written to or swapped out to the target device's memory for caching. The second Transformer layer can be a Transformer layer with a middle priority among the layers stacked in the target large language model. When the key-value data of the second target Transformer layer is needed for current inference, the key-value data of the second target Transformer layer is moved from the second storage area into or swapped into the first storage area.

[0069] The first target Transformer layer discarded from the first memory region is one or more Transformer layers with low positional priority among all Transformer layers. The second target Transformer layer moved from the first memory region into the second memory region has an intermediate positional priority, therefore the positional priority of the second target Transformer layer is higher than that of the first target Transformer layer.

[0070] Optionally, in one embodiment, the method further includes: after moving the key-value data of the second target Transformer layer in the first storage area into the second storage area of ​​the target device memory, if it is detected that the remaining capacity of the second storage area is lower than a second preset capacity threshold, then moving the cached key-value data of the second target Transformer layer in the second storage area into a storage device connected to the target device, and clearing the key-value data of the fourth target Transformer layer in the storage device; if it is detected that the current inference process needs to use the key-value data of the second target Transformer layer, then moving the key-value data of the second target Transformer layer from the storage device into the first storage area; wherein the position priority of the third target Transformer layer is higher than the position priority of the fourth target Transformer layer.

[0071] In this embodiment, if the amount of key-value data of the second target Transformer layer currently cached in the second storage area of ​​the target device memory increases to the point that the second storage area is close to saturation, the key-value data of the third target Transformer layer (i.e., some or all of the Transformer layers in the second target Transformer layer) currently cached in the second storage area can be written to or swapped out to a storage device connected to the target device for caching. The third Transformer layer can be a Transformer layer with a lower priority in the second target Transformer layer. When the key-value data of the third target Transformer layer is needed for the current inference, the key-value data of the third target Transformer layer is moved from the storage device into or swapped into the second storage area, and then moved from the second storage area into or swapped into the first storage area.

[0072] Furthermore, if the remaining capacity of the storage device is lower than a preset capacity threshold, the key-value data of the fourth target Transformer layer with low location priority in the storage device is cleared.

[0073] The third target Transformer layer moved into the storage device from the second storage area is the Transformer layer with the middle position priority among all Transformer layers. The fourth target Transformer layer removed from the storage device is the Transformer layer with the lowest position priority among all Transformer layers. Therefore, the position priority of the third target Transformer layer is higher than that of the fourth target Transformer layer.

[0074] The above multi-level storage area collaborative caching strategy ensures that the large language model system can still run normally in extremely long text scenarios, while giving priority to the key-value data discarding + recalculation strategy to reduce the dependence on storage devices.

[0075] When the first storage area is saturated, high-level key-value data is discarded first (without relying on secondary storage). The second storage area and storage device swapping mechanism are only activated in extreme scenarios, forming a collaborative management mechanism of "prioritizing discarding recomputed data and using swapped-out storage as a fallback".

[0076] In this application embodiment, the scope of storage devices includes, but is not limited to, consumer solid state drives (cSSD), enterprise solid state drives (eSSD), and universal flash storage (UFS).

[0077] Figure 4 This application illustrates the storage area structure for caching key-value data according to one embodiment, such as... Figure 4 As shown, the key-value data cache adopts multi-level storage area collaborative management: including the target device's video memory (VRAM), RAM, and storage devices (such as eSSD / cSSD / UFS). The first storage area of ​​the video memory (VRAM) can be further divided into a corresponding low-level storage area 10 and a high-level storage area 20, which respectively cache the key-value data of the Transformer layer with high location priority and the key-value data of the Transformer layer with low location priority.

[0078] If the video memory (VRAM) used for caching key-value data is saturated, the key-value data in the higher-level storage area 20 is discarded. If the VRAM is still insufficient after discarding, the key-value data in the intermediate layer is swapped out to main memory (RAM). If the main memory (RAM) used for caching key-value data is saturated, the key-value data in main memory (RAM) is swapped out to a storage device (e.g., eSSD / cSSD / UFS). If the storage device is saturated, the key-value data with the lowest location priority in the storage device is cleared or discarded.

[0079] The discarded key-value data, such as the key-value data discarded in the higher-level storage area 20, can be recovered by re-executing the forward computation of the corresponding layer through the recomputation engine.

[0080] Below, in conjunction with Figure 5 The overall process of the key-value cache management method for a large language model according to embodiments of this application is described.

[0081] like Figure 5 As shown, it includes the following steps: Step S1: Initial allocation of KV cache.

[0082] For example, after the pre-filling stage is completed, all KV caches of the Transformer layers are written to the fast area of ​​the video memory according to their positional priority, which is the first storage area including the low-level storage area 10 and the high-level storage area 20. Initially, all layer KV data is written to the low-level storage area 10. When the low-level storage area 10 is saturated, the KV data of the layers with lower positional priority in the low-level storage area 10 are moved to the high-level storage area 20.

[0083] Step S2: Fast Zone Capacity Monitoring, continuously monitor the remaining capacity of the fast zone.

[0084] Step S2b: Determine if the fast region is saturated. If yes, it means that the current generation of a new symbol requires additional KV data and the remaining capacity of the fast region is insufficient, proceed to step S3; otherwise, continue with the inference steps.

[0085] Step S3: Decision on discarding high-level key-value data.

[0086] Select the Transformer layer to be discarded from the higher storage area 20 according to the position priority (starting from the Nth layer LN of the highest layer).

[0087] Specifically, this includes: starting from the highest level LN, moving the high-level KV data into the shared area of ​​memory (if the shared area is available), or discarding it directly (if the shared area is unavailable or not enabled); directly discarding the high-level KV data from the high-level storage area 20 of the fast area without writing it to the shared area, and using the freed high-level storage area 20's fast area space to store the low-level KV data of the new symbol.

[0088] Step S4: Select the discarding mode as either even-level discarding mode or odd-level discarding mode.

[0089] To avoid inference latency spikes caused by concentrated recomputation, a selective discarding mode for even / odd layers can be adopted.

[0090] Even-level discard mode: Discard key-value data cached in even-numbered levels (e.g., L2, L4, L6, ...), and retain key-value data cached in odd-numbered levels (e.g., L1, L3, L5, ...); Odd-level discard mode: Discard key-value data cached in odd-numbered levels (e.g., L3, L5, L7, ...), and retain key-value data cached in even-numbered levels (e.g., L2, L4, L6, ...) (but retain the lowest level L1); Alternating mode: Alternate between even-level and odd-level discard in consecutive inference steps to distribute the recomputation load evenly.

[0091] Choosing either an even-numbered layer discarding mode or an odd-numbered layer discarding mode allows for alternating discarding when key-value data needs to be recalculated during inference for discarded Transformer layers. This ensures that each inference step only requires recalculating about half of the discarded layers, avoiding latency spikes caused by concentrated recalculation of key-value data for all discarded layers in a single inference step. This results in smoother fluctuations in inference speed and improves the user experience.

[0092] Step S5: Perform the drop.

[0093] Step S6: When the decoding stage reaches the inference of a certain discarded layer, determine whether the current inference needs to access the KV data of the discarded layer. If yes, proceed to step S5b; otherwise, proceed to step S7.

[0094] Step S5b: Recalculate and restore KV data.

[0095] The recalculation process does not involve I / O operations on RAM and / or storage devices, and relies entirely on the computing resources of the device's video memory. Therefore, when the video memory computing power is sufficient, the recalculation latency can be lower than the swap-in latency.

[0096] Step S7: Large memory allocation and fragment management.

[0097] A large-block memory allocation strategy is used to manage the fast region, dividing it into fixed-size memory blocks. Allocation and release are performed in blocks to avoid fragmentation caused by fine-grained allocation. When discarding high-level key-value pairs, the entire memory block is released, creating contiguous free areas. New key-value data is preferentially written to these contiguous free areas to reduce fragmentation. Further reasoning will follow.

[0098] In this embodiment, the target device caches the key-value data generated by each Transformer layer stacked during the inference process of the target large language model in the first storage area of ​​the target device's video memory. The key-value data includes the key-value data obtained by each Transformer layer performing forward computation on the target symbol input to each Transformer layer. If the remaining capacity of the first storage area is monitored to be lower than a first preset capacity threshold, the key-value data of the first target Transformer layer cached in the first storage area is discarded according to a preset discarding strategy. If the current inference process needs to use the key-value data of the first target Transformer layer, the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to restore the key-value data of the first target Transformer layer and write it into the first storage area. Therefore, a strategy of discarding key-value data and recompiling is adopted. The key-value data cache of the target Transformer layer is directly deleted from the GPU memory. During inference, it is recomputed and restored using the GPU memory computing resources of the target device. Since the recompiling is completed entirely within the target device and does not involve I / O transmission with the storage device, the recompiling latency is lower than the storage device swap-in latency when the GPU memory computing power of the target device is sufficient. This improves the overall inference speed of large language models and does not require additional storage devices and bandwidth resources, thereby improving the efficiency of KV cache management.

[0099] The key-value caching management method in the above embodiments executes key-value data discarding and recalculation of the target Transformer layer based on whether the storage space of the video memory, main memory, or storage device is saturated or whether the remaining capacity exceeds a preset capacity threshold. In other embodiments, the corresponding key-value data discarding strategy can also be dynamically adjusted based on the load during the inference process. For example, when the load is low during inference, all key-value data generated by each Transformer layer of the target large language model during inference is not discarded and is cached in the video memory of the target device; when the load is medium during inference, the even-numbered layer discarding strategy and the odd-numbered layer discarding strategy can be used alternately to discard Transformer layers with low priority in the video memory of the target device; when the load is high during inference, more Transformer layers with low priority in the video memory of the target device are discarded, and the main memory of the target device is enabled to move the key-value data of the Transformer layers with intermediate priority in the video memory of the target device into the main memory for caching; when the load is extremely high during inference, the main memory is enabled to move some of the key-value data of the Transformer layers in the main memory into the main memory for caching.

[0100] The key-value cache management method in this embodiment can achieve a balance between performance and resource usage through adaptive adjustment.

[0101] This application embodiment also provides a storage controller connected to a target device for caching key-value data generated by the Transformer layer of a target large language model during inference on the target device; the target device, when executed, implements each step of the above-described key-value cache management method embodiment.

[0102] This application also provides a storage device, including the storage controller described in the above embodiments.

[0103] Optionally, such as Figure 6 As shown, this application embodiment also provides a key-value cache management device 2000, including a processor 2400 and a memory 2200. The memory 2200 stores a program or instructions that can run on the processor 2400. When the program or instructions are executed by the processor 2400, they implement the various steps of the above-described key-value cache management method embodiment and can achieve the same technical effect. To avoid repetition, they will not be described again here.

[0104] This application also provides a non-volatile computer-readable storage medium storing a program or instructions. When the program or instructions are executed by a processor, they implement the various processes of any of the above-described key-value cache management method embodiments and achieve the same technical effect. To avoid repetition, further details are omitted here. The non-volatile computer-readable storage medium includes, but is not limited to, consumer solid-state drives (cSSDs), enterprise solid-state drives (eSSDs), universal flash storage (UFS), read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0105] This application also provides a computer program product, which includes a non-volatile computer-readable storage medium storing a computer program. The computer program is operable to enable a computer to execute various processes of any of the above-described key-value cache management method embodiments and achieve the same technical effect. To avoid repetition, it will not be described again here.

[0106] It should be noted that, in this document, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Unless otherwise specified, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes that element.

[0107] Through the above description of the embodiments, those skilled in the art can clearly understand that the methods of the above embodiments can be implemented by means of software plus necessary general-purpose hardware platforms. Of course, they can also be implemented by hardware, but in many cases the former is a better implementation method. Based on this understanding, the technical solution of this application, in essence, or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product is stored in a storage medium (such as ROM / RAM, magnetic disk, optical disk) and includes several instructions to cause a terminal (which may be a mobile phone, computer, server, air conditioner, or network device, etc.) to execute the methods described in the various embodiments of this application.

[0108] The embodiments of this application have been described above with reference to the accompanying drawings. However, this application is not limited to the specific embodiments described above. The specific embodiments described above are merely illustrative and not restrictive. Those skilled in the art can make many other forms under the guidance of this application without departing from the spirit and scope of the claims, and all of these forms are within the protection scope of this application.

Claims

1. A key-value cache management method, characterized in that, Applied to a target device, the method includes: The key-value data generated by each Transformer layer stacked during the inference process of the target large language model is cached in the first storage area of ​​the target device's video memory. The key-value data includes the key-value data obtained by each Transformer layer performing forward computation on the target symbol input to each Transformer layer. If the remaining capacity of the first storage area is found to be lower than the first preset capacity threshold, the key-value data of the first target Transformer layer cached in the first storage area will be discarded according to the preset discarding policy. If the monitoring detects that the current inference process requires the use of the key-value data of the first target Transformer layer, then the forward computation of the target symbol of the first target Transformer layer is re-executed using the computing resources of the video memory to restore the key-value data of the first target Transformer layer and write it to the first storage area.

2. The method according to claim 1, characterized in that, The first storage area is a contiguous storage area comprising multiple fixed-size memory blocks. Each memory block is used to store key-value data obtained by performing forward computation on one or more symbols corresponding to the input of the same Transformer layer. Specifically, when discarding the key-value data of the first target Transformer layer cached in the first storage area, the key-value data stored in the memory block corresponding to the first target Transformer layer is released in its entirety.

3. The method according to claim 1, characterized in that, The first storage area includes a low-level storage area and a high-level storage area. The low-level storage area is used to store key-value data of low-level Transformer layers with high position priority, and the high-level storage area is used to store key-value data of high-level Transformer layers with low position priority. The Transformer layers closer to the input end of the target large language model have higher position priority, and the Transformer layers closer to the output end of the target large language model have lower position priority. According to a preset discarding policy, the key-value data of the first target Transformer layer cached in the first storage area is discarded, including: According to the order of position priority from low to high, the key-value data of the first Transformer layer in the high-level storage area is discarded, and the key-value data of the second Transformer layer in the low-level storage area is moved into the high-level storage area. New key-value data generated by the lower-level third Transformer layer is cached in the lower-level storage area, wherein the position priority of the third Transformer layer is higher than that of the second Transformer layer.

4. The method according to claim 1, characterized in that, According to a preset discarding policy, the key-value data of the first target Transformer layer cached in the first storage area is discarded, including: A discard strategy is adopted for even-numbered layers, discarding key-value data from even-numbered Transformer layers; or An odd-number layer discarding strategy is adopted, discarding the key-value data of the odd-numbered Transformer layers in each Transformer layer, excluding the first Transformer layer at the input of the target large language model; or When the inference process includes multiple consecutive inference steps, each inference step alternates between the even-numbered layer discarding strategy and the odd-numbered layer discarding strategy, correspondingly discarding the key-value data of the even-numbered Transformer layers and discarding the key-value data of the odd-numbered Transformer layers in each Transformer layer.

5. The method according to claim 1, characterized in that, According to a preset discarding policy, the key-value data of the first target Transformer layer cached in the first storage area is discarded, including: Based on the variance of the output of the first target Transformer layer or its contribution to the final output of the target large language model, the importance score of the key-value data in the first target Transformer layer is determined. If the importance score is lower than a preset score threshold, the key-value data of the first target Transformer layer will be discarded.

6. The method according to claim 3, characterized in that, Also includes: After discarding the key-value data of the first target Transformer layer cached in the first storage area, if it is detected that the remaining capacity of the first storage area is lower than the first preset capacity threshold, the key-value data of the second target Transformer layer in the first storage area is moved into the second storage area of ​​the target device memory. If the monitoring indicates that the current inference process requires the use of key-value data of the second target Transformer layer, then the key-value data of the second target Transformer layer is moved from the second storage area to the first storage area; The position priority of the second target Transformer layer is higher than that of the first target Transformer layer.

7. The method according to claim 6, characterized in that, Also includes: After moving the key-value data of the second target Transformer layer in the first storage area into the second storage area of ​​the target device memory, if the remaining capacity of the second storage area is monitored to be lower than the second preset capacity threshold, the key-value data of the third target Transformer layer cached in the second storage area is moved into the storage device connected to the target device, and the key-value data of the fourth target Transformer layer in the storage device is cleared. If the monitoring indicates that the current inference process requires the use of key-value data from the second target Transformer layer, then the key-value data from the third target Transformer layer is moved from the storage device into the first storage area; The position priority of the third target Transformer layer is higher than that of the fourth target Transformer layer.

8. The method according to claim 1, characterized in that, Using the computing resources of the video memory, re-execute the forward computation of the target symbol of the first target Transformer layer to recover the key-value data of the first target Transformer layer, including: Re-execute the forward computation of the target symbol for the key-value data in the first target Transformer layer; Retrieve the key from the key-value data of the Transformer layer adjacent to the first target Transformer layer; Based on the recalculated values ​​and the retrieved keys, the key-value data of the first target Transformer layer is recovered.

9. A storage controller, characterized in that, Connected to the target device, used to cache key-value data generated during inference by the Transformer layer of the target large language model on the target device; The target device, when executed, implements the steps of the method as described in any one of claims 1-8.

10. A storage device, characterized in that, Includes the storage controller as described in claim 9.

11. A key-value cache management device, characterized in that, It includes a processor and a memory, the memory storing a program or instructions that can run on the processor, the program or instructions being executed by the processor to implement the steps of the method as described in any one of claims 1-8.

12. A non-volatile computer-readable storage medium, characterized in that, The non-volatile computer-readable storage medium stores a program or instructions that, when executed by a processor, implement the steps of the method as described in any one of claims 1-8.

13. A computer program product, characterized in that, The computer program product includes a non-volatile computer-readable storage medium storing a computer program operable to cause a computer to perform the steps of the method as described in any one of claims 1-8.

Images