Key-value cache data scheduling method and device, and large model inference method and device

Abstract

The application provides a key-value cache data scheduling method and device and a large model inference method, and relates to the technical field of artificial intelligence. The method comprises the following steps: predicting target key-value cache data required for subsequent token inference of a large model by using a prediction model; the subsequent token refers to a token that has not been inferred by the large model; determining whether the target key-value cache data needs to be scheduled; if the target key-value cache data needs to be scheduled, a scheduling strategy is generated; the target key-value cache data is obtained from a first storage space according to the scheduling strategy, and the target key-value cache data is stored in a second storage space; wherein the target key-value cache data is used to make the large model infer the subsequent token. The application reduces frequent access to the first storage space, and only the required target key-value cache data is transmitted each time, rather than all key-value cache data, thereby reducing the occupation of bandwidth resources caused by the transmission of key-value cache data.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence, in particular to a key-value cache data scheduling method and a large model inference method and device. BACKGROUND

[0002] In recent years, with the wide application of large language models (LLMs) in search engines, intelligent dialogue systems, real-time translation and other scenarios, the high latency and high cost problems of LLM inference services have become increasingly prominent. The inference process of a large model usually includes a prefill stage and a decoding stage:

[0003] The prefill stage is a computationally intensive task that requires the input sequence to be encoded into key-value cache data (KV cache) at one time, occupying a large amount of GPU memory (e.g., a 32B model requires 12GB of KV cache at a context length of 50K).

[0004] The decoding stage is a process of generating tokens in series, and the performance bottleneck is concentrated in the data transfer between the GPU high-bandwidth memory (HBM) and the computing unit. Since the generation of each token depends on the historical KV cache, repeated reading and writing of HBM in a long context scenario will consume a large amount of bandwidth resources, causing waste of computing power. SUMMARY

[0005] The purpose of the embodiments of the present application is to provide a key-value cache data scheduling method and device, electronic equipment and storage medium, which can reduce the consumption of bandwidth resources by scheduling key-value cache data.

[0006] In a first aspect, the embodiments of the present application provide a key-value cache data scheduling method, comprising:

[0007] using a prediction model to predict target key-value cache data required for a subsequent token of a large model inference; the subsequent token refers to a token that has not been inferred by the large model;

[0008] determining whether the target key-value cache data needs to be scheduled;

[0009] if the target key-value cache data needs to be scheduled, generating a scheduling strategy;

[0010] obtaining the target key-value cache data from a first storage space according to the scheduling strategy, and storing the target key-value cache data in a second storage space; wherein the target key-value cache data is used to infer the subsequent token of the large model; the first storage space is used to store full key-value cache data; and the large model reads the required key-value cache data from the second storage space during the inference process.

[0011] The embodiment of the present application predicts the target key-value cache data required by the subsequent token of the large model inference through a prediction model, and performs scheduling judgment, scheduling strategy generation, data migration and the like on the target key-value cache data, reduces the frequent access to the first storage space, and only transmits the required target key-value cache data, rather than the full amount of key-value cache data, greatly reducing the occupation of bandwidth resources caused by the transmission of key-value cache data.

[0012] In any embodiment, the prediction model is used to predict the target key-value cache data required by the subsequent token of the large model inference, including:

[0013] The predicted tokens, the current token, and the position information of the current token are input into the prediction model to obtain the target key-value cache data required by the subsequent token inference output by the prediction model; wherein the current token refers to the token being inferred by the large model.

[0014] The embodiment of the present application can capture context information by analyzing the predicted tokens, the current token, and the position information of the current token using the prediction model, thereby more accurately predicting the target key-value cache data required by the subsequent token.

[0015] In any embodiment, whether the target key-value cache data needs to be scheduled is judged, including:

[0016] Whether the second storage space stores the target key-value cache data is judged.

[0017] If the second storage space does not store the target key-value cache data, it is determined that scheduling is needed; otherwise, scheduling is not needed.

[0018] The embodiment of the present application determines whether scheduling of the target key-value cache data is needed by judging whether the target key-value cache data is stored in the second storage space, thereby reducing unnecessary scheduling operations and reducing the cost and time of data migration.

[0019] In any embodiment, the scheduling strategy is generated, including:

[0020] The target key-value cache data, the key-value cache data already stored in the second storage space, and the environment characteristics are input into the scheduling model to obtain the scheduling strategy output by the scheduling model; wherein the environment characteristics include at least one of high bandwidth memory (HBM) capacity, HBM occupancy, dynamic random access memory (DRAM) capacity, DRAM occupancy, central processing unit (CPU) direct read-write bandwidth, graphics processing unit (GPU) direct read-write bandwidth, CPU occupancy, and GPU occupancy.

[0021] In the embodiments of the present application, the scheduling model can consider the current environmental characteristics, thereby more effectively allocating storage and computing resources and reducing resource waste; and can be applied to a heterogeneous system of different types of memories and processors to fully utilize the overall performance of the system.

[0022] In any embodiment, the scheduling model is obtained by training based on a reinforcement learning algorithm.

[0023] In the embodiments of the present application, the reinforcement learning algorithm can enable the scheduling model to self-adjust and optimize according to the feedback of the environment, and through continuous learning, can consider the occupation and performance characteristics of multiple resources, thereby making more reasonable resource allocation strategies.

[0024] In any embodiment, before storing the target key-value cache data into the second storage space, the method further comprises:

[0025] If the second storage space is insufficient, part of the key-value cache data is deleted from the second storage space.

[0026] In the embodiments of the present application, since the second storage space is limited, when the second storage space is insufficient, in order to be able to store the target key-value cache data required for subsequent reasoning of tokens, part of the key-value cache data in the second storage space is deleted to meet the needs of subsequent reasoning.

[0027] In any embodiment, deleting part of the key-value cache data from the second storage space comprises:

[0028] Excluding the key-value cache data corresponding to the first preset number of tokens in the reasoning token sequence in the second storage space to obtain the excluded key-value cache data;

[0029] From the excluded key-value cache data, the key-value cache data that enters the second storage space earliest is deleted according to the size of the target key-value cache data.

[0030] In the embodiments of the present application, by excluding the key-value cache data corresponding to the first preset number of tokens in the reasoning token sequence, the most relevant key-value cache data in the reasoning process is retained, and the key-value cache data that enters the second storage space earliest is deleted, so that the second storage space is effectively utilized.

