A draft model inference optimization method and system for long context speculation decoding

Abstract

This invention discloses a method and system for optimizing draft model inference for long-context speculative decoding, belonging to the field of electronic information technology. The method includes: dividing the historical context into multiple context units according to a predetermined strategy; evaluating the relevance of the context units to select the most relevant subset of contexts for the current generation step; adapting rotational position encoding to the tree-like draft tokens generated by the draft model, and simultaneously maintaining the position encoding state during sliding window pruning and rollback; jointly constraining attention reachability relationships, outputting results and scores, and generating a candidate token sequence; submitting the candidate token sequence to the target model for verification, and performing a consistent rollback of the generation state, context cache state, and position encoding state of the failed verification token sequences. This invention optimizes the inference process of the draft model by introducing a rollbackable sliding window mechanism and a relevance-based context retrieval mechanism.

Description

Technical Field

[0001] This invention belongs to the field of electronic information technology, and specifically relates to a draft model inference optimization method and system for long context speculation decoding. Background Technology

[0002] With the rapid development of artificial intelligence technology, Large Language Models (LLMs) have been widely applied in various fields such as document generation, multi-turn dialogue, long text analysis, intelligent summarization, data analysis, and knowledge question answering. The length of the context that needs to be processed during the reasoning process is constantly increasing, gradually reaching the scale of millions of tokens. Long context windows have become the core direction for breaking through the application bottleneck of large language models and improving model performance. However, in long context scenarios, model reasoning usually faces the problems of excessive memory consumption and soaring computational complexity, which restricts the application of large language models in complex scenarios.

[0003] To alleviate the problem of low efficiency in autoregressive decoding reasoning of large language models, the industry has proposed speculative decoding technology. This technology adopts a dual-model collaborative reasoning mode of "draft model + main model". The draft model, which has a simple structure and fast reasoning speed, quickly generates candidate token sequences, and then the high-precision main model verifies and corrects them in one go, thereby reducing the number of reasoning steps of the main model and improving the overall reasoning efficiency.

[0004] Existing speculative decoding methods mostly adopt inference schemes for short-context scenarios, without specifically designing for long-context scenarios and the inference requirements of draft models. Although some techniques attempt to optimize through positional encoding extensions and retrieval-enhanced generation, these methods generally have significant drawbacks: they either require additional model training or fine-tuning, increasing training and deployment costs; or they rely on retrieval modules or complex model structures, leading to increased memory consumption and limited applicability, making it difficult to strike a balance between memory usage, inference efficiency, and token acceptance rate. Furthermore, issues such as inconsistent training distribution between long-context and draft models, and excessive memory overhead caused by large-scale loading of KV caches, further restrict the performance of speculative decoding techniques in long-context scenarios.

[0005] In summary, current speculative decoding technology in long-context scenarios urgently needs a plug-and-play, draft model inference optimization solution that requires no additional training or fine-tuning, and can balance inference efficiency, memory usage, and token acceptance rate. This solution should address the problems of high cost, poor adaptability, and insufficient performance of existing technologies without changing the original model architecture and hardware configuration, and fully leverage the speed-up advantages of speculative decoding. Summary of the Invention

[0006] The purpose of this invention is to provide a method and system for optimizing the inference process of a draft model for long-context speculative decoding. This invention addresses the problems faced by draft models in long-context scenarios during speculative decoding, such as context state inflation, high attention computational complexity, and frequent rollbacks due to decreased prediction consistency. Without requiring additional training of the draft model, this invention optimizes the inference process of the draft model by introducing a rollbackable sliding window mechanism and a relevance-based context retrieval mechanism.

[0007] To achieve the above objectives, the present invention adopts the following technical solution: An optimization method for draft model inference for long-context speculative decoding includes: 1) Divide the historical context into multiple context units according to a predetermined strategy, and maintain the representation information of the context units; 2) Based on the query representation corresponding to the current generation position, the relevance of the context units is evaluated, thereby selecting the most relevant subset of context to the current generation step; 3) Adapt rotation position encoding to the tree-like draft tokens generated by the draft model, and maintain the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness; 4) With the support of the tree-like draft tokens after position encoding adaptation and related context subsets, the core computation logic of Attention is encapsulated. Combined with the sliding window mask and the tree structure mask used to describe the draft generation tree structure, the attention reachability relationship is jointly constrained, the results and scores are output, and a candidate token sequence is generated. 5) Submit the candidate token sequence to the target model for verification, and perform a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

[0008] A further improvement of this invention is that step 1) divides the historical context into multiple context units according to a predetermined strategy, and maintains the representation information of the context units, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

[0009] A further improvement of this invention is that step 2) evaluates the relevance of the context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant subset of context units to the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

[0010] A further improvement of this invention is that step 3) adapts the rotation position encoding to the tree-like draft token generated by the draft model, and synchronously maintains the position encoding state during the sliding window clipping and rollback process to ensure Attention position awareness, including: Receive a custom position index corresponding to the tree draft structure, apply rotation position encoding to the query vector and key vector based on the custom position index, and match the actual position of each tree draft token in the tree structure; During the sliding window pruning process, the position index is pruned and remapped synchronously, retaining only the nearest position index within the local sliding window that is relevant to the current inference, and deleting redundant position indexes outside the window; The location encoding status is updated in real time during the inference process, and the corresponding custom location index is automatically appended to the historical location index after each new candidate token is generated; When a rollback operation is triggered, the historical position index records are restored to the state before the draft model generates the tree draft.

