Key value cache multiplexing method for retrieval enhancement generation system

Abstract

The invention provides a key value cache multiplexing method for a retrieval enhancement generation system. The method comprises three stages of knowledge retrieval, prompt construction and reasoning generation. In the knowledge retrieval stage, the system carries out vectorization on user input and retrieves related documents, and retrieval results serve as enhanced information; in the prompt construction stage, user input and enhancement information is coded into a token sequence, and the hash value and length of each part of the token sequence are calculated; in the inference generation stage, whether the key value cache is hit or not is judged through hash comparison, the key value cache is directly reused if the key value cache is hit, differential video memory allocation is executed if the key value cache is not hit, a new key value cache is generated in combination with a partition position coding strategy, and finally a result is generated through large language model inference and returned to a user. According to the method, the key value cache reuse rate of the large language model in multi-round reasoning of the retrieval enhancement generation system can be improved, so that the calculation amount and the video memory overhead are reduced.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence and machine learning technology, specifically relating to a method for key-value cache allocation and reuse in a retrieval enhancement generation system. Background Technology

[0002] With the widespread application of large language models (LLMs) in fields such as intelligent question answering and knowledge retrieval, improving reasoning efficiency and knowledge utilization under limited computing resources has become a significant challenge. Traditional language models rely on internal knowledge for generation, but their coverage of real-time updated knowledge and specialized domains is insufficient, leading to inaccurate generated results. To address this issue, Retrieval-Augmented Generation (RAG) technology has emerged.

[0003] In a typical RAG system, the system first acquires user input and performs semantic encoding to generate a representation that can be used for vector retrieval. Then, it retrieves relevant document fragments based on this representation and concatenates the retrieval results with the user input to form an enhanced prompt sequence as model input. During model inference, the system performs cache detection and key-value cache management to determine whether to reuse existing caches or create new caches. Finally, the model performs forward inference computation to generate the final output result.

[0004] In this generation process, the model typically receives two main types of input: retrieval enhancement information and user input. User input changes most frequently, while the content, quantity, and order of the retrieval enhancement information also dynamically change with fluctuations in the retrieval results. When encoding the complete input sequence, the model generates a corresponding key-value cache to accelerate attention computation in subsequent inference.

[0005] However, existing RAG systems generally employ a global sequence-level caching strategy. This strategy requires that the prefixes of the input sequences be completely identical before the corresponding key-value cache can be reused. While this approach can effectively reduce redundant computation when the task structure is simple, it shows significant shortcomings when dealing with complex scenarios containing multiple retrieval enhancements. For example, during multi-round inference, if only the first of the retrieved enhancements changes, the system still cannot reuse the key-value caches of the remaining content, even if this content is exactly the same as in the previous round of inference. This design leads to the system's inability to fully utilize cache resources, resulting in a large amount of unnecessary redundant computation and affecting inference efficiency. This problem is particularly prominent in scenarios with a large number of multi-round request inferences.

[0006] To address the aforementioned issues, existing RAG systems suffer from shortcomings in cache reuse granularity, cache recognition capabilities, and GPU memory utilization efficiency. Therefore, this invention proposes a key-value cache reuse method for retrieval enhancement generation systems. This method efficiently manages and reuses the cache corresponding to retrieval enhancement content and user input, thereby improving inference efficiency and reducing GPU memory overhead. Summary of the Invention

[0007] This invention aims to address the problems of low cache reuse efficiency and wasted video memory in LLM-based RAG systems. Existing systems typically reuse key-value caches only when the prefixes of the input sequences are completely identical, leading to a large amount of redundant computation when multiple enhancement information segments change. This invention proposes a cache reuse method: encoding and hashing user input and retrieved enhancement information separately to achieve local key-value cache reuse; simultaneously, it employs differentiated video memory allocation and partitioned location encoding.

[0008] The method steps described in this invention are as follows: Figure 1 As shown, the main steps include:

[0009] Knowledge retrieval stage: Receive user input text, convert user input into vector representation through text vectorization encoding, and then perform similarity retrieval in the vector database to obtain the most relevant document fragments;

