A method for implementing sparse attention calculation based on a hybrid analog and digital in-memory computing architecture

Abstract

The application provides a method for realizing sparse attention calculation based on a hybrid digital and analog in-memory computing architecture, and belongs to the technical field of novel storage and computing. The method adopts a KV cache pruning mechanism, and at each decoding step in an autoregressive decoding process, similarity evaluation with high parallelism and high energy efficiency is realized based on an ACIM mode for important token selection, and the attention score of a top-k selected token is accurately calculated based on a DCIM mode, so that efficient sparse attention calculation is realized. In addition, a dynamic top-k selection mechanism is introduced, the reference value is dynamically adjusted according to the accumulated attention score, and then the value of k is adjusted, so as to adapt to different sparse attention distributions of different models, different layers and different attention heads, and then the inference accuracy close to dense attention is ensured while the LLM inference speed and energy efficiency are improved. The method has wide application prospects.

Description

Technical Field

[0001] This invention relates to the field of novel storage and computing technologies, specifically to a method for implementing sparse attention computing based on a hybrid digital and analog in-memory computing architecture. Background Technology

[0002] With the rapid development of artificial intelligence technology, Large Language Models (LLMs) based on the Transformer architecture have made breakthrough progress in the field of Natural Language Processing (NLP) and are widely used in tasks such as dialogue systems, content generation, intelligent search, and document summarization. Typical LLMs, such as GPT and LLaMA, rely on self-attention mechanisms to globally model all tokens in the input sequence to achieve contextual understanding and semantic reasoning. When the sequence length is n, the computational complexity of the self-attention mechanism is O(n^2). 2 To reduce computational overhead, researchers proposed a key-value cache (KV cache) technique. By storing the key and value of historical tokens, it avoids repeatedly calculating the attention weights of historical contexts in each round of generation, thus accelerating multi-round generation. However, this introduces greater storage overhead, especially as model depth and sequence length increase, the capacity requirement of the KV cache grows linearly. The computation of self-attention also becomes a key bottleneck limiting LLM deployment and response speed. Furthermore, KV cache access leads to frequent data transfer overhead and requires numerous matrix-vector multiplication operations to calculate attention scores, resulting in high latency and power consumption, particularly noticeable on edge devices.

[0003] To reduce data migration overhead and improve computational efficiency, researchers have accelerated matrix-vector multiplication operations in attention computation based on a Compute-in-Memory (CIM) architecture. Traditional Digital CIM (DCIM) performs well in terms of accuracy and reliability, but its computational parallelism is limited by the scale of the logic circuit, restricting the inference speed of LLM. Analog CIM (ACIM) has high parallelism and energy efficiency, but its computational accuracy and reliability are lower due to factors such as device non-ideals and noise, limiting the inference accuracy of LLM. Furthermore, in attention computation scenarios, KVcache exhibits dynamic sparsity; given a query vector, only a portion of the tokens significantly contribute to the final output, meaning only a portion of the key vectors have a high similarity to the query. This characteristic provides an opportunity to reduce computational and storage access overhead. However, the existing CIM architecture lacks a computation mechanism for dynamic sparse patterns, and neither DCIM nor ACIM architectures can simultaneously meet the high speed and high accuracy required for LLM inference, resulting in insufficient energy efficiency of the CIM architecture in sparse attention computation. Summary of the Invention

[0004] To address the problems existing in the prior art, this invention proposes a method for implementing sparse attention computation based on a hybrid analog in-memory computation (ACIM) and digital in-memory computation (DCIM) architecture. By combining a dynamic key-value cache pruning algorithm, the method performs fast filtering of important tokens based on ACIM and accurate computation on the selected tokens based on DCIM, thereby significantly reducing power consumption and latency while ensuring inference accuracy.

[0005] The technical solution of the present invention is as follows:

