KVCache multiplexing acceleration model pre-filling method and system

Abstract

The invention discloses a KVCache multiplexing acceleration model pre-filling method and a KVCache multiplexing acceleration model pre-filling system. The method comprises the following steps: an off-line preprocessing stage: for a determined high-frequency text segment, retrieving a set number of similar texts from a knowledge base based on similarity, executing similarity-guided cross-segment preprocessing and cross attention fusion, and generating an enhanced KVCache containing cross attention information between the similar texts of the text segment for constructing a KVCache database; and an online reasoning stage: matching a prefix of a new user request by using a step-by-step matching strategy, and carrying out recalculation on the selected key token by combining a key token dynamic selection method based on an attention score so as to output a result responding to the user request. According to the method, the quality and the speed are both considered, and the method has good universality.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and more specifically, to a method and system for pre-filling a KVCache reuse acceleration model. Background Technology

[0002] Retrieval-Augmented Generation (RAG) is a hybrid reasoning framework that combines external knowledge retrieval with generative models. Unlike traditional pure generation or pure retrieval systems, RAG introduces a collaborative process between a retriever and a generator in the generation phase. First, it uses vector retrieval or sparse retrieval techniques to select several document fragments most relevant to the input query from a large-scale document library. Then, these fragments, along with the original query, are fed into the model as context, providing real-time supplementary external knowledge.

[0003] Attention is a computationally intensive neural network mechanism used to dynamically allocate computational resources based on the relevance between different positions in the input sequence when processing sequential data. It maps the features of each position to three representations: "Query," "Key," and "Value." It calculates the similarity score between the Query and all Keys, and then uses these scores to weightedly sum the Values ​​to generate a new contextual representation. The Attention calculation formula is as follows:

[0004]

[0005] Large Language Models (LLMs) are a class of pre-trained models based on deep neural networks. They utilize self-supervised learning on massive amounts of text data to acquire deep representations of language context. Unlike traditional rule-based or shallow feature engineering-based natural language processing systems, LLMs typically have hundreds of billions or even trillions of parameters, enabling rapid transfer and application across various downstream tasks at the zero-shot, few-shot, or fine-tuning levels. The core architecture of LLMs often employs a Transformer structure, efficiently capturing long-distance dependencies through multi-head self-attention. Leveraging hierarchical encoder-decoder or pure decoder architectures, they possess capabilities such as text generation, summarization, translation, and question answering. During the inference phase, LLMs can be flexibly controlled through "prompt engineering," allowing them to complete new tasks without retraining. Furthermore, through continuous online learning and incremental updates, they can continuously absorb new knowledge, enhancing their adaptability and scalability.

[0006] Before the generation task begins, the model first performs a complete forward computation on the user-provided "context" text (Prompt). This process is called Prefill. The core functions of Prefill are: 1. Calculating the self-attention key-value (KV) cache: For each existing token in the prompt, the Transformer calculates the corresponding Key and Value at each layer using the self-attention module and stores them in the cache (KVCache); 2. Generating context representations: Simultaneously, the model outputs the hidden states of each layer, providing rich contextual information for subsequent generation stages; 3. Completing the most costly part in one go: Since the Prompt is often quite long, Prefill requires performing a complete attention computation at each position, making this the most time-consuming step in the inference process. The processing speed of the Prefill stage affects the time it takes for the model to generate the first token (TimeToFirstToken, TTFT). The Prefill stage processing flow can be summarized by the following formula:

[0007] in For text input, For the input length, For location indexing, the Prefill stage will use the input X Processed as .

[0008] Decode Phase: After Prefilling, the model enters the Decode phase, which is the process of gradually generating new tokens. Each time a new token is generated, the model only needs to perform attention calculations between the query at that position and the previously cached key-value pairs, without repeatedly calculating the entire context, thus reducing the computational cost per step from O(n²) to O(n) (where n is the number of tokens already generated). The newly generated token is converted into its corresponding key-value pair and appended to the cache for use in the next decoding step; simultaneously, the hidden state of the last layer is updated to keep up with context changes. The Decode phase occurs at time steps... The processing procedure is as follows:

[0009] Analysis reveals that in high-knowledge-reuse scenarios such as RAG, user queries often contain repetitive or similar text fragments. Existing technologies attempt to accelerate pre-filling by reusing KVCaches of these fragments, but they still have the following drawbacks: (1) FullCacheRecompute: FullCacheRecompute completely recompiles the user request. In this process, only the Decode reuses the KVCache in the Prefill stage, resulting in the highest quality and the lowest reuse rate.