[0011] A further improvement of this invention is that step 4) encapsulates the core computational logic of Attention and supports combining a sliding window mask and a tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, outputting results and scores for use in driving the draft model to generate candidate token sequences based on context subsets, including: Receive the sliding window mask and the tree structure mask generated by the draft model, and jointly restrict the attention reachability relationship; Based on the mask, the attention calculation within the local sliding window is performed first, and then the supplementary attention calculation of the global relevant context unit is performed to achieve a complementary calculation mode of "sliding window + context retrieval". The calculation results and attention scores are based on the Attention calculation output. The calculation results are used to drive the draft model to generate candidate token sequences, and the attention scores serve as the basis for context relevance evaluation and state maintenance, forming a closed-loop feedback mechanism for inference.

[0012] A further improvement of this invention is that, in step 5), after the target model verifies the candidate token sequence, a consistent rollback is performed on the generation state, context cache state, and position encoding state of the failed token sequence, including: Receive the number of tokens to be rolled back, determine the rollback target state as the state before the most recent candidate token generation, and clear the temporary Attention score and relevance score temporary data generated in the current inference step; Based on the number of rollback tokens, clear the corresponding global cache state in the context block and cache management module; trim the local sliding window cache and delete the local KV Cache corresponding to the verification failure candidate tokens; Based on the number of rollback tokens, the rollback of the position-encoded state is performed through the position encoding module.

[0013] An optimization system for draft model inference for long-context speculative decoding includes: The historical context partitioning unit divides the historical context into multiple context units according to a predetermined strategy and maintains the representation information of the context units. The context unit relevance evaluation unit evaluates the relevance of context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant context subset to the current generation step. The rotation position encoding adaptation unit adapts the rotation position encoding to the tree-shaped draft tokens generated by the draft model, and maintains the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness. The candidate token sequence generation unit, supported by the position-encoded and adapted tree-structured draft tokens and related context subsets, encapsulates the core computational logic of Attention, and combines the sliding window mask and the tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, output the result and score, and generate the candidate token sequence. The verification unit submits the candidate token sequence to the target model for verification and performs a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

[0014] A further improvement of this invention is that the historical context partitioning unit divides the historical context into multiple context units according to a predetermined strategy, and maintains the representation information of the context units, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

[0015] A further improvement of this invention is that the context unit relevance evaluation unit evaluates the relevance of the context unit based on the query representation corresponding to the current generation position, thereby selecting the most relevant subset of context units to the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

[0016] A computer-readable storage medium storing a computer program that, when executed by a processor, implements the steps of the draft model inference optimization method for long context speculative decoding.

[0017] Compared with the prior art, the present invention has at least the following beneficial technical effects: Compared to existing technologies, the technical solution of this invention solves the problems of low inference efficiency of draft models and poor quality of generated tokens leading to reduced acceptance rates in long context speculation decoding through multi-dimensional collaborative optimization, and has the following significant advantages: (1) No need to retrain the draft model, resulting in low deployment costs. This invention improves performance through algorithm and system optimization during the inference phase without changing the network structure and parameters of the draft model. It avoids the high computational costs and data dependencies caused by retraining or fine-tuning the draft model, significantly reduces the cost of algorithm deployment and maintenance, and facilitates rapid implementation in existing large language model inference decoding systems.

[0018] (2) The sliding window limits the context size, which significantly reduces the computational complexity and latency of attention. This invention introduces a sliding window attention mechanism on the draft model side, which limits the local context range by the parameter n_local, thereby reducing the computational complexity of attention. Combined with the computation mode of "local sliding window attention + global supplementary context unit", it effectively avoids redundant computation of global attention while ensuring the semantic integrity of long context.

[0019] (3) A global block selection mechanism based on context retrieval improves the acceptance rate of candidate tokens. This invention divides a long context into multiple units and constructs a compressed representation for each unit. It then uses Attention relevance calculation to retrieve the global context, selecting only context units highly relevant to the current query for inference. This mechanism effectively mitigates the prediction bias caused by the inconsistency between the long context and the training distribution of the draft model, increasing the probability that the token generated by the draft model is accepted by the target model, thereby maintaining the overall acceleration of speculative decoding.

[0020] (4) The reasoning process is collaboratively optimized, with strong versatility and scalability. This invention integrates and coordinates the inference process, including KV Cache initialization, context segmentation, position encoding adaptation, attention calculation, and result generation, reducing data redundancy and inter-module waiting time. Meanwhile, key parameters (such as cell size, number of TopK, and sliding window length) can be flexibly configured according to actual application scenarios without modifying the core structure of the draft model. It can adapt to long context inference requirements of different lengths, is compatible with multiple mainstream draft model architectures, and has good versatility and scalability.

