Medusa model token screening method and device

Abstract

The disclosure provides a medusa model token screening method and device. The method comprises: performing inference calculation on an input sequence by using a target model to generate a plurality of candidate tokens; sampling the plurality of candidate tokens to extract part of the candidate tokens; inputting the part of the candidate tokens into a medusa model for speculative inference to generate a plurality of speculative tokens in a tree structure; screening the plurality of speculative tokens based on the tree structure in a pruning manner to generate a target token; and performing subsequent inference calculation based on the target token to generate an output result. The medusa model token screening method and device disclosed in the application can ensure the quality of the inference result, optimize the token sampling performance, effectively improve the efficiency of speculative sampling during the inference optimization of the large model, and have a wide range of adaptation.

Description

Technical Field

[0001] This disclosure relates to the field of computer information processing, and more specifically, to a method and apparatus for token filtering based on the medusa model. Background Technology

[0002] With the widespread application of large-scale language models (such as GPT and Qwen), text generation tasks hold an important position in the field of natural language processing. However, the high computational cost of these models limits their application in real-time response and resource-constrained scenarios. To reduce inference costs and improve generation efficiency, speculative sampling techniques have been extensively researched and applied.

[0003] Speculative sampling is an optimization method that reduces the computational cost of the target model through prediction. Its core idea is to use a smaller model (such as the Medusa model) to generate candidate tokens (draft tokens), and then use the target model to verify the acceptance probability of these tokens, thereby improving inference efficiency.

[0004] The Medusa model generates multiple draft tokens in parallel using multiple prediction heads, enabling it to generate multiple candidate tokens at once, thus accelerating the inference process. However, due to the lack of correlation between the prediction heads, duplicate or incoherent tokens are easily generated, affecting generation efficiency and quality.

[0005] Therefore, a new token filtering method and device based on the Medusa model is needed.

[0006] The information disclosed in the background section is only intended to enhance the understanding of the background of this application, and therefore may include information that does not constitute prior art known to those skilled in the art. Summary of the Invention

[0007] In view of this, this application provides a token filtering method and apparatus based on the Medusa model, which has the following technical advantages:

[0008] By utilizing a scoring mechanism and path pruning, the likelihood of draft tokens generated by Medusa being accepted by the target model is increased, thereby optimizing generation efficiency;

[0009] Repeated pruning techniques are used to remove duplicate tokens, reducing path complexity and unnecessary computation.

[0010] Low-scoring and disjointed paths are removed to ensure that the final generated token sequence is more coherent and accurate;

[0011] This method is not only applicable to the Medusa model, but can also be extended to other multi-head prediction models and speculative sampling techniques, demonstrating broad applicability.

[0012] Other features and advantages of this application will become apparent from the following detailed description, or may be learned in part from practice of this application.

[0013] According to one aspect of this application, a token filtering method based on a medusa model is proposed. This method includes: performing inference calculations on an input sequence using a target model to generate multiple candidate tokens; sampling the multiple candidate tokens to extract a subset of candidate tokens; inputting the subset of candidate tokens into a medusa model for speculative inference to generate multiple speculative tokens in a tree structure; filtering the multiple speculative tokens using a pruning method based on the tree structure to generate a target token; and performing subsequent inference calculations based on the target token to generate an output result.

[0014] In one exemplary embodiment of this application, the input sequence is inferred and calculated using a target model to generate multiple candidate tokens, including: inputting the input sequence into the target model; the target model calculates multiple candidate tokens based on inference from the input sequence; and caching the key-value cache information corresponding to the multiple candidate tokens.

[0015] In one exemplary embodiment of this application, sampling the plurality of candidate tokens and extracting a portion of the candidate tokens includes: accepting k candidate tokens from the plurality of candidate tokens and caching the kv cache information of p tokens; and discarding the kv cache information of the kp tokens that are not cached from the portion of the candidate tokens.

[0016] In one exemplary embodiment of this application, the candidate tokens are input into the medusa model for speculative inference to generate multiple speculative tokens in a tree structure, including: each prediction head of the medusa model generates a set of speculative tokens and organizes them hierarchically into a tree structure.

[0017] In one exemplary embodiment of this application, the process of inputting the candidate tokens into the medusa model for speculative reasoning to generate a tree-structured plurality of speculative tokens further includes: arranging the plurality of speculative tokens in the order of Depth-First Search (DFS).

[0018] In one exemplary embodiment of this application, filtering the plurality of speculative tokens using a pruning method based on a tree structure to generate a target token includes: filtering the plurality of speculative tokens using a predictive scoring pruning method based on a tree structure to generate a target token; and / or filtering the plurality of speculative tokens using a leaf node pruning method based on a tree structure to generate a target token; and / or filtering the plurality of speculative tokens using a repetition pruning method based on a tree structure to generate a target token.