[0010] (2) vLLM first proposed Prefix Cache: Similar to reusing the KVCache generated in the Prefill stage during the Decode stage in the same set of dialogues, vLLM proposed that when the server receives a new user request, it will look up the longest matching prefix in the global cache and directly reuse the KVCache corresponding to that prefix. It will only re-perform the Prefill calculation for the remaining unmatched suffix, thereby significantly reducing the amount of redundant calculation and speeding up the response. This method does not affect the quality of model generation, but because the requirement for a perfect prefix match is too strict, the cache hit rate is low in practical applications.

[0011] (3) TurboRAG further proposes Full Cache Reuse: To address the issue of low hit rates under diverse prompts in the aforementioned schemes, TurboRAG relaxes the matching rules. That is, as long as any continuous substring in the user request matches a certain text segment in the global prefix cache, all layers of KVCache corresponding to that segment are directly reused, and only the subsequent unmatched parts are recalculated, thereby achieving the highest cache hit rate. However, since this "substring-level" reuse ignores the strict causal order in Transformer self-attention, that is, the Query and Key / Value must be precisely aligned at the time step, partial reuse will lead to distortion in attention weight calculation, which in turn will cause semantic drift and a decrease in generation quality.

[0012] (4) CacheBlend points out that the fundamental reason for the decline in TurboRAG generation quality is that the cross-attention information between different text blocks is ignored when reusing KVCache. Therefore, CacheBlend proposes the CacheFusion scheme: based on Full Cache Reuse, according to an empirically set 15% update ratio, based on the difference in L2 norm between the KVCache generated by Full Cache Reuse and the KVCache generated by Full Recompute on each token, the top 15% of tokens with the largest differences are selected as Critical Tokens, and their corresponding KVCaches are recalculated. The remaining tokens are directly reused from the cache, thus restoring cross-block cross-attention and maximizing the cache reuse rate. However, experiments have shown that the generation quality after CacheBlendKVCache reuse still declines, up to 55%.

[0013] In summary, existing technologies need to be improved to maximize the reuse of KVCache while ensuring that the generation quality is not compromised, so as to achieve high-speed processing in the pre-filling stage. Summary of the Invention

[0014] The purpose of this invention is to overcome the shortcomings of the prior art and provide a method and system for pre-filling a KVCache reuse acceleration model.

[0015] According to a first aspect of the present invention, a method for pre-filling a KVCache reuse acceleration model is provided. The method includes the following steps: Offline preprocessing stage: For identified high-frequency text segments Based on similarity, a set number of similar texts are retrieved from the knowledge base, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate a sequence of texts containing similar information. Enhanced KVCache, which uses cross-attention information between similar texts, is used to build the KVCache database; In the online inference phase: a step-by-step matching strategy is used to match the prefix of the new user request, and a key token dynamic selection method based on attention score is combined to recalculate the selected key tokens, and then output the result in response to the user request.

[0016] According to a second aspect of the present invention, a KVCache multiplexing acceleration model pre-filling system is provided. The system includes: Offline preprocessing module: used for processing specific high-frequency text segments. Based on similarity, a set number of similar texts are retrieved from the knowledge base, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate a sequence of texts containing similar information. Enhanced KVCache, which uses cross-attention information between similar texts, is used to build the KVCache database; The online inference module is used to match the prefix of new user requests using a step-by-step matching strategy, and then combine it with a key token dynamic selection method based on attention scores to recalculate the selected key tokens and output the result in response to the user request.

[0017] Compared with existing technologies, the advantages of this invention are that it increases the reuse probability of KVCache in the knowledge reuse domain while ensuring the quality of model generation, thereby accelerating the response speed of the model prefill stage. This invention improves processing speed while maintaining good generation quality, achieving generation quality at the level of FullCacheRecompute and processing speed exceeding that of CacheBlend. It solves the problems of FullCacheRecompute's good generation quality but slow processing speed, and CacheBlend's fast processing speed but poor generation quality.

[0018] Other features and advantages of the invention will become clear from the following detailed description of exemplary embodiments of the invention with reference to the accompanying drawings. Attached Figure Description

[0019] The accompanying drawings, which are incorporated in and form part of this specification, illustrate embodiments of the invention and, together with their description, serve to explain the principles of the invention.