[0021] (5) The project is simple to implement and the system deployment cost is low. The technical solution of this invention is based on mature deep learning inference technology. It only optimizes the KVCache management and attention calculation process during the inference stage. It does not require reconstructing the encoder or decoder structure of the draft model, nor does it require the introduction of additional hardware resources. It has low engineering implementation difficulty and is easy to integrate into existing speculative decoding systems.

[0022] (6) Strong robustness and high system stability This invention introduces a multi-level state verification and rollback mechanism for draft model rollback during the inference process, and maintains consistency of KV Cache, position encoding state and inference length, effectively avoiding abnormal situations such as memory out-of-bounds and data corruption; at the same time, through pre-allocation of cache space, asynchronous loading and non-blocking copying, the stability and reliability of the system in ultra-long context continuous inference scenarios are further improved.

[0023] In summary, the technical solution of this invention, through multi-dimensional collaborative optimization, effectively solves the core defects of the prior art, significantly improves the inference efficiency of the draft model in long context speculation decoding, reduces memory overhead, can give full play to the speed-up advantage of speculation decoding, and promotes the large-scale application of long context large language models. Attached Figure Description

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

[0025] Figure 1 This is a diagram illustrating the overall architecture of a draft model inference optimization method for long context-based speculative decoding according to the present invention.

[0026] Figure 2 This is a block diagram of a draft model inference optimization system for long context speculation decoding according to the present invention. Detailed Implementation

[0027] In the following description, only certain exemplary embodiments are briefly described. As those skilled in the art will recognize, the described embodiments can be modified in various ways without departing from the spirit or scope of the invention. Therefore, the drawings and description are considered to be exemplary in nature and not restrictive.

[0028] In the description of this invention, it should be understood that, when used in this specification and the appended claims, the terms "comprising" and "including" indicate the presence of the described features, integrals, steps, operations, elements and / or components, but do not exclude the presence or addition of one or more other features, integrals, steps, operations, elements, components and / or collections thereof.

[0029] It should also be understood that the terminology used in this specification is for the purpose of describing particular embodiments only and is not intended to limit the invention. As used in this specification and the appended claims, the singular forms “a,” “an,” and “the” are intended to include the plural forms unless the context clearly indicates otherwise.

[0030] It should also be further understood that the term "and / or" as used in this specification and the appended claims refers to any combination of one or more of the associated listed items and all possible combinations, and includes such combinations.

[0031] The accompanying drawings illustrate various structural schematic diagrams according to embodiments disclosed in this invention. These drawings are not to scale, and some details have been enlarged for clarity, and some details may have been omitted. The shapes of the various regions and layers shown in the drawings, as well as their relative sizes and positional relationships, are merely exemplary and may deviate from reality due to manufacturing tolerances or technical limitations. Furthermore, those skilled in the art can design regions / layers with different shapes, sizes, and relative positions as needed.

[0032] The embodiments of the present invention will now be described in detail with reference to the accompanying drawings.

[0033] Example 1 This invention provides a draft model inference optimization method for long context-based inference decoding, comprising: 1) Divide the historical context into multiple context units according to a predetermined strategy, and maintain the representation information of the context units; 2) Based on the query representation corresponding to the current generation position, the relevance of the context units is evaluated, thereby selecting the most relevant subset of context to the current generation step; 3) Adapt rotation position encoding to the tree-like draft tokens generated by the draft model, and maintain the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness; 4) With the support of the tree-like draft tokens after position encoding adaptation and related context subsets, the core computation logic of Attention is encapsulated. Combined with the sliding window mask and the tree structure mask used to describe the draft generation tree structure, the attention reachability relationship is jointly constrained, the results and scores are output, and a candidate token sequence is generated. 5) Submit the candidate token sequence to the target model for verification, and perform a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

[0034] In this embodiment, step 1) divides the historical context into multiple context units according to a predetermined strategy and maintains the representation information of the context units, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

[0035] In this embodiment, step 2) evaluates the relevance of the context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant subset of contexts to the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

[0036] In this embodiment, step 3) adapts the rotation position encoding to the tree-like draft token generated by the draft model, and synchronously maintains the position encoding state during the sliding window clipping and rollback process to ensure Attention position awareness, including: Receive a custom position index corresponding to the tree draft structure, apply rotation position encoding to the query vector and key vector based on the custom position index, and match the actual position of each tree draft token in the tree structure; During the sliding window pruning process, the position index is pruned and remapped synchronously, retaining only the nearest position index within the local sliding window that is relevant to the current inference, and deleting redundant position indexes outside the window; The location encoding status is updated in real time during the inference process, and the corresponding custom location index is automatically appended to the historical location index after each new candidate token is generated; When a rollback operation is triggered, the historical position index records are restored to the state before the draft model generates the tree draft.

[0037] In this embodiment, step 4) encapsulates the core computational logic of Attention and supports combining a sliding window mask and a tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, outputting results and scores, which are used to drive the draft model to generate candidate token sequences based on context subsets, including: Receive the sliding window mask and the tree structure mask generated by the draft model, and jointly restrict the attention reachability relationship; Based on the mask, the attention calculation within the local sliding window is performed first, and then the supplementary attention calculation of the global relevant context unit is performed to achieve a complementary calculation mode of "sliding window + context retrieval". The calculation results and attention scores are based on the Attention calculation output. The calculation results are used to drive the draft model to generate candidate token sequences, and the attention scores serve as the basis for context relevance evaluation and state maintenance, forming a closed-loop feedback mechanism for inference.