[0019] In one exemplary embodiment of this application, a target token is generated by filtering the plurality of speculative tokens using a prediction score pruning method based on a tree structure. This includes: obtaining a prediction score for each speculative token, wherein the prediction score represents the probability that the speculative token is accepted by the target model; and filtering the plurality of candidate tokens using the prediction score to generate the target token.

[0020] In one exemplary embodiment of this application, the multiple speculative tokens are filtered using a leaf node pruning method based on a tree structure to generate a target token, including: removing speculative tokens located at leaf nodes based on the tree structure to generate the target token.

[0021] In one exemplary embodiment of this application, filtering the multiple speculative tokens based on a tree structure using a repeated pruning method to generate a target token includes: removing duplicate speculative tokens based on the tree structure to generate the target token.

[0022] According to one aspect of this application, a token filtering device based on a medusa model is proposed. The device includes: a target module for performing inference calculations on an input sequence using a target model to generate multiple candidate tokens; a sampling module for sampling the multiple candidate tokens and extracting a subset of candidate tokens; a speculation module for inputting the subset of candidate tokens into a medusa model for speculative inference to generate multiple speculative tokens in a tree structure; a pruning module for filtering the multiple speculative tokens using a pruning method based on the tree structure to generate a target token; and an inference module for performing subsequent inference calculations based on the target token to generate an output result.

[0023] According to one aspect of this application, an electronic device is provided, comprising: one or more processors; a storage device for storing one or more programs; and, when the one or more programs are executed by the one or more processors, causing the one or more processors to implement the method as described above.

[0024] According to one aspect of this application, a computer-readable medium is provided having a computer program stored thereon that, when executed by a processor, implements the method described above.

[0025] According to the token filtering method and apparatus based on the Medusa model of this application, multiple candidate tokens are generated by performing inference calculations on the input sequence using a target model; the multiple candidate tokens are sampled to extract a portion of the candidate tokens; the portion of the candidate tokens are input into the Medusa model for speculative inference to generate multiple speculative tokens in a tree structure; the multiple speculative tokens are filtered using a pruning method based on the tree structure to generate a target token; and subsequent inference calculations are performed based on the target token to generate the output result. This method can optimize token sampling performance while ensuring the quality of the inference result, effectively improve the efficiency of speculative sampling during large model inference optimization, and has a wide range of applications.

[0026] It should be understood that the above general description and the following detailed description are merely exemplary and do not limit this application. Attached Figure Description

[0027] The above and other objects, features, and advantages of this application will become more apparent from the detailed description of exemplary embodiments with reference to the accompanying drawings. The drawings described below are merely some embodiments of this application, and those skilled in the art can obtain other drawings based on these drawings without any inventive effort.

[0028] Figure 1 This is a flowchart illustrating a token filtering method based on the Medusa model according to an exemplary embodiment.

[0029] Figure 2 This is a flowchart illustrating a token filtering method based on the Medusa model according to another exemplary embodiment.

[0030] Figure 3 This is a schematic diagram illustrating a token filtering method based on the Medusa model according to another exemplary embodiment.

[0031] Figure 4 This is a schematic diagram illustrating a token filtering method based on the Medusa model according to another exemplary embodiment.

[0032] Figure 5 This is a schematic diagram illustrating a token filtering method based on the Medusa model according to another exemplary embodiment.

[0033] Figure 6This is a schematic diagram illustrating a token filtering method based on the Medusa model according to another exemplary embodiment.

[0034] Figure 7 This is a block diagram illustrating a token filtering device based on the Medusa model according to an exemplary embodiment.

[0035] Figure 8 This is a block diagram illustrating an electronic device according to an exemplary embodiment.

[0036] Figure 9 This is a block diagram illustrating a computer-readable medium according to an exemplary embodiment. Detailed Implementation

[0037] Exemplary embodiments will now be described more fully with reference to the accompanying drawings. However, these exemplary embodiments can be implemented in many forms and should not be construed as limited to the embodiments set forth herein; rather, they are provided so that this application will be thorough and complete, and will fully convey the concept of the exemplary embodiments to those skilled in the art. The same reference numerals in the drawings denote the same or similar parts, and therefore repeated descriptions of them will be omitted.

[0038] Furthermore, the described features, structures, or characteristics can be combined in any suitable manner in one or more embodiments. Numerous specific details are provided in the following description to give a thorough understanding of embodiments of this application. However, those skilled in the art will recognize that the technical solutions of this application can be practiced without one or more of the specific details, or other methods, components, apparatuses, steps, etc., can be employed. In other instances, well-known methods, apparatuses, implementations, or operations are not shown or described in detail to avoid obscuring various aspects of this application.

