Why Chatbots Are Fast

A chatbot writing a five hundred word answer needs to run its huge network hundreds of times. If each step re-read the entire conversation, it would crawl. But most of that work is exactly the same as the previous step. So what if we just remembered it?

• Hundreds of runs per answer
• Re-reading is wasteful
• Just remember the work

The key value cache

Each new token only needs its own keys and values, because the earlier ones haven't changed. So we store them. In the slow version, every step recomputes all tokens so far. In the fast version, we only compute the new token, and add it to a cache. Run it. A huge difference, and the gap grows with the length. Prompt caching in A P Is uses the same idea, and it saves large amounts of money.

import time
import torch
from torch import nn

dim, T = 256, 400
proj = nn.Linear(dim, dim)
words = torch.rand(1, T, dim)

def slow():
    for t in range(1, T + 1):
        k = proj(words[:, :t])
    return k

def fast():
    cache = []
    for t in range(T):
        cache.append(proj(words[:, t:t + 1]))
    return torch.cat(cache, 1)

with torch.no_grad():
    a = time.time(); slow(); b = time.time(); fast(); c = time.time()
print(f"Without cache: {b - a:.3f}s")
print(f"With cache:    {c - b:.3f}s")

Without cache: 0.019s
With cache:    0.003s

Other speed tricks

Other tricks make serving cheaper still. Continuous batching adds new requests into a running batch. Speculative decoding lets a small model guess ahead, and a large model checks its work. And quantisation shrinks the weights. Together, these have cut the cost of each token by huge factors, over only a few years.

Summary