[0038] In this embodiment, step 5) after the target model verifies the candidate token sequence, performs a consistent rollback on the generation status, context cache status, and position encoding status of the failed token sequence, including: Receive the number of tokens to be rolled back, determine the rollback target state as the state before the most recent candidate token generation, and clear the temporary Attention score and relevance score temporary data generated in the current inference step; Based on the number of rollback tokens, clear the corresponding global cache state in the context block and cache management module; trim the local sliding window cache and delete the local KV Cache corresponding to the verification failure candidate tokens; Based on the number of rollback tokens, the rollback of the position-encoded state is performed through the position encoding module.

[0039] Example 2 This invention provides an optimization method for draft model inference for long context speculative decoding. This invention does not require additional training or modification of the network structure of the draft model. It optimizes the inference process of the draft model as a whole by focusing on context organization, position encoding maintenance and attention participation in the inference stage.

[0040] During the inference decoding process, the KV Cache corresponding to the historical context is managed in blocks. A sliding window mechanism is used to limit the local context scope of the current inference stage. Simultaneously, based on the query representation of the current generation position, the relevance of historical context units is evaluated, and context units relevant to the current generation step are selected for inference. During the attention calculation stage, the model focuses on the context within the sliding window and the most relevant context units, ignoring irrelevant context. This improves the quality of token generation in long-context scenarios and reduces the computational complexity of attention calculations.

[0041] Meanwhile, when generating multi-branch inferences in the draft model, a position encoding maintenance mechanism that evolves synchronously with the sliding window is introduced, so that the position index in the tree draft can be dynamically adjusted with context pruning and rollback operations, ensuring the consistency of positional relationships in different inference branches.

[0042] Furthermore, this solution maintains state information related to context, location encoding, and inference progress in a unified manner during the inference process, and supports reverting the above states to the corresponding states at historical moments when inference fails, thus forming a rollback-capable inference decoding process for sliding window pruning and context retrieval mechanisms.

[0043] (I) System Structure and Functional Modules The overall system architecture of this algorithm can be divided into five collaborative functional modules. Each module dynamically controls the context scope, cache loading, and attention execution path during the inference phase, thereby achieving long-context inference decoding without requiring additional training or structural modifications to the draft model. The overall architecture includes the following functional modules: 1. Context Block Management Module: This module is used to divide historical context into multiple context units according to a predetermined strategy and maintain the representation information of the context units.

[0044] 2. Context relevance assessment and filtering module: It is used to assess the relevance of the context unit based on the query representation corresponding to the current generation position, thereby filtering out the context subset most relevant to the current generation step.

[0045] 3. Position encoding module adapted in conjunction with sliding window: Adapts rotation position encoding (RoPE) to the tree-like draft tokens generated by the draft model, and maintains the position encoding state synchronously during the sliding window clipping and rollback process to ensure position awareness of Attention.

[0046] 4. Attention Calculation and Execution Module: Encapsulates the core calculation logic of Attention and supports combining the sliding window mask and the tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, output the result and score, and use it to drive the draft model to generate candidate token sequences based on the context subset.

[0047] 5. Rollback Awareness State Maintenance Module: After the target model verifies the candidate token, it performs a consistent rollback of the generation state, context cache state, and location encoding state of the failed tokens.

[0048] (ii) Context Block Management Module This module is used to perform structured management of key-value cache corresponding to long contexts during the inference phase. Its core function is to divide the historical context into multiple context units according to a predetermined granularity, and to uniformly manage the cache status, storage location and loading strategy of the context units. In this way, without modifying the draft model network structure, the size of the key-value cache and the memory usage during long context inference are controlled, providing basic support for subsequent context relevance evaluation and dynamic selection.

[0049] 1. Context partitioning mechanism During inference, the system divides the key-value cache corresponding to the historical context into multiple context units of uniform size according to preset chunking rules. For context fragments that do not yet meet the complete chunking conditions, the system maintains them as independent remaining context states and incorporates them into the context unit system once the chunking conditions are met. This approach ensures consistency in representation dimensions and structure among the context units, thus providing a stable foundation for subsequent context relevance assessment and selection.

[0050] 2. State maintenance of the context unit The system maintains corresponding context representation information and related state description information for each context unit. The state description information includes the context unit's identifier, representative token representation, and storage state. Through unified management of this state information, the system achieves controllable maintenance of the entire lifecycle of the context unit.

[0051] 3. Local context supplementation mechanism In addition to global context block management, the system also maintains a local context region directly related to the current inference position. This local context region exists in the form of a sliding window, retaining only the context representation closest to the current generation position to ensure local semantic continuity. The local context region and the global context unit complement each other, effectively limiting the scale of context participation while ensuring inference quality.

