Cascade speculation inference method and system based on hierarchical decline KV cache compression

Abstract

The invention discloses a cascade speculation inference method and system based on hierarchical decline KV cache compression, and the method comprises the steps: firstly inputting a context prompt text into a target model for coding, generating a KV cache, and calculating an attention score between tokens; secondly, descending sorting is carried out based on the attention scores of the final input token, KV cache blocks corresponding to the first k attention scores are selected as cascading middle layer cache, a lightweight large language model is loaded to serve as a draft model, and a hierarchical decline KV cache compression strategy is adopted to maintain draft model cache. Then, based on all the caches, a double-layer cascade speculation reasoning framework is constructed, a target reasoning path is obtained, and the caches are updated; and finally, repeating the operation until target response data corresponding to the context prompt text is output according to the target reasoning path. According to the method, the KV cache proportion is reduced, meanwhile, the draft token acceptance rate of the target model of the full KV cache is improved, and reduction of precision is reduced.

Description

Technical Field

[0001] The present invention relates to the fields of computer artificial intelligence and large-model reasoning acceleration, and specifically to a cascaded speculative reasoning method and system based on hierarchical descending KV cache compression. Background Art

[0002] Large language models are currently booming, achieving remarkable results in natural language processing and numerous cross-domain applications. From simple text generation and question-answering systems to complex code writing and logical reasoning tasks, large language models have demonstrated powerful capabilities. Currently, large language models achieve word-by-word text generation through autoregressive decoding, but as the model size continues to expand, the computational cost increases exponentially, becoming a key bottleneck for improving reasoning efficiency. To balance reasoning speed and accuracy, speculative reasoning methods have emerged. This method draws on the concept of speculative execution in computer architecture and introduces a lightweight draft model to generate preliminary text sequences, which are then verified by the target large language model, effectively reducing reasoning time.

[0003] However, existing speculative reasoning methods have shortcomings in long-context scenarios: the natural divergence between the draft model and the target model increases with increasing context length, which means that the acceptance rate of speculative reasoning decreases with increasing context length. Furthermore, in large language models, temperature is a hyperparameter used to adjust the randomness of the model-generated text. Existing speculative reasoning methods are sensitive to temperature and produce poor results at high temperatures. Furthermore, training a draft model that matches the context length of the target model not only consumes massive computing resources, but it is also unclear whether a small model can achieve the same accuracy as a large model in long-context scenarios.

[0004] At the same time, when large language models perform inference based on autoregressive properties, each time a new token is generated, the complete KV cache and model parameters must be loaded into the GPU's video memory. The KV cache refers to the keys and values ​​calculated by the large language model based on the token sequence. This results in low computing core utilization and a significant increase in latency. Moreover, the memory usage of the KV cache increases linearly with the sequence length, gradually becoming a core bottleneck for large model inference in long-context scenarios. Currently, most KV cache compression research uses a fixed KV cache size for each layer of large language models, ignoring the differences in attention patterns at different layers in the large language model. In the high-level layers where attention is sparse, a large number of unimportant tokens are retained, while in the low-level layers where attention is dense, many key tokens are ignored. This "one-size-fits-all" strategy prevents the model from fully realizing its performance.

[0005] Based on the above status quo, an innovative method is urgently needed that can overcome the shortcomings of speculative reasoning in long context scenarios, and fine-tune the KV cache according to the characteristics of different layers of large language models, thereby significantly improving the reasoning efficiency and performance of large language models. The present application is based on this need. SUMMARY

[0006] The purpose of the present application is to overcome the shortcomings of existing methods and propose a cascading speculative reasoning method and system based on hierarchical decreasing KV cache compression, to effectively reduce the natural divergence of draft models and target models in long context scenarios, and to fully utilize the performance of the model by fine-tuning the KV cache budget of each layer of the large language model.

[0007] To achieve the above purpose, the technical solutions adopted by the present application are as follows:

[0008] On the one hand, the present application provides a cascading speculative reasoning method based on hierarchical decreasing KV cache compression, comprising the following steps:

[0009] Step 1, pre-filling stage: input the long context prompt text into the target model, encode the input context, generate the complete input token sequence, and calculate the attention scores between tokens. Then, generate the KV cache C target based on the input token sequence;