[0020] Figure 1 This is a flowchart of a KVCache reuse acceleration model pre-filling method according to an embodiment of the present invention; Figure 2 This is a schematic diagram of a pre-filling method for a KVCache reuse acceleration model according to an embodiment of the present invention. Detailed Implementation

[0021] Various exemplary embodiments of the present invention will now be described in detail with reference to the accompanying drawings. It should be noted that, unless otherwise specifically stated, the relative arrangement, numerical expressions, and values ​​of the components and steps set forth in these embodiments do not limit the scope of the invention.

[0022] The following description of at least one exemplary embodiment is merely illustrative and is in no way intended to limit the invention or its application or use.

[0023] Techniques, methods, and equipment known to those skilled in the art may not be discussed in detail, but where appropriate, such techniques, methods, and equipment should be considered part of the specification.

[0024] In all the examples shown and discussed herein, any specific values ​​should be interpreted as merely exemplary and not as limitations. Therefore, other examples of exemplary embodiments may have different values.

[0025] It should be noted that similar labels and letters in the following figures indicate similar items; therefore, once an item is defined in one figure, it does not need to be discussed further in subsequent figures.

[0026] This invention improves model generation quality by increasing the CacheBlend recalculation ratio and refining the CriticalTokens selection method, thereby increasing generation efficiency while ensuring model generation quality.

[0027] In summary, the KVCache reuse acceleration model pre-filling method provided by this invention includes an offline preprocessing stage and an online inference stage. Offline preprocessing stage: Knowledge base text segments expected to be frequently retrieved in the RAG system are preprocessed to generate and cache their KVCache. Online inference stage: When actually processing new user requests, the offline enhanced cache is combined with a "two-stage" reuse and correction strategy to achieve a balance between high speed and high quality.

[0028] Specifically, in combination Figure 1 and 2 As shown, the provided KVCache reuse acceleration model pre-filling method includes the following steps: Step S1, Offline Preprocessing Stage: For a given high-frequency text segment, a set number of similar texts are retrieved from the knowledge base based on similarity, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate an enhanced KVCache containing cross-attention information between similar texts of the given text segment.

[0029] In step S1, for the retrieved high-frequency text, KV Caches corresponding to each segment are pre-calculated and cached to obtain cross-text cross-attention. Before inference, for high-frequency text retrieved from online search results or RAG, KV Caches corresponding to each segment can be pre-calculated and cached for reuse during subsequent generation. Furthermore, each segment of the KV Cache undergoes secondary processing.

[0030] 1) Similarity-guided cross-segment preprocessing For each text segment First, use RAG or vector retrieval to retrieve relevant data from the knowledge base. Top k most similar texts And load their corresponding pre-cached KVCache. .

[0031] 2)FullCacheReuse recalculation The system prompt The KVCache of similar segments, and the input of the original segment are concatenated in order to form text X, represented as: (1) This is then used as a one-time input to rerun the pre-filled calculation, generating a new KVCache: (2) (3) in, This indicates that the initial state of the KVCache is empty. Represents a text block KVCache, Indicates the first A recall-related text block.

[0032] 3) Cross-attention fusion Through the above recalculation and splicing, the new The middle already implies Cross-textual attention between similar text segments effectively fills the gaps in contextual interaction information that simple prefix or substring reuse cannot capture. This design achieves an enhanced KVCache with built-in relevant contextual information, laying the foundation for subsequent high-quality reuse.

[0033] Step S2, Online Inference Stage: A step-by-step matching strategy and a key token selection mechanism based on attention scores are adopted to ensure a balance between generation speed and generation quality.

[0034] (1) AlternativePath matching strategy To improve cache hit rate while prioritizing quality, the following order should be followed: Prioritize Prefix Cache matching: First, search the global cache for the longest prefix that exactly matches the current request's input prefix. If a match is found, reuse the high-quality KVCache with that prefix directly (because prefix matching does not violate causality).

[0035] The second priority is to try Full Cache Reuse: If the prefix does not match completely (common in the new question section of a question-and-answer platform), then an offline-enhanced cache is used to find any consecutively matching substring in the input for reuse, thereby achieving a high hit rate.

[0036] (2) Key token selection mechanism based on attention score: calculate attention score and select a set proportion of key tokens based on the score.

[0037] After deciding to enhance the cache through Full Cache Reuse, to ensure the final generated quality, instead of directly trusting all cached tokens, a more scientific dynamic recalculation method based on internal model signals was designed, specifically including: Step S21: Generate the final layer Query matrix.

