Cache management system during large model reasoning

Abstract

The invention provides a cache management system during large model reasoning, and relates to the technical field of large-scale artificial intelligence reasoning optimization. The system comprises a cache module, a cache manager, a retrieval enhancement generation controller, a reasoning time controller, a large language model reasoning engine, a process reward model reasoning engine, a result reward model reasoning engine, a Monte Carlo tree search scheduler and a runtime scheduler, the cache module is used for caching required key value pair data in a pre-filling stage, and the cache manager is used for storing the key value pair data in the pre-filling stage. The key-value pair data is used for storing context information of the preorder lexical elements, and the key-value pair data is also used for being reused in the decoding stage. According to the embodiment of the invention, in the pre-filling stage, the required key value pair data is generated and cached, and the cached key value pair data is used for storing the context information of the preorder lexical elements. In the decoding stage, the caches are reused, repeated calculation is avoided, and therefore the GPU memory is saved.

Description

Technical Field

[0001] The present disclosure relates to the technical field of large-scale artificial intelligence reasoning optimization, and in particular to a cache management system for large model reasoning. Background Art

[0002] With the rapid increase in the number of parameters of large language models (such as GPT and LLaMA), their inference computing and memory requirements have also increased dramatically. The traditional single-device inference method is difficult to meet real-time requirements. Distributed inference architecture is an important method to improve inference efficiency, but it faces the following problems in actual applications:

[0003] The limited GPU memory resources conflict with KV-Cache management during inference: In large-scale model inference (such as large language models), GPU resources (especially memory) are limited, and the intermediate cache data (KV-Cache) generated during inference occupies a large amount of GPU memory. This leads to memory competition between KV-Cache and inference calculations, limiting the efficient use of memory and inference performance, especially in high-load scenarios, prone to memory overflow or cache duplication.

[0004] Repeated generation of KV-Cache data leads to waste of computing resources: In large model reasoning, KV-Cache is a crucial intermediate data that stores the context information of the previous token. Since the decoding process needs to rely on the cache data generated in the Prefill stage, if KV-Cache is regenerated for each reasoning, it will lead to repeated calculations and waste of GPU memory.

[0005] It should be noted that the information disclosed in the above background technology section is only used to enhance the understanding of the background of the present disclosure, and therefore may include information that does not constitute the prior art known to ordinary technicians in the field. Summary of the invention

[0006] The present disclosure provides a cache management system for large model reasoning, which reduces repeated calculations and memory waste of the cache to at least a certain extent, thereby improving the utilization rate of GPU resources.

[0007] Other features and advantages of the present disclosure will become apparent from the following detailed description, or may be learned in part by the practice of the present disclosure.

[0008] According to one aspect of the present disclosure, a cache management system for large model inference is provided, based on a pre-filled instance and a decoded instance separation architecture, including:

[0009] The cache manager includes a cache retriever and a cache controller connected to each other; the cache retriever is connected to a cache module, and the cache module is used to store key-value pairs in large-scale documents; the cache controller is connected to a retrieval enhancement generation controller and an inference time controller;

[0010] The retrieval enhancement generation controller connects the external database and the large language model inference engine;

[0011] The reasoning time controller includes a process reward model strategy and a result reward model strategy; the process reward model strategy is connected to the process reward model reasoning engine, and the result reward model strategy is connected to the result reward model reasoning engine; the reasoning time controller is also connected to the large language model reasoning engine;

[0012] The large language model inference engine, the process reward model inference engine and the result reward model inference engine are connected to the Monte Carlo tree search scheduler;

[0013] The cache module is used to cache the required key-value pair data in the pre-filling stage. The key-value pair data is used to store the context information of the previous word, and the key-value pair data is also used to be reused in the decoding stage.

[0014] In one embodiment of the present disclosure, a cache retriever is used to locate required data in a cache module; a cache controller determines how key-value pair data is processed, whether to be stored in a GPU or removed, by interacting with a retrieval enhancement generation controller and an inference time controller.

[0015] In one embodiment of the present disclosure, the cache controller is also used to call the cache retriever to check whether the data requested by the data request already exists in the cache module after receiving the data request. If the cache hits, the data requested by the data request is returned; if the cache misses, the data requested by the data request is loaded from the computing process or external data source into the cache module.

[0016] In one embodiment of the present disclosure, the cache controller updates the data in the cache module according to the policy; the cache controller determines the data that needs to be removed from the cache module according to the invalidation policy.

[0017] In one embodiment of the present disclosure, when there are multiple cache requests, the cache controller determines the order in which data is stored in the cache module through priority scheduling.

[0018] In one embodiment of the present disclosure, the cache controller also loads the target data into the cache in advance based on a pre-fetch strategy.

[0019] In one embodiment of the present disclosure, a retrieval enhancement generation controller calculates a priority based on the frequency of accessing key-value pair data, the size of the key-value tensor, the last access time, and the prefix-aware recalculation cost; a cache controller determines the data that needs to be removed from the cache module based on the invalidation strategy and priority.

[0020] In one embodiment of the present disclosure, the retrieval enhancement generation controller is further used to combine the prediction and reservation strategies to enable knowledge retrieval and LLM reasoning to be performed in an overlapping manner.

[0021] In one embodiment of the present disclosure, the inference time controller coordinates various stages of model inference by optimizing the process reward model strategy and the result reward model strategy.