[0052] (III) Context Relevance Assessment and Screening Module The core function of this module is to receive the query representation corresponding to the current generation position, score the relevance of all managed context units (Blocks) based on a predetermined evaluation strategy, and finally select the Top K context units most relevant to the current generation step. This provides a decision basis for subsequent dynamic selection and scheduling of contexts. Its core role is to reduce the participation of invalid contexts, enabling the draft model to focus on the part most relevant to the current inference from a large number of historical contexts, thereby reducing the amount of inference computation.

[0053] 1. Query means to retrieve This module first extracts the global query vector for the current inference step. This vector is generated by the query projection layer of the draft model and adapted by the tree-like position encoding of the draft model. It serves as the query representation for the current generated position, ensuring that the query representation carries the positional and semantic information of the current inference.

[0054] 2. Relevance Assessment Mechanism In each inference step, the system calculates the relevance of the maintained context units based on the query representation corresponding to the currently generated position. The relevance calculation is based on the degree of matching between the representative representation of the context unit and the current query representation.

[0055] 3. Filtering relevant context subsets Based on the relevance assessment results, the system selects several context units most relevant to the current generation step from multiple context units, forming a global context participation subset for the current inference step. Through this selection mechanism, the system avoids involving all context units in the inference process, effectively reducing computational complexity.

[0056] 4. Rating updates are dynamic. After each reasoning step is completed, the module calculates the output score based on the attention of the current reasoning, dynamically updates the relevance score of each context unit, incorporates the contribution of the current reasoning into the scoring system, ensures the real-time and accuracy of relevance assessment, and makes the subsequently selected context subsets more suitable for the current reasoning needs.

[0057] (iv) Position encoding module that adapts to sliding windows This module is the core module for solving the problem of inconsistent position encoding in tree-like draft generation scenarios. Its core function is to adapt the rotated position encoding (RoPE) to the tree-like draft tokens generated by the draft model, and to maintain the position encoding state synchronously during sliding window pruning and context rollback, so as to ensure the accuracy of position perception in Attention calculation and ensure that the relative positional relationship between tokens remains consistent in multi-branch speculative decoding scenarios, without the need for additional training of the draft model or modification of the underlying position encoding logic.

[0058] 1. Location encoding adaptation of tree-structured draft tokens During inference, the system employs a positional encoding mechanism corresponding to the tree-like generation structure. This positional encoding does not rely on fixed linearly increasing numbers but is maintained as an evolvable state variable. An extended rotational positional encoding is used during draft model inference, supporting the reception of custom position indices (position_ids) corresponding to the tree-like draft structure. This accurately reflects the hierarchical and relative order relationships between different draft branches. Based on these custom position indices, rotational positional encoding is applied to the query vector and key vector, ensuring that the positional encoding of each tree-like draft token matches its actual position in the tree structure, avoiding semantic confusion caused by traditional linear positional encoding in tree-like unfolding scenarios.

[0059] 2. Synchronous evolution with sliding windows During the sliding window pruning process, the position indexes are pruned and remapped synchronously, retaining only the nearest position indexes within the local sliding window that are relevant to the current inference, and deleting redundant position indexes outside the window. At the same time, the pruned position indexes are remapped to ensure that the relative positional relationships of tokens within the window do not change, thus guaranteeing the continuity and stability of position awareness in Attention calculation. The pruning process is executed synchronously with the dynamic adjustment of the sliding window, and the position encoding state evolves in real time with the coverage of the sliding window.

[0060] 3. State synchronization maintenance during rollback process When the rollback awareness state maintenance module triggers a rollback, this module will respond to the rollback command and restore the position encoding state. It will restore the historical position index records to the state before the draft model generates the tree draft, ensuring that the position encoding state is synchronized with the context cache state and generation state after the rollback, guaranteeing that subsequent position encoding calculations can proceed normally.

[0061] 4. Dynamic updates of location codes As the draft model inference progresses, the position encoding status is updated in real time. Each time a new candidate token is generated, the corresponding custom position index is automatically appended to `past_position_ids` to adapt to the position information of the new token. At the same time, it supports modifying the distance scaling factor parameter (distance_scale) to dynamically adjust the scaling ratio of the position encoding, adapting to context sequences of different lengths and further improving the accuracy of position awareness.

[0062] (v) Attention Calculation and Execution Module This module is used to perform attention calculation operations during the inference process. Its core function is to encapsulate the underlying calculation logic of Attention, receive context data, mask information and control signals passed from other modules, support the combination of sliding window masks and tree structure masks to jointly constrain attention reachability relationships, perform attention calculations and output calculation results and attention scores, provide core calculation support for subsequent candidate token generation and relevance score updates, and fully adapt to the inference process of the draft model.

[0063] 1. Sliding window mask constraint The system receives a sliding window mask from the context dynamic selection and scheduling module. This mask is used to limit the local scope of Attention calculation, allowing only the current token to perform attention calculation on tokens within the local sliding window. It also shields the interference of redundant tokens outside the window, keeping the computational complexity of local Attention at O(n×W) (where n is the current sequence length and W is the sliding window size). This significantly reduces computational redundancy and improves inference efficiency.

