Autoregression text generation acceleration method and system based on dynamic token tree

Through the autoregressive text generation method based on dynamic token tree, the draft model is used to generate candidate token trees and parallel verification sequences, which solves the problem of memory access overhead in large model reasoning and improves reasoning efficiency and throughput.

CN120671850APending Publication Date: 2025-09-19PIO CLOUD COMPUTING (SHANGHAI) CO LTD
View PDF 0 Cites 2 Cited by

Patent Information

Application Number
CN202510862443.X
Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
Filing Date
2025-06-25
Publication Date
2025-09-19

AI Technical Summary

Technical Problem

Existing speculative sampling methods have memory access overhead problems in large model reasoning, resulting in low reasoning efficiency.

Method used

An autoregressive text generation method based on a dynamic token tree is adopted. A draft model is used to generate candidate token trees, calculate path probabilities, retain the highest probability paths, merge them into equal-length matrices for parallel verification sequences, and map token positions for batch copying to reduce memory access overhead.

Benefits of technology

This significantly reduces the number of tree nodes, improves the attention calculation speed, reduces memory access overhead, and improves throughput.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN120671850A_ABST
    Figure CN120671850A_ABST
Patent Text Reader

Abstract

The invention relates to the technical field of large model reasoning acceleration, in particular to an autoregressive text generation acceleration method based on a dynamic token tree, which comprises the following steps of: S1, performing path expansion on a current context by using a draft model, generating a candidate token tree, calculating the path probability from each leaf node to a root node, and reserving the first N leaf nodes with the highest path probability, deleting other paths; s2, combining all the candidate sequences into a combined matrix filled to be equal in length, generating a self-defined attention mask, enabling each token to only pay attention to an ancestor token of a path where the token is located, inputting the combined matrix and the self-defined attention mask into a target model, and verifying all the sequences in parallel through single-time forward propagation; and S3, mapping the position of the accepted token in the candidate token tree to the continuous logic position of the main sequence, and copying the key value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship. According to the method, the token acceptance rate in speculation sampling of the large model acceleration technology is improved, so that the throughput is greatly improved.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the technical field of large model reasoning acceleration, and in particular to an autoregressive text generation acceleration method, system, device and storage medium based on a dynamic token tree. Background Art

[0002] Speculative sampling based on scratch models is a method that can fundamentally decode the computational memory access ratio, ensuring that the sampling distribution is exactly the same as that using the original model. It uses the original target model and an approximate model that is much smaller than the original model. The approximate model is used for autoregressive serial sampling, while the large model is used to evaluate the sampling results. During the decoding process, some tokens are relatively easy to decode, while others are difficult. Therefore, simple token generation can be handled by the small model, while difficult tokens are handled by the large model. The small model here can adopt the same structure as the original target model, but with fewer parameters, or simply use an n-gram model. The small model not only has a smaller computational footprint, but more importantly, reduces the need for memory access. In addition, there is speculative sampling based on additional output heads and speculative sampling based on feature autoregression.

[0003] Current speculative sampling has significant advantages in inference efficiency, but in actual implementation, it still faces the dilemma of memory access overhead. Therefore, a new autoregressive text generation acceleration method based on dynamic token tree is needed to achieve the improvement of token acceptance rate in speculative sampling of large model acceleration technology, thereby significantly improving throughput. Summary of the Invention

[0004] The purpose of the present invention is to solve the shortcomings of the prior art and provide an autoregressive text generation acceleration method based on a dynamic token tree, comprising the following steps: S1: Use the draft model to perform path expansion on the current context, recursively expand to generate a candidate token tree containing multiple branches, calculate the path probability from each leaf node to the root node in the candidate token tree, retain the first N leaf nodes with the highest path probability, and delete the remaining paths; S2: Merge all candidate sequences in the candidate token tree into a merge matrix filled with equal length, generate a custom attention mask so that each token only pays attention to the ancestor token in its path, input the merge matrix and the custom attention mask into the target model, and verify all sequences in parallel with a single forward propagation; S3: Map the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copy the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship.

[0005] Preferably, in step S1, the path expansion of the current context using the draft model and the recursive expansion to generate a candidate token tree containing multiple branches further include: Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold; When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first.

