Data compression method and electronic device

Abstract

The embodiments of the present application relate to the technical field of artificial intelligence, and in particular relate to a data compression method and an electronic device. The method can reduce the memory space occupied by kv cache data, thereby alleviating a memory wall problem. The method comprises: determining a first key tensor and a first value tensor to be data objects to be compressed, wherein the first key tensor comprises m first key vectors corresponding to m tokens, the first value tensor comprises m first value vectors corresponding to the m tokens, and m is a positive integer; and compressing the first key tensor and the first value tensor, so as to obtain a second key tensor and a second value tensor, wherein the second key tensor comprises n second key vectors, the second value tensor comprises n second value vectors, n is a positive integer, and n＜m; and the dimension of a single first key vector is the same as the dimension of a single second key vector, and the dimension of a single first value vector is the same as the dimension of a single second value vector.

Description

Data compression methods and electronic devices

[0001] This application claims priority to Chinese Patent Application No. 202410501112.9, filed on April 24, 2024, entitled “Data Compression Method and Electronic Device”, the entire contents of which are incorporated herein by reference. Technical Field

[0002] This application relates to the field of artificial intelligence technology, and in particular to a data compression method and an electronic device. Background Technology

[0003] In the inference process of a large language model, to avoid redundant calculation of the key-value vectors for each token in the input sequence, the first S-1 key-value vectors can be stored in memory and reused in each subsequent inference step. For example, in the inference task corresponding to the S-th token, the first S-1 key-value vectors can be directly reused. This strategy is called KV-Cache. KV-Cache can avoid redundant calculations and significantly reduce the amount of parameter computation.

[0004] Under the KV-Cache mechanism, the memory space required for the kv cache data (e.g., the first S-1 key-value vectors) increases linearly with the length of the statement to be inferred (input sequence), which leads to the memory wall problem. Summary of the Invention

[0005] This application provides a data compression method and an electronic device that compresses m key vectors and m value vectors corresponding to m tokens into n key vectors and n value vectors corresponding to n tokens, without changing the memory space occupied by a single key vector or value vector, thereby reducing the memory space resources occupied by key-value cache data.

[0006] In a first aspect, embodiments of this application provide a data compression method. First, a first key tensor and a first value tensor are determined as the data objects to be compressed. The first key tensor includes m first key vectors corresponding to m tokens; the first value tensor includes m first value vectors corresponding to m tokens; m is a positive integer. The first key tensor and the first value tensor are compressed to obtain a second key tensor and a second value tensor. The second key tensor includes n second key vectors; the second value tensor includes n second value vectors; n is a positive integer, and n < m. The dimension of a single first key vector is the same as the dimension of a single second key vector; the dimension of a single first value vector is the same as the dimension of a single second value vector. Alternatively, a single first key vector or a single first value vector occupies a first memory space, and a single second key vector or a single second value vector occupies a second memory space; the capacity of the second memory space is equal to or less than the capacity of the first memory space.

[0007] The first key tensor and the first value tensor are the tensors before compression, i.e., historical key-value data. A token represents an element, which is the smallest linguistic unit processed by the large language model, or the basic unit in text processing. This can be a single character, word, phrase, subword, or even a longer text segment. A key represents a key, and the key vector is the key vector. A value represents a value, and the value vector is the value vector.

[0008] The data compression method proposed in this application compresses the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor. This is equivalent to compressing the m first key vectors and m first value vectors corresponding to m tokens to obtain n second key vectors and n second value vectors. Before and after compression, the token dimension corresponding to the tensor is reduced from m to n, or in other words, the m vectors corresponding to m tokens are compressed into n vectors. In some embodiments, the internal dimensions of the vector corresponding to a single token remain unchanged. That is, the dimensions of the second key vector are the same as those of the first key vector, and the dimensions of the second value vector are the same as those of the first value vector. When the same data type is used for storage, the dimensions of the vectors are the same, and the amount of memory space required to store the vectors is also the same. In other words, the amount of memory space required by the first key vector is the same as that required by the second key vector, and the amount of memory space required by the first value vector is also the same as that required by the second value vector. Therefore, if n < m, the amount of memory space (second memory space) required by n second key vectors and n second value vectors is smaller than the amount of memory space (first memory space) required by m first key vectors and m first value vectors, thus reducing the consumption of memory resources. For example, in bytes, the number of bytes occupied by the second key vector or the second value vector in memory is equal to the number of bytes occupied by the first key vector or the first value vector in memory. The number of bytes occupied by the compressed n second key vectors and n second value vectors in memory is less than the number of bytes occupied by the m first key vectors and m first value vectors in memory.

[0009] This method can reduce the token dimension of key-value data from m to n, which not only saves memory space, but also preserves some of the complete feature information of the tokens because it does not require compressing the internal dimension of the vector corresponding to a single token. In contrast, quantization compression methods, such as quantization compression from standard FP16 data type to INT8 or INT4, will lose some feature information within each token during the quantization compression process, and cannot preserve the complete feature information of the tokens. The method proposed in this application, compared with the above-mentioned quantization compression methods that require changing the data storage precision, can reduce the loss of local important feature information, reduce the impact of feature information loss on prediction accuracy, and ensure that the inference accuracy is within the expected range.

[0010] In other embodiments, the compression process not only compresses the m vectors corresponding to m tokens into n vectors, but also reduces the memory space required for a single token vector. Specifically, after compression, the memory space required for the second key vector is less than that required for the first key vector; the memory space required for the second value vector is less than that required for the first value vector. For example, by reducing the token dimension from m to n and then performing quantization compression on the vectors corresponding to a single token, this method can further reduce the storage space required for a single token vector. However, due to the loss of feature information, it may have some impact on prediction accuracy, requiring a balance between saving memory space and ensuring prediction accuracy.

[0011] In some embodiments, compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor can be achieved by performing cluster analysis on the m first key vectors in the first key tensor to determine n categories and their corresponding cluster centers; n < m; n is a positive integer; one of the n categories includes at least one first key vector; then, based on the cluster centers corresponding to the n categories, n second key vectors corresponding to the n categories are obtained; based on the n second key vectors, the second key tensor is obtained; based on at least one first value vector corresponding to at least one first key vector included in each of the n categories, n second value vectors corresponding to the n categories are obtained; based on the n second value vectors, the second value tensor is obtained.

[0012] Clustering algorithms are generally lightweight, requiring minimal computational resources and minimizing their impact on inference task execution speed. Furthermore, cluster centers, as representatives of multiple key vectors within a category, preserve complete feature information, avoiding information loss.

[0013] It should be noted that compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor can be achieved using the clustering analysis method described above, or it can be achieved in other ways. For example, other classification algorithms can be used to classify the m first key vectors and determine the class centers, or a neural network can be used to achieve classification and determination of class centers.

[0014] In some embodiments, based on the cluster centers corresponding to the n categories, n second key vectors corresponding to the n categories are obtained; based on the n second key vectors, a second key tensor is obtained, which can be obtained by taking the cluster center corresponding to the i-th category as the second key vector corresponding to the i-th category; i is an integer, and 1≤i≤n; the n second key vectors corresponding to the n categories are concatenated to obtain the second key tensor.

[0015] In other embodiments, the cluster centers can be appropriately transformed before being used as the second key vector, such as by performing a simple linear transformation or other transformations.

[0016] Based on at least one first value vector corresponding to at least one first key vector included in each of the n categories, n second value vectors corresponding to the n categories are obtained; based on the n second value vectors, a second value tensor is obtained, which can be obtained by using at least one first value vector corresponding to at least one first key vector included in each of the i categories to obtain the second value vector corresponding to the i-th category; the n second value vectors corresponding to the n categories are concatenated to obtain the second value tensor.

[0017] In some embodiments, the second value vector corresponding to the i-th category is obtained based on at least one first value vector corresponding to at least one first key vector included in the i categories. This can be achieved by performing a summation operation on at least one first value vector corresponding to at least one first key vector included in the i categories to obtain the second value vector corresponding to the i categories.

[0018] By accumulating multiple values ​​in a category, the feature information of multiple values ​​in that category can be preserved, and a single category corresponds to a compressed value vector.

[0019] In some embodiments, compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor may be done when the value of m reaches a first threshold, or when the storage space occupied by the first key tensor and / or the first value tensor reaches a second threshold.

[0020] The threshold setting can trigger compression operations periodically. Whenever the generated key-value data reaches a certain trigger condition, the compression operation is triggered.

