Transform model reasoning-based importance perception KV coverage method and system
By calculating the importance score of tokens in the Transformer model and dynamically managing key-value storage addresses, the problem of insufficient ReRAM storage resources is solved, and efficient inference of the Transformer model is achieved.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- ZHEJIANG UNIV
- Filing Date
- 2026-01-08
- Publication Date
- 2026-04-17
AI Technical Summary
In the inference process of the Transformer model, how to effectively store a large number of dynamically generated key-value pairs (KV) in limited ReRAM storage while ensuring the accuracy of model inference and smooth execution becomes a challenge.
By calculating the importance score of each token, dynamically allocating key-value (KV) write addresses, prioritizing the storage of high-importance KV information and discarding low-importance KV, an importance-aware KV overwriting method was designed to achieve KV overwriting updates in ReRAM.
When storage resources are insufficient, the inference accuracy and smooth execution of the Transformer model are guaranteed, and the resource consumption of key-value pairs is effectively controlled.
Smart Images

Figure CN121879685A_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of computer science and artificial intelligence, and in particular to a key-value (KV) coverage method and system based on importance perception inference using a Transformer model. Background Technology
[0002] The Transformer model, with its powerful generation and understanding capabilities, has become the mainstream architecture in many fields such as natural language processing and computer vision. The core mechanism of the Transformer model is self-attention, used to determine the strength of dependencies between input labels (attention score). However, the computational complexity of the attention module increases quadratically with the length of the input sequence. This scalability introduces a severe data movement bottleneck into traditional von Neumann architectures (such as GPUs), resulting in high inference latency. Unlike the von Neumann architecture, Processing-in-Memory (PIM) architecture integrates memory and processor, eliminating the need for weight movement during inference and significantly reducing system energy consumption and computational latency. PIM-based processing technology, using resistive random access memory (ReRAM), stands out due to its high density, non-volatility, fast read speed, and low energy consumption. This technology has achieved efficient in-situ computation in the analog domain, effectively eliminating the large-scale data transfer overhead involved in matrix-vector multiplication (MVM), a core Transformer operation. Therefore, exploring ReRAM-based architectures is a valuable research direction for accelerating Transformer workloads.
[0003] While many previous works have studied ReRAM-based Transformer accelerators by focusing on mitigating constraints such as durability, capacity, and non-idealities, they often overlook a key dynamic characteristic of the attention mechanism: during autoregressive inference, keys (K) and values (V) are newly generated activations, not static weights. Therefore, any ReRAM-based attention computation first requires converting these dynamic KV values into conductance and writing them into ReRAM. During Transformer inference, KV values are continuously generated, and their total amount is positively correlated with model size and inference sequence length, resulting in a very large volume. When implementing this inference process using a ReRAM architecture, storing a large number of KV values in the limited ReRAM while ensuring model inference accuracy and smooth inference execution presents a challenge. Summary of the Invention
[0004] The purpose of this invention is to address the shortcomings of existing technologies by providing an importance-aware key-value coverage method and system based on Transformer model inference.
[0005] The objective of this invention is achieved through the following technical solution: In a first aspect, the present invention provides an importance-aware key-value (KV) coverage method based on Transformer model inference, comprising the following steps: Step 1: Obtain the structural information and input token information of each attention layer in the Transformer model as state information; Step 2: Based on the state information from Step 1, within each head of each Attention layer of the model, calculate the attention score matrix between each token and all historical tokens before the current position. The size of this matrix is... ,in Represents the length of the reasoning sequence (i.e., the number of tokens); Step 3: Based on the attention score matrix from Step 2, evaluate the importance of each token by summing the columns of the attention score matrix to obtain a value of [value missing]. An importance score vector, where each value in the vector represents the importance of a token; Step 4: Regarding the importance of the token in Step 3, dynamically allocate write addresses for the newly generated KV in the ReRAM used to store KV, and write the KV to the corresponding row or column of the ReRAM to achieve KV overwrite and update the cache. Furthermore, the attention score matrix for each token and all historical tokens preceding the current position is calculated using the following formula: In the formula, The matrix represents the attention scores between the 0th and (i-1)th tokens, and its size is [value missing]. ,in This represents the length of the inference sequence, i.e., the number of tokens. Further, step four is implemented through the following sub-steps: When a new key-value pair needs to be written, determine if there are any spare ReRAM resources; if there are spare resources, sequentially write the newly generated key-value pair into the rows or columns of the spare ReRAM; if there are no spare resources, identify the row and column containing the key-value pair of the token with the lowest importance, and write the newly generated key-value pair into the row and column corresponding to that token's key-value pair, thus achieving key-value overwrite and update of the cache.
[0006] Furthermore, the formula for calculating the importance of each token is as follows: In the formula, The matrix represents the attention scores between the 0th and (i-1)th tokens, and its size is [value missing]. ,in This represents the length of the reasoning sequence, i.e., the number of tokens. This represents the importance score of the 0th to the (i-1)th token.
[0007] Secondly, the present invention provides an importance-aware key-value coverage system based on Transformer model inference, for implementing the method described in the first aspect, comprising: The state information acquisition module is used to obtain the structural information and input token information of each attention layer in the Transformer model; the attention calculation module is used to calculate the attention score matrix of the current token and the historical tokens in each attention head. The importance assessment module is used to calculate and output the importance score of each historical token based on the attention score matrix. The storage management module is used to connect to ReRAM and, based on the importance score, execute the write address allocation and overwrite strategy for the newly generated KV.
[0008] The storage management module and the ReRAM are integrated on the same chip to form a model inference acceleration unit; the attention calculation module and the importance evaluation module are implemented by dedicated hardware circuits or programmable logic units.
[0009] Thirdly, the present invention provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the importance-aware key-value coverage method based on Transformer model inference as described in the first aspect.
[0010] The beneficial effects of this invention are: This invention addresses the in-memory computing architecture based on ReRAM, analyzes the importance characteristics of tokens in the Transformer model, calculates the importance of each token based on the attention score matrix, and proposes a method to prioritize saving the key-value (KV) information of more important tokens and discard less important KV information when storage resources are insufficient. Specifically, an importance-aware KV overlay method for ReRAM architecture is designed and implemented. When storage resources are insufficient, the KV information of newly generated tokens is written in the position of the KV information of the existing lowest-importance token, achieving KV overlay. This method can effectively control the total amount of resources occupied by KV information while ensuring the accuracy of model inference, and still ensures smooth inference even when storage resources are insufficient. Attached Figure Description
[0011] Figure 1 This is a ReRAM hardware system diagram, which consists of computing units integrated into a multi-level structure such as Tile and PE.
[0012] Figure 2 This diagram illustrates the matrix multiplication calculation in the XB module of ReRAM, which includes two matrices, Q and K, four DACs (digital-to-analog converters), and three ADCs (analog-to-digital converters). The diagram demonstrates the process of calculating matrix multiplication of Q and K using XB.
[0013] Figure 3 This is a diagram illustrating importance score calculation and KV coverage. Detailed Implementation
[0014] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0015] It should be noted that, unless otherwise specified, the features in the following embodiments and implementations can be combined with each other. This invention provides an importance-aware key-value (KV) coverage method for Transformer model inference, comprising the following steps: Step 1: Obtain the structural information and input token information of each attention layer in the Transformer model as state information, as shown in Table 1, which contains a total of 8 state information items.
[0016] Table 1. Correspondence between Status Information Symbols Step 2: Based on the state information in Step 1, within each head of each Attention layer of the model, calculate the attention score matrix between each token and all tokens before the current position. The calculation formula is as follows: in The matrix represents the attention scores between the 0th and (i-1)th tokens, and its size is [value missing]. ,in This represents the length of the inference sequence (i.e., the number of tokens). ReRAM hardware systems, such as... Figure 1As shown, each level contains input / output buffers for temporary storage of input and output data. XB is the module in the PE structure responsible for matrix multiplication calculations, including basic hardware such as DAC, ADC, and ReRAM. The horizontal circuitry consists of word lines, and the vertical circuitry consists of bit lines. A memristor is integrated at the intersection of the two lines. The word lines are used to apply voltage, and the bit lines are used to generate and accumulate current. The memristor can pre-store parameter weight values. The specific calculation process is as follows... Figure 2 As shown. Figure 2 It includes two matrices, Q and K, four DACs (digital-to-analog converters), and three ADCs (analog-to-digital converters), demonstrating a ReRAM-based computation... Examples. Among them. As input, the voltage is fed into the DAC in three cycles, converted into voltage values (V1-V4) and applied to the ReRAM array. As weights, the conductance values are pre-stored in ReRAM (G1-G12). According to Ohm's law and Kirchhoff's laws, the voltage and conductance on the same circuit are multiplied to obtain the current value, which is then accumulated to form the total current (I1-I3). The current is converted back to the digital domain by the ADC to obtain the matrix multiplication result of Q and K. Thus, the analog domain matrix multiplication operation is completed through the cooperation of DAC, ReRAM and ADC.
[0017] Step 3: Based on the attention score matrix from Step 2, evaluate the importance of each token by summing the attention score matrix column by column. The calculation formula is as follows: in, This vector represents the importance score of tokens 0 to (i-1), and its size is [value missing]. ,in This represents the length of the reasoning sequence, i.e., the number of tokens. Each value in the vector represents the importance of a token.
[0018] Step 4: Regarding the token importance mentioned in Step 3, dynamically allocate write addresses for the newly generated KV in the ReRAM used to store KV, and... Write the corresponding row or column to ReRAM to achieve key-value cache overwrite update, which includes the following sub-steps: Determine if there are any available ReRAM resources; If there are available resources, a new one will be generated. Write sequentially into the rows or columns of the empty ReRAM; If there are no spare resources, find the row and column of the key-value pair of the token with the lowest importance, and then add the newly generated key-value pair. Write the key-value pairs (KV) into the rows and columns corresponding to the token to achieve KV overwriting, such as... Figure 3 As shown.
[0019] Figure 3 This diagram illustrates an example of KV overlay writing. The left image shows a 5x5 attention score matrix, which contains a total of [number missing] values. Five tokens, For newly generated tokens, after calculating the attention score, the score matrix is summed column-wise to obtain the importance of each token. The importance is the lowest at 0.7, therefore a new token will be generated ( KV () Write KV ( In the rows and columns (e.g.) Figure 3 (As shown in the right figure), this achieves KV coverage, ensuring smooth inference when ReRAM capacity is insufficient.
[0020] A key-value (KV) covering system based on Transformer model inference, used to implement the aforementioned KV covering method, includes: The state information acquisition module is used to obtain the structural information and input token information of each attention layer in the Transformer model; the attention calculation module is used to calculate the attention score matrix of the current token and the historical tokens in each attention head. The importance assessment module is used to calculate and output the importance score of each historical token based on the attention score matrix. The storage management module is used to connect to ReRAM and, based on the importance score, execute the write address allocation and overwrite strategy for the newly generated KV.
[0021] The storage management module and the ReRAM are integrated on the same chip to form a model inference acceleration unit; the attention calculation module and the importance evaluation module are implemented by dedicated hardware circuits or programmable logic units.
[0022] The present invention provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the importance-aware key-value coverage method based on Transformer model inference.
[0023] The effects of the present invention will be further explained below with reference to a specific embodiment. The experimental results are shown in Table 2.
[0024] The specific experiment is as follows: (1) ReRAM accelerator simulator: MNSIM; (2) Model: OPT-125M standard code; (3) Dataset: WikiText2.
[0025] Test results: This experiment simulated the process of inferring 128 tokens using the OPT-125M model. ReRAM resources were uniformly limited to storing 64 tokens' key-value pairs (KV), meaning KV covering began from token 65. Three different KV covering methods were compared, and the impact of each method on accuracy was discussed. Kullback-Leibler divergence and cosine similarity were used as accuracy evaluation metrics. Compared to the ideal baseline (no KV covering, retaining all 128 tokens' KV, with no loss of accuracy), the optimal result was achieved with the smallest Kullback-Leibler divergence and the largest cosine similarity. The experimental results are shown in Table 2.
[0026] Table 2 Experimental test results of the present invention Table 2 shows that, compared with object B, which directly discards the newly generated KV method (Ignore), and object C, which sequentially covers the original KV method (Rotary), the proposed method D, which achieves importance-aware KV coverage, achieves the best accuracy performance under the same experimental conditions. Specifically, object D has the best accuracy index, with a KL difference of only 1.03 and a Cosine similarity as high as 0.99 compared with the ideal case A (Baseline).
[0027] Obviously, the above embodiments are merely illustrative examples for clear explanation and are not intended to limit the implementation. Those skilled in the art will recognize that other variations or modifications can be made based on the above description. It is neither necessary nor possible to exhaustively list all possible implementations here. However, obvious variations or modifications derived therefrom are still within the scope of protection of this invention.
Claims
1. A key-value (KV) coverage method based on importance-aware reasoning using a Transformer model, characterized in that, Includes the following steps: Obtain the structural information and input token information of each attention layer in the Transformer model as state information; Based on the state information, within each attention head of the model, calculate the attention score matrix of each token and all historical tokens before the current position; The importance of each token is evaluated based on the attention score matrix; Based on the importance of each token being evaluated, a write address is dynamically allocated in the ReRAM used to store the key-value pairs (KV). The KV is then written to the corresponding row or column of the ReRAM, thus enabling the KV to overwrite and update the cache.
2. The KV coverage method as described in claim 1, characterized in that, The attention score matrix for each token and all historical tokens preceding the current position is calculated using the following formula: In the formula, The matrix represents the attention scores between the 0th and (i-1)th tokens, and its size is [value missing]. ,in This represents the length of the reasoning sequence, i.e., the number of tokens.
3. The KV coverage method as described in claim 1, characterized in that, The method for evaluating the importance of each token is as follows: sum the attention score matrix column by column, and the sum of each column is the importance score of the corresponding historical token.
4. The KV coverage method as described in claim 3, characterized in that, The formula for calculating the importance of each token is as follows: In the formula, The matrix represents the attention scores between the 0th and (i-1)th tokens, and its size is [value missing]. ,in This represents the length of the reasoning sequence, i.e., the number of tokens. This represents the importance score of the 0th to the (i-1)th token.
5. The KV coverage method as described in claim 1, characterized in that, The step of dynamically allocating write addresses for newly generated KV in the ReRAM used to store KV specifically includes: when a new KV needs to be written, determining whether there are any spare ReRAM resources; if there are spare resources, sequentially writing the newly generated KV into the rows or columns of the spare ReRAM; if there are no spare resources, identifying the row and column where the KV of the token with the lowest importance is located, and writing the newly generated KV into the row and column of the KV corresponding to that token, thereby realizing the KV overwrite update of the cache.
6. A key-value (KV) coverage system based on importance-aware reasoning using a Transformer model, characterized in that, For implementing the method as described in any one of claims 1-5, characterized in that it comprises: The state information acquisition module is used to obtain the structural information and input token information of each attention layer in the Transformer model; the attention calculation module is used to calculate the attention score matrix of the current token and the historical tokens in each attention head. The importance assessment module is used to calculate and output the importance score of each historical token based on the attention score matrix. The storage management module is used to connect to ReRAM and, based on the importance score, execute the write address allocation and overwrite strategy for the newly generated KV.
7. The system as described in claim 6, characterized in that, The storage management module and the ReRAM are integrated on the same chip to form a model inference acceleration unit; the attention calculation module and the importance evaluation module are implemented by dedicated hardware circuits or programmable logic units.
8. A computer-readable storage medium, characterized in that, It stores a computer program that, when executed by a processor, implements the importance-aware KV coverage method based on Transformer model reasoning as described in any one of claims 1-5.