[0010] Step 2, initialize the cascading intermediate layer cache C mid : based on the attention scores of the last input token, sort in descending order, and select the KV cache blocks of the top k attention scores as the cascading intermediate layer cache C mid ;

[0011] Step 3, initialize the lightweight draft model C draft : load a lightweight large language model as a draft model, and use a hierarchical decreasing KV cache compression strategy to maintain the draft model cache C draft ;

[0012] Step 4, based on the KV cache, the cascading intermediate layer cache and the draft model cache, construct a two-level cascading speculative reasoning architecture to obtain the target reasoning path, the process is as follows:

[0013] Step 4.1, first layer of cascading speculative reasoning: the draft model generates γ1 draft tokens based on the draft model cache C draft , and inputs the draft tokens into the cascading intermediate layer, which is based on the cache C mid ; Step 4.2, second layer of cascading speculative reasoning: the cascading intermediate layer retains the draft tokens that pass the verification, and rejects the draft tokens based on C midResampling, together with the draft token, constitutes the first layer output sequence, which contains γ2 draft tokens in total;

[0014] Step 4.2 Cascade Speculative Inference Second Layer: The first layer output sequence is input into the target model as the input sequence of the second layer, and the target model is based on the cache C target Verification, and keep the verified token, join the second layer output sequence;

[0015] Step 5, in the two-layer cascade speculative inference architecture, update the cascade intermediate layer KV cache C based on the second layer output sequence mid With the draft model KV cache C draft ;

[0016] Step 6, repeat steps 4 to 5 until the target response data corresponding to the context prompt text is output according to the target inference path, and the speculative inference is completed.

[0017] Further, in step 2, based on the length of the input context, select the appropriate cascade intermediate layer cache budget Budget to initialize the cascade intermediate layer cache C mid , based on the attention score of the last input token in descending order, select the KV cache block corresponding to the token with the top k attention score as the cascade intermediate layer cache C mid . Wherein, the attention score A c The calculation formula is as follows:

[0018]

[0019] Wherein, Q is the query matrix calculated by the model according to the input token sequence, is the transpose of the key matrix of each cache block, and c is the number of tokens in each cache block. After calculating the attention score of each cache block, the attention score Top_K budget Key token cache block.

[0020] Further, in step 3, a hierarchical decreasing KV cache compression strategy is adopted to maintain the draft model cache C draft In the initialization stage, the draft model retains KV cache with Start_size and Recent_size size. Wherein, Start_size parameter is used to control the size of attention pooling window, that is, the number of tokens that attract a large number of attention scores at the head of the input sequence. Recent_size is the size of the sliding window, which retains the locality of tokens.

[0021] Further, in step 4.1, the first layer inference is the speculative inference of the draft model and the cascade intermediate layer, and the draft model Mdraft Cache-based C draft Generate γ1 draft tokens x1, x2, ..., x γ1 ~P draft (x i |x 1:i-1 ),i=1,2,…,γ1, where x1,x2,…,x γ1 The draft model is based on the probability distribution P draft (x i |x 1:i-1 ) generates γ1 draft tokens, where i is the position of the token. Then, the draft token is input into the cascade middle layer and the probability distribution P of the cascade middle layer is calculated. mid (x i |x 1:i-1 ),i=1,2,…,γ1,to verify the draft token. Then the draft token is accepted, otherwise it is rejected, and the acceptance threshold ε~U(0,1). After the probability distribution is corrected by the cascade middle layer, the C mid Resampling generates rejection sampling token x k1 , where the subscript k1 is the position of the rejection token in the cascade middle layer. The correction formula is as follows:

[0022]

[0023] Furthermore, in step 4.2, the second layer of cascaded speculative reasoning: the cascaded middle layer and the draft token generate a total of γ2 draft tokens, and the target model probability distribution P of the draft token is calculated by the target model loaded with all KV caches. target (x i |x 1:i-1 ),i=1,2,…,γ2, to verify the draft token sequence. If Then the draft token is accepted, otherwise it is rejected. The target model modifies the probability distribution and resamples to generate the draft tokenx k2 , where the subscript k2 is the position of the target model rejection token, and the correction formula is as follows:

[0024]

