KV Cache 从 0 到 1：把 nanoGPT 的 generate 加速 100 倍的关键技术

1. 引子：跑 nanoGPT 的 generate() 慢得令人发指

跟着 Karpathy 的 Ep.5 把一个 GPT 训出来之后，第一件事是兴奋地调 generate() 看它"说人话"。然后你会发现一件事：

生成第 1 个 token 很快，生成第 100 个 token 越来越慢，生成第 500 个 token 几乎在爬。

如果做个 profile（或者直接打时间戳），会看到 attention 层吃掉了 80%+ 的时间，而且每生成一个新 token，attention 的耗时会随当前序列长度增长。

2. 复盘 nanoGPT 的 generate()：朴素实现长啥样

Karpathy 在 Ep.5 里的 generate() 是这样的（精简版）：

@torch.no_grad()
def generate(self, idx, max_new_tokens):
    # idx: (B, T) 当前已知的 token 序列
    for _ in range(max_new_tokens):
        # 1) 把上下文截到 block_size（位置编码上限）
        idx_cond = idx[:, -self.block_size:]

        # 2) 整段 forward：每次都从头算
        logits, _ = self(idx_cond)              # (B, T, vocab_size)

        # 3) 只取最后一个位置的预测
        logits = logits[:, -1, :]               # (B, vocab_size)
        probs  = F.softmax(logits, dim=-1)

        # 4) 采样下一个 token
        idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
        idx = torch.cat([idx, idx_next], dim=1)              # 拼到末尾
    return idx

整个函数只有四步，但第 2 步隐藏了所有的性能问题——"整段 forward"意味着：每生成一个新 token，模型把已经算过 N 遍的历史 token 又算了一遍。

3. 哪些计算在重复：把浪费量化出来

假设当前已经生成了 3 个 token：[A, B, C]。在 self-attention 层，每个 token 都会算出自己的 Q / K / V：

然后 attention 输出（causal mask 保证只看自己和过去）：

generate() 只拿 out_2（最后一个位置）去预测 token D，out_0 和 out_1 算了就扔。

下一步生成 D 时，输入变成 [A, B, C, D]，又把 attention 跑一遍：

把浪费量化

设序列长度 $T$、hidden_dim $d$。单步 attention 计算量 $\\mathcal{O}(T^2 \\cdot d)$（Q @ K^T 主导）。

生成 $N$ 个 token 的总计算量：

$$\\sum_{t=1}^{N} \\mathcal{O}(t^2 \\cdot d) = \\mathcal{O}(N^3 \\cdot d)$$

而真正"有用"的计算只有每步最后一个位置的 attention，也就是 $\\mathcal{O}(t \\cdot d)$ 这部分。其余 $\\mathcal{O}(t^2 - t)$ 都是冗余。

4. 关键 Insight：K、V 为什么可缓存（两条事实）

整个 KV Cache 的理论基础是两条几乎显然的事实，但必须同时成立：

合起来："历史 token 的 K、V 是只读的不可变数据"。这跟操作系统里"只读页可以无成本共享"是一个味道。

5. Q 为什么不需要 cache

很多人第一次看 KV Cache 会问：那 Q 为什么不也缓存？答案分两层：

① Q 也是只依赖自身位置的（理论上可以缓存）

从公式看，Q_i = x_i · W_q，跟 K、V 完全对称，也是只依赖自身。所以理论上"可以"缓存。

② 但 decode 阶段每步只用一个 Q，缓存它没意义

生成第 t 个 token 时，attention 输出公式：

只有 Q_t（当前位置的 Q）参与运算。下一步生成 t+1 时，又只用 Q_{t+1}。历史 Q 算了就扔，根本不会再读——所以缓存它纯粹浪费显存。

6. Prefill 与 Decode：两个阶段，两套行为

引入 KV Cache 之后，推理过程被自然切成两个阶段。这两个阶段的计算模式完全不同，是后面理解 vLLM 调度器的关键。

Prefill：处理用户的 prompt（一次性、并行）