[0021] In some embodiments, the operation can be performed based on a heterogeneous computing architecture, which includes a host side and a device side. The host side includes at least one central processing unit (CPU); the device side includes at least one of a graphics processing unit (GPU), a neural network processing unit (NPU), and a tensor processing unit (TPU); wherein the data object to be compressed is stored on the device side; the device side is used to perform inference tasks. Based on this architecture, before compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor, it can be determined that compression is performed on the host side; the data object to be compressed is copied from the device side to the host side; then, compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor can be performed on the host side.

[0022] Offloading the compression operation to the host side can reduce the impact of the compression algorithm on the execution of inference tasks on the device side, make full use of heterogeneous computing resources, and achieve overlap between compression and inference tasks.

[0023] In some embodiments, compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor can be achieved by calling the same kernel function to compress the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor.

[0024] By calling the same kernel function, the compression algorithm is implemented through a fusion operator. Compared to compression achieved by combining multiple small operators, the fusion operator can achieve centralized scheduling and centralized memory access in one go, reducing the overhead of kernel launch and DRAM access.

[0025] In some embodiments, compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor can be based on the first key tensor and the first value tensor corresponding to the first task to obtain the compressed second key tensor and the second value tensor corresponding to the first task.

[0026] Compression can be task-based. In practical applications, a large model can respond to a single task request and perform single-task inference, or it can execute inference for multiple tasks in parallel. The first task can be a single task or one of multiple tasks. The compression operation compresses the key-value data to be compressed for a given task, obtaining the compressed key-value data for that task.

[0027] In some embodiments, after obtaining the compressed second key tensor and the second value tensor, the second key tensor can be concatenated with the q first key vectors corresponding to the q tokens to obtain the third key tensor; q is a positive integer and q≥1; the q tokens are the (m+1)th to (m+q)th tokens in the input sequence; the second value tensor is concatenated with the q first value vectors corresponding to the q tokens to obtain the third value tensor.

[0028] The compression method provided in this application allows the compressed key-value data to be directly concatenated with newly generated key-value data and used in subsequent attention calculations. This method is compatible with large models and eliminates the need to transform n-dimensional key-value data into m-dimensional data. This is a feature that quantization compression methods that reduce data storage accuracy in related technologies do not possess.

[0029] In some embodiments, after obtaining the third key tensor and the third value tensor, it can be determined that the third key tensor and the third value tensor are data objects to be compressed; if the value of n+q reaches the first threshold, or if the size of the storage space occupied by the third key tensor and / or the third value tensor reaches the second threshold, the compressed fourth key tensor and the fourth value tensor are obtained based on the third key tensor and the third value tensor.

[0030] The third key tensor includes n+q key vectors corresponding to n+q tokens; the first value tensor includes n+q value vectors corresponding to n+q tokens.

[0031] In the next round of compression, the third key tensor and the third value tensor can be used as the data objects to be compressed, and the next round of compression operation can be performed.

[0032] Secondly, embodiments of this application also provide an electronic device, the electronic device comprising: a processor, the processor being configured to execute a computer program or instructions in a memory to implement the method as described in any of the foregoing.

[0033] Thirdly, embodiments of this application also provide a computer-readable storage medium comprising a stored program, wherein the program, when executed by a processor, implements the method as described in any of the preceding claims.

[0034] Fourthly, embodiments of this application also provide a computer program product, the program product including a program that, when run by an electronic device, causes the electronic device to perform the method described in any of the above-mentioned embodiments.

[0035] Fifthly, embodiments of this application also provide a chip system, including: a communication interface for inputting and / or outputting data; and a processor for executing a computer-executable program, causing a device equipped with the chip system to perform the method described in any of the preceding claims. Attached Figure Description

[0036] Figure 1(a) is a schematic diagram of the KV-Cache caching mechanism;

[0037] Figure 1(b) is an example diagram of memory fragmentation generated by related technologies;

[0038] Figure 2 is an example diagram of an application scenario for the data compression method provided in the embodiments of this application;

[0039] Figure 3 is a schematic diagram of the hardware system architecture of the data compression method provided in the embodiments of this application;

[0040] Figure 4 is a schematic diagram of the compression process of the data compression method provided in the embodiment of this application;

[0041] Figure 5 is a schematic diagram illustrating the use of cluster analysis to achieve compression in some embodiments of the data compression method provided in this application.

[0042] Figure 6 is a flowchart illustrating the relevant control logic of the data compression method provided in the embodiments of this application;

[0043] Figures 7(a) and 7(b) are schematic diagrams illustrating the specific compression process in some embodiments of the data compression method provided in this application.

[0044] Figure 8 is a schematic diagram of the specific compression process in some embodiments of the data compression method provided in this application, based on a heterogeneous parallel computing architecture and multi-task parallel execution of inference tasks.

[0045] Figure 9 is a schematic diagram showing that the compression process in the data compression method provided in the embodiment of this application is implemented as a fusion operator;

[0046] Figure 10 is a schematic diagram of the fusion operator deployed on a relevant platform or framework in the data compression method provided in the embodiments of this application. Detailed Implementation

[0047] To better understand the technical solutions in this specification, the embodiments of this application will be described in detail below with reference to the accompanying drawings.

[0048] It should be understood that the described embodiments are merely some, not all, of the embodiments in this specification. All other embodiments obtained by those skilled in the art based on the embodiments in this specification without inventive effort are within the scope of protection of this specification.

[0049] The terminology used in the embodiments of this application is for the purpose of describing particular embodiments only and is not intended to be limiting of this specification. The singular forms “a,” “the,” and “the” used in the embodiments of this application and the appended claims are also intended to include the plural forms unless the context clearly indicates otherwise.

[0050] Large models are machine learning models with a large number of parameters and complex computational structures. A key component is large-scale inference. Taking Large Language Models (LLMs) as an example, a large language model can output corresponding content based on prompts input by the user. The core computational unit of current large language models is the attention module, which performs attention calculations. The input corpus generates high-dimensional tensors Q, K, and V, and the results of the attention calculation are used for subsequent computations. Based on the computational mechanism of large language models, the semantic information of historical sentences can be cached in the memory of an accelerator chip (such as a GPU) to reduce the computational workload of large-scale inference.

[0051] Some large language models divide the reasoning process into a prefill (or prefilling) stage and a decode (or decoding) stage. The prefill stage is a preprocessing stage that can calculate the prompt words input by the user in parallel to predict the first word of the answer. Then, the output words and historical sentences are concatenated to form a new sentence to predict subsequent words.

[0052] During the Decode stage of large model inference, the key-value pairs of historical statements can be cached in the memory of the acceleration chip, which can avoid the additional computational overhead caused by repeated calculations. This is the key-value caching strategy.

[0053] The prefill phase allows for a single input of all tokens for model inference. During this phase, inference continues based on user input, generating key-value cache data that supplements the KV-Cache for use in the decode phase. Simultaneously, the first output token can be generated to initiate autoregression during the decode phase. In the decode phase, tokens can be generated one by one.

[0054] For example, as shown in Figure 1(a), Figure 1(a) illustrates the computation logic of an attention unit with a KV-Cache caching strategy. k0,...,k t-1 k represents the key vector corresponding to the first t tokens in the input sequence. t Let v0, ..., vt+1 represent the key vector corresponding to the (t+1)th token.t-1 Let vt represent the value vector corresponding to the first t tokens in the input sequence, and vt represent the value vector corresponding to the (t+1)th token. As shown in Figure 1(a), the key vector and value vector corresponding to the first t tokens are stored in memory (i.e., put in cache as shown in the figure). When performing the inference task corresponding to the (t+1)th token, the values ​​of the key vector and value vector corresponding to the first t tokens can be directly read from memory to perform subsequent attention calculations, and the attention calculation result at the (t+1)th time step (i.e., attention result as shown in the figure) can be obtained without performing the repeated calculations corresponding to the first t tokens.

[0055] The value vectors corresponding to the first t tokens form the value tensor, and the key vectors corresponding to the first t tokens form the key tensor. During the inference process of the (t+1)th token, the (t+1)th value vector and the (t+1)th key vector are generated. The value tensor and the key tensor are incremented by 1 in the token dimension.

[0056] The KV-Cache mechanism can avoid redundant calculations, significantly reduce the amount of parameter calculations, and improve inference efficiency. Its essence is to exchange storage for computation, using memory overhead to exchange for computational overhead.