[0039] The block diagrams shown in the accompanying drawings are merely functional entities and do not necessarily correspond to physically independent entities. That is, these functional entities can be implemented in software, in one or more hardware modules or integrated circuits, or in different network and / or processor devices and / or microcontroller devices.

[0040] The flowcharts shown in the accompanying drawings are merely illustrative and do not necessarily include all content and operations / steps, nor do they necessarily have to be performed in the described order. For example, some operations / steps can be broken down, while others can be combined or partially combined; therefore, the actual execution order may change depending on the specific circumstances.

[0041] It should be understood that although the terms first, second, third, etc., may be used herein to describe various components, these components should not be limited by these terms. These terms are used to distinguish one component from another. Therefore, the first component discussed below may be referred to as the second component without departing from the teachings of this application. As used herein, the term "and / or" includes all combinations of any one and more of the associated listed items.

[0042] Those skilled in the art will understand that the accompanying drawings are merely schematic diagrams of exemplary embodiments, and the modules or processes in the drawings are not necessarily essential for implementing this application, and therefore cannot be used to limit the scope of protection of this application.

[0043] The technical abbreviations used in this application are explained as follows:

[0044] NLP (Natural Language Processing) is a technology for interaction between computers and human language, applied in fields such as text generation, translation, and question answering.

[0045] A token is the smallest unit in natural language processing, usually referring to a word, subword, or character, used as input or output for a language model.

[0046] Draft Model: A language model with a small number of parameters, used to quickly generate candidate tokens as an initial reference for reasoning in the target model.

[0047] Medusa Model, a multi-head prediction model, is a speculative sampling technique that generates multiple candidate tokens through multiple parallel prediction heads, making it suitable for efficient multi-path inference.

[0048] The target model refers to the large-scale language model that actually performs the final inference task, used to verify and generate the final high-quality output.

[0049] EOS (End of Sequence) is a sequence end marker that indicates that the generated text sequence has reached its end point. The model stops generating after encountering an EOS token.

[0050] BFS (Breadth-First Search) is an algorithm that traverses nodes in a tree structure from top to bottom in a hierarchical manner.

[0051] DFS (Depth-First Search) is an algorithm that delves deep into nodes along sub-paths in a tree structure. It is often used to optimize path selection during the reasoning process.

[0052] Speculative Acceptance Rate is the proportion of candidate tokens accepted by the target model, reflecting the efficiency and accuracy of speculative sampling.

[0053] Figure 1 This is a flowchart illustrating a token filtering method based on the Medusa model according to an exemplary embodiment. The token filtering method 10 based on the Medusa model includes at least steps S102 to S110.

[0054] like Figure 1 As shown, in S102, the target model is used to perform inference calculations on the input sequence to generate multiple candidate tokens. The input sequence can be input into the target model; the target model calculates multiple candidate tokens based on inference from the input sequence; and the key-value cache information corresponding to the multiple candidate tokens is cached.

[0055] During inference, the target model generates intermediate computation results, which are stored as a key-value cache. The key-value cache contains crucial data for the attention mechanism, used to reduce redundant computations in subsequent inference and improve inference efficiency. These key-value caches are partially retained or discarded based on the selection of candidate tokens.

[0056] The calculation result can be sampled, and a portion of the token in the calculation result can be accepted to generate a target token; the target token can be used to perform the next inference calculation, and the target token that was not cached can be recalculated in the next inference calculation.

[0057] In S104, the plurality of candidate tokens are sampled, and a subset of candidate tokens are extracted. For example, k candidate tokens from the plurality of candidate tokens can be accepted, and the kv cache information of p tokens can be cached; the kv cache information of the kp tokens that are not cached from the subset of candidate tokens can be discarded.

[0058] In one specific embodiment, k candidate tokens in a tree structure can be accepted, and the kvcache information of p tokens can be cached; the kvcache information of the kp tokens that are not cached among the candidate tokens can be discarded. More specifically, k candidate tokens in the tree structure can be accepted; and the kvcache information of the first p consecutive tokens among the k candidate tokens can be cached.

[0059] Add the kp uncached tokens and a new token to the input sequence to generate the current sequence, and recalculate the kp uncached tokens for the next inference calculation.

[0060] For the kp tokens that are not cached, their key-value cache information is recalculated and added to the inference calculation. This dynamic recalculation mechanism ensures the correctness of the generated sequence while effectively utilizing cache resources and avoiding the storage of redundant information.

[0061] In S106, the candidate tokens are input into the medusa model for speculative inference, generating multiple speculative tokens in a tree structure. Each prediction head of the medusa model generates a set of speculative tokens, which are organized hierarchically into a tree structure.

[0062] More specifically, the sampled candidate tokens are input into the Medusa model for speculative inference. The multi-head structure of the Medusa model allows for the parallel generation of multiple candidate tokens, which are organized into a tree structure. Each prediction head generates a new set of tokens based on the tokens of the previous layer. The speculatively generated tokens are organized in a tree structure, with each layer's tokens corresponding to the parent node of the layer above it.