[0038] After receiving the Prompt, which consists of System, RAG supplementary materials, and user questions, Transformers first generates the final Query matrix separately from the user questions.

[0039] (4) Step S22: Calculate the attention score to locate the key token.

[0040] For each retrieved text block Take the key matrix of its final layer. Attention scores are calculated using the attention mechanism. (5) Where d represents the size of the hidden layer of the model. Indicates the first Attention score of each text block Step S23: Score and select key tokens.

[0041] Summing each column yields a length of CriticalTokens score.

[0042] (6) in, Indicates the first The first text block The key score of each token.

[0043] A certain proportion of CriticalTokens are selected based on scores. For example, tokens with the highest scores are selected based on a preset or dynamically adjusted proportion (such as empirical values) and defined as "Critical Tokens".

[0044] Step S24, selective recalculation.

[0045] When generating the final output responding to the user's request, only the KVCache of the selected "key tokens" can be recalculated to accurately recover the cross-attention information for that part. The majority of other tokens can be directly reused from the offline enhanced cache.

[0046] Accordingly, the present invention also provides a KVCache reuse acceleration model pre-filling system. This system includes: an offline preprocessing module for pre-filling determined high-frequency text segments. Based on similarity, a set number of similar texts are retrieved from the knowledge base, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate a sequence of texts containing similar information. An enhanced KVCache, incorporating cross-attention information between similar texts, is used to construct the KVCache database. The online inference module utilizes a step-by-step matching strategy to match prefixes of new user requests. Combined with a dynamic key token selection method based on attention scores, the selected key tokens are recalculated to output the result in response to the user request. Within the system, each module can be implemented using dedicated processors, general-purpose processors, FPGAs, or a combination of hardware and software.

[0047] In summary, considering that PrefixCache generates high-quality data but has a low hit rate, this invention designs a two-stage improvement scheme for accelerating TTFT in RAG (Retrieval Enhanced Generation) scenarios to maintain both high generation quality and high hit rate. Information aggregation is performed in the Offline stage, significantly reducing the recalculation ratio required in the Online stage. This allows the model to achieve generation quality consistent with complete recalculation at a lower recalculation ratio. Furthermore, the Offline stage can be combined with various Online recalculation methods to further improve generation quality.

[0048] In summary, compared with the prior art, the present invention has the following main advantages: (1) Balancing quality and speed: The pre-built KVCache has higher quality through the "cross-attention fusion" in the offline pre-stage. The "key token selection based on attention score" in the online inference stage can accurately restore the core parts that affect the generation quality with minimal recalculation cost, so that the quality of the final output reaches the same level as full recalculation (FullCacheRecompute), while the processing speed exceeds that of CacheBlend and other solutions.

[0049] (2) Flexible mechanism and strong scalability: Offline information aggregation enhancement can be combined with various online caching matching strategies and recalculation mechanisms, and has good universality.

[0050] (3) Significantly reduce the delay of the first word: By greatly improving the "quality" and "quantity" of the effective cache, the time-consuming pre-filling stage is significantly shortened, thereby reducing the "first word generation time" of the entire reasoning process and ultimately improving the user's real-time experience when facing intelligent assistants or search questions and answers.

[0051] This invention can be a system, method, and / or computer program product. A computer program product may include a computer-readable storage medium having computer-readable program instructions loaded thereon for causing a processor to implement various aspects of the invention.

[0052] Computer-readable storage media can be tangible devices capable of holding and storing instructions for use by an instruction execution device. Computer-readable storage media can be, for example, but not limited to, electrical storage devices, magnetic storage devices, optical storage devices, electromagnetic storage devices, semiconductor storage devices, or any suitable combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), static random access memory (SRAM), portable compact disc read-only memory (CD-ROM), digital multifunction disc (DVD), memory sticks, floppy disks, mechanical encoding devices, such as punch cards or recessed protrusions storing instructions thereon, and any suitable combination thereof. The computer-readable storage media used herein are not to be construed as transient signals themselves, such as radio waves or other freely propagating electromagnetic waves, electromagnetic waves propagating through waveguides or other transmission media (e.g., light pulses through fiber optic cables), or electrical signals transmitted through wires.

[0053] Various aspects of the present invention are described herein with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the invention. It should be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer-readable program instructions.

