Generating text and the KV cache

Module 01, Lesson 3 ran temperature experiments with DeepSeek, when the model was a black box. Now the model is one we trained ourselves, and we can open it up to see what it actually does at each step.

python generate.py

All the model gives is probabilities

After reading a text, the model gives each of the 6,289 characters in its vocabulary a score at the last position, which softmax turns into probabilities:

== 1. 读完一段文字，模型给下一个字的概率（前 8 名）
  「白日依山尽，黄河入海」→ 无 6.9%  间 4.5%  多 4.3%  深 3.7%  边 3.0%  空 2.6%  难 2.3%  遥 2.1%
  「床前明月光，疑是地上」→ 天 3.2%  花 2.9%  楼 2.4%  清 1.8%  看 1.7%  游 1.7%  人 1.6%  枝 1.3%
  「春风」→ 。 14.2%  ， 5.1%  吹 3.3%  满 2.9%  雨 1.9%  起 1.9%  尽 1.4%  落 1.3%

After "黄河入海" ("the Yellow River into the sea"), the original poem has "流" (flows), but "流" isn't in the top 8. After "疑是地上" ("I took it for, on the ground,"), "霜" (frost) isn't either. We checked last lesson that this model doesn't recite: it didn't memorise these two famous poems, and just gives "reasonable" characters according to the patterns it learned.

An interesting discovery along the way: the version of "静夜思" ("Quiet Night Thoughts") in the Complete Tang Poems is "床前看月光，疑是地上霜。举头望山月，低头思故乡", different from the textbook version "床前明月光……举头望明月". The version we know is generally thought to come from Ming and Qing dynasty anthologies. So even if the model had memorised the poem, it would have memorised "看月光".

After "春风" (spring breeze), the most likely next item is a full stop (14.2%). I counted in the training data: "春风" appears 797 times, 111 of them immediately followed by a full stop, 13.9%, almost exactly the probability the model gives. What the model learned is how these characters occur in Tang poetry.

Generating text means doing the same thing over and over: get the probabilities, choose a character, append it, and compute the probabilities for the next one. How you "choose" determines what gets written.

Temperature

Module 01 covered temperature: divide the scores by the temperature, then softmax. Below 1, the gap between high and low scores widens and the model is more "conservative"; above 1, the gap narrows and the model is more "adventurous".

logits = logits[:, -1, :] / temperature
...
next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)

With the same random seed, 3 poems at each of three temperatures:

  温度 0.3：
    白云无限意，白发不如何。何事无人在，无人不可知。
    一望东山路，千行万里愁。夜深人不见，秋色月无声。
    江南春水上，江上月明时。独有清风起，何人别有期。
  温度 1.0：
    白露起春光，知君得舞酲。霜时飏先没，桐死势悠扬。
    白发摩牛放，真关万趾兴。夜寒休绕郡，秋杀杜娘宫。玉镜收新镜，金尊坐晚麕。何言吊鱼道，双锡空余年。
    戍楼映江岸，淼危棹度斜。欲分凝炯娩，归盖乱精冰。
  温度 1.5：
    淡露依春似映兹，襞舞凋年雪路揉筝衬情。死势悠扬双觉处，摩牛放将溉噭胞。……

Temperature 0.3 reads most smoothly; "一望东山路，千行万里愁。夜深人不见，秋色月无声" ("one gaze at the east mountain road, a thousand lines of ten-thousand-li sorrow; deep in the night no one is seen, autumn colours, the moon without a sound") could almost pass for the real thing. The price is very common characters and lots of repetition: "无人" (no one) appears twice, "江" (river) twice, and after reading a few they all feel the same.

Temperature 1.0 chooses characters by the model's raw probabilities. The vocabulary is richer, but rare characters like "麕", "酲" and "炯娩" start appearing, and the meaning is harder to follow.

Temperature 1.5 falls apart completely: the format collapses ("襞舞凋年雪路揉筝衬情" is a nine-character line) and it's full of rare characters. The reason is that at high temperature the probability distribution is flattened, and the characters among the 6,000-odd whose probability was originally tiny together take a sizeable share. Once a strange character is chosen, the following context becomes strange, and the errors pile up.

top-k: choose only among the leading characters

The problem at high temperature is the "long tail": thousands of low-probability characters together take too big a share. top-k is direct: keep only the k most probable characters, drop all the rest, and choose among those k by probability.

if top_k is not None:
    kth = torch.topk(logits, top_k).values[:, -1:]
    logits = logits.masked_fill(logits < kth, float("-inf"))