[0025] Furthermore, in step 5, the KV cache is updated as follows:

[0026] Step 5.1 Cascade the middle-tier KV cache C mid Update: After generating a fixed number of steps (such as every 64 tokens), reselect the attention score size to Top_K based on the attention score of the latest tokenbudget cache blocks.

[0027] Step 5.2 Draft Model KV Cache C draft Update: After two layers of cascaded speculative inference, the draft model updates the cache C based on the token sequence accepted by the target model. drafe , delete the historical token according to the sliding window size Recent_size, and ensure the locality of the token while maintaining the sliding window size. In addition, the present invention caches C according to the characteristics of different layers of the large language model. draft Perform fine-grained allocation to make cache C draft With the characteristic of decreasing levels, the cache budget is pyramid-shaped. draft The expression is as follows:

[0028]

[0029] Among them, R is the sliding window size Recent_size, Layer i is the cache block size of the i-th layer of the draft model, L is the total number of layers of the draft model, and α controls the decreasing gradient of the KV cache budget of each layer of the draft model.

[0030] On the other hand, the present invention also provides a cascaded speculative reasoning system based on hierarchical decremental KV cache compression, the system comprising:

[0031] The input module is used to convert the input long context prompt text into an input token sequence that can be recognized by the large language model and fill the system input sequence.

[0032] The cache initialization module is used to generate the target model KV cache, cascaded middle-layer cache, and draft model cache based on the system input token sequence.

[0033] The cascade reasoning module builds a two-layer cascade speculative reasoning architecture, performs model reasoning, obtains the target reasoning path, and fills the model output into the output sequence. In the two-layer cascade speculative reasoning architecture, the first-layer architecture is the collaboration between the draft model and the cascade middle layer. The draft model uses a lightweight large language model architecture to quickly generate draft tokens. The cascade middle layer uses the KV cache compressed version of the target large language model as the target model of the first-layer architecture to verify the draft token sequence generated by the draft model. The second-layer architecture is the collaboration between the cascade middle layer and the target model: the cascade middle layer serves as the draft model of the second-layer architecture, and generates the second-layer token based on the draft token sequence verified by the first layer. The target model loaded with the full context KV cache verifies the final draft token sequence.

[0034] The cache update module is used to update the cascade middle-layer KV cache and the draft model KV cache based on the model output of the cascade inference module.

[0035] The output module is used to generate the target reasoning path based on the filled output sequence of the cascade reasoning module to output the target response data corresponding to the data to be inferred.

[0036] Furthermore, the cascaded reasoning module is composed of a draft model, a cascaded intermediate layer, and a target model, with each part having a clear division of labor and close coordination. For the draft model, the Llama-1B model was selected. Due to its small number of parameters and low computational complexity, it can quickly generate preliminary text sequences, namely draft tokens, which provides a time advantage for the entire reasoning process. For the cascaded intermediate layer, which serves as a bridge between the draft model and the target model, the target model compressed by the KV cache is used as the cascaded intermediate layer. This can effectively reduce the natural divergence between the draft model and the target model in long-context scenarios, improve the draft token acceptance rate, receive the draft tokens generated by the draft model and perform preliminary verification, and screen out draft tokens with higher reliability. On the other hand, it further generates draft tokens based on its own cache, providing higher-quality candidate tokens for the target model. For the target model, that is, the large language model that users expect to use, the Llama-13B model is selected. It has powerful language understanding and generation capabilities, but the inference computational cost is high. Its responsibility is to perform final verification on the draft token passed from the cascade middle layer, and incorporate the verified token into the final output sequence to ensure the quality and accuracy of the generated text.

[0037] The cascaded speculative reasoning method and system based on hierarchical decremental KV cache compression described above effectively reduces the natural divergence between the draft model and the target model in long-context scenarios, and fully utilizes the model performance by finely allocating the KV cache budget of each layer of the large language model. This has the following beneficial effects:

[0038] (1) Through speculative reasoning, a lightweight draft model is used to generate draft tokens, and the draft token sequence is verified by the target model. This can significantly reduce the delay in token sequence generation caused by autoregression in large language models. At the same time, compared with existing speculative reasoning methods, the method of the present invention introduces a cascaded intermediate layer. By using the target model compressed by KV cache as the cascaded intermediate layer, the natural divergence between the draft model and the target model in long context scenarios is reduced, effectively improving the target model's acceptance rate of draft tokens.