[0063] In one embodiment, the method further includes arranging the plurality of speculative tokens in a Depth-First Search (DFS) order. Normally, the sequence is sorted using a Breadth-First Search (BFS) method; however, in this application, the BFS arrangement of draft tokens can be replaced with a DFS arrangement.

[0064] DFS sorting can organize tokens according to the path order from the root node to the leaf node, so that subsequent reasoning can make full and continuous use of the filtering results.

[0065] The DFS sorting process is as follows:

[0066] Starting from the root node of the Draft Tree, each child node is visited sequentially and recursively. After completing the traversal of a branch, other branches are visited in return, until the entire tree structure has been traversed. The sorted token sequence is then passed as input to the target model to ensure the orderliness and coherence of the reasoning process.

[0067] In S108, the multiple speculative tokens are filtered using a pruning method based on a tree structure to generate a target token.

[0068] In one embodiment, for example, the multiple speculative tokens can be filtered using a prediction score pruning method based on a tree structure to generate a target token; more specifically, for example, a prediction score can be obtained for each speculative token, the prediction score being used to represent the probability that the speculative token is accepted by the target model; the multiple candidate tokens can be filtered using the prediction score to generate a target token.

[0069] In a specific application, during the Tree Decoding process of the draft model, the model originally stored two core member variables: Draft Parents: used to record the parent node relationship of each node in the tree structure; DraftLevels: used to record the level of each node in the tree structure.

[0070] In this application, a new member variable, Draft Credit, can be added. This variable is used to store the predicted score of each Draft Token, which represents the estimated probability that the Token will be accepted in the TargetModel.

[0071] Since the probability of the target model cannot be obtained before sampling, it is impossible to know the true probability of the token being accepted in advance during the calculation process. In this application, the probability calculated by speculation (e.g., draft model) is used as an approximation of the acceptance probability. The approximation error depends on the matching degree.

[0072] In one specific embodiment, the corresponding prediction score is generated based on the estimated probability and speculative probability of the candidate token obtained during speculative inference, including:

[0073]

[0074] in, To predict the score, For the current token, This represents the estimated probability of the parent node of the current token. This represents the speculative probability of the current token.

[0075] Through recursive calculation, the predicted score of a node can be represented as the product of the speculative probabilities of all its parent nodes along its path. This calculation employs the concept of dynamic programming, caching the credit values ​​of parent nodes to avoid redundant calculations, thereby improving computational efficiency.

[0076] Since Draft Parents already record the connection information between nodes, Draft Credit only needs to be mapped one-to-one with DraftTokens. This design avoids repeatedly storing redundant data, making the tree-structured inference more efficient. This improvement lays the foundation for subsequent pruning optimization and sampling strategies. By predicting scores, high-quality DraftTokens can be selected more accurately, improving generation quality and reducing inference computation overhead.

[0077] More specifically, the multiple candidate tokens can be arranged in descending order of the predicted scores, and candidate tokens with a target ranking can be extracted according to the arrangement order; the multiple candidate tokens can be filtered based on the predicted scores of the candidate tokens with the target ranking; and candidate tokens located in leaf nodes can be removed.

[0078] In one embodiment, for example, the multiple speculative tokens can be filtered using a leaf node pruning method based on a tree structure to generate a target token; more specifically, for example, speculative tokens located in leaf nodes can be removed based on a tree structure to generate a target token.

[0079] In this application, a leaf node can refer to a node that has reached the user-specified maximum sequence length or the machine hardware limit. In this application, a leaf node can also refer to a node containing an EOS (End of Sequence) marker, which indicates the end of the sequence and prevents further expansion.

[0080] Tokens from leaf nodes cannot generate subsequent tokens, thus increasing inference costs without improving generation efficiency. The solution is to iterate through all candidate tokens and remove all leaf node tokens. This removal operation reduces the computational burden of inference on the target model while avoiding the generation of redundant and invalid sequences.

[0081] If a leaf node is used as a draft token, then it has two possible outcomes:

[0082] 1) It is accepted, and then the token is returned directly to the user without generating any more tokens.

[0083] 2) Not accepted. Unable to generate subsequent tokens.

[0084] Using a single token as the draft token aims to generate more tokens during a single decoding round. Leaf nodes do not generate subsequent tokens, so their presence or absence does not affect the maximum acceptable length. In fact, leaf nodes increase the computational cost of inference in the target model. Therefore, all leaf nodes can be removed.

[0085] For example:

[0086] If "you ate" is followed by only one draft token ".", then there are three possible outcomes: accept or reject ".". Possible results include:

[0087] 1) Do not accept ".", generate 1 token: ["You ate", "apple"]