[0022] In one embodiment of the present disclosure, the cache management system during large model inference also uses the document identifier as the node identifier and stores the document's key-value tensor as a leaf node to form a multi-branch tree structure; the Monte Carlo tree search scheduler uses a heuristic strategy or a priority-based approach to select nodes in the decision tree that need to be expanded, simulated, or returned.

[0023] The cache management system for large model inference provided by the disclosed embodiment separates the key-value pair cache data during pre-filling and decoding. In the pre-filling stage, the required key-value pair data is generated and cached. The cached key-value pair data is used to store the context information of the previous token. In the decoding stage, these caches are reused to avoid repeated calculations, thereby saving GPU (Graphics Processing Unit) memory. In addition, the disclosed embodiment also coordinates the large language model inference engine, the process reward model inference engine, and the result reward model inference engine through the Monte Carlo tree search scheduler, thereby realizing a dynamic cache management strategy, flexibly releasing and retaining cache data according to computing needs, and optimizing memory utilization.

[0024] It is to be understood that the foregoing general description and the following detailed description are exemplary and explanatory only and are not restrictive of the present disclosure. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] The accompanying drawings, which are incorporated in and constitute a part of this specification, illustrate embodiments consistent with the present disclosure and, together with the description, serve to explain the principles of the present disclosure.

[0026] Obviously, the drawings described below are only some embodiments of the present disclosure, and a person skilled in the art can obtain other drawings based on these drawings without any creative work.

[0027] Figure 1 A schematic diagram of a process reward model in an embodiment of the present disclosure is shown;

[0028] Figure 2 A schematic diagram of a result reward model in an embodiment of the present disclosure is shown;

[0029] Figure 3 A schematic diagram showing the selection steps in the Monte Carlo tree search in an embodiment of the present disclosure is shown;

[0030] Figure 4 A schematic diagram showing the expansion steps in the Monte Carlo tree search in an embodiment of the present disclosure is shown;

[0031] Figure 5 A schematic diagram showing simulation steps in a Monte Carlo tree search in an embodiment of the present disclosure is shown;

[0032] Figure 6 A schematic diagram showing the backpropagation step in the Monte Carlo tree search in an embodiment of the present disclosure is shown;

[0033] Figure 7 A structural block diagram of a cache management system for large model reasoning in an embodiment of the present disclosure is shown. DETAILED DESCRIPTION

[0034] In order to make the purpose, technical scheme and advantages of the embodiments of the present disclosure clearer, the technical scheme in the embodiments of the present disclosure will be clearly and completely described below in conjunction with the drawings in the embodiments of the present disclosure. Obviously, the described embodiments are only part of the embodiments of the present disclosure, rather than all of the embodiments. The components of the embodiments of the present disclosure generally described and shown in the drawings here can be arranged and designed in various different configurations. Therefore, the following detailed description of the embodiments of the present disclosure provided in the drawings is not intended to limit the scope of the present disclosure for protection, but merely represents the selected embodiments of the present disclosure. Based on the embodiments of the present disclosure, all other embodiments obtained by those skilled in the art without making creative work belong to the scope of protection of the present disclosure.

[0035] For ease of understanding, the following first explains the relevant concepts involved in the present disclosure as follows:

[0036] Prefill: The process of generating the first output token based on the input sequence in a large language model generation task

[0037] Decoding: The process by which a large language model gradually generates each token in a generation task. Each new token generated depends on the previously generated token.

[0038] KV Cache: intermediate cache data generated during model inference, including Key (K) and Value (V) vectors, used to store context information of previous tokens

[0039] Separation architecture: The large language model reasoning separation architecture specifically refers to the separation of prefill and decoding, as well as the management strategy of KV Cache under this setting

[0040] Alignment during large model reasoning: Process Reward Model (PRM) and Outcome Reward Model (ORM) in large model reasoning (such as large language model, LLM) are two optimization strategies based on reward signals, which are used to improve the performance and effect of the model during reasoning. They guide the model to learn better behaviors or results in the generation task in different ways.

[0041] Monte Carlo Tree Search (MCTS): is an algorithm used in decision-making processes. It is widely used in tasks that require exploring complex decision spaces. MCTS gradually estimates the pros and cons of a decision node by simulating a large number of randomized future decision paths, and finally selects the best action. This method is particularly suitable for effective search when the decision tree is too large and cannot be exhausted.

[0042] Process Reward Model (PRM): The Process Reward Model (PRM) focuses on each step in the reasoning process and guides the model to learn how to reason more effectively by rewarding or punishing the model's behavior during the generation process. The goal of PRM is to optimize the intermediate process of reasoning through a reward mechanism, rather than just optimizing the final output result.

[0043] Outcome Reward Model (ORM): The Outcome Reward Model (ORM) focuses on the final result of the reasoning task, and guides the model to optimize by evaluating the quality of the final output and giving rewards or penalties. ORM not only evaluates whether the generated results meet the predetermined goals, but also gives feedback based on the quality of the output results, thereby encouraging the model to generate more appropriate outputs.

[0044] RAG (Retrieval-Augmented Generation) is a technology that combines information retrieval (IR) and generative models, aiming to improve model performance in generation tasks, especially in tasks that require rich external knowledge.

[0045] The following introduces the relevant technologies used in the cache management system during large model reasoning provided by the embodiments of the present disclosure.

[0046] Alignment technology for large model reasoning: The alignment problem in large model reasoning usually refers to ensuring that the behavior of the model meets the expected goals, especially in the fields of reinforcement learning and policy optimization. This technology consists of three components:

[0047] 1. Train the reward model: Train the model by designing a suitable reward function so that it can gradually learn and optimize its behavior based on feedback.

[0048] 2. Construct Monte Carlo Tree (MCTS): Explore and optimize decision trees through Monte Carlo tree search to help the model select the optimal strategy in a complex decision space.

[0049] 3. Combined with reinforcement learning (RL) technology: Reinforcement learning continuously adjusts strategies to maximize cumulative rewards through the interaction between the agent and the environment. It includes the process reward model (PRM) and the outcome reward model (ORM); PRM focuses on the behavioral steps in the optimization process, while ORM focuses on the reward feedback of the final task results.

[0050] Process Reward Model (PRM): Figure 1 As shown in Figure 1, the process reward model is a way of rewarding based on the performance of the model in the process of reasoning or behavior. In this model, the reward is given based on the performance of the process itself rather than the final result. PRM focuses on the specific steps and behavior process of the model in completing the task, rather than whether the task is finally completed. PRM aims to provide more refined reward feedback on the behavior process of the intelligent agent, thereby guiding it to continuously optimize each decision and step in the process.

[0051] Outcome Reward Model (ORM): Figure 2 As shown in Figure 1, unlike PRM, the outcome reward model is a mechanism for evaluating rewards based on the final result of the model reasoning process. In ORM, the agent's actions will only receive rewards when the final goal is achieved. ORM focuses on the quality of the final result without considering every detail or intermediate step in the process.

[0052] Retrieval-augmented generation (RAG): The RAG model mainly generates more accurate and contextual answers or text by retrieving relevant information from external documents or databases and combining this information with the generation model. RAG performs well in tasks such as open-domain question answering, dialogue generation, and document summarization. It can make up for the knowledge limitations of traditional generation models and improve the quality and accuracy of generated text. However, the effectiveness of RAG depends on the quality of the retrieval module and the maintenance of the external knowledge base, so it still faces some challenges. Retrieval-augmented generation mainly involves information retrieval and text generation.

[0053] Information Retrieval: The first step of RAG is to retrieve documents or information related to the input question or context from an external knowledge base (such as a document set, web page, database, etc.). Typically, information retrieval uses traditional retrieval methods such as BM25 based on sparse representation or vector retrieval methods based on dense representation (such as through tools such as FAISS or ElasticSearch).

[0054] Input: The input text (such as a question or query) is encoded into a query vector.

[0055] Retrieval: This query vector is similarly matched with the document vectors in the knowledge base to retrieve the documents that are most relevant to the input text.

[0056] Text Generation (Generation) Once the relevant documents or snippets are retrieved, RAG provides this information together with the original input to the generation model (such as GPT, BART, T5, etc.) and lets the generation model generate answers or text based on this information.

[0057] Input to the generative model: The retrieved documents and the original input are used together as the input of the generative model. The model must not only understand the question itself, but also understand the relevant knowledge provided by the external party.

[0058] Output: Generative models generate answers or text based on the retrieved information.

[0059] MCTS (Monte Carlo Tree Search): The core idea of ​​MCTS is to use Monte Carlo simulation to estimate each node in the decision tree, and then select the best decision based on these estimated values. The MCTS algorithm mainly includes the following four steps:

[0060] Selection: Figure 3 As shown in Figure 1, starting from the root node, a node is selected for expansion according to a certain strategy (such as the upper confidence bound (UCB)). The selection process gives priority to nodes that have not been fully expanded or those with higher estimated returns.

[0061] Expansion: Figure 4 As shown in the figure, one or more child nodes are expanded on the currently selected node. The expanded node represents a possible subsequent action or state.

[0062] Simulation: Figure 5As shown, a simulation is performed starting from the newly expanded node. This step is usually performed by randomly selecting subsequent actions until the game ends or the predetermined number of simulation steps is reached. In game applications, simulation is often a process of making random decisions, or heuristic strategies can be used to make decisions based on actual conditions.

[0063] Backpropagation: Figure 6 As shown in Figure 1, after the simulation is completed, the statistical information of all nodes selected along the path is updated according to the simulation results, usually updating the node's winning rate, score or other performance indicators. This information will be sent back to the root node to affect the subsequent decision-making process.

[0064] The above four steps will be executed repeatedly until the computing resource limit (such as time or number of simulations) is reached. Finally, MCTS will select the node with the best performance in all simulations as the final decision.

[0065] The cache management system for large model reasoning proposed in the present invention is a separation architecture for cache during large model reasoning, so as to optimize the utilization of GPU space and efficiently manage KV-Cache in the reasoning process, especially when combined with reasoning alignment technology (such as Monte Carlo Tree Search (MCTS), Process Reward Model (PRM), Result Reward Model (ORM)) and Retrieval Enhanced Generation (RAG) method, to avoid waste of GPU memory space and duplication of calculations.

[0066] In the process of large-scale model reasoning, GPU resources (especially GPU memory) are usually limited. At the same time, both the calculation during reasoning and the management of KV-Cache require GPU space. Specifically, KV-Cache is an indispensable intermediate cache data in the reasoning process. Especially in the reasoning of large language models (such as GPT, etc.), KV-Cache not only needs to store the context information of the model, but also needs to be continuously updated during the generation process. The limited GPU space makes the memory competition between KV-Cache and reasoning calculations a serious problem.