[0057] However, in large model inference, besides computational complexity, the memory wall problem is also a significant bottleneck restricting the development of large models. For example, a typical model with 70 parameters (requiring 140GB of memory) requires approximately 176GB of memory to store the key-value cache data (i.e., the key and value vectors corresponding to the first t-1 tokens) during 8-batch, 8k sequence length inference, which exceeds the memory space required for the model parameters themselves. Therefore, reducing the memory consumption of key-value cache data in large model inference has become one of the urgent problems to be solved.

[0058] Here, 8batch represents the number of task requests that the large model executes in parallel, and 8k sequence length represents the length of the input sequence of the large model, which can be understood as the number of tokens corresponding to the input sequence.

[0059] The following lists three related technologies to illustrate the technical background for improving the embodiments of this application.

[0060] Related Technology 1

[0061] In one related scheme, as shown in Figure 1(b), the kv cache capacity is allocated based on the maximum sentence length or maximum sequence length of the large model during pre-training. For example, a contiguous kv cache space is pre-allocated in the memory of the acceleration chip (e.g., GPU), and the kv cache data generated during the Decode stage is written to the pre-allocated kv cache space in an append-only manner.

[0062] However, in actual reasoning, sentence lengths are often smaller than or even much smaller than the Max Sequence Length, causing internal or external memory fragmentation and wasting memory space. For example, in Figure 2, the actual memory space occupied is much smaller than the maximum sequence length, with 2038 variable slots remaining unused, resulting in internal fragmentation. Furthermore, this is a coarse-grained memory management scheme. Coarse-grained, large-block memory allocation methods may cause external or internal memory fragmentation, meaning that memory space may not be fully occupied by multiple coarse-grained spaces, resulting in external fragmentation and wasting memory space. It is evident that this method cannot effectively utilize memory, a scarce resource.

[0063] Related Technology 2

[0064] To address the memory fragmentation issue of related technology one, an alternative reasoning framework employs a paged attention kv cache memory management method, similar to operating system paging memory management. This is a fine-grained kv cache memory management scheme where physical memory is managed in blocks, which are not necessarily contiguous in physical space. By maintaining a block table, the mapping relationship between logical blocks and physical blocks is reflected, and the address spaces of logical blocks are contiguous. This scheme can significantly mitigate the memory fragmentation impact of coarse-grained memory management.

[0065] However, while this solution alleviates the memory fragmentation problem, the key-value cache memory space still increases linearly with the length of the inference sentence, and the memory wall problem is still faced. This memory management method cannot fundamentally solve the problem of large memory consumption of key-value cache data.

[0066] Related Technology 3

[0067] In addition, another method of kv cache quantization compression in this technical field is to compress storage space by reducing the storage precision of the data. For example, each value in the kv data can be quantized and compressed from the standard FP16 (half-precision floating-point format) data type to INT8 (8-bit signed integer) or INT4 (4-bit signed integer) for storage, which can reduce the amount of memory required for kv cache by a factor of two.

[0068] While this method can significantly reduce the memory required for the key-value cache, the reduction in data storage precision during the quantization and compression of FP16 data types to INT8 inevitably leads to information loss. Furthermore, each token's corresponding feature vector incurs information loss, which impacts end-to-end inference accuracy. The smaller the number of bits after quantization and compression, the greater the impact on end-to-end inference accuracy. Therefore, this technical field generally uses INT8 quantization compression (i.e., the compressed data format is INT8) to preserve accuracy. This also limits the compression space of this method, restricting the key-value cache memory space compression ratio to 2-4 times, preventing further improvements.

[0069] In view of the above problems, this application proposes a data compression method to compress high-dimensional key tensors and high-dimensional value tensors stored in memory to obtain compressed lower-dimensional key tensors (second key tensors) and value tensors (second value tensors), thereby reducing the occupation of memory space resources and alleviating the memory wall problem to a certain extent.

[0070] The data compression method provided in this application embodiment is a method for compressing key-value cache data and reducing memory consumption in incremental inference scenarios using large models. This method can be applied to various application scenarios involving inference or prediction tasks based on large models. For example, it can assist large language models in performing various natural language processing tasks, such as question-and-answer interaction, translation, and text processing. As shown in Figure 2, in one application scenario example, the data compression method proposed in this application embodiment can be executed on a heterogeneous computing architecture to compress key-value cache data when a large language model performs interactive question-and-answer tasks, thereby reducing the consumption of memory resources in the heterogeneous computing architecture.

[0071] Specifically, a heterogeneous computing architecture can be a heterogeneous parallel computing architecture, which can be one or more combinations of the following architectures: Compute Unified Device Architecture (CUDA); and Compute Architecture for Neural Networks (CANN). The CUDA architecture can be used in conjunction with the CUDA Deep Neural Network library (cuDNN).

[0072] Specifically, a heterogeneous computing architecture can include a host side and a device side. The host side includes at least one central processing unit (CPU). The device side includes one or more of the following processors: graphics processing unit (GPU), neural processing unit (NPU), and tensor processing unit (TPU).

[0073] For example, as shown in Figure 3, taking CUDA as an example, the host side includes at least one central processing unit (CPU). The CPU may include a controller as shown in Figure 3, multiple arithmetic and logic units (ALUs), as well as cache and dynamic random access memory (DRAM).

[0074] The device side includes a GPU, which comprises multiple ALUs for performing parallel computing, and DRAM for storing data generated during the parallel computing process. The DRAM may include one or more of the following types of memory: global memory, high-bandwidth memory (HBM), shared memory, or constant memory (not shown in Figure 3). In the GPU, multiple ALUs in a row can share a controller and a cache.

[0075] During large model inference, inference tasks that require parallel computation can be executed on the device side, and the key-value cache data generated by parallel computation can be stored in DRAM or cache on the device side; for example, stored in global memory or HBM.

[0076] Specifically, the data compression method provided in this application uses the key-value cache data generated during the large model inference process as the compressed data object. Specifically, the key-value cache data refers to the key tensor and value tensor generated during the large model inference process. The key-value cache data is hereinafter referred to as key-value data.

[0077] As shown in Figure 4, assuming that in the current state, the length of the input sequence (i.e., the statement) corresponds to m tokens, and the key tensor before compression includes m key vectors corresponding to the m tokens, namely k1, k2…k m The value tensor consists of m value vectors corresponding to m tokens, namely v1, v2, ..., vm, where m is a positive integer. The key tensor or value tensor can be understood as the key tensor or value tensor generated by one head in a multi-head attention mechanism or a single head in a single-head attention mechanism.

[0078] It should be noted that a vector is a one-dimensional tensor. In this embodiment of the application, for clarity, a vector is used to describe the feature value corresponding to a single token, for example, the feature value corresponding to a single token is expressed as a key vector or a value vector; a tensor is used to describe the feature value corresponding to multiple tokens, for example, the feature value corresponding to multiple tokens is expressed as a key tensor or a value tensor.

[0079] The compression method provided in this application embodiment can compress the number of tokens, that is, compress a tensor with m token dimensions into a tensor with n token dimensions, where n < m or n is much smaller than m. The vector dimension corresponding to a single token remains unchanged, that is, the dimension of the key vector or value vector corresponding to a single token is not changed. Therefore, the feature information corresponding to a single token is retained without loss after compression, and the complete feature information of representative tokens is preserved.

[0080] As shown in Figure 4, the compressed key tensor includes n key vectors, namely k'1, k'2, ..., k'. n The compressed value tensor consists of n value vectors, namely v'1, v'2, ..., v' n The number of tokens in the compressed key or value tensor is reduced from m to n, thereby reducing the memory footprint.

[0081] It should be noted that the compressed n key vectors k'1, k'2…k' nThese are the n vectors determined from the m key vectors before compression using a compression algorithm. They may be discretely distributed among the m key vectors, or they may not be. Furthermore, the dimension of each individual vector remains unchanged before and after compression; that is, the feature dimension corresponding to a single token is not reduced. For example, k'1, k'2…k' n The dimension of a vector is related to k1, k2, ... k m The dimension of a vector is the same, and the same applies to the value vector before and after compression.

[0082] In some embodiments, the compression method proposed in this application can achieve online key-value cache compression by clustering and fusion of similar keys and accumulating the corresponding values, effectively reducing the amount of memory required for the key-value cache. Secondly, a periodic scheduling compression method can be used to schedule the inference process and compression computation concurrently, hiding or reducing the overhead caused by compression. Furthermore, considering heterogeneous hardware resources, the compression process is offloaded to the CPU for execution, minimizing the mutual influence between the compression process and inference computation. Finally, the key-value cache compression method is provided to the system platform in the form of a fusion operator, further improving the computational efficiency of compression itself, which will be described in detail below.