[0039] (2) For the target model, the method of block caching is used for compression. By dividing the KV cache into small blocks, the attention between the query and the average key cache of each small block is calculated, which can identify and obtain the key information of the context, realize lossless approximation, effectively reduce the KV cache ratio, and improve the draft token acceptance rate of the target model of the full KV cache.

[0040] (3) For the draft model, a hierarchical decreasing streaming compression method is used. The starting token is reserved, and the token locality is reserved through a sliding window. At the same time, the present application observes that the attention weight distribution of the upper hidden layer is more consistent with the attention pool and locality, and the attention weight distribution of the bottom layer of the model is more uniform. Therefore, compared with the existing fixed mode KV cache compression method of each layer, the present application adopts a hierarchical decreasing pyramid cache budget, which allocates more KV cache budget to the bottom layer and less cache budget to the upper layer. In reducing the KV cache ratio, the present application reduces the decrease in precision by reasonably allocating the KV cache budget between the levels. BRIEF DESCRIPTION OF DRAWINGS

[0041] Other features, objects, and advantages of the present application will become more apparent from the following detailed description of non-limiting embodiments with reference to the following drawings:

[0042] Figure 1 A framework schematic diagram of a hierarchical decreasing KV cache compression based cascading speculative reasoning method is provided for the present application.

[0043] Figure 2 A flowchart of the hierarchical decreasing speculative reasoning method in the present application. DETAILED DESCRIPTION

[0044] The present application will be described in detail below with reference to specific embodiments. The following embodiments will help those skilled in the art to further understand the present application, but do not limit the present application in any form. It should be noted that the embodiments in the present application and the features in the embodiments can be combined with each other without conflict.

[0045] As shown in Figure 1 , the present application provides a hierarchical decreasing KV cache compression based cascading speculative reasoning method, which is composed of a draft model, a cascading intermediate layer and a target model. Specifically, the hierarchical decreasing cascading speculative reasoning method comprises the following steps:

[0046] Step 1, pre-filling stage: inputting a long context prompt text into the target model Llama-13B, pre-processing the input context, and using RoPE rotating position encoding to generate a complete KV cache C target . Further, the specific encoding steps are as follows:

[0047] Step 1.1: Input sequence (B is the batch size, L is the sequence length) converted into word vector E=W e X, where the matrix Word vectors The query matrix, key matrix, and value matrix are calculated through word vectors, which are briefly referred to as Q, K, and V.

[0048] Step 1.2: Convert the input matrix Q into complex form Q complex =Q ...,2i +i·Q ...,2i+1 (0≤i <d / 2),d为模型隐藏层的维数,同理将输入K转化为复数形式K complex .

[0049] Step = 1.3: Frequency initialization formula:

[0050] θ i =10000 -2i / d (0≤i <d / 2)

[0051] freqs=m·θ i

[0052] Among them, m is the position information, θ i To adjust the parameters.

[0053] Step 1.4: Inject the positional encoding into the input vector x, where x is the query and key, and the subscript j is the corresponding position in the vector:

[0054]

[0055] x rot =x⊙cos m +rotate_half(x)⊙sin m

[0056] Where ⊙ represents element-by-element multiplication, cos m and sin m The parameters for the corresponding positions.

[0057] Step 2: Based on the input context length, taking a 32KB context as an example, select 12KB as the cascaded middle layer cache budget to initialize the cascaded middle layer cache C. mid , based on the attention score of the last input token, through Calculate the attention score of each cache block and retain the KV cache blocks of the first 12K tokens in descending order of attention scores as the initial cache C mid .

[0058] Step 3: Initialize the lightweight draft model C draft : Load the lightweight draft model Llama-1B and use the hierarchical descending KV cache compression strategy to maintain the draft model cache C draft ,like Figure 2 As shown. For a 32K length context, a draft model cache budget of 1K token is used, Start_size is set to 16, and Recent size =1K-Start_size.

[0059] In long-context scenarios, the KV cache increases dramatically and the proportion of video memory gradually increases. When video memory resources are scarce, the KV cache occupancy needs to be reduced.

