Large language model reasoning acceleration method and device based on two-stage speculative decoding and storage medium

Abstract

The invention discloses a large language model reasoning acceleration method and device based on two-stage speculative decoding and a storage medium, and the method comprises the steps: constructing and initializing a Trie tree, and inserting a historical corpus, and phrase sequences in a document library or a code library into the Trie tree one by one; in the reasoning process, longest prefix matching is carried out based on a Trie tree, and a candidate draft sequence is generated by adopting branch backtracking and recursive search; performing confidence evaluation on the candidate draft sequence, calculating a joint confidence score of the sequence through probability multiplication and a Top-K screening mechanism, and judging whether the joint confidence score reaches a confidence threshold; if the accumulated confidence of the candidate sequence reaches a threshold value, skipping a small model generation stage, and directly entering large model verification; otherwise, entering a small model draft completion stage; and the final large model takes the replaced and updated draft sequence as final output. According to the method, adaptive acceleration of the decoding process can be realized, and the long text reasoning delay of the large language model is remarkably reduced while the generation quality is ensured.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of natural language processing, and in particular to a large language model inference acceleration method and device based on two-stage speculative decoding and a storage medium. BACKGROUND

[0002] In recent years, large language models have shown excellent performance in natural language processing tasks and are widely used in intelligent customer service, search and question answering, code generation, and other scenarios. However, as the model parameter size continues to grow, especially when processing long text inputs, inference speed has gradually become a key bottleneck limiting the actual deployment of the model.

[0003] Large language model inference mainly includes two stages: prefill and decoding. The prefill stage uses self-attention mechanisms to process the entire input sequence and build key-value caches (KV caches), which has high parallelism and is easy to accelerate. The decoding stage, however, uses autoregressive generation, which can only generate one token at a time and relies on all previous context information, resulting in a lack of parallelism and significant single-token latency. At the same time, the decoding stage's computation for a single batch is limited by memory bandwidth, leading to underutilization of GPU hardware computing resources and becoming the main bottleneck of the overall inference system.

[0004] Speculative decoding, as a recent inference acceleration paradigm, significantly improves the throughput of the decoding stage by generating draft sequences first and then verifying them in parallel by a large model. However, existing methods generally have the following problems: 1. Small model draft sequence generation has high computational overhead. Although generating drafts by a small model is more accurate than non-model methods, it still requires complete prefill, which may offset the acceleration benefits. 2. Non-model draft sequence generation has low quality. Relying on N-gram or retrieval-based splicing methods is extremely fast, but the resulting draft sequences have poor semantic coherence, leading to low acceptance rates by the large model.

[0005] 3. The strategy is fixed and lacks dynamic adaptation. Existing systems usually only use a single draft generation strategy, lacking a mechanism for intelligent switching between "high-quality-high-overhead" and "high-speed-low-quality" strategies based on context or task characteristics, making it difficult to achieve an optimal balance between efficiency and output quality.

[0006] Therefore, developing a multi-stage dynamic decision speculative decoding mechanism that combines efficient draft generation and high-quality output is crucial for improving the practicality of large language models in long text tasks. SUMMARY

[0007] To overcome the problems of unstable draft generation quality, waste of computing resources and limited acceleration efficiency in existing speculation decoding technology, the application provides a large language model inference acceleration method based on two-stage speculation decoding, a device and a storage medium, which realizes adaptive acceleration of the decoding process through the cooperative mechanism of model-free generation and small model completion, significantly reduces the long text inference delay of the large language model while ensuring the generation quality.

