QKV 与 KV Cache：从注意力计算到推理缓存（图解 + PyTorch 源码）

1. QKV 总流程：一次注意力的六步源码

读图 1：注意力前向就三件事——X 分别乘三个矩阵得到 Q（查询）、K（钥匙）、V（内容）；Q 与全部 K 做缩放点积并遮住未来；softmax 得到的权重对 V 加权求和。d_h = d_model / H，多头维度是切出来的，不是新参数。下面六步源码与图一一对应：

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleAttention(nn.Module):
    """对应图 1：一次完整的 QKV 注意力前向。"""
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.dh = d_model // n_heads
        # 三个投影矩阵，分别得到 Q、K、V
        self.Wq = nn.Linear(d_model, d_model, bias=False)
        self.Wk = nn.Linear(d_model, d_model, bias=False)
        self.Wv = nn.Linear(d_model, d_model, bias=False)
        # 输出投影
        self.Wo = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor, causal: bool = True):
        # x: [B, S, d_model]
        B, S, _ = x.shape

        # 1) 线性投影：X -> Q, K, V
        q = self.Wq(x)  # [B, S, d_model]
        k = self.Wk(x)
        v = self.Wv(x)

        # 2) 拆多头：[B, S, d_model] -> [B, H, S, dh]
        q = q.view(B, S, self.h, self.dh).transpose(1, 2)  # [B, H, S, dh]
        k = k.view(B, S, self.h, self.dh).transpose(1, 2)
        v = v.view(B, S, self.h, self.dh).transpose(1, 2)

        # 3) 相似度：Q Kᵀ / √dh -> [B, H, S, S]
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)

        # 4) causal mask：上三角（不含对角线）置 -inf，禁止看未来
        if causal:
            mask = torch.triu(
                torch.ones(S, S, dtype=torch.bool, device=x.device), diagonal=1
            )
            scores = scores.masked_fill(mask, float("-inf"))

        # 5) softmax -> 注意力权重，再乘 V
        attn = F.softmax(scores, dim=-1)        # [B, H, S, S]
        out = attn @ v                          # [B, H, S, dh]

        # 6) 拼回多头 + 输出投影
        out = out.transpose(1, 2).contiguous().view(B, S, -1)  # [B, S, d_model]
        return self.Wo(out)

2. 单个 token 的视角：causal mask 在挡什么

谁在查：生成进行到第 5 个位置，前 4 个 token 是它的历史。当前 token 生成一个查询向量 q₅，拿它去和每个历史 token 的钥匙 k₁、k₂、k₃、k₄ 做点积——点积越大，越需要那个历史 token 的信息。

分数怎么变权重：四个点积先除以 √d_h 缩放，再 softmax 归一化，得到注意力权重 a₁、a₂、a₃、a₄。含义很直白：当前 token 给每个历史位置分配多少注意力。

未来为什么看不见：causal mask 在 softmax 之前，把第 5 个位置及以后的分数直接置为 -∞——softmax 后权重为 0，等价于完全不参与。作用就一条：第 t 个 token 只能依据 1 至 t 的历史做决策，绝不偷看未来。

输出怎么来：o₅ = 用权重对历史内容 v₁、v₂、v₃、v₄ 加权求和——权重越大，对应 v 的影响越强。

一句话收拢：用 Q 去和 K 匹配得到权重，再用权重加权汇总 V。训练时全部位置并行做这件事（第 i 个位置只能看前 i 个）；推理时每个位置只剩「当前一个 q 查全部历史」——这正是 KV Cache 要服务的场景。

3. 多头注意力：拆头不是循环 32 次

读图 3：代码里没有「对每个头循环」——view + transpose 一次把 [B,S,d_model] 重排成 [B,H,S,dh]，32 个头在同一组矩阵乘里并行；concat 是它的逆操作（transpose 回来再 view）。批、头、序列维在张量形状里都有名字，形状对不上时先读形状再读代码。