[0067] In order to effectively meet this challenge, the cache management system during large model reasoning disclosed in the present invention adopts a flexible KV-Cache management strategy. By separating the KV-Cache in the Prefill and Decoding stages, it avoids the repeated generation of the same cache data each time reasoning is performed, and designs a dynamic KV-Cache release and retention mechanism. Through this mechanism, the model can generate and cache the required KV data in the Prefill stage, and reuse these caches in the Decoding stage, reducing repeated calculations and cache occupancy, thereby effectively saving GPU memory space. At the same time, the strategy also flexibly adjusts the release and retention of the cache according to the computing requirements during the reasoning process, ensuring efficient management of computing and cache data within the limited GPU space, thereby ensuring the smoothness and efficiency of the reasoning process.

[0068] Especially when combined with alignment techniques (such as MCTS, PRM, ORM) during reasoning, these techniques usually rely on the cache of intermediate reasoning results (ie, KV-Cache) to optimize the decision-making process and reasoning path. If KV-Cache is not managed properly, it may cause repeated generation of cached data or excessive occupation of GPU memory space, thereby affecting reasoning performance. The separation architecture disclosed in the present invention optimizes the storage and update strategy of the cache, flexibly manages KV-Cache at different reasoning stages, ensures that not only the GPU space is maximized during the reasoning process, but also the previously generated cache data can be reused in the Decoding stage to improve reasoning efficiency. In addition, the present disclosure also takes into account that the RAG model needs to obtain information by retrieving an external knowledge base during the reasoning process and combining it with the generated model. Since the RAG model involves complex cache management and the integration of multi-source data, the architecture disclosed in the present invention can ensure the effective separation and reuse of external data retrieved from the knowledge base and KV-Cache in the model generation process, further improving reasoning efficiency and reducing GPU resource usage.

[0069] The disclosed embodiments adopt an innovative KV-Cache separation management architecture and design flexible cache release and retention strategies for scenarios with limited GPU memory, so as to avoid unnecessary memory waste during the reasoning process and improve the efficiency and resource utilization of large model reasoning, especially when combining the reasoning alignment technology (such as MCTS, PRM, ORM) and RAG model, to ensure efficient use of memory space and maximum reuse of cached data during the reasoning process.

[0070] The defects existing in the above schemes and the proposed solutions are the results obtained by the inventor after practice and careful research. Therefore, the discovery process of the above problems and the solutions proposed by the present disclosure for the above problems below should be the contributions made by the inventor to the present disclosure during the disclosure process.

[0071] It should be noted that the data involved in this technical solution (including but not limited to the data itself, the acquisition or use of the data) shall comply with the requirements of relevant laws, regulations and relevant provisions.

[0072] Figure 7 A cache management system for large model reasoning in an embodiment of the present disclosure is shown, based on a prefilling instance and a decoding instance separation architecture, such as Figure 7 As shown, the cache management system for large model inference provided in the embodiment of the present disclosure includes a cache module, a cache manager (Cache Manager), a retrieval enhancement generation controller (RAG Controller), an inference time controller (Inference Time Controller), a large language model inference engine (LLM Inference Engine), a process reward model inference engine (PRM InferenceEngine), a result reward model inference engine (ORM Inference Engine), a Monte Carlo tree search scheduler (MCTSScheduler) and a runtime scheduler (Runtime Scheduler).

[0073] like Figure 7 As shown, the cache manager includes a cache retriever and a cache controller connected to each other; the cache retriever is connected to the cache module, and the cache module is used to store key-value pairs in large-scale documents; the cache controller is connected to the retrieval enhancement generation controller (RAG Controller) and the inference time controller;

[0074] The retrieval enhancement generation controller (RAG Controller) connects the external database (External Database) and the large language model inference engine (LLM Inference Engine);

[0075] Inference Time Controller, including PRM Strategy and ORM Strategy; PRM Strategy is connected to PRM Inference Engine, ORM Strategy is connected to ORM Inference Engine; Inference Time Controller is also connected to Large Language Model Inference Engine;

[0076] The Large Language Model Inference Engine (LLM Inference Engine), Process Reward Model Inference Engine (PRM Inference Engine) and Result Reward Model Inference Engine (ORM Inference Engine) are connected to the Monte Carlo Tree Search Scheduler (MCTS Scheduler);

[0077] The cache module is used to cache the required key-value pair data in the pre-filling stage. The key-value pair data is used to store the context information of the previous word, and the key-value pair data is also used to be reused in the decoding stage.

[0078] In some embodiments, the retrieval enhancement generation controller (RAG Controller) further includes a strategy controller (Strategy Controller).

[0079] The following is a detailed introduction to each module in the cache management system during the above large model inference.

[0080] In some embodiments, the cache module can use In-Memory KV-Cache, which is a cache mechanism that stores data in system memory (GPU / TPU, etc.), mainly used to store key-value pairs. Compared with disk cache, in-memory cache provides higher access speed. The intermediate cache data generated by the attention mechanism during large model inference includes Key (K) and Value (V) vectors, which are used to store the context information of the previous token.

[0081] In some embodiments, the cache retriever is used to locate required data in the cache module; the cache controller determines how to process the key-value pair data, whether to store it in the GPU or remove it, by interacting with the retrieval enhancement generation controller and the inference time controller.

[0082] In some embodiments, the cache controller is also used to call the cache retriever to check whether the data requested by the data request already exists in the cache module after receiving the data request. If the cache hits, the data requested by the data request is returned; if the cache misses, the data requested by the data request is loaded from the computing process or external data source into the cache module.