[0008] A large language model inference acceleration method based on two-stage speculation decoding is used to continuously generate long text, including the following steps: (1) Construct and initialize a Trie tree, insert phrase sequences in historical corpus, document library or code library into the Trie tree one by one, and define the nodes of the Trie tree as: ; Wherein, represents the current subword; represents the subnode mapping table, which is in the form of ; represents the end of phrase flag; represents the global cumulative frequency; (2) In the inference process, the longest prefix matching is performed based on the Trie tree, and the branch backtracking and recursive search are used to generate candidate draft sequences; (3) The confidence of the candidate draft sequence is evaluated, the joint confidence score of the sequence is calculated through the probability multiplication and Top-K screening mechanism, and it is judged whether the confidence threshold is reached; (4) If the cumulative confidence of the candidate sequence reaches the threshold, the small model generation stage is skipped, and step (6) of large model verification is directly entered; otherwise, step (5) of small model draft completion stage is entered; (5) The small model is used to decode the input prompt in a tree-like attention manner, Top-N candidate branches are expanded in parallel for M steps, and multiple draft sequences are obtained; (6) The large model is called to verify the draft sequence in parallel, the draft sequence is input into the large model, all candidate tokens are verified in parallel based on the tree-like attention mechanism, if the confidence is insufficient or the token is wrong, the error block is located, and the error token is replaced by local regeneration, and the draft sequence after replacement is updated as the final output.

[0009] In step (1), the Trie tree is constructed in the following manner: Each phrase is a sequence composed of several tokens, for each phrase, tokens are inserted from the root node, if the corresponding subnode does not exist, a new node is created ​And link it to the current node; after traversing, mark the last node as... and will Accumulation; among which, Indicates the first in the phrase One token, The symbol represents an empty set, indicating that the node initially has no child nodes.

[0010] The detailed process of step (1) is as follows: (1-1) Segment the historical text to obtain a phrase sequence. Each phrase .

[0011] (1-2) For each phrase Perform the insertion operation. Traverse the sub-words starting from the root node of the Trie. If the current node exist Move to that child node. Otherwise, create a new node. And link it to the current node. After traversal is complete, mark the last node as... and will Accumulate.

[0012] (1-3) During the online phase, the Trie is updated and dynamically optimized in real time based on user input or new text generated. The update method is the same as described in (1-2).

[0013] The specific process of step (2) is as follows: Initialize the matching pointer, path, and score; perform the longest prefix matching stepwise from the starting position of the input sequence; when encountering multiple branches, recursively traverse all child nodes to form multiple candidate paths; output the candidate draft after reaching the predefined sequence length.

[0014] The detailed process of step (2) is as follows: (2-1) Initialization. Let the input sequence be... Initialize the matching pointer position Candidate list Current path And set the current cumulative score. At the same time, let the current node Point to the root node of the Trie tree The path Used to store the currently matched phrase sequence, score Used to accumulate frequency information corresponding to path nodes.

[0015] (2-2) Prefix traversal matching. From Begin by performing a prefix matching search in a loop. Start from the current node. Start by matching the sequence of tokens. , This represents the length of the subwords matched. If a match is successful, continue; otherwise, terminate and output the match length. The matching process continues until... (n is) (length).

[0016] (2-3) Candidate path expansion. If a matching length is found... Extract matching phrase sequences It is obtained directly from the sequence and added to the current path. .

[0017] (2-4) Branch backtracking and recursive exploration. If the current node has multiple feasible matches among its child nodes, for example... It can match different child nodes. Then, the longest match is recursively performed on each child node to generate different candidate branch paths. This process ensures that all possible candidate sequences can be covered in the case of ambiguous or overlapping phrases.

[0018] (2-5) Move the pointer. After completing a match, move the pointer. Updated to and will Move to the end of the matching path to prepare for the next round of matching.

[0019] (2-6) Sequence termination condition. If , This represents the length of the predefined draft sequence. The generation of draft candidates ends, and the resulting draft sequence is... Then add candidates, Add to candidate list Then exit the loop.