[0083] For example, as shown in Figure 5, in some embodiments, the compression process can be implemented based on cluster analysis. Specifically, cluster analysis can be performed on the m key vectors in the key tensor before compression to determine n categories, where each category may include at least one key vector (a first key vector), and each category has one cluster center. Generally, a category may have only one cluster center, but it is not excluded that in some embodiments, a category may have at least one cluster center.

[0084] To prevent confusion, the tensor or vector before compression is defined as the first tensor or the first vector, and the vector or tensor after compression is defined as the second vector or the second tensor.

[0085] Cluster analysis can divide the m first key vectors into n categories, and obtain n cluster centers corresponding to the n categories. For example, as shown in Figure 5, after cluster analysis, n categories are obtained as C1, C2, ..., Cn. Based on the n cluster centers corresponding to C1, C2, ..., Cn, n second key vectors can be obtained.

[0086] For example, in some embodiments, cluster centers can be directly used as the second key vectors. For instance, the cluster center corresponding to the i-th category can be used as the second key vector corresponding to the i-th category; i is an integer, and 1 ≤ i ≤ n. This is equivalent to retaining only the feature information of the cluster centers in a category, with the cluster centers representing the feature information of at least one vector in that category. n cluster centers can directly yield n second key vectors. Then, based on the n second key vectors, a compressed key tensor (second key tensor) can be obtained. For example, the compressed key tensor can be obtained by concatenating the n second key vectors in sequence.

[0087] Cluster analysis can determine the category corresponding to each of the m first key vectors. Since there is a one-to-one correspondence between the first key vector and the first value vector, the category corresponding to each of the m first value vectors can also be determined.

[0088] In some embodiments, n second value vectors corresponding to n categories can be obtained based on at least one first value vector corresponding to each of the n categories. Specifically, the second value vector corresponding to the i-th category can be obtained based on at least one first value vector corresponding to at least one first key vector included in the i-th category. For example, in some embodiments, a summation operation (or accumulation operation) can be performed on the at least one first value vector corresponding to the i-th category to obtain the second value vector corresponding to the i-th category. Thus, after obtaining n second value vectors, the n second value vectors can be concatenated to obtain the second value tensor.

[0089] For example, as shown in Figure 5, suppose the nth category Cn includes the (m-2)th key vector k. m-2 and the m-th key vector k m Therefore, the category Cn corresponds to the (m-2)th value vector v. m-2 and the m-th value vector v m Then the nth second value vector can be determined according to v. m-2 and v m To obtain, specifically, can be v m-2 and v m The vector obtained by summing is used as the nth second value vector v. n .

[0090] It should be noted that the above summation operation is only one example. In other embodiments, a weighted summation operation can also be performed on multiple first value vectors corresponding to the i-th category. For example, in a category, the value vector corresponding to the cluster center has a higher weight, while the value vectors corresponding to other key vectors have a lower weight. Alternatively, the multiple first value vectors corresponding to the i-th category can be linearly transformed before summing. Under the principle of preserving the feature information of the vectors to the greatest extent, there can be various fusion methods, not limited to summation.

[0091] The clustering analysis algorithm can be any algorithm that can achieve the above objectives. For example, it can be one or more of the following algorithms: K-Means algorithm, Mini Batch K-Means algorithm, K-Modes, AP clustering (Affinity Propagation Clustering, AP) algorithm, DBSCAN (Density-Based Spatial Clustering of Applications with Noise) clustering, hierarchical clustering, etc.

[0092] Furthermore, it should be noted that the compression achieved through clustering analysis algorithm described above is merely an example. Based on the above description, other classification algorithms can be used for classification. One can choose a classification algorithm that can determine the class centers or determine the class centers after classification. For example, in other embodiments, the following classification algorithm can be used to classify m key vectors, resulting in n categories:

[0093] K-Nearest Neighbor (KNN) algorithm, Decision Tree algorithm, Naive Bayes algorithm, Support Vector Machine (SVM) algorithm, Logistic Regression algorithm, etc.

[0094] In other embodiments, the classification of m key vectors or value vectors can also be implemented based on a neural network model with classification capabilities. For example, a neural network model can be trained to perform the classification task of m key vectors and the task of determining class centers.

[0095] After determining the n categories and n class centers, the second key vector and the second value vector can be obtained according to the above example, and then the compressed key tensor and value tensor can be obtained.

[0096] It should be noted that the compression method mentioned in the embodiments of this application does not involve decompression. Inference tasks can be performed directly on the compressed key-value data. Therefore, after compression is completed, the compressed key-value data and the newly generated key-value data from the large model inference are concatenated together to continue the inference task. This process is repeated iteratively until the large model inference task is completed.

[0097] In some embodiments, during large model inference, the compression operation of key-value cache data can be performed under predetermined triggering conditions. For example, when the length of the statement or sequence corresponding to the generated key tensor or value tensor reaches a predetermined first threshold, or when the memory space occupied by the key tensor or value tensor reaches a second threshold, the compression operation is triggered. The compressed key-value data is concatenated with the newly generated key-value data during the compression process. If the concatenated key-value data reaches the first threshold or the second threshold again, the compression operation is triggered again, and so on, the compression operation is triggered periodically.

[0098] The statement length or sequence length can be measured in terms of the number of tokens. Compression can also be triggered when the number of tokens corresponding to the generated key or value tensor reaches a first threshold. The first threshold can represent the number of tokens; if the key or value tensor to be compressed corresponds to m tokens, then the first threshold is the value of m.

[0099] Alternatively, the length of the statement or sequence corresponding to the key tensor or value tensor reaches a predetermined first threshold. Specifically, this can be based on the length of the statement or the number of words included in the input sequence. For example, if the length of the input sequence or the statement to be inferred reaches a certain number of words, then a compression operation is triggered. In other words, the first threshold can also represent the number of words in the statement length.

[0100] The second threshold can represent the size of the storage space, for example, in bytes, or in units such as KB, MB, GB, etc.

[0101] At the same time, at the scheduling level, by taking into account heterogeneous hardware resources, compression can be offloaded to the host side for execution, while the device side continues to execute inference tasks. This allows compression to run concurrently with and overlap with the inference process, reducing the impact of compression computation on inference computation while masking the additional overhead brought by compression operations.

[0102] For example, the periodic compression operation and resource scheduling based on the heterogeneous computing architecture can be implemented according to the control logic shown in Figure 6. First, the threshold for key-value cache compression can be determined, and then large-scale model inference computation can begin. During the computation process, key-value data occupies memory space. At this point, the first step of judgment can be executed. If the memory space occupied by the key-value cache data or the length of the corresponding statement exceeds the set threshold, a compression operation is triggered; otherwise, the inference task continues. If the pre-set threshold is reached, then in the second step, it is determined whether to offload the compression operation to the CPU. If so, compression is performed on the CPU; otherwise, compression is performed on the device side. Specifically, the following process can be included:

[0103] S601, determine the kv cache compression threshold.

[0104] For example, a first threshold or a second threshold is determined. In some embodiments, the specific way to determine the compression threshold may be to set a variable (the first variable) and provide an interface for setting the value of this variable. Users can customize the size of the compression threshold before compression (e.g., the first threshold or the second threshold). For example, setting the value of m to one of 12K, 10K, 8K, 6K, etc. indicates that when the number of tokens corresponding to the key tensor or value tensor output by a single head in the multi-head attention mechanism of the current large model reaches one of these (e.g., 8K), the compression operation is triggered.

[0105] In addition, another variable (second variable) and interface can be set. Users can customize the length of the compressed data. For example, the length threshold of the compressed data can be defined as 6K, 4K, 3K, etc., which means that the number of tokens corresponding to the key tensor or value tensor obtained after compression is 6K, 4K, or 3K, etc. In the above compression operation example, the value of n is 6K, 4K, or 3K, etc.

[0106] In other embodiments, the threshold before compression and the threshold after compression can also be written as constants in the code.

[0107] S602, large model inference generates key-value data.

[0108] The key-value data generated by large model inference, namely the key tensors and value tensors generated during the Decode phase of the inference process, is generally stored in the device-side memory, such as in the device-side global memory or HBM.