[0088] 2) Accept "." and generate 2 tokens: ["You ate", ".", "EOS"]

[0089] 3) Accept "." and generate 2 tokens: ["You ate", ".", "I"]

[0090] If "You ate" is followed by two draft tokens "." and "EOS", then the result can still be one of three possibilities:

[0091] 1) Do not accept ".", generate 1 token: ["You ate", "apple"]

[0092] 2) Accept EOS and generate 2 tokens: ["You ate", ".", "EOS"]

[0093] 3) EOS is not accepted; two tokens are generated: ["You ate", ".", "I"]

[0094] In one embodiment, for example, the multiple speculative tokens can be filtered using a repeated pruning method based on a tree structure to generate the target token. More specifically, for example, duplicate speculative tokens can be removed based on the tree structure to generate the target token.

[0095] More specifically, since there is no direct connection between the heads of Medusa, it often happens that different heads infer the same token. Therefore, this application introduces duplicate word pruning. If a Medusa draft token has the current token in its ancestors (i.e., the parent node and the upstream node path of the parent node), the draft token is deleted.

[0096] In S110, subsequent inference calculations are performed based on the target token to generate the output result. The pruned target token is considered the optimal token set and is passed to the target model for subsequent inference calculations. The high quality of the target token ensures the accuracy of inference and the efficiency of generation.

[0097] According to the token filtering method based on the Medusa model in this application, multiple candidate tokens are generated by performing inference calculations on the input sequence using a target model; the multiple candidate tokens are sampled to extract a portion of the candidate tokens; the portion of the candidate tokens are input into the Medusa model for speculative inference to generate multiple speculative tokens in a tree structure; the multiple speculative tokens are filtered using a pruning method based on the tree structure to generate a target token; and subsequent inference calculations are performed based on the target token to generate the output result. This method can optimize token sampling performance while ensuring the quality of the inference result, effectively improve the efficiency of speculative sampling during large model inference optimization, and has a wide range of applications.

[0098] It should be clearly understood that this application describes how specific examples are formed and used, but the principles of this application are not limited to any details of these examples. Rather, based on the teachings of the disclosure of this application, these principles can be applied to many other embodiments.

[0099] Figure 2 This is a schematic diagram illustrating a token filtering method based on the Medusa model according to another exemplary embodiment. Figure 2 The process shown is... Figure 1 The process is described in detail below.

[0100] First, the input sequence is fed into the target model. During the inference process of the target model, the key-value value of each token needs to be calculated and cached.

[0101] Next, the token sequence is sampled. During the sampling process, some tokens are accepted and some tokens are cached. The first P tokens of these tokens are consecutive to the original sequence.

[0102] After sampling, the target model generates multiple candidate tokens. These candidate tokens are then input into the medusa model for further speculation, generating multiple tokens in total.

[0103] The token sorting can also be optimized. More specifically, the DFS method can be used to optimize the token sorting in order to improve the speed of subsequent speculative sampling.

[0104] Next, multiple tokens are filtered, specifically by one or more methods such as predictive scoring pruning, leaf node pruning, and repetition pruning, to generate the target token.

[0105] Finally, cache the tokens not cached in the previous step and the newly received p tokens, add the remaining tokens and a new token to the sequence, and add the newly obtained target token to the sequence to continue the inference.

[0106] In the scenario of tree-shaped speculation, the maximum degree of each layer takes draft tokens according to 3, 2, 1. The general approach is a full connection, that is, all draft tokens of each layer are connected to all tokens of the lower layer to form the tokens of the lower layer. This can also adopt the Credit Tree and the pruning method based on Credit.

[0107] For example:

[0108] "You ate" uses medusa speculation to generate draft tokens.

[0109] The first layer is "he" (0.3), "I" (0.2), "apple" (0.2).

[0110] The second layer is "of" (0.6), "already" (0.4).

[0111] The third layer is "apple" (0.3).

[0112] Then the draft tree and the result selected according to credit are as Figure 3 shown.

[0113] ["You ate", "he" (0.3), "I" (0.2), "apple" (0.2), "of" (0.18), "already" (0.12), "of" (0.12), "of" (0.12)]

[0114] It can be seen that the sequence is not necessarily smooth under the Cartesian product of multiple heads in medusa. For example, "already" in the second layer obviously cannot be connected to "he" and "I" in the first layer. Another example is that "apple" in the third layer obviously cannot be connected to "already" in the second layer.

[0115] "You ate me already apple" is unreasonable. This is the shortcoming of medusa and one of the reasons for the low acceptance rate of multi-head medusa.

[0116] Since there is no direct connection between the heads of Medusa, the situation where different heads infer the same token often occurs.

[0117] Just like the word "apple" appeared in the first layer and the third layer in the previous example.