[0031] In any embodiment, the first storage space is a storage space in the CPU, and the second storage space is a storage space in the GPU.

[0032] In a second aspect, the embodiments of the present application provide a large model reasoning method, comprising:

[0033] obtaining target key-value cache data of the to-be-reasoned token through the large model; the target key-value cache data is obtained by using the scheduling method of the key-value cache data according to the first aspect;

[0034] reasoning the to-be-reasoned token based on the target key-value cache data of the second storage space through the large model.

[0035] The embodiments of the present application predict the target key-value cache data required by the large model to reason the subsequent token through the prediction model, and perform scheduling judgment, scheduling strategy generation, data migration and other operations on the target key-value cache data, thereby reducing frequent access to the first storage space and improving reasoning efficiency.

[0036] In any embodiment, the reasoning process of the large model and the scheduling process of the target key-value cache data are executed in an asynchronous manner.

[0037] The embodiments of the present application execute the reasoning of the large model and the scheduling of the target key-value cache data in an asynchronous manner, predict and schedule the target key-value cache data required by the subsequent token during the reasoning process of the large model, and directly use the target key-value cache data when the large model reasons the subsequent token, without temporarily scheduling the target key-value cache data, thereby improving the reasoning efficiency of the large model.

[0038] In a third aspect, the embodiments of the present application provide a key-value cache data scheduling device, comprising:

[0039] a prediction module configured to predict target key-value cache data required by a large model to reason a subsequent token by using a prediction model; the subsequent token is a token that has not been reasoned by the large model;

[0040] a judgment module configured to judge whether the target key-value cache data needs to be scheduled;

[0041] a scheduling strategy generation module configured to generate a scheduling strategy if the target key-value cache data needs to be scheduled;

[0042] a scheduling module configured to obtain the target key-value cache data from a first storage space according to the scheduling strategy, and store the target key-value cache data to a second storage space; the target key-value cache data is used to make the large model reason the subsequent token; the first storage space is used to store full key-value cache data; and the large model reads required key-value cache data from the second storage space during a reasoning process.

[0043] In a fourth aspect, the embodiments of the present application provide a large model reasoning device, comprising:

[0044] The acquisition module is configured to acquire target key-value cache data of the token to be inferred by using the large model, wherein the target key-value cache data is obtained by using the scheduling method of the key-value cache data according to the first aspect;

[0045] The inference module is configured to infer the token to be inferred based on the target key-value cache data of the second storage space by using the large model.

[0046] In a fifth aspect, an electronic device is provided, including a processor, a memory and a bus, wherein:

[0047] The processor and the memory complete mutual communication through the bus;

[0048] The memory stores program instructions executable by the processor, and the processor calling the program instructions can execute the method of the first aspect or the second aspect.

[0049] In a sixth aspect, an electronic device is provided, including a processor, a memory and a bus, wherein:

[0050] The non-transitory computer readable storage medium stores computer instructions, and the computer instructions make the computer execute the method of the first aspect or the second aspect.

[0051] In a seventh aspect, a computer program product is provided, including computer program instructions, and the computer program instructions are read and run by a processor to execute the method of the first aspect.

[0052] Other features and advantages of the present application will be described in the following description, and some will become apparent from the description, or will be understood through implementation of the embodiments of the present application. The purposes and other advantages of the present application can be achieved and obtained by the structures specifically pointed out in the written description, claims, and drawings. BRIEF DESCRIPTION OF DRAWINGS

[0053] In order to more clearly illustrate the technical solutions of the embodiments of the present application, the following will briefly introduce the drawings needed to be used in the embodiments of the present application. It should be understood that the following drawings only show some embodiments of the present application, and therefore should not be regarded as limiting the scope. For those skilled in the art, other related drawings can also be obtained without creative labor on the basis of these drawings.

[0054] Figure 1 A key-value cache data scheduling method flowchart provided by the embodiments of the present application;

[0055] Figure 2 A large model inference method flowchart provided by the embodiments of the present application;

[0056] Figure 3 A large model inference architecture diagram provided for an embodiment of the present application;

[0057] Figure 4 A prefill stage schematic diagram provided for an embodiment of the present application;

[0058] Figure 5 A decoding stage schematic diagram provided for an embodiment of the present application;

[0059] Figure 6 A key-value cache data scheduling device structure schematic diagram provided for an embodiment of the present application;

[0060] Figure 7 A large model inference device structure schematic diagram provided for an embodiment of the present application;

[0061] Figure 8 An electronic device entity structure schematic diagram provided for an embodiment of the present application. DETAILED DESCRIPTION

[0062] The embodiments of the technical solutions of the present application will be described in detail below with reference to the accompanying drawings. The following embodiments are only used to more clearly illustrate the technical solutions of the present application, and therefore only serve as examples, and cannot limit the protection scope of the present application.

[0063] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as commonly understood by one of ordinary skill in the art to which the present application belongs; the terms used herein are only for the purpose of describing specific embodiments and are not intended to limit the present application; the terms "include" and "have" and any variations thereof in the specification and claims of the present application and the above description of drawings are intended to cover non-exclusive inclusion.

[0064] In the description of the embodiments of the present application, the technical terms "first", "second", etc. are only used to distinguish different objects, and cannot be understood as indicating or implying relative importance or implicitly indicating the number, specific order or primary and secondary relationship of the indicated technical features. In the description of the embodiments of the present application, the meaning of "a plurality of" is two or more, unless otherwise explicitly and specifically limited.

[0065] Reference herein to "an embodiment" means that a particular feature, structure, or characteristic described in connection with the embodiment can be included in at least one embodiment of the present application. The appearance of the phrase in various places in the specification does not necessarily all refer to the same embodiment, nor is it necessarily independent or alternative to other embodiments. It is explicitly and implicitly understood by those skilled in the art that the embodiments described herein can be combined with other embodiments.

[0066] In the description of the embodiments of the present application, the term "and / or" is only to describe the association relationship of the associated objects, which means that there can be three relationships, for example, A and / or B, which can represent the three cases of A alone, A and B together, and B alone. In addition, the character " / " in this paper generally represents that the front and rear associated objects are a "or" relationship.