[0064] 2. Tree structure mask constraints Receive the tree structure mask (tree_mask) generated by the draft model. This mask is used to describe the structural relationship of the draft generation tree, limit the attention reachability between tokens in different branches, allow attention interaction only within the same draft branch, shield invalid attention calculations between tokens in different branches, avoid semantic confusion between different branches under the tree draft structure, and ensure the coherence and rationality of candidate token generation.

[0065] 3. Attention Calculation Mechanism under Joint Constraints When performing attention calculations, the system introduces local contextual constraints and tree structure constraints to jointly restrict attention reachability relationships, allowing only attention calculations that conform to the contextual scope and generative structure constraints to occur. During the calculation process, attention calculations within the local sliding window are performed first, followed by supplementary attention calculations for globally relevant contextual units, achieving a complementary calculation mode of "sliding window + contextual retrieval". 4. Output of calculation results The attention calculation results are used to drive the draft model to generate candidate results, and at the same time serve as the basis for context relevance evaluation and state maintenance, forming a closed-loop feedback mechanism in the reasoning process.

[0066] (vi) Rollback Awareness State Maintenance Module The core function of this module is to perform a consistent rollback of all relevant states of the draft model, such as the generation state, context cache state, and position encoding state, when the candidate result generated by the draft model is not accepted by the target model. This restores the draft model to a stable state before the candidate token was generated, avoiding state confusion during the rollback process and ensuring that subsequent inference can proceed normally. At the same time, no modifications to the draft model structure are required.

[0067] 1. Rollback Trigger Mechanism Rollback triggering mechanism: When the target model returns a signal that the candidate token verification failed, the rollback process is triggered. The module receives the number of tokens to be rolled back and determines the target state of the rollback, that is, the state before the most recent candidate token generation. The rollback triggering process is synchronously driven by the inference process and does not require manual intervention.

[0068] 2. Rollback of generated state Restore the generation state of the draft model, delete the generation records corresponding to the candidate tokens that failed verification, restore the generation progress to the state before the rollback, and clear the temporary data such as the temporary Attention score and relevance score generated in the current inference step to ensure that the generation state is completely consistent with that before the rollback.

[0069] 3. Context cache state rollback To achieve consistent rollback of context caching: ① Global cache snapshot rollback: restore the context block and the global cache state of the cache management module to the cache snapshot before rollback; ② Local cache rollback: trim the local sliding window cache, retain the local context before rollback, and delete the local KV Cache corresponding to the verification failure candidate token; 4. Location coding status rollback The position encoding state is rolled back using the method in the tree-structured position encoding module.

[0070] Example 3 like Figure 2 As shown, the present invention provides a draft model inference optimization system for long context speculation decoding, comprising: The historical context partitioning unit divides the historical context into multiple context units according to a predetermined strategy and maintains the representation information of the context units. The context unit relevance evaluation unit evaluates the relevance of context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant context subset to the current generation step. The rotation position encoding adaptation unit adapts the rotation position encoding to the tree-shaped draft tokens generated by the draft model, and maintains the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness. The candidate token sequence generation unit, supported by the position-encoded and adapted tree-structured draft tokens and related context subsets, encapsulates the core computational logic of Attention, and combines the sliding window mask and the tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, output the result and score, and generate the candidate token sequence. The verification unit submits the candidate token sequence to the target model for verification and performs a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

[0071] In the historical context partitioning unit of this embodiment, the historical context is divided into multiple context units according to a predetermined strategy, and the representation information of the context units is maintained, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

[0072] In the context unit relevance evaluation unit of this embodiment, the relevance of the context unit is evaluated based on the query representation corresponding to the current generation position, thereby filtering out the context subset most relevant to the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

[0073] Example 4 The present invention provides a computer-readable storage medium storing a computer program, which, when executed by a processor, implements the steps of the draft model inference optimization method for long context speculative decoding.

[0074] Those skilled in the art will understand that embodiments of this application can be provided as methods, systems, or computer program products. Therefore, this application can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, this application can take the form of a computer program product embodied on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0075] This application is described with reference to flowchart illustrations and / or block diagrams of methods, systems, and computer program products according to embodiments of this application. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, generate instructions for implementing the flowchart... Figure 1 One or more processes and / or boxes Figure 1 A system that specifies functions in one or more boxes.

[0076] These computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to function in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The function specified in one or more boxes.

[0077] These computer program instructions may also be loaded onto a computer or other programmable data processing equipment to cause a series of operational steps to be performed on the computer or other programmable equipment to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable equipment for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0078] The inventive points protected by this invention are as follows: This invention proposes an optimized draft model inference method for long-context speculative decoding. Addressing the technical problems of significantly reduced inference efficiency, decreased token acceptance rate, and insufficient system stability in existing speculative decoding methods under ultra-long context scenarios, it constructs a draft model inference mechanism with a rollback-capable sliding window control and context retrieval. Its core innovations are summarized as follows: 1. Inference control framework for long context-based inference decoding draft models without training conditions To address the issue that existing speculative decoding methods rely on retraining or positional encoding extension algorithms in long-context scenarios, this invention proposes an inference control framework that eliminates the need for additional training of the draft model. This framework dynamically controls the context scope, KVCache loading, and attention computation path during the inference phase, combining the sliding window concept and context retrieval mechanism. This enables the draft model, which was originally trained only under short-context conditions, to stably and efficiently participate in the speculative decoding process for long contexts with millions of tokens, thereby significantly reducing algorithm deployment costs and engineering complexity.