[0109] The decoding stage of a large model typically includes multiple layers and employs a multi-head attention mechanism. In this embodiment, a key tensor or a value tensor can be understood as a tensor generated by one head in a layer. For example, a key tensor or a value tensor can have four dimensions of information: B, N, S, and D. Here, B represents the batch size, i.e., the number of tasks processed by the large model at one time; N represents the number of heads, i.e., the number of heads included in the multi-head attention in the same attention layer; S represents the statement length, which can specifically be the number of tokens; and D is the hidden dimension, representing the dimension of a single token in the large model. For example, if D is 128, then the vector corresponding to a single token includes 128 elements, meaning that the key vector or value vector is a one-dimensional vector containing 128 elements.

[0110] The compression method proposed in this application does not change B, N, and D, but only changes the dimension of S.

[0111] S603, determine whether the key-value data has reached the compression threshold. If yes, proceed to S604; otherwise, return to S602.

[0112] It should be noted that in some embodiments, the monitoring object can be the key-value data volume of a single task, or it can be the total key-value data volume of a group of at least two tasks. If the monitoring object is a single task, when the key-value data volume of that task reaches the compression threshold corresponding to that task, a compression operation on the key-value data of that task is triggered. If the monitoring object is a group of tasks, the total data volume of the group of tasks is monitored, and when it reaches the threshold, a compression operation on the key-value data corresponding to the group of tasks is triggered.

[0113] For example, in a multi-task scenario, i.e. when the batch size is greater than 1, determining whether the key-value data has reached the compression threshold can be done by determining whether the key-value data corresponding to a single task has reached the compression threshold. More specifically, it can be determined whether the token dimension (i.e., S dimension) of the key tensor or value tensor generated by a single head in the large model attention mechanism has reached the compression threshold.

[0114] For example, read the value of the first variable in S601, compare whether the number of tokens corresponding to the key tensor or value tensor generated by the current single header reaches the value of the first variable, or compare whether the memory space required by the key tensor or value tensor generated by the current single header reaches the value of the first variable. If yes, proceed to S604.

[0115] In other embodiments, the compression threshold can also be calculated using multi-task statistics or multi-head statistics. The first or second threshold can be multiplied by a corresponding factor. For example, if x tasks are executed in parallel, the first or second threshold is multiplied by x to obtain the compression threshold. Alternatively, if the multi-head attention mechanism contains z heads, the first or second threshold is multiplied by z to obtain the compression threshold. There are various ways to set the compression threshold, and it is not limited to the examples above.

[0116] For example, in some embodiments, using multi-task statistics, assuming batchsize = 3, that is, the number of inference tasks currently executed in parallel by the large model is 3, it is possible to detect whether the storage space occupied by the current 3 tasks has reached 3 times the second threshold, or whether the total amount of key-value data generated by the current 3 tasks has reached 3 times the first threshold. If so, a compression operation is triggered. For example, the key-value data in one task can be compressed, and the compressed data is concatenated with the newly generated key-value data corresponding to the same task.

[0117] Similarly, using multi-head statistics, assuming that in the current large model's multi-head attention mechanism, an attention layer contains Y heads, then we can detect whether the total memory space occupied by the key-value data jointly generated by the Y heads reaches Y times the second threshold; or detect whether the total amount of key-value data generated by the Y heads reaches Y times the first threshold. If so, a compression operation is triggered. For example, the key-value data in a task can be compressed, and the compressed data can be concatenated with the newly generated key-value data corresponding to the same task.

[0118] S604, Determine whether to offload the compression to the host side? If yes, proceed to S605; otherwise, proceed to S606.

[0119] Since the generated key-value data is generally stored on the device side, some compression operations can be performed on the device side. The device side also needs to perform parallel computation in the inference task. Considering the impact of compression operations on the inference task, this application proposes to offload the compression operations to the host side for execution.

[0120] In some embodiments, whether to perform compression unloading on the host side can be set by a corresponding switch, allowing the user to specify whether to enable unloading to the host side. Alternatively, in other embodiments, it can be set to unload to the host side by default.

[0121] Alternatively, in some other embodiments, the current status parameters on the Device side and / or the Host side can be detected. For example, the system can automatically determine whether to offload the compression operation to the Host side based on the status parameters on the Device side. That is, without user specification, the system automatically determines whether to offload to the Host side based on the resource status under the current heterogeneous computing architecture.

[0122] Of course, the compression operation proposed in this application embodiment can also be performed on the device side.

[0123] S605, compression is performed by the CPU on the host side.

[0124] When the CPU performs compression, it needs to copy the data to be compressed from the Device side to the Host side. For example, it can copy the data from the Device side's main memory (DRAM) global memory or from HBM used as main memory to the Host side's main memory (DRAM).

[0125] S606, compression is performed on the device side.

[0126] If compression is performed on the device side, there is no need to copy the data to be compressed; it can be read directly from main memory.

[0127] The following is a specific example.

[0128] In this embodiment, by way of example, it is performed based on the heterogeneous computing architecture shown in Figure 3, and the Kmeans clustering algorithm is used in the compression operation.

[0129] Clustering analysis algorithms such as K-means require less computation and can perform compression tasks without consuming a large amount of computing resources on the device.

[0130] For example, as shown in Figure 7(a), in this embodiment, the compression process based on the K-means clustering algorithm may include:

[0131] S701, if m is detected to have reached the compression threshold (e.g., the first threshold), clustering is triggered.

[0132] In this embodiment, the number of tokens corresponding to the key tensor or value tensor to be compressed is counted, with a first threshold of 8K. If the currently generated key tensor or value tensor corresponds to m tokens, and m ≥ 8K, a compression operation is triggered. In this embodiment, the compression operation is implemented using the K-means clustering analysis algorithm, thus triggering a clustering operation.

[0133] S702 performs a compression operation based on K-means clustering.

[0134] Before compression, the S-dimensionality of the key or value tensor is m. Using the K-means clustering algorithm, the m keys are clustered into n categories, each with a cluster center, resulting in n cluster centers. These n cluster centers are the compressed n key vectors. Concatenating these n key vectors yields the compressed key tensor. For value vector compression, multiple value vectors corresponding to each category are summed, and the sum becomes the compressed value vector. This process yields n value vectors, which are then concatenated in chronological order to obtain the compressed value tensor. The S-dimensionality of the compressed key or value tensor is reduced to n.

[0135] S703: The compressed key-value data is concatenated with the newly generated key-value data, and the concatenated tensor participates in the subsequent attention calculation.

[0136] In this embodiment, data compression can be performed in parallel with the large model inference process. During the data compression process, the large model can continue to perform the inference process and generate new key-value data. The compressed key-value data can be concatenated with the generated new key-value data, and the concatenated tensor can continue to participate in the subsequent attention calculation of the large model.

[0137] It should be noted that during the splicing, the key-value data corresponding to the same task is spliced, that is, the compressed key-value data is spliced ​​with the new key-value data generated by the same task during the compression process. For example, the data splicing example corresponding to a single task shown in Figure 7(a) is as follows: the newly generated key-value data and the compressed key-value data are key-value data generated based on the same task request.

[0138] For example, as shown in Figure 7(a), assume that the data generated by the large model continuing to perform inference tasks during the compression process corresponds to q tokens, where q is a positive integer. The compressed n tokens are concatenated with the q tokens to obtain a key tensor or value tensor with dimension S of n+q. Here, q is a positive integer, and q≥1, and the q tokens are the (m+1)th to (m+q)th tokens in the input sequence.

[0139] Specifically, for the q tokens corresponding to the q first key vectors and q first value vectors, the compressed key tensor (second key tensor) is concatenated with the q first key vectors to obtain the concatenated third key tensor. The compressed value tensor (second value tensor) is concatenated with the q first value vectors to obtain the concatenated third value tensor.

[0140] In this embodiment, the compression operation can be triggered periodically. For example, after obtaining the third key tensor and the third value tensor, the third key tensor and the third value tensor can continue to be used as data objects to be compressed. When the value of n+q reaches a first threshold, or when the storage space occupied by the third key tensor and / or the third value tensor reaches a second threshold, the compressed fourth key tensor and the fourth value tensor are obtained based on the third key tensor and the third value tensor. The compressed fourth key tensor and the fourth value tensor can continue to be concatenated with newly generated key-value data to obtain the fifth key tensor and the fifth value tensor, and so on. In this way, the compression operation can be triggered periodically, and compression can be performed whenever the key-value data meets the triggering conditions.