[0083] In some embodiments, the cache controller updates the data in the cache module according to the policy; the cache controller determines the data that needs to be removed from the cache module according to the invalidation policy.

[0084] In some embodiments, when there are multiple cache requests, the cache controller determines the order in which data is stored in the cache module through priority scheduling.

[0085] In some embodiments, the cache controller also loads the target data into the cache in advance based on a prefetch strategy.

[0086] In some embodiments, the Cache Retriever efficiently locates the key-value tensor of documents stored in the memory cache by leveraging a knowledge tree. The tree structure is built as a prefix tree based on the document ID, matching the sensitivity of the large language model (LLM) to the order of documents. Each path in the tree represents a specific sequence of documents that are referenced in the request, and each node holds a key-value tensor of the referenced document. Different paths may share the same nodes, indicating that these documents are shared between different requests. This structure enables the retriever to quickly access the key-value tensors of documents in a specified order.

[0087] In some embodiments, the cache controller is the core component responsible for managing and scheduling the operation of the cache system (such as In-Memory KV-Cache). Its main functions include cache allocation, access, update, eviction, and cache data optimization strategy. By interacting with the RAG Controller and Inference Time Controller, the cache controller decides which KV Cache should be stored in the GPU, which should be removed, and how to efficiently use the limited cache space when processing KV Cache access requests to maximize system performance.

[0088] The cache controller manages the lifecycle of cached data and ensures the validity and performance of cached data by working closely with the application or model's reasoning process. Its main functions include:

[0089] Cache access: When the system requests a piece of data, the cache controller first checks whether the data already exists in the cache. If the cache hits, the cache controller returns the data directly; if the cache misses, the data is loaded into the cache from the computing process or external data source.

[0090] Cache update: The cache controller can automatically update the data in the cache according to policies (such as LRU, LFU, etc.). For example, when the cache space is full, the cache controller will select the least frequently used data or the data that has not been accessed recently for eviction.

[0091] New data or intermediate results generated during the inference process are also stored in the cache for subsequent use. The cache controller decides how to update and manage these newly generated cache items.

[0092] Cache Eviction: The cache controller decides which cached data needs to be removed to make room for new data based on the invalidation policy (such as LRU (Least Recently Used), LFU (Least Frequently Used), etc.).

[0093] Invalidation policies typically make decisions based on how often data is accessed, how long it has been used, or its priority, ensuring that the most frequently used or most valuable data always remains in the cache.

[0094] Cache priority management: When there are multiple cache requests, the cache controller can decide which data to store in the cache first and which data can be loaded later through priority scheduling (such as based on the type, importance or time requirement of the request). This is particularly important in complex reasoning tasks, which can improve cache hit rate and reduce computing latency.

[0095] Cache prefetching and preloading (Prefetching): To further improve performance, the cache controller can also adopt a prefetching strategy to load the data that may be needed into the cache in advance. For example, in a large model inference task, the controller can predict and preload the input data that may be needed in the next step based on the current inference progress.

[0096] Cache consistency: When multiple computing threads or processes access the cache at the same time, the cache controller must ensure the consistency and integrity of the cached data. It usually achieves this goal through a lock mechanism (such as a read-write lock) or a distributed consistency protocol to avoid data errors or inconsistencies during concurrent access.

[0097] In some embodiments, the retrieval enhancement generation controller calculates the priority based on the frequency of accessing the key-value pair data, the size of the key-value tensor, the last access time, and the prefix-aware recalculation cost; the cache controller determines the data that needs to be removed from the cache module based on the invalidation strategy and priority.

[0098] In some embodiments, the retrieval enhancement generation controller is further configured to combine the prediction and reservation strategies so that knowledge retrieval and LLM reasoning are performed in an overlapping manner.

[0099] In some embodiments, the Retrieval Enhanced Generation Controller (RAG Controller) is responsible for coordinating some system interactions related to RAG optimization. Different strategies can be combined to minimize cache miss rates. Priority is calculated based on the frequency of access to KV-Cache content, the size of the key-value tensor, the last access time, and the prefix-aware recalculation cost. Cache eviction is performed based on priority to ensure that the most valuable tensors are retained. Cache-aware reordering schedules requests to improve cache hit rates and prevent cache jitter, while ensuring the fairness of requests and reducing starvation problems. At the same time, prediction and reservation strategies can be combined to enable knowledge retrieval and LLM reasoning to overlap, thereby minimizing latency. This optimization uses the mid-term generation of retrieval results to start LLM reasoning in advance.

[0100] In some embodiments, the inference time controller coordinates the various stages of model inference by optimizing the process reward model strategy and the result reward model strategy.

[0101] In some embodiments, the Inference Time Controller is a core component that manages and optimizes the KV-Cache during inference. It coordinates the various stages of model inference by optimizing PRM Strategy and ORM Strategies to ensure that the inference process proceeds smoothly, while maximizing performance, reducing latency, optimizing computing resource usage, and dynamically adjusting the inference strategy based on demand.

[0102] In some embodiments, the LLM Inference Engine is an interface for optimizing LLM inference performance, represented by vllm; the PRM Inference Engine is an interface for optimizing LLM inference performance, represented by vllm, using Process Reward Model; and the ORM Inference Engine is an interface for optimizing LLM inference performance, represented by vllm, using Process Reward Model. In the above, "vllm" is a specific library or tool for optimizing the inference performance of a large language model (LLM, Large LanguageModel).