[0118] Repeated tokens that appear later are very likely not to be accepted. ("You ate an apple" is possible, but "You ate an apple apple" is impossible.) We have made a statistics, and the probability of consecutive identical tokens being accepted is extremely low:

[0119] The following table is a table showing the situation of medusa generating 2 draft tokens:

[0120] Number of times No duplicate draft tokens Duplicate draft token The second draft token was accepted. 712 7 The second draft token is not accepted. 1407 228

[0121] Using the chi-square distribution to verify the correlation between the two random variables of whether there are repeated draft tokens and whether the second draft token is accepted, there is a 99.9% certainty that these two variables are correlated.

[0122] When there are no repeated draft tokens, the probability that the second draft token is accepted is 33.6%. However, when there are repeated draft tokens, the probability that the second draft token is accepted drops to 3.0%. This is sufficient to show that repeated draft tokens will significantly reduce the acceptance rate.

[0123] Pruning of repeated words can be carried out. If a draft token of medusa has the current token in its ancestors (that is, the path of the parent node and the upstream nodes of the parent node), then this draft token is not added.

[0124] In the above example, the repeated "apple" will be crossed out, resulting in Figure 4 the result.

[0125] In another specific embodiment, medusa often has a similar distribution for certain tokens in different heads.

[0126] Using medusa to speculatively generate draft tokens for "You ate", as Figure 5 shown:

[0127] The first layer is "le" (0.3), "apple" (0.2), "de" (0.2),

[0128] The second layer is "le" (0.6), "de" (0.4),

[0129] The third layer is "le" (0.3)

[0130] The result obtained after repeated pruning is as Figure 6As shown, under this repeated pruning, the acceptance rate of medusa was improved to some extent (around 6.5%, with the acceptance rate increasing from 11.87% to 12.65% under 7 draft tokens). Due to the reduction of many useless token combinations, throughput was improved (around 4%).

[0131] Those skilled in the art will understand that all or part of the steps of the above embodiments are implemented as a computer program executed by a CPU. When the computer program is executed by the CPU, it performs the functions defined by the method provided in this application. The program can be stored in a computer-readable storage medium, such as a read-only memory, a magnetic disk, or an optical disk.

[0132] Furthermore, it should be noted that the above figures are merely illustrative representations of the processes included in the method according to exemplary embodiments of this application, and are not intended to be limiting. It is readily understood that the processes shown in the above figures do not indicate or limit the temporal order of these processes. Additionally, it is readily understood that these processes may be executed synchronously or asynchronously, for example, in multiple modules.

[0133] The following are embodiments of the apparatus described in this application, which can be used to execute the embodiments of the method described in this application. For details not disclosed in the apparatus embodiments of this application, please refer to the embodiments of the method described in this application.

[0134] Figure 7 This is a block diagram illustrating a token filtering device based on the Medusa model, according to another exemplary embodiment. Figure 7 As shown, the token screening device 70 based on the Medusa model includes: a target module 702, a sampling module 704, a speculation module 706, a pruning module 708, and an inference module 710.

[0135] The target module 702 is used to perform inference calculations on the input sequence using the target model to generate multiple candidate tokens; the target module 702 is also used to input the input sequence into the target model; the target model calculates multiple candidate tokens based on inference from the input sequence; and caches the kvcache information corresponding to the multiple candidate tokens.

[0136] The sampling module 704 is used to sample the plurality of candidate tokens and extract a portion of the candidate tokens; the sampling module 704 is also used to accept k candidate tokens from the plurality of candidate tokens and cache the kv cache information of p tokens; and discard the kv cache information of the kp tokens that are not cached from the portion of the candidate tokens.

[0137] The speculation module 706 is used to input the candidate tokens into the medusa model for speculative inference, generating multiple speculative tokens in a tree structure; each prediction head of the medusa model generates a set of speculative tokens, which are organized hierarchically into a tree structure. The speculation module 706 is also used to arrange the multiple speculative tokens in the order of DFS.

[0138] The pruning module 708 is used to filter the multiple speculative tokens based on a tree structure using a pruning method to generate a target token; the pruning module 708 is also used to filter the multiple speculative tokens based on a tree structure using a predictive scoring pruning method to generate a target token; the pruning module 708 is also used to filter the multiple speculative tokens based on a tree structure using a leaf node pruning method to generate a target token; the pruning module 708 is also used to filter the multiple speculative tokens based on a tree structure using a repeating pruning method to generate a target token.

[0139] The reasoning module 710 is used to perform subsequent reasoning calculations based on the target token to generate output results.