[0006] Preferably, calculating the path probability from each leaf node to the root node in the candidate token tree further comprises: Calculate the path probability from each leaf node to the root node as follows; in Indicates the first tokens, represents the draft model's predicted acceptance rate for this token, represents the confidence score derived from the draft model, The path from each leaf node to the root node.

[0007] Preferably, in step S2, all candidate sequences in the candidate token tree are merged into a merge matrix filled to equal length to generate a custom attention mask, further comprising: Flattening the candidate token tree into the candidate sequence by pre-order traversal; Merge the candidate sequences to form a merge matrix; Generate a custom attention mask with a block-diagonal structure, where each diagonal block corresponds to the causal mask of a path; Dependencies are stored using a compressed bitmap, where a mask matrix element M[i,j] = 1 if and only if token j is a path ancestor of token i.

[0008] Preferably, in step S2, after generating the custom attention mask, the following steps are included: The following masked attention calculation formula is used: in represents element-wise multiplication, is the key vector dimension, Q is the query matrix, K is the key matrix, V is the value matrix, and M is the merge matrix; Use FP16 precision to load the query matrix Q, key matrix K, and value matrix V, use FP32 precision to calculate the dot product and softmax to obtain the calculation matrix, and store the calculation matrix in FP16 precision; Use the GPU's Tensor Core to accelerate matrix multiplication operations.

[0009] Preferably, mapping the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copying the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship further includes: Mapping the positions of accepted tokens in the candidate token tree to consecutive logical positions in the main sequence; Use CUDA scatter operations to perform batch copies of non-contiguous to contiguous addresses; Select the copy strategy based on the hardware bandwidth. When the bandwidth reaches the threshold, direct parallel copy is used; otherwise, phased aggregate copy is used.

[0010] Based on the same concept, the present invention also provides an autoregressive text generation acceleration system based on a dynamic token tree, comprising: A candidate token tree construction module uses a scratch model to perform path expansion on the current context, recursively expands and generates a candidate token tree containing multiple branches, calculates the path probability from each leaf node to the root node in the candidate token tree, retains the first N leaf nodes with the highest path probability, and deletes the remaining paths; A candidate token verification module merges all candidate sequences in the candidate token tree into a merge matrix padded to equal length, generates a custom attention mask so that each token only pays attention to the ancestor token in its path, inputs the merge matrix and the custom attention mask into the target model, and verifies all sequences in parallel with a single forward propagation; The cache adjustment module maps the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copies the key value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship.

[0011] Preferably, the method of performing path expansion on the current context using the scratch model and recursively expanding to generate a candidate token tree containing multiple branches further comprises: Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold; When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first.

[0012] Based on the same concept, the present invention also provides a computer device, including a memory and one or more processors, wherein the memory stores computer code, and when the computer code is executed by the one or more processors, the one or more processors execute the autoregressive text generation acceleration method based on the dynamic token tree as described in any one of the embodiments.

[0013] Based on the same concept, the present invention also provides a computer-readable storage medium, which stores computer code. When the computer code is executed, the autoregressive text generation acceleration method based on a dynamic token tree as described in any one of the embodiments is executed.

[0014] Compared with the prior art, the present invention has the following beneficial effects: (1) The present invention uses the draft model to perform path expansion on the current context, recursively expands and generates a candidate token tree containing multiple branches, calculates the path probability from each leaf node to the root node, retains the first N leaf nodes with the highest path probability, and deletes the remaining paths, thereby significantly reducing the total number of tree nodes; (2) The present invention merges all candidate sequences in the candidate token tree into a merge matrix filled with equal length, generates a custom attention mask so that each token only pays attention to the ancestor token in its path, inputs the merge matrix and the custom attention mask into the target model, and verifies all sequences in parallel with a single forward propagation to achieve high-speed attention calculation.

[0015] (3) The present invention maps the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copies the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship, realizing a one-time copy to a reasonable position, which greatly reduces the memory access overhead. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Various other advantages and benefits will become apparent to those skilled in the art by reading the following detailed description of the preferred embodiment.The accompanying drawings are only for the purpose of illustrating the preferred embodiments and are not to be considered as limiting the invention.