[0103] In some embodiments, the cache management system during large model inference also uses the document identifier as the node identifier and stores the document's key-value tensor as a leaf node to form a multi-branch tree structure; the Monte Carlo tree search scheduler uses heuristic strategies or priority-based methods to select nodes in the decision tree that need to be expanded, simulated, or returned.

[0104] In some embodiments, the Monte Carlo Tree Search Scheduler (MCTS Scheduler) may also be referred to as a knowledge tree and Monte Carlo Tree Search Scheduling Module.

[0105] The knowledge tree is a data organization method based on a prefix tree structure, which is used to store and retrieve key-value pairs in large-scale document collections. It forms a multi-branch tree structure by using the document ID as the node identifier and storing the document's key-value tensor as a leaf node. The structure of this tree is consistent with the sensitivity of the large language model (LLM) to the order of documents, making it easy to quickly locate and access the key-value tensor of a specific document.

[0106] Each path in the knowledge tree represents a series of document IDs, which are organized according to the order of requests.

[0107] Each node stores a key-value tensor for a document, and different paths may share the same nodes, representing documents shared across different requests.

[0108] Since the tree structure is prefix-based, requests with the same prefix will share part of the path, thereby reducing memory and computational redundancy and improving cache reuse efficiency.

[0109] The knowledge tree enables the cache system to efficiently organize and access the key-value tensors of documents according to the request order, thereby improving data retrieval speed and reducing document access latency. It is particularly suitable for scenarios that require sequential access based on document IDs, such as LLM reasoning and external knowledge retrieval tasks in the RAG (Retrieval-Augmented Generation) model.

[0110] The Monte Carlo Tree Search Scheduler (MCTS Scheduler) is a task scheduling system for managing and optimizing the Monte Carlo Tree Search (MCTS) algorithm. It is responsible for scheduling and controlling the execution of MCTS in large model reasoning, especially when multiple decision paths need to be explored and evaluated. The MCTS Scheduler coordinates the execution order and resource allocation of different computing tasks to improve the overall efficiency and reasoning speed of the algorithm. It needs to be used in conjunction with PRM Strategy and ORM Strategy.

[0111] The MCTS algorithm iteratively selects, expands, simulates, and returns in the decision tree to find the optimal decision path. The MCTS scheduler is responsible for scheduling these steps according to priority to ensure that computing resources are used reasonably.

[0112] The scheduler uses heuristic strategies or priority-based approaches to select which nodes need to be expanded, simulated, and backhauled to minimize the amount of computation and improve decision-making efficiency.

[0113] In complex reasoning tasks, the MCTS scheduler can also dynamically adjust the expansion strategy of the decision tree according to the different stages of the reasoning process, thereby optimizing the calculation of each step and maximizing the reasoning performance.

[0114] The MCTS scheduler enables Monte Carlo tree search to efficiently explore and optimize in large-scale decision spaces, especially when resources are limited, so that computing resources can be reasonably allocated and reasoning latency can be reduced. In addition, it can be combined with other optimization strategies (such as dynamic cache management, reward models, etc.) to further improve the performance and accuracy of reasoning tasks.

[0115] The disclosed embodiment designs a unified separation architecture for KV-Cache from different sources during large model reasoning, which facilitates unified management and optimization of KV-Cache, which is specifically reflected in the following three aspects:

[0116] Unified cache control system: A standardized Inference Time Controller component is proposed to manage the KV-Cache generated during large model inference;

[0117] Unified dynamic cache management mechanism: The Knowledge Tree / MCTS Scheduler component is proposed to uniformly manage and build the KV-Cache Tree generated during RAG and large model reasoning;

[0118] Algorithm compatibility for optimizing model inference efficiency: Based on the resource utilization characteristics of different usage scenarios, the algorithm compatibility is strong. The RAG Controller and Inference Time Controller components are compatible with different scheduling strategies.

[0119] The disclosed embodiment solves the GPU memory limitation problem of KV-Cache management in the large model RAG system and the large model o1 reasoning system through an innovative KV-Cache separation architecture and dynamic management mechanism, and improves reasoning efficiency and GPU utilization. In particular, when combining MCTS, PRM, ORM and RAG technologies, it can ensure efficient execution of reasoning tasks and optimal utilization of resources.

[0120] The separation architecture and dynamic cache management mechanism of the disclosed embodiment can reduce the generation of duplicate cache data and improve the utilization rate of GPU memory, especially when complex computing tasks are performed in parallel during the inference process, which can avoid cache overflow.

[0121] The disclosed embodiment introduces components such as Inference Time Controller and RAG Controller to achieve multi-task collaborative reasoning and enhance the system's adaptability to diverse reasoning tasks. In the reasoning process combining the RAG model, MCTS search, and process reward and result reward models, cache data can be effectively managed while maintaining high efficiency and accuracy.

[0122] The system architecture of the disclosed embodiment reduces cache misses and resource waste by optimizing cache access and release strategies, thereby improving the overall efficiency of reasoning, and ensuring high performance output of the system, especially under large-scale reasoning and high-load tasks.

[0123] The disclosed embodiments are compatible with different reasoning stages and optimization algorithms, provide flexible configuration and expansion capabilities, and meet the complex requirements for resource and task scheduling in large model reasoning.