[0079] 2. The collaborative mechanism between tree-based positional encoding and sliding window strategy To address the issue of nonlinear unfolding of contextual positional relationships in speculative decoding where the draft model generates multiple candidate tokens using a tree structure, this invention proposes a tree-based positional encoding method that is compatible with sliding windows. During draft model inference, a custom positional index generated based on the tree-based draft structure is used to apply rotational positional encoding to the query vector and key vector. This positional index is synchronously maintained and remapped during sliding window pruning, allowing the positional encoding state to dynamically evolve with the sliding window's coverage and draft rollback operations. This maintains consistency in the relative positional information between tokens in multi-branch speculative decoding scenarios. This invention ensures the continuity and consistency of positional encoding calculation in multi-branch speculative decoding scenarios combined with sliding windows without requiring additional training or modification of the draft model's network structure.

[0080] 3. Attention-Related Context Unit Construction and Context-Aware Participation Methods This invention proposes a method for constructing and participating in global context units for long-context speculative decoding draft models. During the draft model inference process, the historical key-value cache is divided into multiple units according to a predetermined size, and a TopK search is performed based on the attention correlation between the query vector and the representation of each unit to determine candidate units for global attention computation. In the attention computation stage, a tree structure mask is introduced to describe the tree-like draft tokens generated by the draft model during speculative decoding, and this mask, in conjunction with a sliding window mask, constrains the attention reachability relationship between the query position and the key position. This achieves context-aware token-level attention control within the selected context units, reducing interference from irrelevant context on the inference results.

[0081] 4. A rollbackable context state maintenance mechanism for scenarios where speculative decoding fails. To address the issue that draft model-generated results may be rejected by the target model during speculative decoding, necessitating rollback, this invention proposes a rollbackable context state maintenance mechanism for speculative decoding. This mechanism uniformly maintains and controls the consistency of inference-related context states during draft model inference, including context states within the local sliding window, historical context states not yet fixed as context units, and position-encoded states related to the inference process. When a speculation failure occurs and triggers a rollback operation, the mechanism can restore the entire context state to its predetermined historical state without re-initializing the model, thereby avoiding redundant calculations and state inconsistencies, and significantly reducing the additional computational and system overhead caused by speculative decoding failures.

[0082] 5. System architecture and device implementation for collaborative design of inference process Based on the above method, this invention constructs a system architecture for collaborative optimization of the inference process. Each module works collaboratively under a unified control strategy, realizing end-to-end optimization from KVCache management and context selection to attention calculation and result generation. Without modifying the draft model network structure, it significantly improves the overall performance of long context inference decoding.

[0083] The foregoing has shown and described the basic principles, main features, and advantages of the present invention. It will be apparent to those skilled in the art that the invention is not limited to the details of the exemplary embodiments described above, and that the invention can be implemented in other specific forms without departing from its spirit or essential characteristics. Therefore, the embodiments should be considered illustrative and non-limiting in all respects, and the scope of the invention is defined by the appended claims rather than the foregoing description. Thus, all variations falling within the meaning and scope of equivalents of the claims are intended to be included within the scope of the invention. No reference numerals in the claims should be construed as limiting the scope of the claims.

[0084] Furthermore, it should be understood that although this specification describes embodiments, not every embodiment contains only one independent technical solution. This narrative style is merely for clarity. Those skilled in the art should consider the specification as a whole, and the technical solutions in each embodiment can be appropriately combined to form other embodiments that can be understood by those skilled in the art. The above content is only for illustrating the technical concept of the present invention and should not be construed as limiting the scope of protection of the present invention. Any modifications made based on the technical concept proposed in this invention shall fall within the scope of protection of the claims of this invention.

Claims

1. A method for optimizing draft model inference for long-context speculative decoding, characterized in that, include: 1) Divide the historical context into multiple context units according to a predetermined strategy, and maintain the representation information of the context units; 2) Based on the query representation corresponding to the current generation position, the relevance of the context units is evaluated, thereby selecting the most relevant subset of context to the current generation step; 3) Adapt rotation position encoding to the tree-like draft tokens generated by the draft model, and maintain the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness; 4) With the support of the tree-like draft tokens after position encoding adaptation and related context subsets, the core computation logic of Attention is encapsulated. Combined with the sliding window mask and the tree structure mask used to describe the draft generation tree structure, the attention reachability relationship is jointly constrained, the results and scores are output, and a candidate token sequence is generated. 5) Submit the candidate token sequence to the target model for verification, and perform a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

2. The draft model inference optimization method for long context-based speculative decoding according to claim 1, characterized in that, Step 1) Divide the historical context into multiple context units according to a predetermined strategy, and maintain the representation information of the context units, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

3. The draft model inference optimization method for long context-based speculative decoding according to claim 1, characterized in that, Step 2) Based on the query representation corresponding to the current generation position, evaluate the relevance of the context units to select the most relevant subset of contexts for the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