[0060] For the draft model, a hierarchical decreasing KV cache compression method is adopted to retain the attention pool, the token at the beginning of the input sequence, and the token locality is retained through a sliding window. At the same time, because the attention weight distribution of the upper hidden layer of the model is more in line with the attention pool and locality, the bottom attention weight distribution of the model is more uniform. Therefore, compared with the existing KV cache compression method with a fixed pattern per layer, the present invention adopts a hierarchical decreasing pyramid cache budget, allocating more KV cache budget to the bottom layer and less cache budget to the upper layer. While reducing the KV cache ratio, the accuracy drop is reduced by reasonably allocating the KV cache budget between levels.

[0061] Because the attention score of the draft model exhibits locality and convergence, that is, when calculating the attention score of a token, it pays more attention to the starting position and a small number of adjacent tokens in the nearby window. These two characteristics are more significant in the upper hidden layers of the model. Therefore, the hierarchical compression strategy adopted in this paper adjusts the size of Recent_size layer by layer, reducing the Recent_size of the hidden layer layer by layer by decreasing the gradient coefficient α:

[0062]

[0063] R i Recent_size is the size of the sliding window of layer i, Layer i is the cache block size of the i-th layer of the draft model, L is the total number of layers of the draft model, and α controls the decreasing gradient of the KV cache budget of each layer of the draft model.

[0064] Step 4: Based on the KV cache, cascaded middle-layer cache, and draft model cache, build a two-layer cascaded speculative reasoning architecture to obtain the target reasoning path:

[0065] In long-context scenario generation tasks, loading the complete target model weights for inference will introduce higher latency, and in long-context scenarios, the natural divergence between the draft model and the target model will increase, which means that the acceptance rate of speculative inference will decrease as the context length increases.

[0066] The present invention adopts a two-layer cascade speculative reasoning architecture consisting of a draft model, a cascaded intermediate layer, and a target model. Speculative reasoning is used to introduce a lightweight draft model to reduce the delay in generating inference tokens. At the same time, compared with existing speculative reasoning methods, this method introduces a cascaded intermediate layer and uses the target model compressed by KV cache as the cascaded intermediate layer, which reduces the natural divergence between the draft model and the target model in long-context scenarios and effectively improves the target model's acceptance rate of draft tokens.

[0067] Step 4.1: Cascade speculative reasoning first layer:

[0068] Step 4.1.1: Draft model Llama-1B based on cache C draft Use greedy sampling autoregression to generate 6 draft tokens [x t+1 ,x t+2 ,...,x t+6 ]

[0069] Step 4.1.2: Input the draft token sequence into the cascade middle layer, token[x t+1 ,x t+2 ,...,x t+6 ]~P draft (x i |x 1:i-1 ), i = 1, 2, ..., 6, where t is the starting position of the input token. Then, the probability distribution P of the cascaded intermediate layer is calculated. mid (x i |x 1:i-1 ),i=1,2,…,6

[0070] Step 4.1.3: If If all 6 tokens are accepted, it will be directly advanced to the t+6 position. If it is rejected at position t+k (k≤6), the subsequent tokens will be discarded. At t+k, the probability distribution is corrected by the cascaded intermediate layer and based on C mid Resample to generate draft token x k , the correction formula is as follows:

[0071]

[0072] The first level of reasoning ends when the draft model accepts 6 tokens.

[0073] Step 4.2, Cascade Speculative Inference Second Layer:

[0074] Step 4.2.1: Cascade Intermediate Layer and Draft Token Together Generate γ2 Draft Tokens [x t+1 ,x t+2 ,...,x t+γ2 ]

[0075] Step 4.2.2: Input the Draft Token Sequence into the Target Model Llama-13B, Token [x t+1 ,x t+2 ,...,x t+γ2 ] ~ P mid (x i |x 1:i-1 ), i = 1, 2,..., γ2, and calculate the probability distribution P target (x i |x 1:i-1 ), i = 1, 2,..., γ2 of the target model

[0076] Step 4.2.3: If , accept the draft token, otherwise reject the draft token, ε ~ U(0, 1). If all γ2 tokens are accepted, directly advance to position t+γ2, if rejected at position t+k (k≤γ2), discard subsequent tokens, and at t+k, generate draft token x target based on C k by re-sampling after the target model Llama-13B corrects the probability distribution, the correction formula is as follows:

[0077]

[0078] Step 5.1: Update the Cascade Intermediate Layer KV Cache C mid : After two-layer cascade speculative inference, the cascade intermediate layer will reselect Top_K budget key tokens based on the latest tokens by calculating the attention score of each cache block .

[0079] Step 5.2: Update the Draft Model KV Cache C draft : After two-layer cascade speculative inference, the draft model updates the cache C draft based on the accepted token sequence from the target model, deletes historical tokens according to the sliding window size Recent_size, maintains the sliding window size while ensuring token locality. And, allocate different budgets for each layer of the draft model KV cache, so that the cache C draftWith the characteristic of decreasing levels, the cache budget is pyramid-shaped. draft The expression is as follows:

[0080]

[0081] Among them, R is the sliding window size Recent_size, Layer i is the cache block size of the i-th layer of the draft model, L is the total number of layers of the draft model, and α controls the decreasing gradient of the KV cache budget of each layer of the draft model.

[0082] Step 6: Repeat steps 4 to 5 until the cascaded speculative reasoning method outputs target response data corresponding to the data to be inferred according to the target reasoning path;

[0083] Experimental setup:

[0084] The draft model used in the experiment of this invention is Llama-1B, and the target models are Llama-7B and Llama-13B. The datasets used are PG-19, NarrativeQA and InfiniteBench. Among them, PG-19 is a language modeling dataset launched by DeepMind containing 28,752 books published before 1919, which aims to promote the research of long context sequence modeling; NarrativeQA is a reading comprehension dataset launched by DeepMind, which contains the complete stories of 1,567 books and film scripts, as well as 46,765 complex questions that require understanding of the overall narrative to answer; InfiniteBench is a long text evaluation dataset jointly launched by Tsinghua University and Mianbi Intelligence, focusing on evaluating the ability of large models to handle 100K-2000K ultra-long contexts, covering 12 tasks in 5 major fields such as retrieval, mathematics, code, dialogue and books.

[0085] Result Analysis

[0086] In an environment with a temperature of 0.6, a Top_p of 0.9, and a γ1 of 6, the draft model is Llama-1B and the target model is Llama-7B. On a 32GB GPU-accelerator, we tested the speedup ratio of the autoregressive token generation rate and the draft token acceptance rate at different context lengths. The data in the table shows the speedup ratio of each method for autoregressive and the acceptance rate of the target model for draft tokens. The experimental results are shown in Table 1:

[0087] Table 1 Draft speedup ratios under different context lengths when the target model is Llama-7B

[0088]

[0089] In an environment with a temperature of 0.6, a Top_p of 0.9, and a γ1 of 6, the draft model is Llama-1B, the target model is Llama-13B, and a 2×GPU-accelerator 48G graphics card is used to test the acceptance rate of draft tokens under different context lengths. The data in the table are the speedup ratio of each method for autoregression and the acceptance rate of the target model for draft tokens. The experimental results are shown in Table 2:

[0090] Table 2 Draft token acceptance rate at different context lengths when the target model is Llama-13B

[0091]

[0092] In order to evaluate the stability of the proposed method, the performance of the method was tested on multiple datasets under an environment with Top_p of 0.9 and γ1 of 6, with the draft model being Llama-1B and the target model being Llama-13B, a 2×GPU-accelerator, and a 32G graphics card. The data in the table are the speedup ratio of each method for autoregression and the acceptance rate of the target model for the draft token. The experimental results are shown in Table 3:

[0093] Table 3 Speedup and acceptance rate of different context lengths tested on multiple datasets

[0094]

[0095] The results show that the speedup of this method compared to autoregressive inference increases with increasing context length and is applicable to multiple draft model-target model pairs, achieving a maximum speedup of 4.809x on a single GPU-accelerator, 32GB, and 4.314x on a dual GPU-accelerator, 48GB. Furthermore, the present invention has demonstrated its effectiveness across multiple datasets, reducing the natural divergence between the draft model and the target model in long contexts across various text scenarios. Compared to naive speculative inference methods, the method employed by this invention maintains a high draft token acceptance rate in long contexts, reaching over 80%.