[0017] Figure 1 Flowchart of the autoregressive text generation acceleration method based on dynamic token tree of the present invention; Figure 2A schematic diagram of candidate token tree construction for the autoregressive text generation acceleration method based on a dynamic token tree of the present invention; Figure 3 This is a schematic diagram of the tree-shaped verification mask construction of the autoregressive text generation acceleration method based on the dynamic token tree of the present invention. DETAILED DESCRIPTION

[0018] In order to make the purpose, technical solutions and advantages of the present invention more clear, the present invention is further described in detail below with reference to the accompanying drawings and examples. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention. Obviously, the embodiments described are part of the embodiments of this application, rather than all of the embodiments. Based on the embodiments in this application, all other embodiments obtained by those of ordinary skill in the art without making creative work are within the scope of protection of this application.

[0019] Those skilled in the art will understand that, unless otherwise specified, the singular forms "a," "an," and "the" used herein may also include plural forms. It should be further understood that the term "comprising" used in the specification of the present invention refers to the presence of the stated features, integers, steps, operations, elements, and / or components, but does not preclude the presence or addition of one or more other features, integers, steps, operations, elements, components, and / or groups thereof.

[0020] First embodiment Speculative sampling (Medusa) with additional output heads is similar to speculative sampling in that it reduces data handling costs by adding additional computation, thereby improving overall inference speed. The difference between them is that multi-head Medusa utilizes multiple prediction heads for multi-word prediction. These additional prediction heads are called Medusa heads. Multi-head Medusa can be applied to pre-trained models without changing the original model structure; simply add the Medusa heads. The training process is independent: the parameters of the original model are fixed and not trained; only the parameters of the additional Medusa heads are trained. This significantly reduces computational effort and greatly eases the difficulty of training convergence.

[0021] EAGLE, a speculative sampling algorithm based on feature autoregression, combines the advantages of speculative sampling based on the draft model and Medusa, using the Token-Level & Feature-Level method to sample draft tokens. This not only improves the accuracy of each draft token (compared with speculative sampling based on the draft model), but also increases the correlation between each node (compared with Medusa), greatly improving the acceptance rate. In addition, the autoregressive sampling method also avoids generating draft trees through the Cartesian product method, greatly reducing the length of the verification sequence.

[0022] PagedAttention: PagedAttention is an attention algorithm inspired by the classic concepts of virtual memory and paging in operating systems. Unlike traditional attention algorithms, PagedAttention allows for the storage of consecutive keys and values ​​in non-contiguous memory spaces. Specifically, PagedAttention partitions the KV cache of each sequence into blocks, each containing keys and values ​​for a fixed number of tokens. During attention computation, the PagedAttention kernel can efficiently identify and retrieve these blocks.

[0023] However, existing technologies still face the problem of memory access overhead. Figure 1 As shown, the autoregressive text generation acceleration method based on a dynamic token tree provided in this embodiment can significantly improve the response delay and concurrent processing capability of online services based on large models without compromising accuracy, including: S1: Use the draft model to perform path expansion on the current context, recursively expand to generate a candidate token tree containing multiple branches, calculate the path probability from each leaf node to the root node in the candidate token tree, retain the top N leaf nodes with the highest path probability, and delete the remaining paths. Specifically, in this embodiment, the draft model is a network whose dictionary is consistent with the target model, which may be in the form of an independent neural network or an additional output head of the target model. Given a single verified candidate token tree, use the draft model to forward reason N times to obtain N candidate tokens or reason once to obtain N candidate tokens from N output heads. The deletion operation includes marking all nodes on the retained path, depth-first traversal of the tree structure, deleting unmarked nodes and their subtrees, and updating the child node pointer of the parent node.

[0024] Preferably, in step S1, the path expansion of the current context using the draft model and the recursive expansion to generate a candidate token tree containing multiple branches further include: Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold. Specifically, in this embodiment, when the probability value of the expanded node is less than the dynamic threshold, the node is stopped from expanding. Otherwise, it continues to expand. Figure 2 As shown in Figure 2, when the dynamic threshold is 0.1 and K is 2, each node takes the first two child nodes for expansion. When the node probability is less than 0.1, the next token is no longer sampled for that token. When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first. Specifically, in this embodiment, when the probabilities of different nodes are equal, the shallow nodes are retained in order to ensure the structure of the tree, such as Figure 2 As shown, when K is 8, the first 8 tokens with the highest probability in the tree are taken to form a new draft tree (candidate token tree). In the standard causal inference sequence, the calculation of each token depends on all the tokens before it, and the sequence dependency graph presented will be a downward triangle. In the tree structure, the dependency path of the token is traced back from the leaf node to the root node. After flattening the tree structure, the dependency relationship of the token is presented as an additional triangle. Figure 3 As shown, a check mark indicates that the token with the vertical axis is dependent on the generated token. The vertical axis is from top to bottom, and the horizontal axis is from left to right for the order of token generation.