class MultiHeadAttention(nn.Module):
    """对应图 3：显式展示按头拆分 / 并行 / 拼接。"""
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.h, self.dh = n_heads, d_model // n_heads
        self.Wq = nn.Linear(d_model, d_model, bias=False)
        self.Wk = nn.Linear(d_model, d_model, bias=False)
        self.Wv = nn.Linear(d_model, d_model, bias=False)
        self.Wo = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, causal=True):
        B, S, _ = x.shape
        q, k, v = self.Wq(x), self.Wk(x), self.Wv(x)   # [B,S,d_model]

        # 按头拆分
        def split(t):
            return t.view(B, S, self.h, self.dh).transpose(1, 2)
        q, k, v = split(q), split(k), split(v)         # [B,H,S,dh]

        # 每个头并行计算（矩阵一次搞定，等价于 for 循环）
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)  # [B,H,S,S]
        if causal:
            m = torch.triu(torch.ones(S, S, dtype=torch.bool, device=x.device), 1)
            scores = scores.masked_fill(m, float("-inf"))
        heads = F.softmax(scores, dim=-1) @ v          # [B,H,S,dh]

        # concat 多头：转置回 [B,S,H,dh] 再 view
        concat = heads.transpose(1, 2).contiguous().view(B, S, self.h * self.dh)
        return self.Wo(concat)

4. 没有 KV Cache：每步都在重算

读图 4：浪费发生在 K/V 投影上——历史的钥匙和内容明明没变，每生成一个新 token 却要全部重算一遍（generate_naive 每步 forward 整段序列）。注意点积本身每步本来就要做（新的 q 要扫全部 k），缓存省的是「投影重算」这部分。repeat_kv_demo 把这个重复显式写了出来：

@torch.no_grad()
def generate_naive(model, prompt_ids, max_new_tokens):
    """
    对应图 4：每生成一个 token，都把整段历史重新 forward 一遍。
    """
    ids = prompt_ids
    for _ in range(max_new_tokens):
        logits = model(ids)                # [B, S, vocab]
        next_tok = logits[:, -1].argmax(-1, keepdim=True)  # 贪心
        ids = torch.cat([ids, next_tok], dim=1)
    return ids


def repeat_kv_demo(x_full: torch.Tensor,
                   Wk: nn.Linear, Wv: nn.Linear,
                   steps: int = 3):
    """
    演示“重复计算 K/V”：每一步都算一遍 k = Wk(x)。
    """
    all_k = []
    for t in range(1, steps + 1):
        x_t = x_full[:, :t, :]        # 取前 t 个 token
        k = Wk(x_t)                   # ← 重复计算历史 K
        v = Wv(x_t)                   # ← 重复计算历史 V
        all_k.append(k[:, -1, :])
    return all_k

5. 有 KV Cache：Prefill 一次写满，Decode 每步追加

读图 5：缓存的写入分两段——prefill 把整段 prompt 一次算完写入（并行、吃算力），decode 每步只处理 1 个新 token、算它的 q/k/v 后把 k、v 追加进缓存（串行、吃访存）。代码里 prefill 与 decode 共用同一个 attention，差别只在输入长度 S：S 大于 1 时需要 causal mask，S=1 时天然无未来。下面三段是缓存容器、带缓存的注意力和生成主循环：

class KVCache:
    """对应图 5：一个极简 KV Cache 容器，按层存 K/V。"""
    def __init__(self):
        self.k = None   # [B, H, t, dh]
        self.v = None

    def append(self, k_new: torch.Tensor, v_new: torch.Tensor):
        # k_new/v_new: [B, H, 1, dh]
        if self.k is None:
            self.k, self.v = k_new, v_new
        else:
            self.k = torch.cat([self.k, k_new], dim=2)  # 沿序列维拼接
            self.v = torch.cat([self.v, v_new], dim=2)

    @property
    def seq_len(self) -> int:
        return 0 if self.k is None else self.k.size(2)