[0006] A method for implementing sparse attention computation based on a hybrid analog in-memory computation (ACIM) and digital in-memory computation (DCIM) architecture is characterized by storing the key and value vectors of the token sequence in a CIM array, using the query vector as input to the CIM array, and operating the CIM array in either ACIM or DCIM mode. In ACIM mode, the CIM array is used as an ACIM array; in DCIM mode, the CIM array is used as a DCIM array. A dynamic key-value cache (KV cache) pruning mechanism is employed. During autoregressive decoding, sparse attention computation is implemented based on both CIM modes. The specific steps of sparse attention computation in each decoding step include:

[0007] 1) Selecting important tokens based on ACIM

[0008] The query of the current decoding step is input into the ACIM array storing the key. All rows of the ACIM array are activated in parallel. The similarity between the input query and the stored key in each row is evaluated. The similarity calculation result of each row of the ACIM array is compared with the reference value. The row with the higher similarity to the reference value indicates that the key stored in that row has a higher similarity to the query. The k tokens corresponding to the k keys with the highest similarity are selected for subsequent precise attention calculation. This is called top-k selection. The unselected tokens are still kept in KVcache and will still participate in the calculation in the subsequent decoding process.

[0009] 2) Accurate attention score calculation based on DCIM

[0010] After selecting important tokens based on ACIM, precise attention calculation is further performed based on DCIM. The top-k selected rows in the DCIM array are selected and activated in turn. Each row is activated, and the attention score of the key and query stored in that row is precisely calculated. The multiply accumulate (MAC) calculation is performed by digital circuits through an adder tree to obtain the attention score of the token.

[0011] Then, the Softmax operation is calculated using digital circuitry, and the attention score, which incorporates attention weights, is calculated using the attention-weighted MAC of the value array.

[0012] Furthermore, in step 1), the similarity score uses cosine similarity as the attention score to evaluate the importance of the KV cache to the current query, as shown in the formula:

[0013] Cosine similarity=query·key T

[0014] Furthermore, the method introduces a dynamic top-k selection mechanism. First, a target value and an initial reference value are set for top-k selection. Based on the attention scores with integrated attention weights, an accumulated attention score is obtained. If the accumulated attention score exceeds the target value, it indicates that the top-k selected by the reference value is sufficient to represent the attention distribution of the current query, and subsequent autoregressive decoding calculations continue. If the accumulated attention score is less than or equal to the target value, it indicates that the number of selected top-k is too small. By reducing the initial reference value, more keys and queries are compared and selected for precise attention calculation based on DCIM, until the accumulated attention score of the tokens selected by the top-k exceeds the target value. Then, dynamic KV cache pruning and sparse attention calculation are completed.

[0015] Furthermore, the method for implementing sparse attention computation based on a hybrid ACIM and DCIM architecture, wherein ACIM and DCIM reuse the CIM array, and the CIM cells can be CIM cells based on traditional SRAM or CIM cells based on non-volatile memory (NVM), and different peripheral input circuits and output circuits are used to implement the functions of ACIM and DCIM, wherein the ACIM computation mode can adopt voltage domain CIM, current domain CIM and time domain CIM.

[0016] The technical effects of this invention are as follows:

[0017] 1. This invention proposes a method for sparse attention computation based on a hybrid ACIM and DCIM architecture. By combining a dynamic KV cache pruning algorithm, in each decoding step, a high-parallelism and high-energy-efficiency similarity evaluation is implemented based on the ACIM mode for the selection of important tokens, and the attention score of the top-k selected tokens is accurately calculated based on the DCIM mode, which can achieve efficient sparse attention computation.

[0018] 2. This invention proposes a method for sparse attention computation based on a hybrid ACIM and DCIM architecture. On the basis of dynamic KV cache pruning algorithm, a dynamic top-k selection mechanism is further proposed. The reference value is dynamically adjusted according to the accumulated attention score, and then the value of k is adjusted to adapt to different sparse attention distributions of different models, different layers and different attention heads. Thus, while improving the inference speed and energy efficiency of LLM, the inference accuracy is guaranteed to be close to that of dense attention. Attached Figure Description