[0010] Tip building phase: The user input and retrieved document fragments are concatenated to build enhanced tips, and the concatenated text is encoded into a token sequence using a tokenizer, while the hash value of the token sequence is calculated.

[0011] Inference generation stage: The key-value cache hit status is checked based on the hash value. If there is a miss, new video memory is allocated and a corresponding key-value cache is created. Differentiated video memory allocation strategy and partition position encoding are used. If there is a hit, the existing key-value cache is directly loaded for reuse. Finally, the model forward inference calculation is performed and the output result is generated.

[0012] The main parts of this method will be explained in detail below:

[0013] 1. Knowledge Retrieval Stage

[0014] 1.1 User Input Acquisition

[0015] The system receives a text query request submitted by the user and obtains the original user input text. This step serves as the entry point for the entire RAG system.

[0016] 1.2 Text Vectorization Encoding

[0017] The acquired user input text is converted into a high-dimensional vector representation using a pre-trained text encoding model. This process utilizes a deep learning model to extract semantic features of the text and generate a numerical representation suitable for vector similarity calculation.

[0018] 1.3 Vector similarity retrieval and acquisition of relevant document fragments

[0019] The vectorized user query is matched for similarity in a pre-built vector database. By calculating the distance metric between the query vector and the document vector, the document fragments most relevant to the user query are selected from the knowledge base.

[0020] 2. Tips for the construction phase

[0021] 2.1 Building Enhancement Hints

[0022] After the retrieval process is complete, the system integrates the user's original input with the retrieved relevant document fragments (i.e., enhanced information). The data is organized using a key-value pair structure, with information type as the key and specific content as the value. User input is stored as a string, while enhanced information is saved as a list, with each item in the list corresponding to an independent text chunk. This structure clearly distinguishes different content and provides clear boundaries for subsequent partitioning, encoding, and caching management. The storage structure can be represented as: {"enhanced information":["chunk1","chunk2",...],"user input":"..."}.

[0023] 2.2 Token Sequence Encoding

[0024] After constructing the input structure, the system parses and preprocesses each part. First, the system applies a tokenizer to each text block in the user input and augmented information, converting the raw text into a token sequence that the model can process. Then, the resulting multiple token sequences are concatenated into a single continuous token sequence. During the encoding process, the system records the number of tokens in the user input and each augmented text block, as well as their interval positions within the complete sequence, thus providing structural information for subsequent memory allocation and positional encoding processing.

[0025] 2.3 Hash Calculation

[0026] After encoding the token sequence, the system performs independent hash calculations on each part of the input based on its segmented structure. For the user input, the system divides it into segments according to a preset granularity and generates a corresponding hash value for each segment using cumulative hashing. For the augmented information, the system performs hash calculations directly on the entire content in chunks. These segmented hash values ​​will then correspond one-to-one with the subsequently allocated physical memory blocks.

[0027] 3. Reasoning Generation Stage

[0028] 3.1 Key-value cache hit and miss handling

[0029] Before LLM memory allocation, the system first performs cache status checks on user input and enhancement information based on the segmented hash results generated in the preceding stages. The check process compares the hashes of user input and each text block of enhancement information in the current request with the cached hash indices in the system to determine if each content segment has changed. If no corresponding hash value is found, memory allocation and key-value cache creation are performed; if a corresponding hash value is detected, the key-value cache is reused.

[0030] 3.1.1 Video Memory Allocation and Key-Value Cache Creation

[0031] When the system detects a hash miss for any segment of user input or enhanced information, it indicates that the content in that segment has changed and cannot be reused from the historical cache. At this point, the system initiates the memory allocation and key-value cache construction process. The system allocates contiguous physical memory blocks for the key-value cache, establishes a correspondence based on the hash values ​​of each segment, and places them into the physical memory block pool created during system initialization. Simultaneously, the system creates contiguous virtual address ranges for all cached information in the virtual address space and maps the requested physical memory blocks to these virtual address ranges through a mapping interface. To improve the reuse rate of the key-value cache, the system adopts a differentiated memory allocation strategy: allocating contiguous physical memory blocks matching the number of tokens for each enhanced information chunk, enabling overall reuse across multiple rounds of inference; while for more dynamic user input, it segments the data according to a preset allocation granularity and requests physical memory pages for each segment separately.