Again at temperature 1.5, now with top-k=20:

== 3. 温度 1.5，但只从概率最高的 20 个字里选（top-k=20）
    白头白马在西陵，自是西林下路尘。今日不辞青史士，可堪回首白苹州。
    万国风尘里，千金雨色深。不知归未到，应似到人中。
    东北山川尽不愁，三千春尽一枝愁。春寒未入千门色，雨尽空闻九夜寒。……

The format is back and the rare characters are gone, but there's more variety than at temperature 0.3. High temperature supplies the "adventure" and top-k keeps it from going "too far"; together they work best.

Another common method is top-p (nucleus sampling): instead of a fixed number, add characters from the top down until their cumulative probability reaches p (say 0.9). When probability is concentrated only a few characters remain; when it's spread out, more are kept. We saw this parameter in Module 01 when using the API.

Continuation and acrostic poems

Generation doesn't have to start from nothing. Give it the first half of a line and the model carries on:

== 4. 给出开头，让它续写
    床前明月光，已有四邻情。知君得舞袖，何奈在山期。
    大漠孤烟直，高楼楚水清。雁雕摩露下，鸿劒鼓歌中。夜雨初分郡，秋风又过巴。此才应不遇，空与钓鱼舟。
    人工智能好，言蠹不见人。徒思一岁月，独在岘生春。

"大漠孤烟直" ("in the vast desert a lone column of smoke rises straight") was continued into a complete eight-line five-character regulated verse. It can even carry on from "人工智能" (artificial intelligence), a word the Tang dynasty never had, because each of those four characters appears in Tang poetry.

Going further, we can intervene during generation. An acrostic poem requires the first character of each line to be specified: put the specified character in first, then let the model write the rest of the line one character at a time, forbidding it from outputting punctuation or newlines while it does, and add the punctuation ourselves at the end:

@torch.no_grad()
def acrostic(heads, n_char=7, temperature=0.8):
    ids = encode("\n")
    for i, head in enumerate(heads):
        ids = torch.cat([ids, encode(head)], dim=1)  # 先把指定的字放进去
        for _ in range(n_char - 1):  # 再让模型一个字一个字地写完这一句
            logits = model(ids)[0][0, -1] / temperature
            logits[BANNED] = float("-inf")
            next_id = torch.multinomial(F.softmax(logits, dim=-1), 1)
            ids = torch.cat([ids, next_id[None]], dim=1)
        ids = torch.cat([ids, encode("，。"[i % 2])], dim=1)  # 标点由我们加，保证格式
    return decode(ids[0].tolist()).strip()

== 5. 藏头诗：每句开头的字由我们指定，其余由模型写
    春色欲依春似春，眠花先得舞花年。不知诗句先酬情，觉后还随楚水人。
    学得重雕出，无情忘却回。止行空向兴，境静更闻闻。

The first characters of the lines read "春眠不觉" ("spring sleep, unaware") and "学无止境" ("learning has no end").

This technique has a formal name: constrained decoding. Before choosing a character, set the scores of options that don't meet the requirements to negative infinity. When Module 02, Lesson 4 had DeepSeek output JSON, the server guaranteed the output format was correct using the same principle: at every step, exclude the tokens that would make the JSON invalid.

Repeated computation

Now look at generation efficiency. Back in the generate method, the simplest approach feeds the whole text into the model at every step:

logits, _ = self(idx)  # 不用缓存：每一步都把整段序列从头算一遍

When writing the 50th character, the model processes all 49 characters before it; when writing the 51st, it processes those 49 all over again, plus the 50th.

But think about it: the results for the first 49 characters haven't changed. The causal mask ensures each position looks only backwards, so appending a new character doesn't affect the computation at earlier positions at all. The only thing that's truly new is the last position.

What does the last position need to compute attention? Its own q, plus the k and v of every earlier position. So we just store each layer's k and v, compute only the new character's q, k and v at each step, and append the new k and v to the stored ones. That's the KV cache.

Implementation

In gpt.py's attention, the cache adds only a few lines:

if cache is not None:  # KV 缓存：把之前算过的 K、V 接在前面，不用重算
    if "k" in cache:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v

One detail needs care: the causal mask. Without the cache, q and k have the same length and the mask is a square lower triangle. With the cache, q has only the 1 new character while k has all of them, and the new character comes last, so it can see every k. So the mask has to be offset:

total = k.size(2)
mask = torch.ones(T, total, dtype=torch.bool, device=x.device).tril(diagonal=total - T)

The position embedding has to follow suit: the new character isn't at position 0 but at position start. That's what the start parameter in forward is for.

During generation, the first step feeds in the whole prompt and stores the k and v of every position in the cache; after that, each step feeds in only the latest character:

if use_cache:
    # 第一步把整段提示词送进去；之后每步只送最新的一个词元，前面的 K、V 从缓存里取
    logits, _ = self(idx[:, start:], caches=caches, start=start)
    start = idx.size(1)

Measuring it

== 6. KV 缓存：结果一样吗？快了多少？
  同样的随机种子，写 120 个字：用缓存和不用缓存的结果完全一样
  写  32 个字：不用缓存     23 毫秒，用缓存    12 毫秒，快了 1.8 倍
  写  64 个字：不用缓存     52 毫秒，用缓存    26 毫秒，快了 2.0 倍
  写 120 个字：不用缓存    130 毫秒，用缓存    49 毫秒，快了 2.6 倍
  这个模型每个词元的缓存：2 × 4 层 × 128 维 × 4 字节 = 4 KB，128 个词元共 512 KB

First confirm correctness: with the same random seed, the 120 characters written with and without the cache are exactly the same. The cache only saves repeated computation; the result doesn't change.

Then the speed. The more characters written, the more is saved: 1.8 times faster for 32 characters, 2.6 times for 120. Without the cache, writing the n-th character processes n characters, so total computation grows with the square of the length; with the cache, each step processes 1 character.

On our model the speed-up is only a little over two times, nothing dramatic. The reason is that the model is too small and the text too short: the real computation at each step takes a fraction of a millisecond, and Python's own overhead of calling functions and concatenating tensors is a large share, which the cache can't save. On an LLM, each step computes far more, and texts run to thousands or tens of thousands of tokens, so the KV cache saves orders of magnitude.

The cost of the cache: memory

The KV cache trades memory for time. Each token stores one k and one v in every layer. Our model needs only 4 KB per token, 512 KB for 128 tokens, negligible.

LLMs are different. For a model with dozens of layers and thousands of dimensions per layer, each token's cache may take several hundred KB to over 1 MB. A long conversation of a hundred-odd thousand tokens needs tens of GB of GPU memory for the cache, often more than the model itself. When serving many users at once, it's the main consumer of GPU memory.

So many of the improvements mentioned at the end of last lesson target the KV cache: several heads sharing k and v (grouped-query attention), or compressing k and v into smaller vectors before storing them (DeepSeek's multi-head latent attention, MLA). When Module 10 covers vLLM, you'll also see how servers manage these caches.

The "cache-hit input is cheaper" from Module 06, Lesson 4 is a cache at another level: the provider keeps the KV cache for common prefixes (such as a fixed system prompt), and when the next request uses the same prefix, retrieves it directly without recomputing, which is why it can be so much cheaper.

Exercises

1. Implement top-p sampling: sort the probabilities from high to low and keep the characters whose cumulative probability reaches p. Write a few poems at temperature 1.5 with p=0.9 and compare with top-k=20.
2. Write a "fixed format" constrained decoder: allow only seven-character quatrains, where the 8th, 16th and 24th characters must be a comma or full stop and no other position may contain punctuation.
3. Change gpt.py's block_size to 512 and retrain (or just initialise randomly), and time writing 500 characters with and without the cache.

Self-check

1. What do temperature and top-k each change, and why are they often used together?

Temperature scales all the scores before softmax, changing how "sharp" the probability distribution is: at low temperature it's concentrated on a few high-scoring characters, at high temperature it's more even. top-k keeps only the k most probable characters and drops the rest. High temperature brings variety but lets a mass of low-probability characters take a share; top-k removes that long tail, so together they give variety without going too far.

2. Why doesn't the KV cache change what's generated?

Because of the causal mask, each position's computation depends only on itself and earlier positions, so appending new characters doesn't change earlier positions' k and v. The cache just stores these unchanging results for reuse, saving repeated computation; mathematically it's exactly equivalent.

3. What does the KV cache cost, and why does that cost matter for LLMs?

It costs memory: each token stores a k and a v in every layer. LLMs have many layers and large dimensions, so each token's cache is large, and with long contexts and many simultaneous users the cache can take more GPU memory than the model itself. That's why many architectural improvements (such as grouped-query attention and multi-head latent attention) aim to shrink it.