[0096] In addition, temperature is also an important parameter in large-scale model inference, and is a hyperparameter used to adjust the randomness of the model-generated text. To test the draft token acceptance rate of the draft model and the target model under different temperature parameters, the proposed method was tested with a Top_p of 0.9 and a γ1 of 6. The draft model was Llama-1B, the target model was Llama-13B, and a 2×GPU-accelerator, 32G graphics card, to compare the speedup ratio of the autoregressive token generation rate and the draft token acceptance rate under different temperatures. The experimental data is shown in Table 4:

[0097] Table 4 Speedup and acceptance rate at different temperatures

[0098]

[0099] From the results, we can see that the draft token acceptance rate of this method only decreases slightly when the temperature increases, which ensures that the natural divergence between the draft model and the target model will not increase sharply in an environment with increasing temperature.

[0100] The cascaded speculative reasoning method based on hierarchical decremental KV cache compression proposed in this paper overcomes the problem of increased natural divergence between the draft model and the target model in speculative reasoning in long context scenarios, and refines the KV cache according to the characteristics of different layers of the large language model, thereby significantly improving the reasoning efficiency and performance of the large language model.

[0101] The above shows and describes the basic principles, main features, and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the above embodiments. The above embodiments and descriptions are merely preferred examples of the present invention and are not intended to limit the present invention. Various changes and improvements may be made to the present invention without departing from the spirit and scope of the present invention. Such changes and improvements fall within the scope of the present invention. The scope of protection claimed in the present invention is defined by the appended claims and their equivalents.

Claims

1. A cascaded speculative reasoning method based on hierarchical decremental KV cache compression, characterized in that: The steps include: Step 1: Input the contextual prompt text into the target model for encoding, generate a complete input token sequence, calculate the attention score between tokens, and generate a KV cache based on the input token sequence; Step 2: Sort the attention scores of the last input token in descending order and select the KV cache blocks of the tokens corresponding to the first k attention scores as the cascade middle layer cache; Step 3: Load the lightweight large language model as the draft model and use the hierarchical descending KV cache compression strategy to maintain the draft model cache; Step 4: Based on the KV cache, cascaded middle-layer cache, and draft model cache, a two-layer cascaded speculative reasoning architecture is constructed to obtain the target reasoning path; Step 5: In the two-layer cascaded speculative reasoning architecture, update the cascaded middle layer cache and the draft model cache based on the second layer output sequence; Step 6: Repeat steps 4 to 5 until the target response data corresponding to the context prompt text is output according to the target reasoning path, completing the speculative reasoning.

2. The cascaded speculative reasoning method based on hierarchical decremental KV cache compression according to claim 1 is characterized in that: The specific implementation of step 2 is as follows: based on the input context length, select the cascade middle layer cache budget Budget to initialize the cascade middle layer cache C mid , sort in descending order based on the attention score of the last input token, and select the KV cache blocks corresponding to the tokens with the first k attention scores as the cascaded intermediate layer cache C mid .

3. The cascaded speculative reasoning method based on hierarchical decremental KV cache compression according to claim 2 is characterized in that: In step 3, a hierarchical descending KV cache compression strategy is used to maintain the draft model cache C draft In the initialization phase, the draft model retains the KV cache of Start_size and Recent_size; among them, the Start_size parameter is used to control the size of the attention sink window; Recent_size is the sliding window size, which retains the locality of the token and reduces the Recent_size of the hidden layer layer by layer through the set decreasing gradient coefficient α.

4. The cascaded speculative reasoning method based on hierarchical decremental KV cache compression according to claim 3 is characterized in that: The specific implementation process of step 4 is as follows: Step 4.1 Cascaded speculative reasoning layer 1: Draft model based on the draft model cache C draft Generate γ1 draft tokens and input the draft tokens into the cascade middle layer, which is based on cache C mid Keep the draft token that passes the verification and reject it based on C mid Resample and form the first layer output sequence together with the draft token, which contains γ2 draft tokens in total; Step 4.2 Cascade speculative reasoning second layer: The output sequence of the first layer is passed into the target model as the input sequence of the second layer. The target model is based on the KV cache C target Verify and retain the verified token and add it to the second layer output sequence.