class CachedAttention(nn.Module):
    """Prefill 和 Decode 共用同一个 attention，只是输入长度不同。"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.dh = n_heads, d_model // n_heads
        self.Wq = nn.Linear(d_model, d_model, bias=False)
        self.Wk = nn.Linear(d_model, d_model, bias=False)
        self.Wv = nn.Linear(d_model, d_model, bias=False)
        self.Wo = nn.Linear(d_model, d_model, bias=False)
        self.cache = KVCache()

    def _split(self, t):
        B, S, _ = t.shape
        return t.view(B, S, self.h, self.dh).transpose(1, 2)  # [B,H,S,dh]

    def forward(self, x, is_prefill: bool):
        B, S, _ = x.shape
        q = self._split(self.Wq(x))         # [B,H,S,dh]
        k = self._split(self.Wk(x))
        v = self._split(self.Wv(x))

        # 把新 k/v 追加进 cache
        self.cache.append(k, v)
        K_all, V_all = self.cache.k, self.cache.v   # [B,H,total,dh]

        scores = q @ K_all.transpose(-2, -1) / math.sqrt(self.dh)  # [B,H,S,total]

        # Prefill 且 S>1 才需要 causal mask；Decode 时 S=1 天然无未来
        if is_prefill and S > 1:
            offset = K_all.size(2) - S  # cache 里已有的历史长度（本例 prefill 时为 0）
            row = torch.arange(S, device=x.device)[:, None] + offset
            col = torch.arange(K_all.size(2), device=x.device)[None, :]
            mask = col > row
            scores = scores.masked_fill(mask, float("-inf"))

        attn = F.softmax(scores, dim=-1)
        out = attn @ V_all                                   # [B,H,S,dh]
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.Wo(out)


@torch.no_grad()
def generate_with_cache(model, prompt_ids, max_new_tokens):
    """第一次 Prefill，之后每步只喂 1 个 token。"""
    logits = model(prompt_ids, is_prefill=True)
    next_tok = logits[:, -1].argmax(-1, keepdim=True)

    for _ in range(max_new_tokens - 1):
        logits = model(next_tok, is_prefill=False)
        next_tok = logits[:, -1].argmax(-1, keepdim=True)

    return next_tok

6. Decode 单步的内部数据流

读图 6：一步的形状账——输入 [B,1,d_model]，投影拆头后 q/k/v 各 [B,H,1,dh]；缓存里的 K/V 是 [B,H,t,dh]；点积 scores 为 [B,H,1,t]，softmax 后乘 V_cache 得 [B,H,1,dh]。t 每步加 1，这就是「缓存随生成线性长高」的单步微观图。

@torch.no_grad()
def decode_one_step(x_t: torch.Tensor,
                    cache: KVCache,
                    Wq: nn.Linear, Wk: nn.Linear,
                    Wv: nn.Linear, Wo: nn.Linear,
                    n_heads: int, dh: int):
    """
    对应图 6：Decode 单步。
    x_t: [B, 1, d_model]
    """
    B = x_t.size(0)

    def split(t):
        return t.view(B, 1, n_heads, dh).transpose(1, 2)  # [B,H,1,dh]

    # 只算当前 token 的 q/k/v
    q_t = split(Wq(x_t))    # [B,H,1,dh]
    k_t = split(Wk(x_t))
    v_t = split(Wv(x_t))

    # 追加进 cache：t 从 1 变长
    cache.append(k_t, v_t)

    # 用当前 q 去查全部历史 K
    scores = q_t @ cache.k.transpose(-2, -1) / math.sqrt(dh)  # [B,H,1,t]
    attn = F.softmax(scores, dim=-1)                          # [B,H,1,t]

    # 加权求和全部 V
    out = attn @ cache.v                                      # [B,H,1,dh]
    out = out.transpose(1, 2).contiguous().view(B, 1, n_heads * dh)
    return Wo(out)                                            # [B,1,d_model]

7. 缓存是逐层存在的

读图 7：KV Cache 不是一个大数组，而是每层各存各的——输入每穿过一个 Block，就把当前 token 在这一层的 k、v 追加进这一层自己的缓存；一个 32 层模型就是 32 对抽屉。CachedTransformer 把这个结构直接写成代码，inspect_cache_shapes 可以打印每层的实际形状。

class CachedTransformer(nn.Module):
    """每层各有一份 KV Cache。"""
    def __init__(self, n_layers, d_model, n_heads):
        super().__init__()
        self.layers = nn.ModuleList([
            CachedAttention(d_model, n_heads) for _ in range(n_layers)
        ])

    def forward(self, x, is_prefill: bool):
        for layer in self.layers:
            x = layer(x, is_prefill=is_prefill)
        return x

    def reset_cache(self):
        for layer in self.layers:
            layer.cache.k = None
            layer.cache.v = None


def inspect_cache_shapes(model: CachedTransformer):
    for i, layer in enumerate(model.layers):
        k, v = layer.cache.k, layer.cache.v
        if k is not None:
            print(f"layer {i}: K={tuple(k.shape)}, V={tuple(v.shape)}")

8. Cache 随序列增长与显存账

读图 8：增长是严格的线性——没有压缩、没有淘汰，历史全部保留。显存公式：2 × 层数 × batch × 序列长 × KV 头数 × 每头维数 × 每元素字节数（乘 2 是 K、V 各一份）。按本篇基线（32 层、8 个 KV 头、每头 128、fp16 即 2 字节）代入：每 token 每层 4KB，全模型每 token 128KB。代码里的 kv_cache_memory_bytes(32, 1, 4096, 8, 128) 算出 0.50GB，即下表第二行：

def kv_cache_memory_bytes(n_layers, batch, seq_len,
                          n_kv_heads, head_dim,
                          dtype_bytes=2):
    """
    KV Cache 显存公式：
      2 * L * B * S * H_kv * d_h * bytes
    """
    return 2 * n_layers * batch * seq_len * n_kv_heads * head_dim * dtype_bytes


def demo_growth():
    """演示 cache 每步线性增长的形状与显存。"""
    cache = KVCache()
    B, H, dh = 1, 8, 64
    for step in range(1, 5):
        k_new = torch.randn(B, H, 1, dh)
        v_new = torch.randn(B, H, 1, dh)
        cache.append(k_new, v_new)
        print(f"step {step}: K={tuple(cache.k.shape)}, "
              f"V={tuple(cache.v.shape)}, "
              f"bytes={cache.k.numel() * cache.k.element_size() * 2}")


if __name__ == "__main__":
    print("KV Cache ≈ "
          f"{kv_cache_memory_bytes(32, 1, 4096, 8, 128) / 1024**3:.2f} GB")
    demo_growth()

与不分组对照：末列是同样长度下 32 个 KV 头（MHA）的账——本文基线的 8 头分组已经把显存压到四分之一，第 10 章把这条压缩路线补全。

换算到生成过程：decode 每生成 1 个 token，缓存就多 128KB，一千步就是 128MB——同一张卡能并发几条请求，先数缓存。

9. 两条推理路径对照：平方重算 vs 线性追加

def no_cache_characteristics(prompt_len: int, gen_len: int):
    """
    图 9 的定量描述：
      每步 forward 的 token 数 = prompt_len + step
      总 token-forward 数 ≈ Σ (prompt_len + i)，i=1..gen_len
    """
    total = sum(prompt_len + i for i in range(1, gen_len + 1))
    return {
        "特点": ["每步全量重算 K/V", "计算量随步数平方增长", "显存占用恒定（只有激活）", "实现简单"],
        "总 token-forward 数": total,
        "适用场景": "训练 / 简单原型 / 短序列生成",
    }


def no_cache_peak_memory(prompt_len, gen_len, d_model, n_layers, bytes_per=2):
    """
    无 cache 时显存主要来自激活；峰值与最长一次 forward 长度近似成正比，
    并不随生成长度累积（因为每步都释放）。
    """
    max_len = prompt_len + gen_len
    # 粗略估计：每层激活 ≈ 数个 [B, S, d_model]
    return max_len * d_model * n_layers * 4 * bytes_per

def with_cache_characteristics(prompt_len: int, gen_len: int):
    """
    图 10 的定量描述：
      Prefill 一次算 prompt_len 个 token；
      每个 Decode step 只算 1 个 token。
      总 token-forward 数 = prompt_len + gen_len
    """
    total = prompt_len + gen_len
    return {
        "特点": ["Prefill 一次建 cache", "Decode 只算新 token", "计算量随生成长度线性增长",
                 "显存随生成长度线性增长"],
        "总 token-forward 数": total,
        "适用场景": "自回归推理 / 长文本生成 / 生产部署",
    }


def cache_memory_at_step(prompt_len, gen_len, n_layers, n_kv_heads,
                         head_dim, batch=1, dtype_bytes=2):
    """
    有 cache 时显存的“增长轨迹”：
      step 0 = Prefill 后，长度 = prompt_len
      step i = 长度 = prompt_len + i
    """
    trajectory = []
    for i in range(gen_len + 1):
        S = prompt_len + i
        bytes_ = 2 * n_layers * batch * S * n_kv_heads * head_dim * dtype_bytes
        trajectory.append((f"step {i}", S, bytes_ / 1024**2))
    return trajectory


if __name__ == "__main__":
    print("=== 图 9 无 Cache ===")
    print(no_cache_characteristics(128, 64))

    print("\n=== 图 10 有 Cache ===")
    print(with_cache_characteristics(128, 64))
    for name, S, mb in cache_memory_at_step(128, 64, 32, 8, 128):
        print(f"{name}: seq={S}, KV≈{mb:.1f} MB")

10. 优化路线与一页记忆卡

读图 11：记忆卡的四根分支正好是本篇的四条线——QKV 的计算、缓存的存在形式、两条推理路径的代价、以及压缩缓存的优化。优化分支的共同思路：Q 头一个不少（32 个全保留），压的是 K/V 的份数或精度。同一每头 128 维口径下的量级对照：

PagedAttention（块池）是另一维度的优化——不减每 token 体积，而是把缓存切块按需分配、物理散放、块表拼回，让预留显存可共享（站内《DeepSeek V4.1 全链路主时序》步骤 17 的块表就是它的工程现场）。GroupedQueryAttention 的完整源码：

class GroupedQueryAttention(nn.Module):
    """
    优化分支：Q 有 H 个头，K/V 只有 G 个头（G < H），
    多个 query 共享一个 KV 头，降低 KV Cache 显存。
    """
    def __init__(self, d_model, n_q_heads, n_kv_heads):
        super().__init__()
        assert n_q_heads % n_kv_heads == 0
        self.hq, self.hkv = n_q_heads, n_kv_heads
        self.dh = d_model // n_q_heads
        self.group = n_q_heads // n_kv_heads

        self.Wq = nn.Linear(d_model, self.hq * self.dh, bias=False)
        self.Wk = nn.Linear(d_model, self.hkv * self.dh, bias=False)
        self.Wv = nn.Linear(d_model, self.hkv * self.dh, bias=False)
        self.Wo = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, causal=True):
        B, S, _ = x.shape
        q = self.Wq(x).view(B, S, self.hq, self.dh).transpose(1, 2)   # [B,hq,S,dh]
        k = self.Wk(x).view(B, S, self.hkv, self.dh).transpose(1, 2)  # [B,hkv,S,dh]
        v = self.Wv(x).view(B, S, self.hkv, self.dh).transpose(1, 2)

        # 把每个 KV 头复制 group 次，以对齐 Q 头
        k = k.repeat_interleave(self.group, dim=1)   # [B,hq,S,dh]
        v = v.repeat_interleave(self.group, dim=1)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)
        if causal:
            m = torch.triu(torch.ones(S, S, dtype=torch.bool, device=x.device), 1)
            scores = scores.masked_fill(m, float("-inf"))
        out = F.softmax(scores, dim=-1) @ v
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.Wo(out)

11. 与站内文章的关系

12. 术语速查

来源：11 图与 10 段代码收录于 2026-10-18 底稿（QKV 与 KV Cache 图解，全部可运行）；显存账按底稿基线（32 层、8 KV 头、每头 128、fp16）推得，MHA/MLA 对照口径取自《多头注意力机制精读》