[0067] In the description of the embodiments of the present application, the term "a plurality of" refers to two or more (including two), and similarly, "a plurality of groups" refers to two or more groups (including two groups), and "a plurality of pieces" refers to two or more pieces (including two pieces).

[0068] In the description of the embodiments of the present application, unless otherwise explicitly specified and limited, the technical terms "mounting", "connection", "connection", "fixing" and the like should be understood in a broad sense, for example, it can be fixedly connected, or it can be detachably connected, or it can be integrated; it can be mechanical connection, or it can be electrical connection; it can be directly connected, or it can be indirectly connected through an intermediate medium; it can be the internal communication of two elements or the interaction relationship between two elements. For ordinary skilled in the art, the specific meaning of the above terms in the embodiments of the present application can be understood according to the specific circumstances.

[0069] With the wide application of large models in online services such as search engines, chat robots and virtual assistants, the high operating cost of large models has become a significant obstacle. Large model inference service includes prefill stage and decoding stage. Among them, the prefill stage is a power-intensive task, and the decoding stage is used to generate tokens in series. The performance bottleneck is the bandwidth between HBM and GPU. The decoding stage affects the inference delay between generated tokens, which is a key indicator of inference performance TBT (time between tokens). Especially in the case of long context, such as 32B model single user inference scenario, 50K context length requires 12GB key value cache data (KV cache) demand. If multiple users need to use large models for inference at the same time, the required KV cache will increase greatly.

[0070] Because in the decoding stage of large model inference, all KV cache needs to be read from CPU to GPU for each token inference, the hardware computing power and bandwidth are relatively large, causing waste of computing resources. In subsequent research, it is found that in the decoding process of large model inference service, not all token information is needed. Microsoft's latest work now has a certain pattern for context token information, such as A-shape, block-shape, etc.

[0071] A-shape (A shape) represents a specific sparse pattern, in which the distribution of attention presents a shape similar to the letter "A". Specifically, this pattern may exhibit a trend in some heads, that is, the attention is mainly concentrated in a central area and gradually weakens from the central area to both sides. This reflects that when processing long context, the model may pay more attention to certain key information or specific parts in the context.

[0072] block-shape (block shape) represents another sparse pattern, in which attention is concentrated in several block-shaped areas. In this mode, the model may focus attention on certain specific blocks or paragraphs of text when processing the text.

[0073] Based on this, the embodiments of the present application provide a key value cache data scheduling method and a large model inference method and device. The prediction model is used to predict the target key value cache data (hereinafter referred to as target KV cache) required by the large model for subsequent token inference. If the target KV cache is not in the second storage space, the target KV cache is scheduled from the first storage space storing the full KV cache, and the target KV cache is stored to the second storage space, so that the large model can read the target KV cache from the second storage space when inferring the subsequent token. In the scheduling process, only the target KV cache is scheduled, not the full KV cache, reducing the bandwidth resources occupied by data scheduling.

[0074] The specific embodiments will be described below:

[0075] It should be noted that the key value cache data scheduling method provided by the embodiments of the present application can be applied to an electronic device, which includes a terminal and a server; the terminal can be specifically a smartphone, a tablet computer, a computer, a personal digital assistant (Personal Digital Assitant, PDA), etc.; the server can be specifically an application server or a Web server. The electronic device runs the large model, the prediction model and the scheduling model mentioned in the following embodiments.

[0076] Figure 1 A flowchart of a key value cache data scheduling method provided by an embodiment of the present application is shown in FIG. 1, which includes the following steps: Figure 1

[0077] Step 101: predicting the target key value cache data required by the large model for subsequent token inference by using the prediction model; the subsequent token refers to the token that has not been inferred by the large model.

[0078] ​The prediction model is obtained by pre-training, and its main function is to predict the target KV cache required by the token to be inferred by the large model. The prediction model can include a pooling layer and a fully connected layer, etc. In the training process, the data of the commonly used KV cache sparsification scheme can be used for small model pre-training; for example, the token sequence data required by the large model inference service process, such as Minference, is generated. It can be understood that the prediction model can be deployed on the CPU of the electronic device, or can be deployed on the GPU of the electronic device.

[0079] A large model (LLM) refers to a deep learning model trained using a large amount of text data, so that the model can generate natural language text or understand the meaning of language text. These models can provide in-depth knowledge and language production on various topics by training on massive datasets. The core idea is to learn the patterns and structures of natural language through large-scale unsupervised training, to a certain extent, to simulate the language cognition and generation process of human beings. Commonly used large models include: GPT series, BERT series, and general-purpose large models. The embodiments of the present application do not limit the specific type of large model.

[0080] The subsequent token refers to the token that has not been inferred by the large model at the current time, for example: the current token is being inferred by the large model, then the subsequent token can be the next token to be inferred by the large model, and can also be the next two tokens of the current token. For ease of description, the embodiments of the present application take the subsequent token as the next token of the current inferred token.

[0081] Step 102: Determine whether the target key-value cache data needs to be scheduled.

[0082] After obtaining the target KV cache, it is determined whether the target KV cache needs to be scheduled. The main determination method is: determining whether the second storage space contains the target KV cache. If it contains, it is determined that scheduling is not needed; otherwise, it is determined that scheduling is needed. Among them, for the case that the second storage space only contains part of the target KV cache, it is also considered not to contain the target KV cache, and scheduling is needed. Only in the subsequent scheduling process, only the part of the target KV cache that the second storage space does not contain is stored in the second storage space. For example: if the target KV cache includes the KV cache corresponding to the first 10 tokens, the KV cache corresponding to the 30th-50th token, and the KV cache corresponding to the 100th-110th token. The second storage space currently stores the KV cache corresponding to the first 10 tokens and the KV cache corresponding to the 30th-50th token, and does not store the KV cache corresponding to the 100th-110th token. It is determined that scheduling is needed. When scheduling, the KV cache corresponding to the 100th-110th token is stored in the second storage space.

[0083] It can be understood that the second storage space is a storage space of the target KV cache that needs to be read by the large model in the inference process. Specifically, it can be HBM in the GPU, or other storage spaces in the GPU. The embodiments of the present application do not make specific limitations thereto. In addition, the prediction model can output an index value of the target KV cache, and use the index value to find whether there is a KV cache corresponding to the same index value in the second storage space, so as to determine whether scheduling is needed. And when scheduling is needed, the corresponding KV cache can also be obtained from the first storage space based on the index value, and then stored in the second storage space.