[0019] Figure 1 This is a flowchart of the sparse attention computation based on the hybrid ACIM and DCIM architecture proposed in this invention;

[0020] Figure 2 This is a schematic diagram of the method for implementing sparse attention computation based on a hybrid ACIM and DCIM architecture proposed in this invention. Detailed Implementation

[0021] The present invention will be further clearly and completely described below with reference to the accompanying drawings and specific embodiments.

[0022] The flowchart of the method for sparse attention computation based on a hybrid ACIM and DCIM architecture proposed in this invention is as follows: Figure 1As shown, in the autoregressive decoding process, important tokens are first selected based on ACIM. The query of the current decoding step is input into the ACIM array storing the key, and all rows of the ACIM array are activated in parallel. The cosine similarity between the input query and the stored key in each row is evaluated and compared with the reference value. Then, top-k selection is performed to select the k tokens corresponding to the k keys with the highest similarity. The cosine similarity is used as the attention score to evaluate the importance of the KV cache to the current query, as shown in formula (1).

[0023] Cosine similarity=query·key T (1)

[0024] Further, based on DCIM, precise attention calculation is performed. The top-k selected rows in the corresponding DCIM array are selected and activated sequentially. One row is activated at a time, and the attention score of the key and query stored in that row is precisely calculated. Since DCIM only calculates the attention score of one key and query at a time, digital circuits are used to calculate the MAC result. The attention score is accurately calculated. The calculation of the multiply-accumulate (MAC) calculation result can be completed by the adder tree. Then, the calculation of the Softmax operation and the attention weighted MAC calculation with the value array are completed by the digital circuit to obtain the attention score with attention weight.

[0025] Furthermore, considering the differences in sparsity characteristics across different models, layers, and attention heads, a dynamic top-k selection mechanism is employed. First, a target value and an initial reference value are set for top-k selection. An accumulated attention score is obtained by accumulating the attention scores based on the attention weights. The value of k is dynamically adjusted according to this accumulated attention score to adapt to the different sparse attention distributions across different models, layers, and attention heads. If the accumulated attention score of the selected token is less than or equal to the target value, it indicates that the number of selected top-k tokens is insufficient. The reference value for top-k selection is then decreased until the accumulated attention score exceeds the target value, thus obtaining the attention distribution for the current query. If the accumulated attention score of the selected token exceeds the target value, it indicates that the top-k tokens selected at that reference value are sufficient to represent the attention distribution for the current query, and subsequent autoregressive decoding calculations continue.

[0026] The schematic diagram of the method for sparse attention computation based on a hybrid ACIM and DCIM architecture proposed in this invention is shown below. Figure 2As shown, this invention employs a dynamic KV cache pruning mechanism. During the autoregressive decoding process, all rows of the ACIM array are first activated in parallel to evaluate the similarity between the input query and the stored key. The k tokens corresponding to the k keys with the highest similarity are selected for precise attention calculation. Figure 2 Taking top-1 as an example, in each autoregressive decoding step, the token with the highest similarity is selected for precise attention calculation. Pruned (unselected) tokens remain in the KV cache and can still participate in calculations during subsequent decoding processes; this is called dynamic KV cache pruning, thus ensuring inference accuracy. On the hardware side, the above dynamic KV cache pruning and sparse attention calculation are implemented based on a hybrid ACIM and DCIM architecture. First, important tokens are selected based on ACIM. The query of the current decoding step is input into the ACIM array storing the keys, and all rows of the ACIM array are activated in parallel to evaluate the cosine similarity between the query and the key. Precise calculation of the MAC result is not required; only the k tokens with the largest MAC results need to be selected. After selecting important tokens based on ACIM, precise attention calculation is further performed based on DCIM. Only the top-k selected rows need to be activated. Each time a row of the DCIM array is activated, a digital circuit is used to calculate the MAC result through an adder tree, thereby accurately calculating the attention score of that token. Then, a Softmax operation is performed through digital circuitry to obtain the attention distribution. The attention distribution and the MAC calculation of the value array yield the final result of contextual information incorporating attention weights.