4. The draft model inference optimization method for long context-based speculative decoding according to claim 1, characterized in that, Step 3) Adapt rotational position encoding to the tree-like draft tokens generated by the draft model, and simultaneously maintain the position encoding state during sliding window clipping and rollback to ensure Attention position awareness, including: Receive a custom position index corresponding to the tree draft structure, apply rotation position encoding to the query vector and key vector based on the custom position index, and match the actual position of each tree draft token in the tree structure; During the sliding window pruning process, the position index is pruned and remapped synchronously, retaining only the nearest position index within the local sliding window that is relevant to the current inference, and deleting redundant position indexes outside the window; The location encoding status is updated in real time during the inference process, and the corresponding custom location index is automatically appended to the historical location index after each new candidate token is generated; When a rollback operation is triggered, the historical position index records are restored to the state before the draft model generates the tree draft.

5. The draft model inference optimization method for long context-based speculative decoding according to claim 1, characterized in that, Step 4) Encapsulate the core computational logic of Attention, and support the combination of a sliding window mask and a tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, outputting results and scores, which are used to drive the draft model to generate candidate token sequences based on context subsets, including: Receive the sliding window mask and the tree structure mask generated by the draft model, and jointly restrict the attention reachability relationship; Based on the mask, the attention calculation within the local sliding window is performed first, and then the supplementary attention calculation of the global relevant context unit is performed to achieve a complementary calculation mode of "sliding window + context retrieval". The calculation results and attention scores are based on the Attention calculation output. The calculation results are used to drive the draft model to generate candidate token sequences, and the attention scores serve as the basis for context relevance evaluation and state maintenance, forming a closed-loop feedback mechanism for inference.

6. The draft model inference optimization method for long context-based speculative decoding according to claim 1, characterized in that, Step 5) After validating the candidate token sequence in the target model, perform a consistent rollback of the generation state, context cache state, and position encoding state of the failed token sequences, including: Receive the number of tokens to be rolled back, determine the rollback target state as the state before the most recent candidate token generation, and clear the temporary Attention score and relevance score temporary data generated in the current inference step; Based on the number of rollback tokens, clear the corresponding global cache state in the context block and cache management module; trim the local sliding window cache and delete the local KV Cache corresponding to the verification failure candidate tokens; Based on the number of rollback tokens, the rollback of the position-encoded state is performed through the position encoding module.

7. A draft model inference optimization system for long context-based speculative decoding, characterized in that, include: The historical context partitioning unit divides the historical context into multiple context units according to a predetermined strategy and maintains the representation information of the context units. The context unit relevance evaluation unit evaluates the relevance of context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant context subset to the current generation step. The rotation position encoding adaptation unit adapts the rotation position encoding to the tree-shaped draft tokens generated by the draft model, and maintains the position encoding state synchronously during the sliding window clipping and rollback process to ensure Attention position awareness. The candidate token sequence generation unit encapsulates the core computational logic of Attention, supported by the tree-like draft tokens after position encoding adaptation and related context subsets. It also combines the sliding window mask and the tree structure mask used to describe the draft generation tree structure to jointly constrain the attention reachability relationship, output the results and scores, and generate the candidate token sequence. The verification unit submits the candidate token sequence to the target model for verification and performs a consistent rollback on the generation status, context cache status, and position encoding status of the failed verification token sequence.

8. The draft model inference optimization system for long context speculation decoding according to claim 7, characterized in that, The historical context partitioning unit divides the historical context into multiple context units according to a predetermined strategy, and maintains the representation information of the context units, including: During the inference process of the draft model, the key-value cache (KV Cache) corresponding to the historical context is divided into multiple context units of the same size. For context fragments that do not yet meet the complete block division conditions, they are maintained as independent remaining context states and incorporated into the context unit system after the block division conditions are met. Maintain corresponding context representation information and state description information for the context unit. The state description information includes the context unit's identification information, representative token representation, and storage state. A local context region directly related to the current inference position is maintained synchronously. This local context region exists in the form of a sliding window, and only the context representation closest to the current generation position is retained.

9. The draft model inference optimization system for long context speculation decoding according to claim 7, characterized in that, The context unit relevance evaluation unit evaluates the relevance of context units based on the query representation corresponding to the current generation position, thereby selecting the most relevant subset of context to the current generation step, including: Extract the global query vector for the current inference step. This vector is generated by the Query projection layer of the draft model and adapted by the tree-like position encoding of the draft model to serve as the query representation of the current generated position. Based on the query representation corresponding to the generated location, and combined with the degree of matching between the representative representation of the context unit and the current query representation, the relevance of the maintained context units is evaluated. Based on the relevance assessment results, several context units most relevant to the current generation step are selected from multiple context units to form the global context participation subset of the current inference step; After the reasoning step is completed, the output score is calculated based on the attention of the current reasoning, the relevance score of each context unit is dynamically updated, and the contribution of the current reasoning is incorporated into the scoring system.

10. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed by a processor, implements the steps of the draft model inference optimization method for long context speculative decoding according to any one of claims 1-6.

Images