[0141] For example, as shown in Figure 7(b), in the Tth compression, the key vectors corresponding to m tokens are compressed into n key vectors. During the Tth compression, the large model continues to perform inference tasks to generate q key vectors corresponding to q tokens. The n key vectors are concatenated with the q key vectors. When the concatenated tensor reaches the compression threshold, the (T+1)th compression is triggered, compressing the n+q key vectors into s key vectors corresponding to s tokens. During the (T+1)th compression, the large model generates p key vectors corresponding to p tokens. The s key vectors are concatenated with the p key vectors to obtain a key tensor with dimensions s+p tokens. This cycle repeats, allowing compression operations to be triggered periodically while the large model performs inference tasks. The above example uses key vectors; the same principle applies to value vectors.

[0142] It should be noted that, according to the compression method proposed in the embodiments of this application, the compressed tensor can be directly concatenated with the newly generated tensor (i.e., the new kv data generated during the compression process). After concatenation, no other transformation is required, and it can directly participate in attention calculation.

[0143] The specific implementation of the concatenation operation can be done in the memory space storing key-value data, such as in DRAM. The key-value data corresponding to the compressed n tokens replaces the key-value data corresponding to the uncompressed m tokens. The key-value data corresponding to the n tokens is then concatenated with the key-value data corresponding to the newly generated q tokens. During subsequent inference, the large model can read the key-value data corresponding to the concatenated n+q tokens to perform subsequent attention calculations.

[0144] In contrast, related technologies that use compression methods to change data storage precision, such as compressing FP16 data format to INT8 data format, have another drawback: after storing data in INT8 format, when participating in attention calculations, the stored INT8 data needs to be inversely transformed back to FP16 format. INT8 data cannot be directly used in subsequent calculations. Thus, compressing from FP16 to INT8 consumes computational resources, and inversely transforming from INT8 to FP16 consumes even more, resulting in significant computational resource consumption. The compression method proposed in this application, however, allows the compressed key-value data to be directly concatenated with newly generated key-value data without further transformation, reducing computational resource consumption.

[0145] For example, the compression operation proposed in this application is an online key-value cache compression method based on similar key fusion. In the attention calculation of large models, the role of the key is mainly to determine the weighting of the value. During long sequence inference, the key states corresponding to many tokens in the key-value cache are quite similar. By clustering similar key vectors and retaining only the class centers, and accumulating similar values, the key-value states (i.e., key states and value states) can be effectively compressed, thus ensuring that information is not lost to the greatest extent possible. States represent the feature values ​​corresponding to a time step or a token, and key states are the key vectors.

[0146] For example, in the compression operation, the core calculation formula is:

[0147] Where Attn(q,K,V) represents attention computation, q is the query, K is the query vector, and k is the query vector. i Let c represent the i-th key vector. i v represents the cluster center of the i-th category after key vector clustering. τi i μ represents the value vector corresponding to the τ-th key vector in the i-th category. i Let S represent the value vector obtained by summing the τ value vectors in the i-th category. S represents softmax.

[0148] As can be seen, this method can directly perform attention calculations on the clustered keys, and the values ​​corresponding to all keys in a category represented by the cluster center can be accumulated to ensure that semantic information is not lost. Furthermore, subsequent inference tasks can be performed directly on the compressed key-value data.

[0149] After the above compression operation, m key states are clustered into n categories (n << m). Correspondingly, the values corresponding to the keys in the same category are accumulated to obtain n new values, thus achieving compression by a factor of m / n.

[0150] It can be proved that when the distance between the key and the class center is less than δ, the error between the compressed inference result and the original inference result will be less than σ(δ). The within-class distance can be controlled by adjusting the number of clusters according to the error requirement.

[0151] It should be noted that in other embodiments, other clustering analysis algorithms can be used to replace the kmeans clustering method.

[0152] As shown in FIG. 7(a), it can be an example of performing compression on the Device side in a single-task scenario. In other embodiments, in a single-task scenario, compression can be performed on the Host side. When performing on the Host side, the kv data to be compressed needs to be copied from the Device side to the Host side, specifically, from the DRAM on the Device side to the DRAM on the Host side. The CPU on the Host side performs compression to obtain the compressed kv data, and then copies it from the DRAM on the Host side back to the DRAM on the Device side, and performs a splicing operation on the Device side.

[0153] FIG. 7(a) shows the process of performing compression on the tensor output by a single head for a single task. Taking batchsize = 3 as an example, the following lists embodiments of parallel execution of 3 inference tasks.

[0154] In this embodiment, it can be executed based on the hardware architecture of CPU (Host side) + GPU (Device side). In order to reduce the load of the Device during the inference process, the compression process can be offloaded to the CPU on the Host side for execution. This process is an optional item and can be specified by the customer whether to enable it.

[0155] Specifically, as shown in FIG. 8, the large model executes 3 inference task requests in parallel, r1, r2, and r3.

[0156] S801, the kv data to be compressed reaches the compression threshold, triggering the compression operation.

[0157] For example, the kv data to be compressed corresponding to at least one task among multiple tasks reaches the compression threshold, triggering the compression operation corresponding to the above at least one task.

[0158] For example, if the key-value data corresponding to tasks r1 and r2 reaches the compression threshold, compression operations for tasks r1 and r2 are triggered; or, if the key-value data corresponding to task r3 reaches the compression threshold, compression operations for task r3 are triggered. Alternatively, if the key-value data corresponding to all three tasks reaches the compression threshold, compression operations for all three tasks are triggered.

[0159] Alternatively, monitoring can be performed on a task-by-task basis, using the key-value data corresponding to a single task as the object. If the key-value data of a single task reaches a threshold, the compression operation corresponding to that task can be triggered. For example, if the key-value data corresponding to task r1 reaches the compression threshold, the compression operation for task r1 can be executed.

[0160] S802, the key-value data to be compressed is copied from the Device side to the Host side.

[0161] If the key-value data generated by the large model inference reaches the key-value cache limit, the data in the current key-value cache is copied from the device side to the host side. Specifically, this can be done by copying the data from the device-side DRAM (such as global memory or HBM) to the host-side DRAM. Meanwhile, on the device side, the inference task continues, generating new key-value data.

[0162] S803, the CPU performs compression operations based on Kmeans clustering.

[0163] After the data copy from the device to the host is completed, the compression operation is performed on the CPU. Since the computing resources of the CPU and the device are heterogeneous and do not conflict with each other, they can be executed in parallel.

[0164] It should be noted that in a multi-tasking scenario, compression operations can be performed separately for each task, with different tasks executing their respective compression operations in sequence. Alternatively, multiple tasks can execute compression operations in parallel, meaning that compression operations for multiple tasks can also be executed concurrently in sequence, and a compression operation for one task can correspond to one thread.

[0165] S804, the compressed key-value data is copied from the host side to the device side.

[0166] After the compression operation is completed on the CPU, the compressed key-value data is copied from the host to the device. Specifically, it can be copied from the DRAM on the host side to the DRAM on the device side (such as global memory or HBM).

[0167] S805: The key-value data generated after compression is concatenated with the key-value data newly generated during the compression process, and the concatenated tensor is used for attention calculation.

[0168] In multi-task scenarios, key-value data concatenation is performed only on key-value data corresponding to the same task. For example, key-value data with n tokens compressed in response to task request r1 is concatenated with key-value data with q tokens generated based on task r1. Concatenation across tasks is not allowed.

[0169] After the data copy is completed, the compressed data is concatenated with the newly generated key-value data on the Device side to continue the subsequent inference task.

[0170] It should be noted that Figure 8 shows an example of compression performed on the host side in a multi-tasking scenario. In other embodiments, compression can be performed on the device side in a multi-tasking scenario. When performing compression on the device side, there is no need to copy the key-value data to be compressed from the device side to the host side. Instead, the key-value data to be compressed is read directly from the DRAM on the device side, and the compressed key-value data is written into the DRAM after the compression operation is performed by the fusion operator. This compressed key-value data is then concatenated with the new key-value data generated by the final inference of the big data model during the compression process.

[0171] As shown in Figure 9, in some embodiments, when performing compression operations, the compression process can be implemented in the form of a fusion operator to reduce the time of small operator kernel launch and the overhead of accessing global memory, thereby improving computational execution efficiency.