[0124] The disclosed embodiment aims at the limited GPU memory and cache management problems in the large-scale model reasoning process. By separating the KV-Cache in the Prefill and Decoding stages and designing a flexible dynamic management mechanism, it can effectively avoid repeated cache calculations and memory waste, thereby improving the utilization of GPU resources. This architecture is not only compatible with various optimization algorithms (such as MCTS, PRM, ORM), but can also be combined with complex reasoning models such as RAG to achieve efficient scheduling and management of memory and computing resources, further improving the efficiency and performance of large model reasoning. The core components and core steps of the invention design involve the optimization of KV-Cache management in the large model reasoning process, especially the effective use of GPU memory and cache reuse, and how to support efficient reasoning under different reasoning technologies (such as RAG, MCTS, PRM, ORM).

[0125] Inference Time Controller: Responsible for coordinating cache management during the inference process, dynamically adjusting cache release and retention, and optimizing the utilization of computing resources for inference. It adjusts cache usage during the inference phase by combining different inference technologies (PRM, ORM, etc.), reduces GPU memory waste, and improves inference efficiency.

[0126] RAG Controller (Retrieval Augmented Generation Controller): Mainly responsible for managing and optimizing the cache strategy in the Retrieval Augmented Generation (RAG) process. According to factors such as cache access frequency, key-value tensor size, and cache access time, it drives cache storage and release with priority to reduce cache misses. It supports the parallelization of knowledge retrieval and reasoning tasks, and optimizes the combination of retrieval and generation processes.

[0127] Cache Retriever: Efficiently locates and retrieves key-value tensors in the memory cache through the knowledge tree, reducing memory and computational redundancy during the retrieval process.

[0128] Cache Controller: manages cache allocation, access, update, and invalidation operations. It uses invalidation policies (such as LRU and LFU) to determine which cache data should be removed to make room for new data. During the inference process, it determines which data should be cached and which should be removed to ensure the efficient operation of the cache system.

[0129] Knowledge Tree: Stores and retrieves key-value tensors of documents based on a prefix tree structure, ensuring fast location of cached data for specific documents.

[0130] MCTS Scheduler: Responsible for scheduling the execution of the Monte Carlo Tree Search (MCTS) algorithm, improving decision efficiency by controlling the expansion, simulation, and backpropagation process of the decision tree, especially for optimizing decision paths in complex reasoning tasks.

[0131] PRM / ORM Inference Engines (PRM / ORM Inference Engines) PRM (Process Reward Model) and ORM (Outcome Reward Model): The inference engine is responsible for optimizing the behavior and results of the inference process and improving the quality of inference through a reward mechanism.

[0132] LLM Inference Engine (Large Language Model Inference Engine): Specially designed to optimize the inference performance of large language models (such as GPT, etc.) and handle computing tasks during the inference process.

[0133] In some embodiments, the embodiments of the present disclosure generate and cache the required KV data (i.e., Key-Value pairs) in the Prefill stage. These caches are used to store the context information of the previous token. In the Decoding stage, these caches are reused to avoid repeated calculations, thereby saving GPU memory. Through dynamic cache management strategies, cache data can be flexibly released and retained according to computing needs to optimize memory utilization.

[0134] In some embodiments, the cache data of the disclosed embodiments is stored in In-Memory KV-Cache, and the required data is quickly located in the cache using Cache Retriever. The Cache Controller schedules the update, access, and invalidation of the cache to ensure that the most commonly used data is retained in the cache, reducing the probability of misses.

[0135] In some embodiments, in the RAG framework of the disclosed embodiments, the model first performs information retrieval, then combines the retrieved relevant information with the original input and hands it over to the large model for generation. The RAG Controller optimizes cache access, reduces cache conflicts and misses, and ensures efficient integration of information retrieval and generation processes.

[0136] In some embodiments, the MCTS scheduler of the disclosed embodiment is responsible for managing the calculation of the Monte Carlo tree search, optimizing the search path in the decision tree, and improving the decision efficiency of the reasoning process. The scheduler coordinates the execution order and resource allocation of different computing tasks to reduce the amount of calculation and improve the reasoning efficiency.

[0137] In some embodiments, the disclosed embodiments use PRM (Process Reward Model) and ORM (Outcome Reward Model) to gradually optimize the decision steps and final output results during the reasoning process, thereby improving the quality of the generated tasks.

[0138] In some embodiments, the Inference Time Controller and RAG Controller of the disclosed embodiments are combined to dynamically adjust the cache management strategy and the inference algorithm at different stages of the inference process to ensure efficient use of memory and computing resources and avoid unnecessary memory waste.

[0139] In some embodiments, the embodiments of the present disclosure optimize KV-Cache management to avoid repeated cache generation and improve GPU memory utilization efficiency, especially in large-scale reasoning tasks, to avoid memory overflow and waste of computing resources.

[0140] The management system proposed in the embodiment of the present disclosure can be applied to KV-Cache management in both the large model RAG system and the large model o1 reasoning system, which is specifically reflected in the following aspects: The system is based on a flexible KV-Cache management mechanism implemented based on a separation architecture, which can provide efficient memory management and computing resource scheduling strategies for a variety of reasoning tasks, ensuring efficient operation of different tasks and models on hardware resources such as GPU or TPU. The system is designed with an adjustable cache management strategy, which supports dynamic adjustment of cache access, update and release strategies according to the needs of reasoning tasks.

[0141] The disclosed embodiment combines advanced technologies such as MCTS, PRM, and ORM, which not only optimizes the reasoning performance, but also enhances the transparency and explainability of the reasoning process. In complex decision-making processes, the system can provide traceability of reasoning steps and decision paths, helping developers better understand model behavior and make adjustments.