[0140] The token filtering device based on the Medusa model of this application generates multiple candidate tokens by performing inference calculations on the input sequence using a target model; samples the multiple candidate tokens to extract a portion of them; inputs the portion of candidate tokens into the Medusa model for speculative inference to generate multiple speculative tokens in a tree structure; filters the multiple speculative tokens based on the tree structure using a pruning method to generate a target token; and performs subsequent inference calculations based on the target token to generate output results. This method can optimize token sampling performance while ensuring the quality of inference results, effectively improving the efficiency of speculative sampling during large model inference optimization, and has a wide range of applications.

[0141] Figure 8 This is a block diagram illustrating an electronic device according to an exemplary embodiment.

[0142] The following reference Figure 8 To describe an electronic device 800 according to this embodiment of the present application. Figure 8 The electronic device 800 shown is merely an example and should not impose any limitations on the functionality and scope of use of the embodiments of this application.

[0143] like Figure 8As shown, the electronic device 800 is presented in the form of a general-purpose computing device. The components of the electronic device 800 may include, but are not limited to: at least one processing unit 810, at least one storage unit 820, a bus 830 connecting different system components (including storage unit 820 and processing unit 810), a display unit 840, etc.

[0144] The storage unit stores program code that can be executed by the processing unit 810, causing the processing unit 810 to perform the steps described in this specification according to various exemplary embodiments of this application. For example, the processing unit 810 can perform actions such as... Figure 1 , Figure 2 The steps are shown in the figure.

[0145] The storage unit 820 may include a readable medium in the form of a volatile storage unit, such as a random access memory unit (RAM) 8201 and / or a cache storage unit 8202, and may further include a read-only memory unit (ROM) 8203.

[0146] The storage unit 820 may also include a program / utility 8204 having a set (at least one) program module 8205, such program module 8205 including but not limited to: an operating system, one or more application programs, other program modules and program data, each or some combination of these examples may include an implementation of a network environment.

[0147] Bus 830 can represent one or more of several types of bus structures, including a memory cell bus or memory cell controller, a peripheral bus, a graphics acceleration port, a processing unit, or a local bus using any of the various bus structures.

[0148] Electronic device 800 can also communicate with one or more external devices 800' (e.g., keyboard, pointing device, Bluetooth device, etc.), enabling users to communicate with devices that interact with electronic device 800, and / or any device (e.g., router, modem, etc.) that allows electronic device 800 to communicate with one or more other computing devices. This communication can be performed via input / output (I / O) interface 850. Furthermore, electronic device 800 can also communicate with one or more networks (e.g., local area network (LAN), wide area network (WAN), and / or public networks, such as the Internet) via network adapter 860. Network adapter 860 can communicate with other modules of electronic device 800 via bus 830. It should be understood that, although not shown in the figures, other hardware and / or software modules can be used in conjunction with electronic device 800, including but not limited to: microcode, device drivers, redundant processing units, external disk drive arrays, RAID systems, tape drives, and data backup storage systems.

[0149] From the above description of the embodiments, those skilled in the art will readily understand that the exemplary embodiments described herein can be implemented by software, or by combining software with necessary hardware. Therefore, as... Figure 9 As shown, the technical solution according to the embodiments of this application can be embodied in the form of a software product. The software product can be stored in a non-volatile storage medium (such as a CD-ROM, USB flash drive, mobile hard drive, etc.) or on a network, and includes several instructions to cause a computing device (such as a personal computer, server, or network device, etc.) to execute the above-described method according to the embodiments of this application.

[0150] The software product may employ any combination of one or more readable media. A readable medium may be a readable signal medium or a readable storage medium. A readable storage medium may be, for example, but not limited to, an electrical, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any combination thereof. More specific examples of readable storage media (a non-exhaustive list) include: an electrical connection having one or more wires, a portable disk, a hard disk, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fiber, portable compact disk read-only memory (CD-ROM), optical storage devices, magnetic storage devices, or any suitable combination thereof.

[0151] The computer-readable storage medium may include data signals propagated in baseband or as part of a carrier wave, carrying readable program code. Such propagated data signals may take various 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, capable of transmitting, propagating, or transmitting programs for use by or in connection with an instruction execution system, apparatus, or device. The program code contained on the readable storage medium may be transmitted using any suitable medium, including but not limited to wireless, wired, optical fiber, RF, etc., or any suitable combination thereof.

[0152] Program code for performing the operations of this application can be written in any combination of one or more programming languages, including object-oriented programming languages ​​such as Java and C++, and conventional procedural programming languages ​​such as C or similar languages. The program code can execute entirely on the user's computing device, partially on the user's computing device, as a standalone software package, partially on the user's computing device and partially on a remote computing device, or entirely on a remote computing device or server. In cases involving remote computing devices, the remote computing device can be connected to the user's computing device via any type of network, including a local area network (LAN) or a wide area network (WAN), or it can be connected to an external computing device (e.g., via the Internet using an Internet service provider).