[0172] The compression algorithm provided in this application can achieve the corresponding function by calling multiple interfaces provided by PyTorch multiple times and combining small operators. However, multiple calls increase the overhead of kernel launch and require multiple accesses to global memory, increasing memory access overhead. The fusion operator, that is, implements the compression operation through a single kernel function, can achieve the compression operation by calling a single kernel function, eliminating the need for multiple calls to small operators and reducing launch overhead and memory access overhead.

[0173] The fusion operator can be understood as implementing the clustering compression algorithm of this scheme into an operator and integrating it into the end-to-end inference framework. This fusion operator can be executed on the device side (GPU / NPU) or on the CPU side (in the scenario where compression is offloaded to the CPU side). It takes historical key-value data as input and outputs clustered and compressed key-value data.

[0174] As shown in Figure 10, at the operator and scheduling level of the large model inference system software stack, the fusion operator obtained based on the kv cache compression algorithm can be registered with commonly used inference frameworks in the industry, such as the Huggingface platform, to be compatible with the checkpointing of the pre-trained models provided by the Huggingface platform; or it can be registered with the large model inference deployment framework (Text Generation Inference, TGI) launched by the HuggingFace inference platform. Alternatively, it can be registered with the open-source LLM inference and service engine vLLM.

[0175] This fusion operator can run on heterogeneous computing architectures or acceleration libraries such as CUDA / CUDNN / CANN. Specifically, the server running this compression algorithm can be at least one of GPU servers, NPU servers, or run on a server cluster.

[0176] For example, in a large model inference scenario, a PyTorch+huggingface platform is set up on a server with 8 A100 GPUs (each with 40GB of video memory), using the checkpoint of the pre-trained model llama-7B provided by huggingface. The training sentence length of this model is 32K.

[0177] During inference, k-means is used as the clustering algorithm, and the compression threshold is set to 4k. That is, during inference, as long as the cumulative length of the kv cache data stored in memory or cache exceeds 4k and reaches 8k, compression will be triggered and compressed to 4k.

[0178] To verify accuracy, the batch size was set to 1, meaning one task request was processed at a time. In the evaluation dataset, the sentence length of the user question prompt was 2048, and the maximum sentence length for inference was 32k, meaning it would not exceed the sentence length of the pre-trained model.

[0179] Through experiments, statistical analysis was conducted on inference accuracy, memory usage, and inference speed before and after compression. Under the premise that the accuracy loss meets the requirements (0.153 < 0.2), the algorithm can achieve an 8x compression ratio in inference of a 32k sentence-length model. Furthermore, the algorithm can achieve even higher compression ratios in inference of longer or even infinitely long sequences.

[0180] The compression method proposed in this application, in addition to reducing memory overhead, can also reduce inference time by reducing attention computation.

[0181] K-means computation is not yet highly efficient on Ascend. When the sequence length is short, the computation time of the original attention mechanism is short, so the speed improvement of this embodiment is limited in short sequence scenarios. In long sequence scenarios, the computation time of the original attention mechanism is long, while the computation time of K-means can be diluted by longer periods. Therefore, in long sequence scenarios, the compression method proposed in this embodiment can significantly reduce memory overhead and inference time.

[0182] Specifically, at a length of 8k, even with a simple operator implemented using the torch API in PyTorch, a 2X compression and a 1.16X end-to-end inference speed improvement can be achieved; at a length of 32k, an 8X compression and a 2X end-to-end inference speed improvement can be achieved. When the length exceeds 100K, a compression improvement of more than 10X and an end-to-end inference speed improvement of more than 5X are expected.

[0183] As shown in Table 1, comparing the inference accuracy before and after using the kv cache compression method, the Perplexity (PPL) metric improved to 5.0481 after compression. PPL represents inference accuracy and is an indicator for evaluating the performance of a language model. Peak memory represents the peak memory usage during large model execution; after compression, the peak memory usage was reduced to 20.35GB. KV max length represents the maximum sentence length, or the maximum length of the input sequence; before compression, it was 32k tokens, and after compression, it was 4k tokens.

[0184] Compared to related technologies, the embodiments of this application use the kv cache compression method, which can achieve an 8X compression ratio in 32k sentence-length model inference while meeting the accuracy loss requirement (0.153<0.2). Moreover, the algorithm can achieve a higher compression ratio in longer or even infinitely long sequence inference. In terms of inference speed performance, the inference speed can achieve a 2X speedup.

[0185] It should be noted that the data compression method provided in this application compresses the token dimension corresponding to the tensor from m dimensions to n dimensions. This compression process can be implemented using the above-mentioned compression algorithm based on clustering analysis, or it can be implemented using other compression methods. For example, a compressed vector can be predicted based on X vectors by a trained neural network, achieving a compression ratio of X times. For example, if the input is two vectors, predicting one vector through the neural network achieves a compression ratio of 2 times.

[0186] In summary, this application proposes an online key-value cache compression method based on similar key fusion. The method clusters keys and accumulates corresponding values. By clustering similar keys and retaining only the cluster centers and accumulating similar values, key-value states can be effectively compressed while ensuring minimal information loss. It can be proven that when the distance between the key and the cluster center is less than δ, the error between the compressed inference result and the original inference result will be less than σ(δ). σ(δ) is the mapping function between the distance between the key vector and the cluster center and the error. That is, both the error and the clustering accuracy are controllable, and δ can be adjusted according to the error requirements.

[0187] The method proposed in this application is a scheme in which the compression process and inference computation can overlap. A compression threshold is set in advance. When the key-value data generated by inference reaches the key-value cache limit, compression is triggered. Simultaneously, existing key-value data is used to perform inference tasks and generate new key-value data. After compression, the compressed key-value data and the newly generated key-value data are concatenated for subsequent inference tasks. When the key-value data continues to grow and exceeds the threshold, compression is triggered again. Periodically triggering clustering reduces the overhead of the clustering algorithm while ensuring that the key-value cache size is controlled within a specified range, theoretically supporting inference of infinitely long sequences.

[0188] In the method proposed in this application embodiment, the compression operation can be offloaded to the host CPU for execution, while inference computation is generally performed on the device side (e.g., GPU / Ascend, etc.). The clustering compression method involved in this application embodiment can naturally support overlap with the inference process. Therefore, the compression process can be offloaded to the host for execution, avoiding competition for computing resources on the device side and reducing the impact on the speed of the inference task.

[0189] In addition, the kv cache clustering algorithm can be made into a fusion operator and executed on the device side or CPU side. The fusion operator can reduce the kernel launch time of small operators and the overhead of accessing global memory, thereby improving the computational execution efficiency.

[0190] It should be noted that the compression method proposed in this application is equivalent to aggregating information after cluster compression, and retains some typical features. For example, the features of the class centers are completely unaffected and not completely discarded. Inference calculations can be performed directly on the compressed kv states, avoiding the additional overhead caused by decompression operations.

[0191] During inference, compression is triggered whenever the key-value states reach a pre-set threshold. Theoretically, it supports sentences of unlimited length and has no limit on the length of the input sequence.

[0192] The method proposed in this application is essentially a semantically aware compression method that does not modify numerical precision. Compared with traditional quantization compression methods (which reduce data storage precision), it can retain local feature information and reduce the loss of feature information of representative elements. Furthermore, the method provided in this application can be executed orthogonally with traditional quantization compression methods. That is, a quantization compression method can be executed on the basis of the method provided in this application, or a compression method proposed in this application can be executed on the basis of a quantization compression method.

[0193] As mentioned above, in the method proposed in this application embodiment, the clustering compression process can overlap with the inference process, and the periodic compression operation also reduces the overhead caused by compression. Furthermore, it supports offloading the compression process to the host side for execution, avoiding the compression process from competing for device computing resources and improving end-to-end inference speed.

[0194] It should be noted that the compression method proposed in this application has significant advantages when the sequence is long. Therefore, in practical application scenarios, some pre-trained models in this technical field experience a sharp drop in inference accuracy when the inference sequence length exceeds the pre-training sentence length. In this case, advanced algorithms for training short and inferring long sequences are needed to solve the accuracy problem. Then, the kv cache compression method involved in this application is used to perform compression to save memory. That is, with the increase of sequence length, it is necessary to first solve the problem of the decrease in accuracy of the pre-trained model, and then execute the compression method proposed in this application while ensuring a certain level of accuracy.

[0195] This application also provides an electronic device, the electronic device comprising: a processor, the processor being configured to execute a computer program or instructions in a memory to implement the method as described in any of the above embodiments.