用户输入 100 个 token 的 prompt，这 100 个 token 一起塞进模型：

# x: (1, 100, d) ← 整段 prompt 的 embedding
Q = x @ W_q                 # (1, 100, d)  ← 100 个位置的 Q 同时算
K = x @ W_k                 # (1, 100, d)  ← 100 个位置的 K 同时算
V = x @ W_v                 # (1, 100, d)  ← 100 个位置的 V 同时算

cache_K = K                 # 全部 K 一次性塞进 cache
cache_V = V                 # 全部 V 一次性塞进 cache

# 标准 causal attention
out = softmax(Q @ K.T / √d + causal_mask) @ V

# 只用最后一个位置的输出预测第 101 个 token

Decode：一个一个吐 token（串行、每步只算 1 个）

# x_new: (1, 1, d) ← 只有 1 个新 token 的 embedding
q     = x_new @ W_q          # (1, 1, d)  ← 只算 1 个 Q
k_new = x_new @ W_k          # (1, 1, d)  ← 只算 1 个 K
v_new = x_new @ W_v          # (1, 1, d)  ← 只算 1 个 V

# 把新 K、V 追加到 cache 尾部
cache_K = torch.cat([cache_K, k_new], dim=1)   # (1, T+1, d)
cache_V = torch.cat([cache_V, v_new], dim=1)   # (1, T+1, d)

# 注意：因为只有 1 个 Q，不需要 causal mask（它"过去"的全是有效的）
out = softmax(q @ cache_K.transpose(-1,-2) / √d) @ cache_V   # (1, 1, d)

两阶段对比

7. 改 nanoGPT：手写一个 KV Cache 版 generate()

把上面的 idea 落到 Karpathy 的 nanoGPT 上。我们需要做三处改动：

1. 改 CausalSelfAttention.forward，让它接受外部 cache 并能追加
2. 改 Block.forward，把 cache 透传
3. 改 GPT.generate，引入 prefill / decode 两阶段

7.1 改 Attention 层

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        # 训练时还要 causal mask，这里略

    def forward(self, x, kv_cache=None):
        """
        x: (B, T_in, C)  Prefill 时 T_in = prompt 长度；Decode 时 T_in = 1
        kv_cache: None 表示 prefill；否则是 (cache_k, cache_v)，每个 (B, n_head, T_past, head_dim)
        """
        B, T_in, C = x.shape
        qkv = self.qkv(x)                                 # (B, T_in, 3C)
        q, k, v = qkv.split(C, dim=2)                     # 各 (B, T_in, C)

        # reshape 成多头：(B, n_head, T_in, head_dim)
        q = q.view(B, T_in, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T_in, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T_in, self.n_head, self.head_dim).transpose(1, 2)

        # ====== KV Cache 核心：拼接历史 ======
        if kv_cache is not None:
            cache_k, cache_v = kv_cache                   # 每个 (B, n_head, T_past, head_dim)
            k = torch.cat([cache_k, k], dim=2)            # (B, n_head, T_past + T_in, head_dim)
            v = torch.cat([cache_v, v], dim=2)
        new_cache = (k, v)                                # 返回新的 cache

        # ====== Attention 计算 ======
        # 推理时 T_in 在 decode 阶段为 1，q 对全部 k 做注意力——自然只看到过去和当前，不需要 mask
        # Prefill 阶段 T_in > 1 且 cache 为 None，需要 causal mask
        att = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)   # (B, nh, T_in, T_total)
        if kv_cache is None and T_in > 1:
            mask = torch.tril(torch.ones(T_in, T_in, device=x.device))
            att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v                                                 # (B, nh, T_in, head_dim)
        y = y.transpose(1, 2).contiguous().view(B, T_in, C)
        return self.proj(y), new_cache

7.2 改 Block

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1  = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln2  = nn.LayerNorm(n_embd)
        self.ffn  = FeedForward(n_embd)   # 沿用 Ep.5

    def forward(self, x, kv_cache=None):
        attn_out, new_cache = self.attn(self.ln1(x), kv_cache=kv_cache)
        x = x + attn_out
        x = x + self.ffn(self.ln2(x))
        return x, new_cache