[0084] Step 103: If scheduling is needed, a scheduling strategy is generated.

[0085] The scheduling strategy is used to indicate the way in which the electronic device schedules the target KV cache from the first storage space to the second storage space, for example: how much data block is scheduled, whether the KV cache in the second storage space needs to be deleted, and how many KV caches are deleted, etc. When generating the scheduling strategy, the current bandwidth and other environmental information on the electronic device can be considered.

[0086] Step 104: According to the scheduling strategy, the target key value cache data is obtained from the first storage space, and the target key value cache data is stored in the second storage space; wherein the target key value cache data is used to make the large model infer the subsequent token.

[0087] Wherein, after determining the scheduling strategy, the target KV cache is obtained from the first storage space according to the scheduling strategy, and the target KV cache is stored to the second storage space. The first storage space can be DRAM on the CPU, and can also be other memories, and the embodiments of the present application do not make specific limitation thereto. In the prefll stage of the large model, the full amount of KV cache is stored in the first storage space, and when the target KV cache is scheduled, the target KV cache can be migrated from the first storage space to the second storage space, and after migration, the first storage space will no longer contain the target KV cache; or the target KV cache can be obtained from the first storage space, and then stored in the second storage space, at this time, the first storage space still contains the target KV cache.

[0088] The embodiments of the present application predict the target key value cache data required by the subsequent token of the large model inference through the prediction model, and perform scheduling judgment, scheduling strategy generation and data migration and other operations on the target KV cache, reduce the frequent access to the first storage space, and only transmit the required target key value cache data each time, rather than the full amount of key value cache data, greatly reducing the occupation of bandwidth resources caused by transmitting key value cache data.

[0089] On the basis of the above embodiments, the prediction model is used to predict the target key value cache data required by the subsequent token of the large model inference, comprising:

[0090] The already-inferred tokens, the current token and the position information of the current token are input into the prediction model to obtain the target key value cache data required by the subsequent token inference output by the prediction model; wherein the current token refers to the token being inferred by the large model.

[0091] In the specific implementation process, in the large model inference process, in order to improve the inference efficiency, KV cache is usually used to store the hidden state or intermediate result that has been calculated, so as to avoid repeated calculation.

[0092] The already-inferred tokens refer to the tokens that have been inferred by the large model, which can be in the form of token sequence, and the hidden state or intermediate result of these tokens has been stored in the KV cache.

[0093] The current token refers to the token being inferred by the large model, and the position information of the current token refers to the position of the token being inferred by the large model in the inference service sequence, which can be represented by an index value index, used for the prediction model to understand the order relationship of the token.

[0094] The reasoned tokens, the current token, and the position information of the current token are combined into a complete input sequence, and input into a prediction model, and the prediction model generates target KV cache required by the large model to reason for the next token according to the input content.

[0095] By analyzing the reasoned tokens, the current token, and the position information of the current token by using the prediction model, the context information can be captured, so that the target key value cache data required by the subsequent token can be more accurately predicted.

[0096] On the basis of the above-mentioned embodiments, a scheduling strategy is generated, including:

[0097] The target key value cache data, the key value cache data stored in the second storage space, and the environment characteristics are input into a scheduling model to obtain a scheduling strategy output by the scheduling model; wherein the environment characteristics include at least one of high bandwidth memory (HBM) capacity, HBM occupancy, dynamic random access memory (DRAM) capacity, DRAM occupancy, central processing unit (CPU) direct read-write bandwidth, graphics processing unit (GPU) direct read-write bandwidth, CPU occupancy, and GPU occupancy.

[0098] In the specific implementation process, the embodiments of the present application introduce a scheduling model, take the target KV cache and the environment characteristics as inputs, and output a more optimal scheduling strategy to guide the reading, writing, migration, and elimination of the KV cache and the like.

[0099] The target KV cache is obtained by the prediction model in the above-mentioned embodiments. The environment characteristics refer to the hardware and software environment characteristics of the electronic device, including at least one of the following:

[0100] High bandwidth memory (HBM) capacity: the total capacity of the HBM, reflecting how much data the system can store.

[0101] HBM occupancy: the current usage of the HBM, used to evaluate the remaining space of the HBM.

[0102] Dynamic random access memory (DRAM) capacity: the total capacity of the DRAM, also reflecting the storage capacity.

[0103] DRAM occupancy: the current usage of the DRAM, used to evaluate the remaining space of the DRAM.

[0104] CPU direct read-write bandwidth: the rate of CPU accessing memory (including HBM and DRAM), affecting the data processing speed.

[0105] GPU Direct Read / Write Bandwidth: The rate at which the GPU accesses memory, which is particularly important for tasks such as graphics processing and deep learning.

[0106] CPU Utilization: The current workload of the CPU, reflecting the use of its processing capacity.

[0107] GPU Utilization: The current workload of the GPU, also reflecting the use of its processing capacity.

[0108] The target key-value cache data refers to the KV cache required for the next token of large model inference. The stored tokens in the second storage space are used to determine whether the target KV cache has been stored in the second storage space, thereby determining whether KV cache scheduling is needed.

[0109] The prepared target KV cache, the stored KV cache in the second storage space, and the environmental characteristics are input into the scheduling model. The scheduling model can be a machine learning, deep learning, or reinforcement learning-based model that learns from the input data and outputs the optimal scheduling strategy. That is, it first determines whether KV cache scheduling is needed, and if so, how to schedule the KV cache.

[0110] The scheduling model can include multiple processing layers, such as input layers, hidden layers, and output layers, to process input data and generate output through complex network structures and algorithms.

[0111] The scheduling strategy output by the scheduling model can include the read-write order of key-value cache data, migration strategies, and eviction rules. The scheduling strategy aims to optimize system performance, such as reducing access latency, improving data throughput, or balancing resource usage. In the embodiments of the present application, the scheduling model can consider the current environmental characteristics, thereby more effectively allocating storage and computing resources and reducing resource waste. It can also be applied to heterogeneous systems with different types of storage and processors to fully utilize the overall performance of the system.