[0025] Preferably, calculating the path probability from each leaf node to the root node in the candidate token tree further comprises: Calculate the path probability from each leaf node to the root node as follows; in Indicates the first tokens, represents the draft model's predicted acceptance rate for this token, represents the confidence score derived from the draft model, For the path from each leaf node to the root node, specifically, in this embodiment, a GPU is used to parallelly calculate the path probabilities of all leaf nodes, and the top N paths are found through a parallel reduction algorithm.

[0026] A context-aware draft token tree is constructed, the acceptance rate is estimated by the joint probability of the tokens, and the draft tree is dynamically expanded using the acceptance rate. Under the premise of ensuring that the number of tree nodes remains basically unchanged, that is, the verification overhead remains basically unchanged, the acceptance length of the draft token sequence is improved.

[0027] S2: Merge all candidate sequences in the candidate token tree into a merge matrix filled with equal length, generate a custom attention mask so that each token only pays attention to the ancestor token in its path, input the merge matrix and the custom attention mask into the target model, and verify all sequences in parallel with a single forward propagation. Specifically, in this embodiment, for the candidate tokens generated by the tree structure, multiple possible text sequences are included. In order to improve the verification efficiency, multiple text sequences are merged into one sequence. Please refer to Figure 3 As shown, the mask is a check mark, representing the dependency between tokens. According to the description of the previous draft tree, a token on a path is dependent on the token before it. The mask between tokens is no longer in the form of overall causal inference, but is customized, and the correctness of the computational logic is guaranteed by the redesigned attention kernel.

[0028] Preferably, in step S2, all candidate sequences in the candidate token tree are merged into a merge matrix filled to equal length to generate a custom attention mask, further comprising: Flatten the candidate token tree into the candidate sequence by pre-order traversal. Specifically, in this embodiment, a depth-first traversal is performed starting from the root node. When a node is visited, its token is added to the sequence. The subtree is recursively traversed in descending order of child node probability, and the global index and path information of each node in the sequence are recorded. Merge the candidate sequences to form a merge matrix; Generate a custom attention mask with a block-diagonal structure, where each diagonal block corresponds to the causal mask of a path; Dependencies are stored using a compressed bitmap, where a mask matrix element M[i,j] = 1 if and only if token j is a path ancestor of token i. Specifically, in this embodiment, i is the query token index and j is the key token index.

[0029] Preferably, in step S2, after generating the custom attention mask, the following steps are included: The following masked attention calculation formula is used: in represents element-wise multiplication, is the key vector dimension, Q is the query matrix, K is the key matrix, V is the value matrix, and M is the merge matrix; Use FP16 precision to load the query matrix Q, key matrix K, and value matrix V, use FP32 precision to calculate the dot product and softmax to obtain the calculation matrix, and store the calculation matrix in FP16 precision; The GPU's Tensor Core is used to accelerate matrix multiplication operations. Specifically, in this embodiment, Use the WMMA (War Matrix Multiply Accumulate) API, configure Tensor Core calculations to 16×16×16 blocks, and the kernel function contains inline PTX assembly.

[0030] S3: Map the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copy the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship. Specifically, in this embodiment, the logical sequence numbers of the N accepted tokens on the page cache are usually non-continuous. The newly added N tokens are the self-incrementing sequence numbers on the main sequence, and the sequence numbers of the two sequences are matched one to one. The prompt is Let's introduce a toy. After prefill, the kvcache with sequence numbers 0-5 is used. The draft sequence to be verified, a sequence stored in a tree structure is It is has a the to goodbe, the sequence number corresponding to each token is 6-12, and the corresponding original sequence is It is a good It is the It has tobe Assuming that the final accepted sequence is It is the kvcache with sequence numbers 6, 7, and 10 will be copied together to sequence numbers 6, 7, and 8 of the kvcachepool, that is, the sequence map is {6:6, 7:7, 10:8}.

