KV Cache - From Neuron to GPT

In the previous two chapters, we explored how to control the generation of the model: sampling strategies and stopping criteria. Next, let us look at how to improve the performance of model inference.

Autoregressive generation in a GPT model has a performance bottleneck: every time a new token is generated, the attention of all tokens has to be recomputed. Suppose 100 tokens have been generated; to generate the 101st token we have to compute the attention of 101 tokens, and the time complexity of generating each token is $O(n^2)$ .

KV cache¶

The core idea of the KV cache (key-value cache): cache the keys and values of the tokens that have already been generated. When generating a new token, compute only its query, and then do attention together with the cached keys and values.

How it works¶

In self-attention, the key and value of each token depend only on the token itself and the tokens before it, and have nothing to do with later tokens. Therefore:

the keys and values of generated tokens can be cached and need not be recomputed;
a new token only needs to compute its own query and do attention together with the cached keys and values;
the cache keeps growing as generation proceeds, but each token is computed only once.

Implementation points¶

Cache structure: each Transformer block stores one set of keys and values, where the key is already transposed, with the shape (batch size, number of heads, per-head embedding size, generated length); the value has the shape (batch size, number of heads, generated length, per-head embedding size);
Incremental computation: the key and value of the new token are concatenated with the cache before computing attention;
Position encoding: the position encoding of the new token needs an offset (the generated length).

import time

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

Model¶

GPT cached embedding layer¶

The GPT cached embedding layer (GPTEmbeddingWithCache) needs to support a position offset: the position encoding of the new token starts from offset instead of from 0.

class GPTEmbeddingWithCache(GPTEmbedding):

    def forward(self, x: Tensor, offset=0):
        token = self.embedding(x)
        # Position offset
        position = self.positional_embedding(Tensor(range(offset, offset + x.shape[1])))
        return self.dropout(token + position)

GPT cached attention layer¶

The GPT cached attention layer (GPTAttentionWithCache) is the core of the KV cache:

Without cache: the input is the whole context window before the current position, computed normally;
With cache: the input is only the new token at the current position, and its key and value must be concatenated with the cache before computing attention.

With a cache, no causal mask is needed, because the new token computes attention only with the tokens already generated, which naturally satisfies causality.

class GPTAttentionWithCache(GPTAttention):

    def forward(self, x: Tensor, use_cache=False, cache=None):
        norm = self.normalize(x)
        head_embedding_size = self.embedding_size // self.heads
        multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
        shape = (-1, x.shape[1], self.embedding_size)

        query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
        key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
        value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))

        # With a cache, the input is only the newest token and must be concatenated with the cache
        if use_cache and cache is not None:
            key = Tensor(np.concatenate([cache[0].data, key.data], axis=3))
            value = Tensor(np.concatenate([cache[1].data, value.data], axis=2))
        # Update the cache
        new_cache = (key, value) if use_cache else None

        scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
        scores = query @ key * scale
        # Without a cache, the causal mask is needed
        if cache is None:
            scores = self.mask(scores)
        weights = self.softmax(scores)
        output = self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape))
        return x + self.dropout(output), new_cache

GPT cached Transformer block¶

The GPT cached Transformer block (GPTTransformerWithCache) passes the cache to the attention layer as a parameter and returns the new cache produced by the attention layer.

class GPTTransformerWithCache(Composite):

    def __init__(self, embedding_size, heads):
        self.attention = GPTAttentionWithCache(embedding_size, heads)
        self.feed_forward = GPTFeedForward(embedding_size)

        super().__init__([self.attention,
                          self.feed_forward])

    # Inference function, passing the cache along
    def forward(self, x: Tensor, use_cache=False, cache=None):
        x, new_cache = self.attention(x, use_cache, cache)
        return self.feed_forward(x), new_cache

GPT cached model¶

The GPT cached model (GPTWithCache) manages the caches of all Transformer blocks:

each block has one cache, stored in a list;
during the forward pass the cache is passed block by block and the new caches are collected.

class GPTWithCache(Composite):

    def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
        self.embedding = GPTEmbeddingWithCache(vocab_size, context_size, embedding_size)
        self.transformers = [GPTTransformerWithCache(embedding_size, heads) for _ in range(blocks)]
        self.output = GPTOutput(embedding_size, vocab_size)

        super().__init__([self.embedding] + self.transformers + [self.output])

    # Inference function, passing the cache along
    def forward(self, x: Tensor, use_cache=False, caches=None, offset=0):
        x = self.embedding(x, offset)
        # New cache list, each Transformer block has its own cache
        new_caches = []
        for i, layer in enumerate(self.transformers):
            # Cache of the current block
            cache = caches[i] if caches is not None else None
            x, new_cache = layer(x, use_cache, cache)
            # Update the cache
            new_caches.append(new_cache)
        return self.output(x), new_caches

GPT cached trainer¶

The generate function of the GPT cached trainer GPTTrainerWithCache:

First generation: no cache, compute all tokens;
Later generations: with cache, compute only the new token;
Cache reset: if the sequence length exceeds the context size, reset the cache.

class GPTTrainerWithCache(GPTTrainer):

    def generate(self, dataset, prompt, steps=256, use_cache=True):
        self.model.eval()
        tokens = dataset.encode(prompt)
        # Cache
        caches = None

        with Tensor.no_grad():
            for _ in range(steps):
                if use_cache and caches is not None and len(tokens) <= dataset.context_size:
                    # With a cache and within the context size, use only the newest token
                    feature = Tensor([tokens[-1:]])
                    offset = len(tokens) - 1
                else:
                    # Without a cache, use the whole context window
                    feature = Tensor([tokens[-dataset.context_size:]])
                    offset = 0
                    caches = None

                # Inference, updating the cache at the same time
                logits, caches = self.model(feature, use_cache, caches, offset)

                # Sampling
                last_logits = logits.data[0, -1]
                exp = np.exp(last_logits - np.max(last_logits))
                probs = exp / np.sum(exp)
                token = np.random.choice(len(probs), p=probs)
                tokens.append(token)

        return dataset.decode(tokens)

Hyperparameters¶

Model parameters¶

The model parameters must stay consistent with the pre-training stage.

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

Modeling¶

Loading the model¶

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPTWithCache(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters)
trainer = GPTTrainerWithCache(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)

Without cache¶

np.random.seed(42)

start = time.time()
print(trainer.generate(dataset, prompt="ROMEO:", use_cache=False))
print(f"{time.time() - start:.8f} seconds")

ROMEO:
When I pray.

CAMILLO:
In my liest, I have done drunk.

CAMILLO:
Ay, since the speak, my doubt, fainners!
I am too, as to a kind in honourity,
Can a very disentent till take from here other
out, strike than 'Which children choistones
That would how have m
4.86582589 seconds

With cache¶

np.random.seed(42)

start = time.time()
print(trainer.generate(dataset, prompt="ROMEO:", use_cache=True))
print(f"{time.time() - start:.8f} seconds")

ROMEO:
When I pray.

CAMILLO:
In my liest, I have done drunk.

CAMILLO:
Ay, since the speak, my doubt, fainners!
I am too, as to a kind in honourity,
Can a very disentent till take from here other
out, strike than 'Which children choistones
That would how have m
0.62065578 seconds

In this chapter, we implemented the KV cache: cache the keys and values of the tokens already generated, so that a new token computes attention only against the cache. The time complexity of generating each token drops from $O(n^2)$ to $O(n)$ , and generation is much faster.

Exercises¶

Adjust the maximum generation length steps and observe the difference in speed between no cache and with cache. Does the advantage of the cache grow more obvious the longer the generation?