5. The cascaded speculative reasoning method based on hierarchical decremental KV cache compression according to claim 4 is characterized in that: In step 4.1, the first layer of reasoning is the speculative reasoning of the draft model and the cascaded intermediate layer. draft Cache-based C draft Generate γ1 draft tokens: x1, x2, ..., x γ1 ~P draft (x i |x 1:i-1 ),i=1,2,…,γ1, where x1,x2,…,x γ1 The draft model is based on the probability distribution P draft (x i |x 1:i-1 ) generates γ1 draft tokens by sampling, where i is the position of the token; then, the draft token is input into the cascade middle layer and the probability distribution P of the cascade middle layer is calculated. mid (x i |x 1:i-1 ),i=1,2,…,γ1,to verify the draft token; if Then accept the draft token, otherwise reject the draft token, accept the threshold ε~U(0,1); and the probability distribution is corrected by the cascade middle layer based on C mid Resampling generates rejection sampling token represented as x k1 , where the subscript k1 is the position of the rejection token in the cascade middle layer; In step 4.2, the second layer of cascaded speculative reasoning: the cascaded intermediate layer and the draft token generate a total of γ2 draft tokens, and the target model probability distribution P of the draft token is calculated by the target model loaded with all KV caches. target (x i |x 1:i-1 ), i = 1, 2, ..., γ2, to verify the draft token sequence; if The draft token is accepted, otherwise it is rejected; and the target model modifies the probability distribution and resamples the draft token to generate the draft token represented by x k2 , where the subscript k2 is the position of the target model rejection token.

6. The cascaded speculative reasoning method based on hierarchical decremental KV cache compression according to claim 5, characterized in that: The specific implementation process of step 5 is as follows: Step 5.1 Cascade the middle-tier KV cache C mid Update: After generating a fixed number of steps, reselect the attention score size to Top_K based on the attention score of the latest token budget cache blocks; Step 5.2 Draft Model KV Cache C draft Update: After two layers of cascaded speculative inference, the draft model updates the cache C based on the token sequence accepted by the target model. draft , delete historical tokens according to the sliding window size Recent_size, while maintaining the sliding window size to ensure the locality of tokens; and, based on the characteristics of different layers of the large language model, cache C draft Perform fine-grained allocation to make cache C draft With the characteristics of decreasing levels, the cache budget is pyramid-shaped. The cache C after fine-grained allocation draft The expression is as follows: Among them, R is the sliding window size Recent_size, Layer i is the cache block size of the i-th layer of the draft model, L is the total number of layers of the draft model, and α controls the decreasing gradient of the KV cache budget of each layer of the draft model.

7. A cascaded speculative reasoning system based on hierarchical decremental KV cache compression, used to implement the cascaded speculative reasoning described in any one of claims 1 to 6, characterized in that: Includes the following modules: The input module is used to convert the input long context prompt text into an input token sequence that can be recognized by the large language model; The cache initialization module is used to generate the target model KV cache, cascaded middle-layer cache, and draft model cache based on the input token sequence; The cascade inference module is used to build a two-layer cascade speculative inference architecture based on each model cache and input token sequence, perform model inference, obtain the target inference path, and fill the model output into the output sequence; The cache update module is used to update the cascade middle layer cache and the draft model cache based on the model output of the cascade inference module; The output module is used to generate the target reasoning path based on the filled output sequence of the cascade reasoning module to output the target response data corresponding to the data to be inferred.

8. The cascaded speculative reasoning system based on hierarchical decremental KV cache compression according to claim 7, characterized in that: In the two-layer cascaded speculative reasoning architecture, the first layer architecture is a collaboration between the draft model and the cascaded middle layer; the draft model adopts a lightweight large language model architecture to generate draft tokens; the cascaded middle layer adopts a KV cache compressed version of the target large language model as the target model of the first layer architecture to verify the draft token sequence generated by the draft model; the second layer architecture is a collaboration between the cascaded middle layer and the target model: the cascaded middle layer serves as the draft model of the second layer architecture, generates the second layer token according to the draft token sequence verified by the first layer, and the target model loaded with the complete context KV cache verifies the final draft token sequence.

Images