[0112] Based on the above embodiments, the scheduling model can use an adaptive reinforcement learning framework, where adaptability refers to the reinforcement learning scheme that can adapt to different hardware resources and different large models. In the embodiments of the present application, the scheduling model can be fine-tuned after a reinforcement learning to adapt to different hardware resources. When deploying the scheduling model on different hardware, the model can be fine-tuned to optimize its performance.

[0113] The core idea of reinforcement learning is that an agent learns an optimal policy through interaction with the environment. The agent selects an action in a given state and observes the reward or punishment given by the environment, so as to adjust its policy in order to obtain greater cumulative reward in the future. This process can be summarized as follows: the agent selects an action according to the current state, the environment returns the next state and the reward according to the action, and the agent updates the policy according to the reward.

[0114] The reinforcement learning in the embodiments of the present application is not limited to the classical value-based or policy-based reinforcement learning methods, such as Deep Q-Network, referred to as DQN; Asynchronous Advantage Actor Critic, referred to as A3C; Proximal Policy Optimization (PPO) and the like.

[0115] The reinforcement learning focuses on the representation of state space, reward function and action. The state space includes: the category of large model framework adopted, the tokens that have been reasoned, the current token, the position information and storage location of the current token (including the storage location of the tokens that have been reasoned and the current token), the HBM capacity, the HBM occupancy, the DRAM capacity, the DRAM occupancy, the CPU and GPU direct read-write bandwidth and occupancy, and the output tokens_needs (i.e. target KV cache) of the pre-trained pre-fetch small model, etc.; not limited to HBM, DRAM and other heterogeneous storage media, not limited to CPU, GPU and other heterogeneous computing accelerators. The state composition features are described as follows:

[0116] S=[Token_position,Token,whole_pre_tokens,whole_DRAM,DRAM_Usag e,whole_HBM,HBM_Usage,Read_bandwidth,Write_bandwidth,Read_Usage_bandwidth,Write_Usage_bandwidth;tokens_needs]。

[0117] wherein Token_position is position information of the current token; Token is the current token; whole_pre_tokens is the reasoned tokens; whole_DRAM is the DRAM capacity; DRAM_Usage is the DRAM occupancy rate; whole_HBM is the HBM capacity; HBM_Usage is the HBM occupancy rate; Read_bandwidth is the read bandwidth; Write_bandwidth is the write bandwidth; Read_Usage_bandwidth is the read occupancy rate; Write_Usage_bandwidth is the write occupancy rate; and tokens_needs is the target KV cache.

[0118] The action is a communication or a keeping action of data required for next token reasoning. The reward function includes: the inverse number of the loss of the current reasoning, the HBM occupancy rate, the bandwidth occupancy rate; the reward function can be expressed as:

[0119] Reward(s,index) = -loss + HBM_Usag + Read_Usag_bandwidth + Write_Usag_bandwidth

[0120] wherein Reward(s,index) is the reward value; loss is the loss value of the current reasoning; HBM_Usag is the HBM occupancy rate; Read_Usag_bandwidth is the read bandwidth occupancy rate; and Write_Usag_bandwidth is the write bandwidth occupancy rate.

[0121] The action space of the reinforcement learning defines the operations that can be performed, mainly including: a migration operation and a keeping original state action; the migration operation includes: migrating a certain data block from the HBM to the DRAM; and migrating a certain data block from the DRAM back to the HBM.

[0122] In the embodiment of the application, the reinforcement learning algorithm can enable the scheduling model to self-adjust and optimize according to the feedback of the environment, and through continuous learning, the occupancy conditions and performance characteristics of multiple resources can be considered, so that a more reasonable resource allocation strategy can be made.

[0123] On the basis of the above-mentioned embodiment, before storing the target key value cache data into the second storage space, the method further includes:

[0124] If the second storage space is insufficient, part of the key value cache data is deleted from the second storage space.

[0125] In the specific implementation process, the capacity of the second storage space is limited, and it is impossible to store data unlimitedly, therefore, in the case of needing to schedule the target key value cache data to the second storage space, if it is judged that the second storage space is insufficient, then deleting part of the KV cache from the second storage space is an effective solution. It should be explained that the judgment of whether the second storage space is insufficient can be pre-set judgment basis, for example: when the second storage space occupancy rate reaches 80%, it can be considered that the storage space is insufficient; or the remaining capacity of the second storage space is less than the size of the target KV cache to be scheduled.

[0126] When deleting part of the KV cache, the KV cache with the lowest access frequency can be deleted, which can be realized by maintaining an access counter or using the LRU (Least Recently Used) cache eviction policy. It can also be deleted based on importance or priority, for example: the KV cache in the second storage space all have their own priorities, and the KV cache with low priority is deleted when deleting.

[0127] The number of deleted KV cache also has various strategies, for example: it can be deleted according to the preset size, for example: 5K of KV cache is deleted each time, if after deletion, it still does not meet the requirement of storing the target KV cache to the second storage space, then 5K of KV cache is deleted again, until it meets the requirement of storing the target KV cache to the second storage space. It can also delete the KV cache with a corresponding size from the second storage space according to the size of the target KV cache to be scheduled.

[0128] In the embodiment of the application, since the second storage space is limited, when the second storage space is insufficient, in order to be able to store the target key value cache data required for subsequent reasoning of token, part of the key value cache data in the second storage space is deleted to meet the needs of subsequent reasoning.

[0129] On the basis of the above embodiment, when part of the key value cache data is deleted from the second storage space, the key value cache data corresponding to the first preset number of tokens in the reasoning token sequence can be excluded to obtain the excluded key value cache data; then, from the excluded key value cache data, the key value cache data that enters the second storage space earliest is deleted according to the size of the target key value cache data.