[0054] These computer-readable program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, or other programmable data processing apparatus to produce a machine such that, when executed by the processor of the computer or other programmable data processing apparatus, they create means for implementing the functions / actions specified in one or more blocks of the flowchart and / or block diagram. These computer-readable program instructions can also be stored in a computer-readable storage medium that causes a computer, programmable data processing apparatus, and / or other device to operate in a particular manner; thus, the computer-readable medium storing the instructions comprises an article of manufacture that includes instructions for implementing aspects of the functions / actions specified in one or more blocks of the flowchart and / or block diagram.

[0055] Computer-readable program instructions may also be loaded onto a computer, other programmable data processing apparatus, or other device to cause a series of operational steps to be performed on the computer, other programmable data processing apparatus, or other device to produce a computer-implemented process, thereby causing the instructions executed on the computer, other programmable data processing apparatus, or other device to perform the functions / actions specified in one or more boxes of a flowchart and / or block diagram.

[0056] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of the present invention. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of an instruction containing one or more executable instructions for implementing a specified logical function. In some alternative implementations, the functions marked in the blocks may occur in a different order than those marked in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, can be implemented using a dedicated hardware-based system that performs the specified function or action, or using a combination of dedicated hardware and computer instructions. It will be known to those skilled in the art that implementation in hardware, implementation in software, and implementation using a combination of software and hardware are equivalent.

[0057] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, and are not limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or technical improvements to the embodiments in the market, or to enable others skilled in the art to understand the embodiments disclosed herein. The scope of the invention is defined by the appended claims.

Claims

1. A pre-filling method for a KVCache reuse acceleration model, comprising the following steps: Offline preprocessing stage: For identified high-frequency text segments Based on similarity, a set number of similar texts are retrieved from the knowledge base, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate a sequence of texts containing similar information. Enhanced KVCache, which uses cross-attention information between similar texts, is used to build the KVCache database; In the online inference phase: a step-by-step matching strategy is used to match the prefix of the new user request, and a key token dynamic selection method based on attention score is combined to recalculate the selected key tokens, and then output the result in response to the user request.

2. The method according to claim 1, characterized in that, The enhanced KVCache is calculated according to the following steps: For each retrieved text segment Retrieve from the knowledge base The Top k Most Similar Segments And load the corresponding pre-cached KVCache. ; System prompt The KVCache of similar segments and the input of the original segments are concatenated in order: Rerun the pre-fill calculation according to the following formula to generate a new KVCache as an enhanced KVCache: in, This indicates that the initial state of the KVCache is empty. Represents a text block KVCache, Indicates the first A recall-related text block.

3. The method according to claim 1, characterized in that, The key token is obtained according to the following steps: The Transformers model generates a final-level Query matrix separately based on the user's query: For each text segment Take the key matrix of its final layer. Attention scores are calculated using attention mechanisms and are expressed as follows: Where d represents the size of the hidden layer of the model. Indicates the first Attention score for each text block; After summing each column, we get a length of Key Tokens Score: in, Indicates the first The first text block The key score of each token; Key tokens are selected based on a set percentage threshold according to the score.

4. The method according to claim 2, characterized in that, Use search enhancement to generate RAGs or vector search to retrieve relevant text segments from the knowledge base. The top k most similar texts.

5. The method according to claim 1, characterized in that, The step-by-step matching strategy includes: First, find the longest prefix in the global cache that completely matches the prefix of the current user's request input. If a match is found, the KVCache of that prefix is ​​directly reused. If the prefix does not match completely, an offline-enhanced KVCache is used to find and reuse consecutively matching substrings in the input request.

6. The method according to claim 3, characterized in that, The ratio threshold is a preset fixed ratio or a dynamically adjusted ratio threshold.

7. A pre-filling system for a KVCache reuse acceleration model, comprising: Offline preprocessing module: used for processing specific high-frequency text segments. Based on similarity, a set number of similar texts are retrieved from the knowledge base, and similarity-guided cross-segment preprocessing and cross-attention fusion are performed to generate a sequence of texts containing similar information. Enhanced KVCache, which uses cross-attention information between similar texts, is used to build the KVCache database; The online inference module is used to match the prefix of new user requests using a step-by-step matching strategy, and then combines a key token dynamic selection method based on attention scores to recalculate the selected key tokens, thereby outputting the result in response to the user request.

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

9. A computer device comprising a memory and a processor, wherein a computer program capable of running on the processor is stored in the memory, characterized in that, When the processor executes the computer program, it implements the steps of the method according to any one of claims 1 to 6.

Images