7.3 改 GPT.generate（核心）

@torch.no_grad()
def generate(self, idx, max_new_tokens):
    """
    idx: (B, T_prompt)  用户输入
    """
    B, T_prompt = idx.shape
    device = idx.device

    # ============ Prefill：一次性把整段 prompt 灌进去 ============
    tok_emb = self.tok_emb(idx)                         # (B, T_prompt, C)
    pos_emb = self.pos_emb(torch.arange(T_prompt, device=device))  # (T_prompt, C)
    x = tok_emb + pos_emb

    kv_caches = [None] * len(self.blocks)
    for i, block in enumerate(self.blocks):
        x, kv_caches[i] = block(x, kv_cache=None)        # 注意：prefill 时 cache 是 None

    x = self.ln_f(x)
    logits = self.lm_head(x[:, -1, :])                   # 只用最后一个位置预测
    probs = F.softmax(logits, dim=-1)
    idx_next = torch.multinomial(probs, num_samples=1)   # (B, 1)
    idx = torch.cat([idx, idx_next], dim=1)

    # ============ Decode：循环吐 token，每步只算 1 个 ============
    for step in range(max_new_tokens - 1):
        T_past = idx.shape[1] - 1                        # cache 里已有的长度

        tok_emb = self.tok_emb(idx_next)                 # (B, 1, C)
        pos_emb = self.pos_emb(torch.tensor([T_past], device=device))  # (1, C)
        x = tok_emb + pos_emb

        # 透传 cache，每个 block 更新自己的 cache
        for i, block in enumerate(self.blocks):
            x, kv_caches[i] = block(x, kv_cache=kv_caches[i])

        x = self.ln_f(x)
        logits = self.lm_head(x[:, -1, :])
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, idx_next], dim=1)

    return idx

8. 显存代价：每 token 到底吃多少 KV

省时间不是白来的。每生成一个 token，所有层都要往 cache 写一份 K 和一份 V。

单 token 的 KV Cache 大小：

$$\\text{KV per token} = 2 \\times L \\times H \\times d_h \\times \\text{bytes}$$

其中 2 = K + V，L = 层数，H = 头数，$d_h$ = head_dim，bytes = 单元素字节数（fp16 是 2）。

用 LLaMA-7B 算一遍

所以"减少 KV 显存"本身就是一个独立的大方向，分支出几条优化路线：

9. 引出 PagedAttention（下一篇的钩子）

我们刚刚的实现有一个隐藏问题：每个请求要预留多少 cache 显存？

• 给每个请求预分配 max_len？——大部分请求根本生成不到 max_len，浪费
• 按需 torch.cat 增长？——频繁分配释放，显存碎片化，长时间运行后即使总显存够也分不出来连续大块
• 多请求并发时，每个请求 cache 长度不一样，没法批量算 attention（shape 对不上）

这跟 1960 年代操作系统遇到的问题完全一样：每个进程要多少内存？连续分配浪费、按需分配碎片。OS 的解法是虚拟内存 + 分页。

vLLM 论文（PagedAttention，SOSP'23）做的事情：

10. ⭐ 学习过程高频卡点答疑（10 个 Q）

11. 参考资料与 Colab

• Karpathy nanoGPT 仓库：https://github.com/karpathy/nanoGPT
• vLLM 论文：Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP'23)
• vLLM 仓库：https://github.com/vllm-project/vllm（重点看 vllm/v1/core/）
• Hugging Face transformers 的 KV Cache 实现：transformers/cache_utils.py（生产级参考）
• 本篇 Colab Notebook：（待补：把上面 7.1–7.3 的代码塞进一个可跑的 nanoGPT 改造版）

这是 Stage 1 / Week 1 的第一篇，对应"AI Infra 转型完整学习计划"的推理引擎核心阶段开端。
下一篇：博客 07 · PagedAttention 源码解读 - 块管理篇（逐行读 vllm/v1/core/block_pool.py）。