[0130] In a specific implementation process, the inference token sequence refers to index information corresponding to the context tokens, such as the 2nd token and the 9th token. The inference token sequence is [2, 9]. For example, the KV cache corresponding to the first preset number (for example, 20) of tokens obtained after tokenization of the large model is the KV cache that needs to be excluded in the embodiment of the application. Because the KV cache corresponding to the first preset number of tokens in the inference token sequence is relatively important in the inference process of the large model, it has a great influence on the inference of subsequent tokens, and therefore, this part of the KV cache needs to be retained in the second storage space. The specific value of the preset number can be 5, 10, 20, or 25, etc. It can also be a percentage value, for example, the first 1% of tokens in the token sequence, etc. The specific value can be determined according to experience. When deleting part of the KV cache in the second storage space, the KV cache corresponding to the first preset number of tokens in the inference token sequence can be excluded to obtain the excluded KV cache. Then, according to the size of the target KV cache to be scheduled, the KV cache that enters the second storage space earliest is deleted. That is, the size of the deleted KV cache is greater than or equal to the target KV cache to be scheduled. It can be understood that the target KV cache to be scheduled refers to the KV cache that needs to be stored in the second storage space, which can be all target KV caches or part of the target KV caches. This depends on whether the second storage space contains part of the target KV cache. If it contains, the part of the target KV cache that is not contained in the second storage space is stored in the second storage space.

[0131] In the embodiment of the application, by excluding the key-value cache data corresponding to the first preset number of tokens in the inference token sequence, the most relevant key-value cache data in the inference process is retained, and the key-value cache data that enters the second storage space earliest is deleted, so that the second storage space is effectively utilized.

[0132] Figure 2 A large model inference method provided in the embodiment of the application is shown in a flowchart as shown in Figure 2 The method comprises the following steps.

[0133] Step 201: obtaining target key-value cache data of tokens to be inferred through a large model; wherein the target key-value cache data is obtained by using the key-value cache data scheduling method provided in the above embodiments.

[0134] Step 202: inferring the to-be-reasoned token based on the target key-value cache data of the second storage space by the large model.

[0135] In the specific implementation process, Figure 3 The large model inference architecture diagram provided by the embodiment of the present application is shown in Figure 3 . The large model is deployed on an electronic device, which includes a CPU and a GPU. The CPU includes CPU computing cores and DRAM, wherein the full amount of KV cache is stored in the DRAM, and the related parameters of the data prefetching module can also be stored in the DRAM. It can be understood that the data prefetching module can include a prediction model and a scheduling model. The GPU includes GPU computing cores and HBM, and the weights of the large model LLM and the hot KV cache are stored in the HBM. It can be understood that the hot KV cache refers to the KV cache frequently used in the large model inference process and the KV cache required for the subsequent token inference.

[0136] The prediction model is used to predict the target KV cache required for the subsequent token inference of the large model, and the scheduling model is used to schedule the target KV cache, that is, to schedule the target KV cache from the CPU to the GPU. It should be noted that the prediction method of the prediction model can refer to the above embodiments, which will not be described here. Similarly, the scheduling strategy generated by the scheduling model and the specific method of scheduling according to the scheduling strategy can also refer to the above embodiments, which will not be described here.

[0137] The embodiment of the present application predicts the target key-value cache data required for the subsequent token inference of the large model through the prediction model, and performs scheduling judgment, scheduling strategy generation, data migration and other operations on the target key-value cache data, thereby reducing frequent access to the first storage space and improving inference efficiency.

[0138] On the basis of the above embodiments, the large model includes a prefill stage and a decoding stage in the inference process. Figure 4 The schematic diagram of the prefill stage provided by the embodiment of the present application is Figure 5 The schematic diagram of the decoding stage provided by the embodiment of the present application is shown in Figure 4 and Figure 5 . In the prefill stage, the GPU computing cores of the GPU calculate and generate full amount of KV cache based on the input of the large model, and store the full amount of KV cache in the DRAM of the CPU. In the decoding stage, the HBM of the GPU stores the hot KV cache scheduled from the CPU. It can be understood that the hot KV cache includes the target KV cache required for the subsequent token inference of the large model.

[0139] Furthermore, the scheduling of the target KV cache by the large model inference, prediction model, and scheduling model can be executed asynchronously. For example, while the large model is inferring the current token, the prediction model can predict the target KV cache required for the large model's inference of the next token. If the target KV cache is already stored in the second storage space (e.g., GPU HBM), no scheduling is needed; if the target KV cache is not in the second storage space, the scheduling model will schedule the target KV cache from the first storage space to the second storage space. When the large model infers the next token, it can directly read the target KV cache from the second storage space. At this time, the prediction model can predict the target KV cache required by the large model for inference of the token after that, and so on. Through asynchronous execution, when the large model infers a token, it can directly read the target KV cache required for inference of that token from the GPU, saving the time of reading the entire KV cache from the CPU, improving inference efficiency; and reducing the number of KV cache reads, thus reducing bandwidth consumption.

[0140] Figure 6 This is a schematic diagram of a scheduling device for key-value cache data provided in an embodiment of this application. The device can be a module, program segment, or code on an electronic device. It should be understood that this device is similar to the one described above. Figure 1 The method implementation corresponds to this and can be executed. Figure 1 The specific functions of the device involved in the method embodiment can be found in the description above; to avoid repetition, detailed descriptions are omitted here. The device includes: a prediction module 601, a judgment module 602, a scheduling strategy generation module 603, and a scheduling module 604, wherein:

[0141] Prediction module 601 is used to predict the target key-value cache data required for the large model to infer subsequent tokens using the prediction model; subsequent tokens refer to tokens that the large model has not yet inferred.

[0142] The judgment module 602 is used to determine whether the target key-value cache data needs to be scheduled;

[0143] The scheduling policy generation module 603 is used to generate a scheduling policy if scheduling is required;

[0144] The scheduling module 604 is used to obtain target key-value cache data from the first storage space according to the scheduling strategy, and store the target key-value cache data in the second storage space; wherein, the target key-value cache data is used to enable the large model to infer subsequent tokens; the first storage space is used to store the full amount of key-value cache data; during the inference process, the large model reads the required key-value cache data from the second storage space.

[0145] On the basis of the above-mentioned embodiments, the prediction module 601 is specifically configured to:

[0146] input the reasoned tokens, the current token, and the position information of the current token into the prediction model, to obtain target key-value cache data required for reasoning the subsequent token output by the prediction model; wherein the current token refers to a token being reasoned by the large model.

[0147] On the basis of the above-mentioned embodiments, the judgment module 602 is specifically configured to:

[0148] determine whether the second storage space stores the target key-value cache data;

[0149] if the second storage space does not store the target key-value cache data, it is determined that scheduling is required; otherwise, scheduling is not required.

[0150] On the basis of the above-mentioned embodiments, the scheduling strategy generation module 603 is specifically configured to:

[0151] input the target key-value cache data, the key-value cache data already stored in the second storage space, and environment features into a scheduling model, to obtain a scheduling strategy output by the scheduling model; wherein the environment features include at least one of high bandwidth memory (HBM) capacity, HBM occupancy, dynamic random access memory (DRAM) capacity, DRAM occupancy, central processing unit (CPU) direct read-write bandwidth, graphics processing unit (GPU) direct read-write bandwidth, CPU occupancy, and GPU occupancy.

[0152] On the basis of the above-mentioned embodiments, the scheduling model is obtained based on a reinforcement learning algorithm.

[0153] On the basis of the above-mentioned embodiments, the device further includes a deletion module configured to:

[0154] if the second storage space is insufficient, delete part of the key-value cache data from the second storage space.

[0155] On the basis of the above-mentioned embodiments, the deletion module is specifically configured to:

[0156] exclude, from the second storage space, key-value cache data corresponding to a pre-set number of tokens in the token sequence to obtain excluded key-value cache data;

[0157] from the excluded key-value cache data, delete key-value cache data that entered the second storage space earliest according to the size of the target key-value cache data.

[0158] Based on the above embodiments, the first storage space is the storage space in the CPU, and the second storage space is the storage space in the GPU.

[0159] Figure 7 This is a schematic diagram of a large model inference device provided in an embodiment of this application. The device can be a module, program segment, or code on an electronic device. It should be understood that this device is similar to the one described above. Figure 2 The method implementation corresponds to this and can be executed. Figure 2 The various steps involved in the method embodiment, and the specific functions of the device, can be found in the description above. To avoid repetition, detailed descriptions are appropriately omitted here. The device includes: an acquisition module 701 and an inference module 702, wherein:

[0160] The acquisition module is used to acquire the target key-value cache data of the token to be inferred through the large model; wherein, the target key-value cache data is obtained by using the key-value cache data scheduling method as described in the first aspect;

[0161] The inference module is used to perform inference on the token to be inferred by using the target key-value cache data in the second storage space based on the large model.

[0162] Based on the above embodiments, the inference process of the large model and the scheduling process of the target key-value cache data are executed asynchronously.

[0163] Figure 8 This is a schematic diagram of the physical structure of the electronic device provided in the embodiments of this application, such as... Figure 8 As shown, the electronic device includes: a processor 801, a memory 802, and a bus 803; wherein:

[0164] The processor 801 and the memory 802 communicate with each other through the bus 803;

[0165] The processor 801 is configured to invoke program instructions in the memory 802 to perform the method provided by each of the above method embodiments, for example, including: predicting, by using a prediction model, target key-value cache data required for subsequent token inference of a large model; the subsequent token refers to a token that has not been inferred by the large model; determining whether the target key-value cache data needs to be scheduled; if the target key-value cache data needs to be scheduled, generating a scheduling strategy; obtaining the target key-value cache data from a first storage space according to the scheduling strategy, and storing the target key-value cache data to a second storage space; wherein the target key-value cache data is used to make the large model infer the subsequent token; the first storage space is used to store full key-value cache data; and the large model reads required key-value cache data from the second storage space in an inference process. Alternatively, obtaining, by a large model, target key-value cache data of a token to be inferred; wherein the target key-value cache data is obtained by using the scheduling method of the key-value cache data provided in the above embodiments; and inferring, by the large model, the token to be inferred based on the target key-value cache data of the second storage space.