[0020] In step (3), the confidence scores of the candidate draft sequences are calculated and sorted, and the Top-K high-confidence sequences are selected as candidate sequences. The specific process is as follows: First, perform joint probability calculation, the formula is: ; in, This is a candidate draft sequence. ,Depend on It consists of several phrases; The formula is obtained by querying and calculating from the Trie tree: ; in, It is a smoothing term to avoid division by zero. It is the number of all possible candidate phrases. Indicates the number of times a phrase appears. String concatenation operations representing phrases; After joint probability calculation, the candidate draft sequences are sorted in descending order of cumulative probability score and then normalized to obtain the final confidence score, in the following form: ; in, This represents the temperature coefficient, used to control the smoothness of the normalized distribution; Represents the set of candidate sequences Any candidate sequence in the sequence; The cumulative confidence score of Top-K candidate sequences is calculated as follows: ; If satisfied If the draft is deemed reliable, it is directly input into the large model for verification; otherwise, it is regenerated from the small model draft sequence.

[0021] In step (5), the small model adopts a tree-like attention mechanism, selects Top-N candidate tokens in each decoding step and performs M-step predictions in parallel to form a shallow tree-like candidate structure, and selects the highest-scoring candidate sequences as the completion result based on the cumulative log probability.

[0022] In step (5), the user inputs text during the initialization phase. Input the small model and obtain the initialized token distribution and context information.

[0023] (5-1) Parallel Candidate Expansion Based on Tree-like Attention. During the decoding process of the small model, we expand multiple candidate tokens using a parallel tree-like attention mechanism. Specifically, the small model uses the attention mechanism of the Transformer architecture to generate the probability distribution of candidate tokens for each decoding step. In each decoding step, the top N most probable tokens are selected.

[0024] (5-2) Tree-structured parallel branch path construction. Each Top-N token is considered the start of a new branch, generating a corresponding path. The model will compute the next few tokens in parallel on each branch (total expansion). (Step). This allows for simultaneous prediction on multiple paths in each decoding step, forming a shallow tree-like candidate structure. The generation of each branch is parallel, thus the expansion process can be performed simultaneously on multiple computational resources, avoiding excessive computational redundancy.

[0025] (5-3) Candidate sequence formation. Traverse all branches of this tree structure and extract the path from the root (input prompt) to the end of each branch, thus obtaining the candidate sequence. A candidate sequence of alternative tokens.

[0026] (5-4) Path scoring. Calculate the log probability of the tokens generated by the small model for each draft candidate path, and sum or average the log probabilities of all tokens on that path to obtain the score for that candidate path. .

