How is KV-cache implemented in autoregressive transformer inference?

In autoregressive transformers during inference, it is important to efficiently utilize computations for each new token so as not to recalculate the entire sequence from scratch. For this, a KV cache is used — a cache of keys (Key) and values (Value) from the attention layers (self-attention).

The idea is that when generating the next token, we have already computed the keys and values for all previous tokens. Instead of recomputing them, we store them in memory (cache), and when computing attention for the new token, we use this cache.

It roughly works like this:

• For each transformer layer, keys and values are stored in separate buffers when processing a token.
• When generating the next token, only the key and value for it are computed, then combined with the previously stored ones.
• The attention mechanism uses the combined set of keys and values to compute the context.

This significantly speeds up inference, reducing the computational complexity from quadratic in sequence length to linear.

A pseudocode example for one layer:

# kv_cache stores keys and values for previous tokens
new_key, new_value = compute_kv(new_token)
kv_cache.keys = concatenate(kv_cache.keys, new_key)
kv_cache.values = concatenate(kv_cache.values, new_value)
output = attention(query=new_key, keys=kv_cache.keys, values=kv_cache.values)

Thus, the KV cache is a mechanism for storing intermediate keys and values to accelerate sequential generation in transformers.