KVCacheOperation

Description

Implements KVCache processing. It is supported only by the Atlas A2 inference products .

Operator Context

Stores k and v in kcache and vcache. This function is used in the selfattention operator.

Figure 1 Position of kvCache in SelfAttention

Operator Function Implementation

For each batch:

Figure 2 Functions of the kvCache operator

Compute logic:

prefix_ntokens = 0
for i in range(batch):
    for j in range(seqlen[i]):
        cache_out[layer_id[0]][i][token_offset[i] - seqlen[i] + j][:] = newkv[prefix_ntokens + j][:] # Modifies only the cache content of the layer specified by layer_id.
    prefix_ntokens += seqlen[i]

Application Scenarios

Used in the transformer inference phase.

For example, if a user enters "the world's highest peak" and the model continues to write "is Mount Everest," the computation process of each step of KV Cache is as follows:

• Step 1: The cache K and cache V are empty, and the input is "the world's highest peak". The model performs parallel computing in the common way.
  1. Obtain k, v, and an attention representation corresponding to each token through parallel computing.
  2. Use to predict the next token and obtain "is".
  3. Update the cache by setting and .

• Step 2:
  1. Input "is" into the model, and map the word vector to obtain , , and .
  2. Update the cache by setting and .
  3. Perform and predict the next token and obtain "Mount".

Definition

struct KVCacheParam {
    uint8_t rsv[8] = {0};
};

Parameters

Input

Output

Restrictions

• The number of elements in seqLen array is greater than 0, and the sum of arrays is nTokens.
• The number of elements in the tokenOffset array is greater than 0 and less than the value of maxSeqLen.
• The value of layerId must be smaller than the first dimension of past.
• Atlas A2 inference products : newKv can be 4D [batch, seq_len, head_num, head_size]. In this case, head_num x head_size must be equal to hiddenSize.