[0027] (5-5) Final draft output. Select score. highest strip( Candidate sequences, used as fine-grained draft sequences generated by the small model. .

[0028] In step (6), the large model validation includes: forward parallel validation of each token in the draft sequence; if the token is not in the Top-K prediction distribution, it is marked as an error; consecutive error tokens are aggregated into error blocks; for the error blocks, the large model autoregressive regeneration is triggered based on the local context, and only the content of the error blocks is replaced; the updated draft sequence is used as the final output.

[0029] In step (6), the draft sequence generated by the small model The data will be fed into a large language model for further verification and correction. This stage mainly consists of the following key operations: parallel sequence verification, verification checking, error block identification, local regeneration, and replacement update.

[0030] (6-1) Forward Parallel Sequence Verification. The large language model performs a forward propagation on the concatenated sequence. The position of each draft token is computed in parallel. The real next token on the model output (at position in the input sequence) The probability distribution of the next token is given at each location. For each token in the draft sequence D Check if it is among the top-K (K=10) high-probability candidate tokens predicted by the model. If it is not in the top-K, mark that position. If it's an incorrect token, record its location. and context.

[0031] (6-2) Error Block Identification: Analyze consecutive error-marked locations. Aggregate consecutive error token fragments into error block locations, let... arrive A series of consecutive error tokens constitute an error block. .

[0032] (6-3) Local Regeneration: Constructed from the input, for each erroneous block Construct a locally regenerated input context. This refers to the correct portion of the original hint + draft sequence up to the point before the erroneous block. Using a large language model on this context, a new token sequence is regenerated in an autoregressive manner to replace the erroneous block. It is important to note the generation limits. You can set the maximum length to be generated, or the generation process will automatically stop when a certain natural termination point (such as a period or newline character) is reached.

[0033] (6-4) Replacement Update: Replace the positions of erroneous blocks in the original draft sequence with the correct token sequence generated by the large language model. The erroneous blocks are replaced by the new correct sequence, and the entire draft sequence is updated to the final result. The corrected draft sequence will be used as the final output of the large language model.

[0034] A large language model inference acceleration device based on two-level speculative decoding includes a memory and one or more processors. The memory stores executable code, and when the one or more processors execute the executable code, they are used to implement the above-mentioned large language model inference acceleration method based on two-level speculative decoding.

[0035] A computer-readable storage medium having a program stored thereon, which, when executed by a processor, implements the above-described method for accelerating inference of large language models based on two-level speculative decoding.

[0036] Compared with the prior art, the present invention has the following beneficial effects: 1. This invention proposes a two-level speculative decoding algorithm, which significantly reduces the inference latency of long texts and improves throughput while maintaining the original generation quality.

[0037] 2. This invention designs a lightweight draft generation module based on a high-frequency phrase Trie tree, which can quickly generate multiple preliminary draft sequences based on context and retrieval information with extremely low computational cost.

[0038] 3. This invention designs a confidence evaluation module based on a probability model, which can quickly determine whether the quality of the currently generated draft sequence can skip the generation of draft sequences of small models, thereby optimizing the overall inference process. Attached Figure Description

[0039] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0040] Figure 1 This is a flowchart of a method for accelerating reasoning in a large language model based on two-level speculative decoding, according to an embodiment of the present invention.

[0041] Figure 2 This is a schematic diagram of Trie tree candidate path construction according to an embodiment of the present invention.

[0042] Figure 3 This is a schematic diagram of a large language model inference acceleration device based on two-level speculative decoding according to an embodiment of the present invention. Detailed Implementation

[0043] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0044] It should be noted that, unless otherwise specified, the features in the following embodiments and implementation methods can be combined with each other.

[0045] like Figure 1 As shown, a method for accelerating inference in large language models based on two-level speculative decoding is presented. This method achieves a balance between generation efficiency and result reliability by constructing a Trie tree index, generating drafts without a model, completing small model drafts, and validating and regenerating the large model locally. Its specific implementation is as follows: In the Trie tree construction and initialization (S1) phase, phrase sequences from historical corpora, document repositories, or code repositories are uniformly encoded and inserted into the Trie tree structure to form the basic index. During construction, for each phrase, word-by-word insertion begins from the root node. If the current node does not have a corresponding subword, a new node is created and connected. When inserting the last node, it is marked as the end node and its frequency value is accumulated. During online runtime, the Trie tree can be dynamically updated based on user input or newly generated text to ensure gradual expansion of coverage.

[0046] (S2) Model-free draft pre-generation: In this stage, the longest prefix match is performed between the input token sequence and the Trie tree, and the scores of candidate paths are accumulated by combining node frequency information. For cases with ambiguous or overlapping phrases, the system generates multiple candidate sequences through branch backtracking and recursive search, and outputs a draft set after reaching a preset length. This process achieves lightweight and low-cost draft generation without calling a language model, such as... Figure 2 As shown.

[0047] (S3) Perform confidence evaluation on the generated candidate sequences. The evaluation calculates the conditional probability of the candidates using frequency statistics and obtains the confidence distribution using a normalization method. For example... Figure 2 As shown, assuming K=1, then only the path with the highest probability will be selected in the end. According to its score and threshold The results are compared. If the cumulative confidence of the Top-K candidates exceeds the set threshold, the draft result is deemed credible and can be directly input into the large model for verification; if the threshold is not reached, the process proceeds to the small model generation stage.