[0031] Preferably, mapping the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copying the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship further includes: Mapping the positions of accepted tokens in the candidate token tree to consecutive logical positions in the main sequence; Use CUDA scatter operations to perform batch copies of non-contiguous to contiguous addresses; Select the copy strategy based on the hardware bandwidth. When the bandwidth reaches the threshold, direct parallel copy is used; otherwise, phased aggregate copy is used.

[0032] Second embodiment Based on the same concept, the present invention also provides an autoregressive text generation acceleration system based on a dynamic token tree, comprising: A candidate token tree construction module uses a scratch model to perform path expansion on the current context, recursively expands the path to generate a candidate token tree containing multiple branches, calculates the path probability from each leaf node to the root node in the candidate token tree, retains the first N leaf nodes with the highest path probability, and deletes the remaining paths; A candidate token verification module merges all candidate sequences in the candidate token tree into a merge matrix padded to equal length, generates a custom attention mask so that each token only pays attention to the ancestor token in its path, inputs the merge matrix and the custom attention mask into the target model, and verifies all sequences in parallel with a single forward propagation; The cache adjustment module maps the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copies the key value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship.

[0033] Preferably, the method of performing path expansion on the current context using the scratch model and recursively expanding to generate a candidate token tree containing multiple branches further comprises: Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold; When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first.

[0034] Third embodiment In this embodiment, a computer device is provided, including a memory and one or more processors. The memory stores computer code. When the computer code is executed by the one or more processors, the one or more processors execute the steps of the autoregressive text generation acceleration method based on the dynamic token tree in the first embodiment.

[0035] In some embodiments of the present application, a computer-readable storage medium is also provided. When the computer-readable instructions are executed by one or more processors, the one or more processors execute the steps of the autoregressive text generation acceleration method based on a dynamic token tree as described in any one of the first embodiments.

[0036] It is understandable that, for the aforementioned autoregressive text generation acceleration method based on a dynamic token tree, if it is implemented in the form of a software function module and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention is essentially or the part that contributes to the prior art or all or part of the technical solution can be embodied in the form of a software product, and the computer software product is stored in a storage medium, including several instructions for enabling a computer device (which can be a personal computer server, or a network device, etc.) to execute all or part of the steps of the methods of various embodiments of the present invention. The aforementioned storage medium includes: various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk or an optical disk.

[0037] Computer-readable storage media may include a data signal propagated in baseband or as part of a carrier wave, which carries readable program code. Such propagated data signals may take a variety of forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination thereof. The readable storage medium may also be any readable medium other than a readable storage medium, which may send, propagate, or transmit a program for use by or in conjunction with an instruction execution system, apparatus, or device. The program code contained on the readable storage medium may be transmitted using any appropriate medium, including but not limited to wireless, wired, optical cable, RF, etc., or any suitable combination thereof.

[0038] The above description is merely a preferred embodiment of the present invention. The scope of protection of the present invention is not limited to the above embodiment. All technical solutions based on the concept of the present invention are within the scope of protection of the present invention. It should be noted that for those skilled in the art, various improvements and modifications that do not depart from the principles of the present invention should also be considered within the scope of protection of the present invention.

Claims

1. An autoregressive text generation acceleration method based on dynamic token tree, characterized in that: The following steps are involved: S1: Use the draft model to perform path expansion on the current context, recursively expand to generate a candidate token tree containing multiple branches, calculate the path probability from each leaf node to the root node in the candidate token tree, retain the first N leaf nodes with the highest path probability, and delete the remaining paths; S2: Merge all candidate sequences in the candidate token tree into a merge matrix filled with equal length, generate a custom attention mask so that each token only pays attention to the ancestor token in its path, input the merge matrix and the custom attention mask into the target model, and verify all sequences in parallel with a single forward propagation; S3: Map the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copy the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship.

2. The autoregressive text generation acceleration method based on dynamic token tree according to claim 1, characterized in that In step S1, the path expansion of the current context is performed using the draft model, and the recursive expansion generates a candidate token tree containing multiple branches, further comprising: Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold; When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first.