[0166] The processor 801 can be an integrated circuit chip having a signal processing capability. The processor 801 can be a general-purpose processor, including a central processing unit (CPU), a network processor (NP), and the like; or can be a digital signal processor (DSP), an application-specific integrated circuit (ASIC), a field programmable gate array (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components. The processor 801 can implement or execute the various methods, steps, and logical block diagrams disclosed in the embodiments of the present application. The general-purpose processor can be a microprocessor, or the processor can also be any conventional processor.

[0167] The memory 802 can include, but is not limited to, a random access memory (RAM), a read-only memory (ROM), a programmable read-only memory (PROM), an erasable programmable read-only memory (EPROM), an electrically erasable programmable read-only memory (EEPROM), and the like.

[0168] The embodiment discloses a computer program product, the computer program product comprises a computer program stored on a non-transitory computer readable storage medium, the computer program comprises program instructions, when the program instructions are executed by a computer, the computer can execute the method provided by each method embodiment, for example, comprising: predicting target key-value cache data required for large model inference of subsequent tokens by using a prediction model; the subsequent token refers to a token that has not been inferred by the large model; determining whether the target key-value cache data needs to be scheduled; if scheduling is needed, a scheduling strategy is generated; the target key-value cache data is obtained from a first storage space according to the scheduling strategy, and the target key-value cache data is stored to a second storage space; wherein the target key-value cache data is used to make the large model infer the subsequent token; the first storage space is used to store full key-value cache data; the large model reads the required key-value cache data from the second storage space in the inference process. Or, obtaining target key-value cache data of a token to be inferred by a large model; wherein the target key-value cache data is obtained by using the scheduling method of the key-value cache data provided in the above embodiment; inferring the token to be inferred by the large model based on the target key-value cache data of the second storage space.

[0169] The embodiment provides a non-transitory computer readable storage medium, the non-transitory computer readable storage medium stores computer instructions, the computer instructions make the computer execute the method provided by each method embodiment, for example, comprising: predicting target key-value cache data required for large model inference of subsequent tokens by using a prediction model; the subsequent token refers to a token that has not been inferred by the large model; determining whether the target key-value cache data needs to be scheduled; if scheduling is needed, a scheduling strategy is generated; the target key-value cache data is obtained from a first storage space according to the scheduling strategy, and the target key-value cache data is stored to a second storage space; wherein the target key-value cache data is used to make the large model infer the subsequent token; the first storage space is used to store full key-value cache data; the large model reads the required key-value cache data from the second storage space in the inference process. Or, obtaining target key-value cache data of a token to be inferred by a large model; wherein the target key-value cache data is obtained by using the scheduling method of the key-value cache data provided in the above embodiment; inferring the token to be inferred by the large model based on the target key-value cache data of the second storage space.

[0170] In the embodiments of the present application, it should be understood that the disclosed apparatus and method can be implemented in other manners. The embodiments described above are merely exemplary, for example, the division of the units is only a logical function division, and there can be another division manner in actual implementation; for example, a plurality of units or components can be combined or integrated into another system, or some features can be ignored or not executed. In addition, the displayed or discussed mutual couplings or direct couplings or communication connections can be indirect couplings or communication connections through some interfaces, devices or units, and can be in electrical, mechanical or other forms.

[0171] In addition, the units described as separate components can or can not be physically separate, and the components displayed as units can or can not be physical units, i.e., can be located in one place, or can be distributed on a plurality of network units. Some or all of the units can be selected according to actual needs to achieve the purposes of the embodiments of the present application.

[0172] In addition, the functional modules in each of the embodiments of the present application can be integrated together to form an independent part, or each module can exist independently, or two or more modules can be integrated to form an independent part.

[0173] In this document, the terms such as first and second are used only to distinguish one entity or operation from another, and do not necessarily require or imply that there is any such actual relationship or order between these entities or operations.

[0174] The above only describes the embodiments of the present application, and is not used to limit the protection scope of the present application. For those skilled in the art, various modifications and changes can be made to the present application. Any modification, equivalent replacement, improvement, etc. made within the spirit and principles of the present application shall be included in the protection scope of the present application.

Claims

1. A scheduling method for key-value cached data, characterized in that, include: The predictive model is used to predict the target key-value cache data required for the large model to infer subsequent tokens; the subsequent tokens refer to the tokens that the large model has not yet inferred. Determine whether the target key-value cache data needs to be scheduled; If scheduling is required, a scheduling policy is generated; According to the scheduling strategy, the target key-value cache data is obtained from the first storage space and stored in the second storage space; wherein, the target key-value cache data is used to enable the large model to infer the subsequent token; the first storage space is used to store the full amount of key-value cache data; during the inference process, the large model reads the required key-value cache data from the second storage space; The generation scheduling strategy includes: The target key-value cache data, the key-value cache data already stored in the second storage space, and environmental features are input into the scheduling model to obtain the scheduling strategy output by the scheduling model; wherein, the environmental features include at least one of the following: high bandwidth memory (HBM) capacity, HBM occupancy rate, dynamic random access memory (DRAM) capacity, DRAM occupancy rate, CPU direct read / write bandwidth, GPU direct read / write bandwidth, CPU occupancy rate, and GPU occupancy rate.

2. The method according to claim 1, characterized in that, The target key-value cache data required for predicting subsequent tokens using a predictive model includes: The already inferred tokens, the current token, and the location information of the current token are input into the prediction model to obtain the target key-value cache data required for the subsequent token inference output by the prediction model; wherein, the current token refers to the token that the large model is inferring.

3. The method according to claim 1, characterized in that, The determination of whether the target key-value cache data needs to be scheduled includes: Determine whether the second storage space stores the target key-value cache data; If the second storage space does not store the target key-value cache data, then scheduling is required; otherwise, scheduling is not required.

4. The method according to claim 1, characterized in that, The scheduling model was obtained by training based on a reinforcement learning algorithm.

5. The method according to claim 1, characterized in that, Before storing the target key-value cache data in the second storage space, the method further includes: If the second storage space is insufficient, some key-value cache data will be deleted from the second storage space.

6. The method according to claim 5, characterized in that, The step of deleting a portion of the key-value cache data from the second storage space includes: In the second storage space, exclude the key-value cache data corresponding to the first preset number of tokens in the reasoning token sequence to obtain the excluded key-value cache data; From the excluded key-value cache data, delete the key-value cache data that entered the second storage space earliest, according to the size of the target key-value cache data.

7. The method according to any one of claims 1-6, characterized in that, The first storage space is the storage space in the CPU, and the second storage space is the storage space in the GPU.

8. A large-scale model reasoning method, characterized in that, include: The target key-value cache data of the token to be inferred is obtained through a large model; wherein, the target key-value cache data is obtained by the key-value cache data scheduling method as described in any one of claims 1-7; The large model performs inference on the token to be inferred based on the target key-value cache data in the second storage space.

9. The method according to claim 8, characterized in that, The reasoning process of the large model and the scheduling process of the target key-value cache data are executed asynchronously.

10. A scheduling device for key-value cached data, characterized in that, include: The prediction module is used to predict the target key-value cache data required for large model inference of subsequent tokens using the prediction model; The subsequent token refers to the token that the large model has not yet reasoned about; The judgment module is used to determine whether the target key-value cache data needs to be scheduled; The scheduling policy generation module is used to generate scheduling policies if scheduling is required. A scheduling module is used to obtain the target key-value cache data from a first storage space according to the scheduling strategy, and store the target key-value cache data in a second storage space; wherein, the target key-value cache data is used to enable the large model to infer the subsequent token; the first storage space is used to store the full amount of key-value cache data; during the inference process, the large model reads the required key-value cache data from the second storage space; The scheduling strategy generation module is specifically used for: The target key-value cache data, the key-value cache data already stored in the second storage space, and environmental features are input into the scheduling model to obtain the scheduling strategy output by the scheduling model; wherein, the environmental features include at least one of the following: high bandwidth memory (HBM) capacity, HBM occupancy rate, dynamic random access memory (DRAM) capacity, DRAM occupancy rate, CPU direct read / write bandwidth, GPU direct read / write bandwidth, CPU occupancy rate, and GPU occupancy rate.

11. A large-scale model reasoning device, characterized in that, include: The acquisition module is used to acquire target key-value cache data of the token to be inferred through a large model; wherein, the target key-value cache data is obtained by using the key-value cache data scheduling method as described in any one of claims 1-7; The inference module is used to infer the token to be inferred based on the target key-value cache data in the second storage space using the large model.

12. An electronic device, characterized in that, include: Processor, memory, and bus, among which: The processor and the memory communicate with each other via the bus; The memory stores program instructions that can be executed by the processor, and the processor can execute the method as described in any one of claims 1-9 by calling the program instructions.

13. A non-transitory computer-readable storage medium, characterized in that, The non-transitory computer-readable storage medium stores computer instructions that, when executed by a computer, cause the computer to perform the method as described in any one of claims 1-9.

14. A computer program product, characterized in that, It includes computer program instructions, which, when read and executed by a processor, perform the method as described in any one of claims 1-9.

Images