[0196] It should be noted that a processor can be any chip with computing capabilities, and is not limited to a central processing unit (CPU). For example, a processor can be a chip that includes one or more transistors, resistors, capacitors, and other circuit elements to perform a certain function; or it can be an integrated circuit in various packages that can implement the above methods.

[0197] For example, a processor may include one or more processing units, such as a neural network processing unit (NPU), an application processor (AP), a modem processor, a graphics processing unit (GPU), an image signal processor (ISP), a controller, a digital signal processor (DSP), a baseband processor, etc. The different processing units may be independent devices or integrated into one or more processors. The controller can generate operation control signals based on the instruction opcode and timing signals to control instruction fetching and execution.

[0198] The memory can be used to store executable program code, including instructions. Internal memory may include a program storage area and a data storage area. The program storage area may store the operating system, at least one application program required for a function, etc. The data storage area may store data created during the use of the electronic device (such as input data, output data, etc.). Furthermore, internal memory may include high-speed random access memory, and may also include non-volatile memory, such as at least one disk storage device, flash memory device, universal flash storage (UFS), etc. The processor executes various functional applications and data processing of the electronic device by running instructions stored in the internal memory and / or instructions stored in memory located within the processor.

[0199] It is understood that the structures illustrated in the embodiments of this application are merely examples and do not constitute a limitation on the electronic device. The electronic device in the embodiments of this application may include more or fewer components than illustrated, or combine some components, or split some components, or have different component arrangements. The illustrated components may be implemented in hardware, software, or a combination of software and hardware.

[0200] This application also provides a computer-readable storage medium comprising a stored program, wherein the program, when executed by a processor, implements the method described in any of the above embodiments.

[0201] This application also provides a computer program product, which includes a program that, when run by an electronic device, causes the electronic device to perform the method described in any of the above embodiments.

[0202] This application also provides a chip system, including: a communication interface for inputting and / or outputting data; and a processor for executing a computer-executable program, causing a device equipped with the chip system to perform the methods described in any of the above embodiments.

[0203] The aforementioned computer-readable storage medium may be any combination of one or more computer-readable media. A computer-readable medium may be a computer-readable signal medium or a computer-readable storage medium. A computer-readable storage medium may be, for example, but not limited to, an electrical, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: an electrical connection having one or more wires, a portable computer disk, a hard disk, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM), or flash memory, optical fiber, portable compact disk read-only memory (CD-ROM), optical storage device, magnetic storage device, or any suitable combination thereof. In this document, a computer-readable storage medium may be any tangible medium that contains or stores a program that may be used by or in connection with an instruction execution system, apparatus, or device.

[0204] Computer-readable signal media may include data signals propagated in baseband or as part of a carrier wave, carrying computer-readable program code. Such propagated data signals may take various forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination thereof. Computer-readable signal media may also be any computer-readable medium other than computer-readable storage media, capable of sending, propagating, or transmitting programs for use by or in connection with an instruction execution system, apparatus, or device.

[0205] Program code contained on a computer-readable medium may be transmitted using any suitable medium, including but not limited to wireless, wire, optical fiber, RF, etc., or any suitable combination thereof.

[0206] In the description of this specification, the references to terms such as "one embodiment," "some embodiments," "example," "specific example," or "some examples," etc., refer to specific features, structures, materials, or characteristics described in connection with that embodiment or example, which are included in at least one embodiment or example of this application. In this specification, the illustrative expressions of the above terms do not necessarily refer to the same embodiment or example. Furthermore, the specific features, structures, materials, or characteristics described may be combined in any suitable manner in one or more embodiments or examples. Moreover, without contradiction, those skilled in the art can combine and integrate the different embodiments or examples described in this specification, as well as the features of different embodiments or examples.

[0207] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include at least one of that feature. In the description of this application, "multiple" means at least two, such as two, three, etc., unless otherwise explicitly specified.

[0208] Any process or method description in the flowchart or otherwise herein can be understood as representing a module, segment, or portion of code comprising one or more executable instructions for implementing custom logic functions or processes, and the scope of the preferred embodiments of this application includes additional implementations in which functions may be performed not in the order shown or discussed, including substantially simultaneously or in reverse order depending on the functions involved, as should be understood by those skilled in the art to which embodiments of this application pertain.

Claims

1. A data compression method, characterized in that, The method includes: The first key tensor and the first value tensor are determined to be the data objects to be compressed; the first key tensor includes m first key vectors corresponding to m elements token; the first value tensor includes m first value vectors corresponding to m tokens; m is a positive integer; Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor; the second key tensor includes n second key vectors; the second value tensor includes n second value vectors; n is a positive integer and n < m; Wherein, the dimension of a single first key vector is the same as the dimension of a single second key vector; the dimension of a single first value vector is the same as the dimension of a single second value vector.

2. The method as described in claim 1, characterized in that, Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor, including: Using the m first key vectors in the first key tensor as objects, perform cluster analysis to determine n categories and the cluster centers corresponding to the n categories; n < m; n is a positive integer; one of the n categories includes at least one first key vector; Based on the cluster centers corresponding to the n categories, obtain n second key vectors corresponding to the n categories; based on the n second key vectors, obtain the second key tensor; Based on at least one first value vector corresponding to at least one first key vector included in each of the n categories, n second value vectors corresponding to the n categories are obtained; based on the n second value vectors, a second value tensor is obtained.

3. The method as described in claim 2, characterized in that, Based on the cluster centers corresponding to the n categories, n second key vectors corresponding to the n categories are obtained; Based on the n second key vectors, the second key tensor is obtained, including: Use the cluster center corresponding to the i-th category as the second key vector corresponding to the i-th category; i is an integer, and 1≤i≤n; The n second key vectors corresponding to the n categories are concatenated to obtain the second key tensor; Based on at least one first value vector corresponding to at least one first key vector included in each of the n categories, n second value vectors corresponding to the n categories are obtained; based on the n second value vectors, a second value tensor is obtained, including: Based on at least one first value vector corresponding to at least one first key vector included in the i categories, the second value vector corresponding to the i-th category is obtained; The n second value vectors corresponding to the n categories are concatenated to obtain the second value tensor.

4. The method as described in claim 3, characterized in that, Based on at least one first value vector corresponding to at least one first key vector included in the i categories, the second value vector corresponding to the i-th category is obtained, including: For at least one first value vector corresponding to at least one first key vector included in the i categories, perform a summation operation to obtain the second value vector corresponding to the i categories.

5. The method according to any one of claims 1-4, characterized in that, Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor, including: If the value of m reaches a first threshold, or if the storage space occupied by the first key tensor and / or the first value tensor reaches a second threshold, compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor.

6. The method according to any one of claims 1-5, characterized in that, The method is executed based on a heterogeneous computing architecture, which includes a host side and a device side; wherein the data object to be compressed is stored on the device side. Before compressing the first key tensor and the first value tensor to obtain the second key tensor and the second value tensor, the method further includes: Determine whether compression will be performed on the host side; Copy the data object to be compressed from the device side to the host side; Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor, including: On the host side, the first key tensor and the first value tensor are compressed to obtain the second key tensor and the second value tensor.

7. The method according to any one of claims 1-6, characterized in that, Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor, including: By calling the same kernel function, the first key tensor and the first value tensor are compressed to obtain the second key tensor and the second value tensor.

8. The method according to any one of claims 1-7, characterized in that, Compress the first key tensor and the first value tensor to obtain a second key tensor and a second value tensor, including: Based on the first key tensor and the first value tensor corresponding to the first task, obtain the compressed second key tensor and the second value tensor corresponding to the first task.

9. The method according to any one of claims 1-8, characterized in that, After obtaining the compressed second key tensor and the second value tensor, the method further includes: The second key tensor is concatenated with the q first key vectors corresponding to the q tokens to obtain the third key tensor; q is a positive integer and q≥1; the q tokens are the (m+1)th to (m+q)th tokens in the input sequence; The second value tensor is concatenated with the q first value vectors corresponding to the q tokens to obtain the third value tensor.

10. The method as described in claim 9, characterized in that, After obtaining the third key tensor and the third value tensor, the method further includes: The third key tensor and the third value tensor are determined to be data objects to be compressed; If the value of n+q reaches a first threshold, or if the storage space occupied by the third key tensor and / or the third value tensor reaches a second threshold, then based on the third key tensor and the third value tensor, a compressed fourth key tensor and a fourth value tensor are obtained.

11. An electronic device, characterized in that, The electronic device includes: A processor for executing a computer program or instructions in memory to implement the method as described in any one of claims 1-10.

Images