3. The autoregressive text generation acceleration method based on dynamic token tree according to claim 2, characterized in that Calculating the path probability from each leaf node to the root node in the candidate token tree further includes: Calculate the path probability from each leaf node to the root node as follows; in Indicates the first tokens, represents the draft model's predicted acceptance rate for this token, represents the confidence score derived from the draft model, The path from each leaf node to the root node.

4. The autoregressive text generation acceleration method based on dynamic token tree according to claim 3, characterized in that In step S2, all candidate sequences in the candidate token tree are merged into a merge matrix filled to the same length to generate a custom attention mask, further comprising: Flattening the candidate token tree into the candidate sequence by pre-order traversal; Merge the candidate sequences to form a merge matrix; Generate a custom attention mask with a block-diagonal structure, where each diagonal block corresponds to the causal mask of a path; Dependencies are stored using a compressed bitmap, where a mask matrix element M[i,j] = 1 if and only if token j is a path ancestor of token i.

5. The autoregressive text generation acceleration method based on dynamic token tree according to claim 4 is characterized in that, In step S2, after generating the custom attention mask, it includes: The following masked attention calculation formula is used: in represents element-wise multiplication, is the key vector dimension, Q is the query matrix, K is the key matrix, V is the value matrix, and M is the merge matrix; Use FP16 precision to load the query matrix Q, key matrix K, and value matrix V, use FP32 precision to calculate the dot product and softmax to obtain the calculation matrix, and store the calculation matrix in FP16 precision; Use the GPU's Tensor Core to accelerate matrix multiplication operations.

6. The method for accelerating autoregressive text generation based on a dynamic token tree according to claim 5, wherein: Mapping the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copying the key-value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship, further including: Mapping the positions of accepted tokens in the candidate token tree to consecutive logical positions in the main sequence; Use CUDA scatter operations to perform batch copies of non-contiguous to contiguous addresses; Select the copy strategy based on the hardware bandwidth. When the bandwidth reaches the threshold, direct parallel copy is used; otherwise, phased aggregate copy is used.

7. An autoregressive text generation acceleration system based on a dynamic token tree, characterized in that: include: A candidate token tree construction module uses a scratch model to perform path expansion on the current context, recursively expands the path to generate a candidate token tree containing multiple branches, calculates the path probability from each leaf node to the root node in the candidate token tree, retains the first N leaf nodes with the highest path probability, and deletes the remaining paths; A candidate token verification module merges all candidate sequences in the candidate token tree into a merge matrix padded to equal length, generates a custom attention mask so that each token only pays attention to the ancestor token in its path, inputs the merge matrix and the custom attention mask into the target model, and verifies all sequences in parallel with a single forward propagation; The cache adjustment module maps the position of the accepted token in the candidate token tree to the continuous logical position of the main sequence, and copies the key value cache of the target model to the corresponding position of the main sequence in batches according to the mapping relationship.

8. The method for accelerating autoregressive text generation based on a dynamic token tree according to claim 7, wherein: The method further includes: performing path expansion on the current context using the draft model, and recursively expanding to generate a candidate token tree containing multiple branches. Take the last token output by the base model as the root node; Starting from the root node, each layer is fixedly expanded to k×k child nodes, where k is a preset expansion factor; For each parent node, use the draft model to predict the probability distribution of the next token and select the top k candidates with the highest probability as child nodes; When the path probability product is lower than the dynamic threshold Stop expansion when , d is the current node depth, λ is the attenuation factor, is the initial expansion threshold; When the path probabilities of multiple leaf nodes are the same, the nodes with shallower depth are retained first.

9. A computer device comprising a memory and one or more processors, wherein the memory stores computer code, and when the computer code is executed by the one or more processors, the one or more processors execute the steps of the autoregressive text generation acceleration method based on a dynamic token tree as described in any one of claims 1 to 8.

10. A computer-readable storage medium storing computer code, wherein when the computer code is executed, the steps of the autoregressive text generation acceleration method based on a dynamic token tree according to any one of claims 1 to 8 are executed.

Citation Information

Cited By

  • Training method of file fragment classification model and file fragment classification method

    CN122045149A

  • File fragment classification model training method and file fragment classification method

    CN122045149B