[0142] With the continuous development of large model reasoning technology and hardware architecture, the management system proposed in the embodiment of the present disclosure has good scalability. In the future, the system can be upgraded according to new hardware platforms or optimization algorithms to further improve reasoning efficiency and system stability. In addition, the system can also support the integration of new optimization algorithms (such as quantization, distillation, speculative sampling, etc.) with the reasoning framework to provide users with more diverse reasoning solutions.

[0143] In application scenarios involving sensitive data or privacy protection, the cache management system of the disclosed embodiment can be compatible with existing data encryption and privacy protection mechanisms to ensure the security of data during the reasoning process. By encrypting and controlling access to cached data, data leakage and unauthorized access can be effectively prevented.

[0144] The term "and / or" in this disclosure is only a description of the association relationship of associated objects, indicating that there may be three relationships. For example, A and / or B can represent: A exists alone, A and B exist at the same time, and B exists alone. In addition, the character " / " in this article generally indicates that the associated objects before and after are in an "or" relationship.

[0145] Furthermore, although the steps of the method in the present disclosure are described in a particular order in the drawings, this does not require or imply that the steps must be performed in this particular order or that all the steps shown must be performed to achieve the desired results.

[0146] It should be noted that, although several modules or units of the device for action execution are mentioned in the above detailed description, such division is not mandatory.

[0147] In fact, according to the embodiments of the present disclosure, the features and functions of two or more modules or units described above can be embodied in one module or unit. Conversely, the features and functions of one module or unit described above can be further divided into multiple modules or units to be embodied.

[0148] Some of the blocks shown in the accompanying drawings are functional entities that do not necessarily correspond to physically or logically independent entities. These functional entities can be implemented in software form, or in one or more hardware modules or integrated circuits, or in different networks and / or processor devices and / or microcontroller devices.

[0149] Those skilled in the art will appreciate that all or part of the steps for implementing the above embodiments may be implemented in the following forms, namely: a complete hardware implementation, a complete software implementation (including firmware, microcode, etc.), or a combination of hardware and software implementations, which may be collectively referred to herein as "circuits", "modules" or "systems".

[0150] Other embodiments of the disclosure will be apparent to those skilled in the art from consideration of the specification and practice of the invention disclosed herein.

[0151] The present disclosure is intended to cover any variation, use or adaptation of the present disclosure, which follows the general principles of the present disclosure and includes common knowledge or customary technical means in the technical field not disclosed in the present disclosure. The description and examples are to be regarded as exemplary only, and the true scope and spirit of the present disclosure are indicated by the appended claims.

Claims

1. A cache management system for large model inference, based on a pre-filled instance and a decoded instance separation architecture, characterized in that: include: a cache manager, including a cache retriever and a cache controller connected to each other; The cache retriever is connected to a cache module, and the cache module is used to store key-value pairs in large-scale documents; The cache controller is connected to the retrieval enhancement generation controller and the reasoning time controller; The retrieval enhancement generation controller is connected to an external database and a large language model reasoning engine; The reasoning time controller includes a process reward model strategy and a result reward model strategy; The process reward model strategy is connected to the process reward model inference engine, and the result reward model strategy is connected to the result reward model inference engine; the inference time controller is also connected to the large language model inference engine; The large language model inference engine, the process reward model inference engine and the result reward model inference engine are connected to a Monte Carlo tree search scheduler; The cache module is used to cache required key-value pair data in the pre-filling stage, the key-value pair data is used to store context information of the preceding word element, and the key-value pair data is also used to be reused in the decoding stage.

2. The system according to claim 1, characterized in that The cache retriever is used to locate required data in the cache module; The cache controller interacts with the retrieval enhancement generation controller and the reasoning time controller to determine the processing mode of the key-value pair data, that is, storing it in the GPU or removing it.

3. The system according to claim 2, characterized in that The cache controller is also used to, after receiving a data request, call the cache retriever to check whether the data requested by the data request already exists in the cache module. If the cache hits, the data requested by the data request is returned; if the cache misses, the data requested by the data request is loaded from a computing process or an external data source into the cache module.

4. The system according to claim 2, characterized in that The cache controller updates the data in the cache module according to the policy; The cache controller determines the data that needs to be removed from the cache module according to the invalidation strategy.

5. The system according to claim 2, characterized in that When there are multiple cache requests, the cache controller determines the order in which data is stored in the cache module through priority scheduling.

6. The system according to claim 2, characterized in that The cache controller also loads the target data into the cache in advance based on a pre-fetch strategy.

7. The system according to claim 1, characterized in that The search enhancement generation controller calculates the priority according to the frequency of accessing the key-value pair data, the size of the key-value tensor, the last access time, and the prefix-aware recalculation cost; The cache controller determines the data to be removed from the cache module according to the invalidation strategy and the priority.

8. The system according to claim 7, characterized in that The retrieval enhancement generation controller is also used to combine the prediction and reservation strategies to enable knowledge retrieval and LLM reasoning to be performed in an overlapping manner.

9. The system according to claim 1, characterized in that The inference time controller coordinates the various stages of model inference by optimizing the process reward model strategy and the result reward model strategy.

10. The system according to claim 1, characterized in that The cache management system during the large model reasoning also uses the document identifier as the node identifier and stores the document key value tensor as a leaf node to form a multi-branch tree structure; The Monte Carlo tree search scheduler selects nodes in the decision tree that need to be expanded, simulated, or returned using a heuristic strategy or a priority-based approach.

Images