KVCacheOperation

Applicable Products

Description

Implements KVCache processing.

For each batch:

Figure 1 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]

Operator Context

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

Figure 2 Position of kvCache in SelfAttention

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", which is considered as three words. 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

1
2
3

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

Parameters

Input

Output

Constraints

• 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.
• The value of hiddenSize must be 32-byte aligned.