[0032] 3.1.1.1 Differentiated Memory Allocation Strategy

[0033] (1) Enhanced Information Processing: Enhanced information typically contains multiple chunks returned by the retrieval system. Its content remains relatively stable across multiple rounds of inference, and its overall length is relatively large. To improve the construction and access efficiency of this key-value cache, the system allocates complete and contiguous physical memory space for the token sequence of each chunk in the enhanced information. During memory allocation, the system requests a corresponding number of contiguous physical memory blocks based on the number of chunks recorded during the encoding phase. The size of each memory block is consistent with the number of tokens in the corresponding chunk. Subsequently, the system maps the contiguous physical memory blocks to the virtual address range corresponding to each chunk through mapping operations. This method ensures that when reusing the physical memory blocks related to the enhanced information in subsequent iterations, each chunk only requires one mapping operation to complete reuse, thereby effectively improving cache reuse efficiency.

[0034] (2) User Input Processing: The user input section employs a block-based video memory allocation strategy. The system first divides the user input into fixed-size blocks according to a preset allocation granularity and allocates corresponding physical video memory pages for each block. Subsequently, the system flexibly maps these potentially non-contiguous physical pages to corresponding contiguous virtual address ranges.

[0035] A differentiated video memory allocation strategy is adopted. On the one hand, by implementing contiguous physical video memory allocation for chunks of enhanced information, chunk-level key-value cache reuse is achieved, reducing the overhead of video memory allocation and mapping. On the other hand, block-based video memory allocation is used for user input, reducing cache fragmentation and providing a basis for prefix reuse of subsequent inputs, further optimizing the use of video memory resources and system performance.

[0036] 3.1.1.2 Partition Location Encoding

[0037] After memory allocation, the system needs to perform positional encoding on both the augmented information region and the user input region. Since the attention mechanism of large language models relies on positional encoding to determine the temporal structure of tokens, once the positional offset changes, even if the text content is completely identical, the model cannot reuse the corresponding key-value cache, leading to unnecessary duplicate computations. If the previously unified positional encoding method is used for augmented information and user input, then when the augmented information consists of multiple chunks and their overall order is adjusted due to retrieval fluctuations, even if some chunks have identical content, the cache reuse capability will be lost due to changes in positional offset.

[0038] To address this issue, this invention employs a partitioned positional encoding strategy: the enhanced information area and the user input area undergo independent positional encoding. For the enhanced information portion, its content is relatively stable, and the token length and structure within a single chunk typically remain fixed. Therefore, the system uses an independent and continuous encoding interval for each chunk, ensuring its position remains consistent across multiple requests. This guarantees that the relevant key cache can be reliably reused when the chunk content remains unchanged. For the user input portion, since its token count and content vary with actual requests, its positional encoding is recalculated incrementally from the starting position of the user input area to ensure correct temporal order.

[0039] 3.1.2 Key-value cache reuse

[0040] When the hash value of the enhanced information or user input remains unchanged, the system reuses the allocated video memory resources. For enhanced information, the system directly locates the original contiguous physical video memory block through the hash value and remaps it to the current virtual address space, completely avoiding redundant allocation and mapping operations. For user input, the system uses a block matching mechanism, comparing the hash values ​​of each block sequentially to achieve gradual reuse of the prefix cache until all prefix matches are completed or the first mismatch is encountered. Specifically, the system divides the user input into multiple blocks and compares their hash values ​​sequentially with the prefix blocks in the cache. If the hash value of the current block matches the cache, the corresponding cache block is reused, and the comparison continues to the next block; if a mismatch occurs, the reuse process is terminated, and only the matched portion is retained.

[0041] 3.2 Forward Reasoning Calculation

[0042] After handling key-value cache hits and misses, the system sends the missing input data to the LLM for forward inference computation. During this process, the data passes through the model layer by layer, using a self-attention mechanism to generate predicted tokens, while the calculated key-value pairs are stored in pre-allocated GPU memory. Through multiple forward computations, the system gradually generates tokens, ultimately obtaining a complete token sequence.