[0027] This embodiment fully and in detail describes a method for implementing sparse attention computation based on a hybrid ACIM and DCIM architecture. By combining a dynamic KV cache pruning algorithm, in each decoding step, a high-parallelism and high-energy-efficiency similarity evaluation is implemented based on the ACIM mode for the selection of important tokens. The attention score of the top-k selected tokens is accurately calculated based on the DCIM mode. This can achieve efficient sparse attention computation. Furthermore, by introducing a dynamic top-k selection mechanism, the inference speed and energy efficiency of LLM are improved while ensuring inference accuracy close to that of dense attention.

[0028] Finally, it should be noted that the purpose of disclosing the embodiments is to help further understand the present invention. However, those skilled in the art will understand that various substitutions and modifications are possible without departing from the spirit and scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the content disclosed in the embodiments, and the scope of protection of the present invention is defined by the scope of the claims.

Claims

1. A method for implementing sparse attention computation based on a hybrid ACIM and DCIM architecture, characterized in that, The key vector and value vector of the token sequence are stored in the CIM array. The query vector is used as the input to the CIM array. The CIM array operates in either ACIM mode or DCIM mode. When operating in ACIM mode, the CIM array is used as an ACIM array; when operating in DCIM mode, the CIM array is used as a DCIM array. A dynamic key-value cache (KV) pruning mechanism is adopted. During the autoregressive decoding process, sparse attention calculation is implemented based on the two CIM modes. The specific steps of sparse attention calculation in each decoding step include: 1) Selecting important tokens based on ACIM The query of the current decoding step is input into the ACIM array storing the key. All rows of the ACIM array are activated in parallel. The similarity between the input query and the stored key in each row is evaluated. The similarity calculation result of each row of the ACIM array is compared with the reference value. The row with the higher similarity to the reference value indicates that the key stored in that row has a higher similarity to the query. The k tokens corresponding to the k keys with the highest similarity are selected for subsequent precise attention calculation. This is called top-k selection. The unselected tokens are still kept in the KV cache and will still participate in the calculation in the subsequent decoding process. 2) Accurate attention score calculation based on DCIM After selecting important tokens based on ACIM, precise attention calculation is further performed based on DCIM. The top-k selected rows in the DCIM array are selected and activated sequentially, one row at a time. The attention score of the key and query stored in that row is precisely calculated. Digital circuits are used to perform the multiplication-accumulation calculation MAC result through an adder tree to obtain the attention score of the token. Then, the softmax operation is calculated and the attention weighted MAC with the value array is calculated through digital circuits to obtain the attention score with attention weights.

2. The method as described in claim 1, characterized in that, In step 1), the similarity score uses cosine similarity as the attention score to evaluate the importance of the KV cache to the current query. The formula is: 。 3. The method as described in claim 1 or 2, characterized in that, The method introduces a dynamic top-k selection mechanism. First, a target value and an initial reference value are set for top-k selection. Based on the attention scores with integrated attention weights, an accumulated attention score is obtained. If the accumulated attention score exceeds the target value, it indicates that the top-k selected by the reference value is sufficient to represent the attention distribution of the current query, and subsequent autoregressive decoding calculations continue. If the accumulated attention score is less than or equal to the target value, it indicates that the number of selected top-k is too small. By reducing the initial reference value, more keys and queries are compared and selected for precise attention calculation based on DCIM, until the accumulated attention score of the tokens selected by the top-k exceeds the target value. Then, dynamic KVcache pruning and sparse attention calculation are completed.

4. The method as described in claim 1 or 2, characterized in that, The ACIM and DCIM multiplexed CIM array uses CIM cells based on traditional SRAM and CIM cells based on non-volatile memory (NVM). Different peripheral input and output circuits are used to realize the functions of ACIM and DCIM. The calculation mode of ACIM is voltage domain CIM, current domain CIM and time domain CIM.

Images