[0048] (S4) In the small model draft completion stage, a lightweight model is used for completion in cases of insufficient Trie tree coverage or long-tail phrases in the professional domain. This small model utilizes a tree-like attention mechanism to achieve parallel expansion of multiple tokens: in each decoding step, Top-N candidate tokens are selected, and M-step predictions are further carried out in parallel to form a shallow tree-like candidate structure. All candidate paths are sorted by log probability accumulation, and finally, several high-scoring candidates are selected as the completion result. In this embodiment, the small model uses the Eagle series model Eagle-Qwen2.5-Instruct, with approximately 0.33B parameters.

[0049] (S5) Large Model Validation and Local Regeneration Stage. The large model performs parallel validation on the input draft sequence. If a token is detected not appearing in the Top-K candidates of the predicted distribution, it is marked as an erroneous token; and consecutive erroneous positions are aggregated into erroneous blocks. For each erroneous block, this invention constructs a local context input to trigger the large model to perform autoregressive regeneration, replacing only the erroneous parts and avoiding full repetitive computation. In this embodiment, the large model uses Qwen2.5-14B-Instruct, with approximately 14 bytes of parameters.

[0050] (S6) Finally, the erroneous block is replaced and updated by the correct token sequence generated by the large model, resulting in a corrected complete output sequence, and the context Trie tree is incrementally updated by outputting the new sequence.

[0051] In summary, the embodiments of the present invention achieve a balance between computational efficiency and generation accuracy in sequence generation tasks through a hierarchical processing flow of "Trie tree construction - generation of model-free drafts - completion of small model drafts - verification and local regeneration of large models". It is particularly suitable for scenarios requiring efficient reasoning, such as large model retrieval agents and intelligent code generation.

[0052] Based on the same inventive principle, this invention also provides a large language model text generation device based on two-level speculative decoding. The device includes a memory, one or more processors, and a dedicated database module. The memory stores executable code, which, when executed by the one or more processors, is used to implement the large language model inference acceleration method based on two-level speculative decoding mentioned in the above embodiments. The dedicated database module is used to store and manage historical corpora, document libraries, code libraries, and other specialized data resources, providing underlying data support for the initialization of Trie trees.

[0053] The embodiments of the large language model text generation device based on two-level speculative decoding of the present invention can be applied to any device with data processing capabilities, such as a computer. The device embodiments can be implemented in software, hardware, or a combination of both. Taking software implementation as an example, as a logical device, it is formed by the processor of any data processing device loading the corresponding computer program instructions from non-volatile memory into memory for execution. From a hardware perspective, such as... Figure 3 The diagram shown is a hardware structure diagram of any device with data processing capabilities, including the large language model text generation device based on two-level speculative decoding according to the present invention. (Except for...) Figure 3 In addition to the processor, memory, database, network interface, and non-volatile memory shown, any data processing device in the embodiment may also include other hardware depending on the actual function of the data processing device, which will not be described in detail here.

[0054] Based on the same inventive principle, embodiments of the present invention also provide a computer-readable storage medium having a program stored thereon, which, when executed by a processor, implements the large language model inference acceleration method based on two-level speculative decoding mentioned in the above embodiments.

[0055] The embodiments described above provide a detailed explanation of the technical solutions and beneficial effects of the present invention. It should be understood that the above descriptions are merely specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, additions, and equivalent substitutions made within the scope of the principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A large language model inference acceleration method based on two-stage speculative decoding for continuous generation of long text, characterized in that, Comprising the following steps: (1) Construct and initialize a Trie tree, insert phrase sequences in the historical corpus, document library or code library into the Trie tree one by one, and the nodes of the Trie tree are defined as: Definition: ； wherein, represents the current subword; represents a subnode mapping table, in the form ; represents an end-of-phrase marker; represents a global cumulative frequency; (2) In the reasoning process, the longest prefix matching is performed based on the Trie tree, and branch backtracking and recursive search are adopted to generate candidate draft sequences; (3) Confidence evaluation is performed on the candidate draft sequences, the joint confidence score of the sequence is calculated through probability multiplication and Top-K screening mechanism, and it is judged whether the confidence threshold is reached; (4) If the cumulative confidence of the candidate sequence reaches the threshold, the small model generation stage is skipped, and the large model verification in step (6) is directly entered; otherwise, the small model draft completion stage in step (5) is entered; (5) The small model is used to decode the input prompt in a tree attention manner, Top-N candidate branches are expanded in parallel for M steps, and multiple draft sequences are obtained; (6) The large model is called to verify the draft sequence in parallel, the draft sequence is input into the large model, and all candidate tokens are verified in parallel based on the tree attention mechanism, if the confidence is insufficient or the token is wrong, the error block is located, and the error token is replaced by local regeneration, and the draft sequence after replacement is updated as the final output.