[0043] 3.3 Result Generation

[0044] After the model inference is completed, the generated token sequence will be decoded into readable text and finally output to the user. Attached Figure Description

[0045] Figure 1 System Flowchart Detailed Implementation

[0046] The hardware environment of this invention is primarily a PC host. The host configuration is as follows: CPU is Intel(R) Core(TM) i7-4790, memory is 32GB DDR4, and operating system is 64-bit Ubuntu 24.04.

[0047] The software implementation of this invention is based on the Ubuntu 24.04 platform and developed using Python and C / C++ languages. The system uses Python version 3.12, CUDA version 12.8, GCC version 13.3, an NVIDIA GeForce RTX 4090 GPU, and Qwen-1_8B-Chat LLM.

[0048] The operation process of this invention is mainly divided into three parts: the first part is token sequence encoding; the second part is key-value cache hit and miss processing.

[0049] 1. Token Sequence Encoding

[0050] (1) Algorithm Description

[0051] Algorithm input: InputData, tokenizer

[0052] Algorithm output: InputTokens, TokenInfo

[0053] Description: InputData: User input and augmentation information. Tokenizer: A tokenizer used to encode text into tokens. InputTokens: The final continuous sequence of tokens. TokenInfo: Stores the length of each token component.

[0054] The pseudocode is as follows:

[0055]

[0056] 2. Key-value cache hit and miss handling

[0057] (1) Algorithm Description

[0058] Algorithm input: TokenInfo, HashInfo

[0059] Algorithm output: VirtualSpace

[0060] Notes: TokenInfo: Length of each token component. HashInfo: Hash value corresponding to each component.

[0061] VirtualSpace: Caches the virtual address space.

[0062] The pseudocode is as follows:

[0063]

[0064]

Claims

1. A key-value cache reuse method for retrieval enhancement generation systems, characterized in that... The implementation steps are as follows: (1) Knowledge retrieval stage: This stage retrieves document fragments (i.e., enhanced information) related to the current user input from an external knowledge base through a retrieval system, and integrates them with the user input to provide background information for subsequent generation tasks; (2) Prompt construction phase: In this phase, the system integrates user input and retrieved document fragments into a unified input structure and organizes the data in the form of key-value pairs. User input is stored in the form of strings, while the enhancement information is saved in the form of a list of text chunks. (3) Inference generation stage: After completing the key-value cache hit and miss processing, the input data is sent to the model to perform forward inference calculation. The predicted token is generated through the self-attention mechanism. During the inference process, the calculated key-value pairs will be stored in the pre-allocated video memory space. Through multiple forward calculations, the system gradually generates a complete token sequence and finally outputs it to the user.

2. The key-value caching reuse method for a retrieval enhancement generation system according to claim 1, characterized in that... This method provides suggestions during the build phase: (1) Integrate user input text and search enhancement information into a dictionary-formatted data structure with information type as the key and specific content as the value; (2) Save the user information content as a string and save the enhanced information content as a list of multiple string-formatted chunks.

3. The key-value cache reuse method for a retrieval enhancement generation system according to claim 1, characterized in that... This method provides suggestions during the build phase: (1) Perform word segmentation and encoding on multiple chunks of user input text and enhanced information to generate corresponding token sequences; (2) Use a key-value data structure to record the length and total number of each part of the token.

4. A key-value cache reuse method for a retrieval enhancement generation system according to claim 1, characterized in that... Inference generation stage: (1) For the enhanced information part, the system allocates a contiguous physical video memory block for each chunk and allocates it to the virtual address space through mapping; (2) For the user input part, the system divides the input into multiple blocks according to the preset granularity, allocates physical video memory pages to each block, and then flexibly maps them to the virtual address space.

5. A key-value cache reuse method for a retrieval enhancement generation system according to claim 1, characterized in that... Inference generation stage: (1) For the enhanced information part, the system uses an independent continuous position encoding range for each chunk, thereby ensuring that the position encoding remains consistent in multiple rounds of inference and ensuring the reuse of key-value cache; (2) For the user input part, the system dynamically calculates the position code based on the input content and increments it from the starting position of the input area to ensure the accuracy of the timing relationship.

Images