[0153] The aforementioned computer-readable medium carries one or more programs. When these programs are executed by a device, the computer-readable medium performs the following functions: inference calculations on an input sequence using a target model to generate multiple candidate tokens; sampling the multiple candidate tokens to extract a subset of candidate tokens; inputting the subset of candidate tokens into a medusa model for speculative inference to generate multiple speculative tokens in a tree structure; filtering the multiple speculative tokens using a pruning method based on the tree structure to generate a target token; and performing subsequent inference calculations based on the target token to generate an output result.

[0154] Those skilled in the art will understand that the above modules can be distributed in the device as described in the embodiments, or they can be modified accordingly and placed in one or more devices that are unique to this embodiment. The modules in the above embodiments can be combined into one module, or they can be further divided into multiple sub-modules.

[0155] Through the description of the above embodiments, those skilled in the art will readily understand that the exemplary embodiments described herein can be implemented by software or by combining software with necessary hardware. Therefore, the technical solutions according to the embodiments of this application can be embodied in the form of a software product, which can be stored in a non-volatile storage medium (such as a CD-ROM, USB flash drive, external hard drive, etc.) or on a network, including several instructions to cause a computing device (such as a personal computer, server, mobile terminal, or network device, etc.) to execute the methods according to the embodiments of this application.

[0156] Exemplary embodiments of this application have been specifically shown and described above. It should be understood that this application is not limited to the detailed structures, arrangements, or implementation methods described herein; rather, this application is intended to cover various modifications and equivalent arrangements contained within the spirit and scope of the appended claims.

Claims

1. A medusa model token screening method based on, characterized in that, The method comprises the following steps: performing inference calculation on an input sequence by using a target model to generate a plurality of candidate tokens, wherein the tokens are words, subwords or characters; sampling the plurality of candidate tokens to extract part of the candidate tokens; inputting the part of the candidate tokens into a medusa model to perform speculative inference to generate a plurality of speculative tokens in a tree structure; screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a pruning method; performing subsequent inference calculation based on the target tokens to generate an output result.

2. The method of claim 1, wherein, The method comprises the following steps: performing inference calculation on an input sequence by using a target model to generate a plurality of candidate tokens, comprising: inputting the input sequence into the target model; the target model calculates a plurality of candidate tokens based on the input sequence; 3. The method of claim 1, wherein, caching kv cache information corresponding to the plurality of candidate tokens. The method comprises the following steps: sampling the plurality of candidate tokens to extract part of the candidate tokens, comprising:

4. The method of claim 1, wherein, accepting k candidate tokens in the plurality of candidate tokens, and caching kv cache information of p tokens; discarding the kv cache information of k-p tokens in the part of the candidate tokens that are not cached.

5. The method of claim 4, wherein, The method comprises the following steps: inputting the part of the candidate tokens into a medusa model to perform speculative inference to generate a plurality of speculative tokens in a tree structure, comprising:

6. The method of claim 1, wherein, Each prediction head of the medusa model generates a group of speculative tokens, and organizes the group of speculative tokens into a tree structure in a hierarchical manner. The method comprises the following steps: inputting the part of the candidate tokens into a medusa model to perform speculative inference to generate a plurality of speculative tokens in a tree structure, further comprising: arranging the plurality of speculative tokens in a DFS order.

7. The method of claim 6, wherein, The method comprises the following steps: screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a pruning method, comprising: screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a prediction score pruning method; and / or 8. The method of claim 6, wherein, screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a leaf node pruning method; and / or screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a duplicate pruning method.

9. The method of claim 6, wherein, The method comprises the following steps: screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a prediction score pruning method, comprising: obtaining a prediction score of each speculative token, wherein the prediction score is used to represent a probability that the speculative token is accepted by the target model; screening the plurality of candidate tokens by using the prediction score to generate target tokens. The method comprises the following steps: screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a leaf node pruning method, comprising: eliminating speculative tokens located at leaf nodes based on the tree structure to generate target tokens. The method comprises the following steps: screening the plurality of speculative tokens based on the tree structure to generate target tokens by using a duplicate pruning method, comprising: The repeated speculative tokens are pruned based on a tree structure to generate target tokens.

10. A medusa model token based screening device, characterized in that, The method comprises the following steps: a target module is configured to perform inference calculation on an input sequence by using a target model to generate a plurality of candidate tokens, wherein the tokens are words, subwords or characters; a sampling module is configured to sample the plurality of candidate tokens to extract part of the candidate tokens; a speculative module is configured to input the part of the candidate tokens into a medusa model to perform speculative inference to generate a plurality of speculative tokens in a tree structure; a pruning module is configured to filter the plurality of speculative tokens based on the tree structure by using a pruning method to generate target tokens; an inference module is configured to perform subsequent inference calculation based on the target tokens to generate an output result.

Images