2. The two-stage speculative decoding-based large language model inference acceleration method according to claim 1, characterized in that, In step (1), the construction of the Trie tree is as follows: Each phrase is a sequence of tokens, and for each phrase, we insert a new node for each token starting from the root node, and if the corresponding child node does not exist, we create it , and link to the current node; after the traversal is completed, we mark the last node as , and accumulate ; where denotes the th token in the phrase, denotes the empty set, and represents that the node initially has no child nodes.

3. The two-level speculation-based decoding large language model inference acceleration method according to claim 1, characterized in that, The specific process of step (2) is: Initialize the matching pointer, path and score; perform longest prefix matching step by step from the starting position of the input sequence; when multiple branches are encountered, recursively traverse all child nodes to form multiple candidate paths; output the candidate draft after reaching the predefined sequence length.

4. The two-stage speculative decoding based large language model inference acceleration method according to claim 1, characterized in that, In step (3), the confidence of the candidate draft sequence is calculated and sorted, and Top-K high confidence sequences are selected as candidate sequences, the specific process is: First, the joint probability is calculated, the formula is: ； wherein, is a candidate draft sequence, consists of phrases; is calculated from the Trie tree, the formula is: ； wherein, is a smoothing term to avoid division by zero, is the number of all possible candidate phrases, denotes the number of occurrences of a phrase, denotes the string concatenation operation of a phrase; After the joint probability calculation of the candidate draft sequence, the cumulative probability score is sorted in descending order, and the final confidence score is obtained after normalization, which is: ； wherein, denotes a temperature coefficient, for controlling the smoothness of the normalized distribution, denotes any one of the candidate sequence set in the candidate sequence set The cumulative confidence calculation method of Top-K candidate sequences is: ； If the following conditions are met , it is considered that the draft quality is reliable, and the large model verification is directly input; otherwise, the small model draft sequence is regenerated.

5. The two-stage speculative decoding based large language model inference acceleration method according to claim 1, characterized in that, In step (5), the small model adopts a tree attention mechanism, selects Top-N candidate tokens in each decoding step, and expands M steps in parallel to form a shallow tree candidate structure, and selects the highest scoring candidate sequences based on log probability accumulation as the completion result.

6. The two-stage speculative decoding-based large language model inference acceleration method according to claim 1, characterized in that, In step (6), the large model verification includes: forward parallel verification of each token in the draft sequence, if the token is not in the Top-K prediction distribution, it is marked as an error; the continuous error tokens are aggregated into error blocks; for the error block, the large model autoregressive regeneration is triggered based on the local context, and only the error block content is replaced; the draft sequence after replacement is updated as the final output.

7. A large language model inference acceleration device based on two-level speculative decoding, characterized in that, A memory and one or more processors, the memory has executable code stored therein, and the one or more processors execute the executable code to implement the two-level speculative decoding-based large language model reasoning acceleration method of any one of claims 1-6.

8. A computer-readable storage medium, characterized in that, A program is stored thereon, which is executed by a processor to implement the two-level speculative decoding-based large language model reasoning acceleration method of any one of claims 1-6.

Images