KV 缓存

先把它放回推理链路

KV Cache 位于每个 Transformer attention 层内部，生命周期跨越同一请求的 prefill 与后续 decode step。它缓存的是已由该层投影并完成位置编码处理后、可直接参与 attention 的历史 K/V；不是模型权重，不是 token ids，也不是 attention score、softmax 概率或 Query。

一条请求通常先经过 prefill：prompt 的 $T_p$ 个 token 并行进入模型，各层一次生成 $T_p$ 个 K/V 并写入缓存。随后进入 decode：每步只有一个新 token，每层只计算新的 $q_t,k_t,v_t$，追加 $k_t,v_t$，再让 $q_t$ 读取历史缓存。标准自回归 attention 的这一层数据流是：

其中 $x_t\in\mathbb{R}^{B\times1\times H}$，$q_t\in\mathbb{R}^{B\times H_q\times1\times D}$，而标准 MHA/GQA/MQA 缓存中的 $K,V\in\mathbb{R}^{B\times H_{kv}\times T\times D}$。单卡教学示例中它们都在 cuda:0；张量并行生产实现会按 KV head 或相关投影维度把缓存分给各 TP rank，容量公式应使用该 rank 实际拥有的 $H_{kv}$。

两条数学等价的 decode 路径

设 prompt 长度为 $T_p$，随后给定 $G$ 个 continuation hidden states。为了隔离缓存机制，下面先不讨论 sampling，而是让两条路径使用同一组固定随机权重和同一组输入。

路径 A：每一步完整重算前缀

第 $j$ 个 decode step 把长度 $T_p+j$ 的完整前缀再次送进 causal attention，只取最后一个位置的输出。历史 token 的所有层激活和 K/V 会重复生成；若用 dense reference attention，当前调用还会物化 $(T_p+j)\times(T_p+j)$ 的 score。

路径 B：prefill 一次，随后增量追加

先对 prompt 做一次 prefill 并保存各层 $K_{0:T_p-1},V_{0:T_p-1}$。第 $j$ 步只投影一个新位置，缓存长度从 $T_p+j-1$ 增加为 $T_p+j$，然后计算 shape 为 $[B,H_q,1,T_p+j]$ 的 score。

在相同权重、相同位置编码、相同 causal 可见范围和相同数值精度下，两条路径最后位置的输出应逐 token 相等：

“约等于”来自浮点运算次序和 backend 差异，不是缓存允许改变模型语义。若误差很大，优先检查 position id、causal mask、GQA head 映射、cache append 维度和 dtype。

图中把每个 token 的 $[H_{kv},D]=[2,4]$ 切片画成一个 slot，并省略了 batch 调度、RoPE 的具体写入点、allocator、page table、kernel fusion 和 TP 分片。逻辑上的“追加”也不保证生产实现每步执行 torch.cat；高性能 runtime 会预分配或分页管理存储，再就地写入新位置。物理 block 如何避免必须连续分配，由 PagedAttention 处理，不是 KV Cache 本身的语义。

shape 与所有权：一个可检查的小例子

取单层 GQA：$B=1,H=16,H_q=4,H_{kv}=2,D=4$，prompt 有 8 个 token，再 decode 5 个 token。prefill 后：

这里 $H_q/H_{kv}=2$，所以每个 KV head 服务两个 Query heads。这个 repeat 是逻辑广播关系；教学代码用 repeat_interleave 显式展开，生产 GQA kernel 通常不真的复制 K/V storage。

可直接运行的对照实验

实验问题：完整重算与增量缓存能否逐 token 对齐？prefill 与纯 decode 分别花多长时间？缓存理论容量与 PyTorch allocator 观测有什么区别？

预期结果：所有断言通过，最终小例子 cache shape 为 [1, 2, 13, 4]；给定教学实现与当前测试 shape，缓存 decode 通常少于完整重算，但具体倍数不能外推到生产 kernel 或其他硬件。CPU 也能运行正确性和计时；只有 CUDA 分支报告显存。

"""Full-prefix recomputation and incremental KV cache, using identical weights."""

import gc, math, platform, statistics, time
import torch

torch.manual_seed(7)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float32

def make_weights(hidden, q_heads, kv_heads, head_dim):
    shapes = [(hidden, q_heads * head_dim), (hidden, kv_heads * head_dim)]
    shapes += [(hidden, kv_heads * head_dim), (q_heads * head_dim, hidden)]
    return tuple(torch.randn(s, device=DEVICE, dtype=DTYPE) / hidden**0.5 for s in shapes)

def project(x, weights, q_heads, kv_heads, head_dim):
    batch, tokens, _ = x.shape
    heads = (q_heads, kv_heads, kv_heads)
    return tuple((x @ w).view(batch, tokens, h, head_dim).transpose(1, 2)
                 for w, h in zip(weights[:3], heads))

def expand_kv(x, q_heads):
    return x.repeat_interleave(q_heads // x.shape[1], dim=1)

def prefill(x, weights, q_heads, kv_heads, head_dim):
    q, k, v = project(x, weights, q_heads, kv_heads, head_dim)
    scores = q @ expand_kv(k, q_heads).transpose(-2, -1) / math.sqrt(head_dim)
    tokens = x.shape[1]
    mask = torch.triu(torch.full((tokens, tokens), -torch.inf, device=DEVICE), 1)
    context = torch.softmax(scores + mask, dim=-1) @ expand_kv(v, q_heads)
    output = context.transpose(1, 2).reshape(x.shape[0], tokens, -1) @ weights[3]
    return output, (k, v)

def decode_one(x_new, cache, weights, q_heads, kv_heads, head_dim):
    q, k_new, v_new = project(x_new, weights, q_heads, kv_heads, head_dim)
    k = torch.cat((cache[0], k_new), dim=2)
    v = torch.cat((cache[1], v_new), dim=2)
    scores = q @ expand_kv(k, q_heads).transpose(-2, -1) / math.sqrt(head_dim)
    context = torch.softmax(scores, dim=-1) @ expand_kv(v, q_heads)
    output = context.transpose(1, 2).reshape(x_new.shape[0], 1, -1) @ weights[3]
    return output, (k, v)

def recompute(prompt, continuation, weights, q_heads, kv_heads, head_dim):
    prefix, outputs = prompt, []
    for x_new in continuation.split(1, dim=1):
        prefix = torch.cat((prefix, x_new), dim=1)
        full_output, _ = prefill(prefix, weights, q_heads, kv_heads, head_dim)
        outputs.append(full_output[:, -1:].clone())
    return torch.cat(outputs, dim=1)

def cached_decode(continuation, cache, weights, q_heads, kv_heads, head_dim):
    outputs = []
    for x_new in continuation.split(1, dim=1):
        output, cache = decode_one(x_new, cache, weights, q_heads, kv_heads, head_dim)
        outputs.append(output)
    return torch.cat(outputs, dim=1), cache

def benchmark(fn, warmup=3, rounds=9):
    for _ in range(warmup): fn()
    if DEVICE.type == "cuda": torch.cuda.synchronize()
    samples = []
    for _ in range(rounds):
        if DEVICE.type == "cuda": torch.cuda.synchronize()
        start = time.perf_counter(); fn()
        if DEVICE.type == "cuda": torch.cuda.synchronize()
        samples.append((time.perf_counter() - start) * 1e3)
    return statistics.median(samples)

def peak_memory(fn):
    gc.collect(); torch.cuda.empty_cache()
    base_a, base_r = torch.cuda.memory_allocated(), torch.cuda.memory_reserved()
    torch.cuda.reset_peak_memory_stats(); fn(); torch.cuda.synchronize()
    return (torch.cuda.max_memory_allocated() - base_a,
            torch.cuda.max_memory_reserved() - base_r)

def main():
    b, pt, dt, hidden, hq, hkv, d = 1, 8, 5, 16, 4, 2, 4
    weights = make_weights(hidden, hq, hkv, d)
    x = torch.randn(b, pt + dt, hidden, device=DEVICE, dtype=DTYPE)
    prompt, continuation = x[:, :pt], x[:, pt:]
    full = recompute(prompt, continuation, weights, hq, hkv, d)
    prompt_output, cache = prefill(prompt, weights, hq, hkv, d)
    cached, final_cache = cached_decode(continuation, cache, weights, hq, hkv, d)
    torch.testing.assert_close(prompt_output, prefill(prompt, weights, hq, hkv, d)[0])
    for step in range(dt):
        torch.testing.assert_close(cached[:, step], full[:, step], rtol=1e-4, atol=1e-5)
    assert final_cache[0].shape == final_cache[1].shape == (b, hkv, pt + dt, d)

    pt, dt = ((256, 32) if DEVICE.type == "cuda" else (64, 8))
    b, hidden, hq, hkv, d = 2, 512, 8, 2, 64
    weights = make_weights(hidden, hq, hkv, d)
    x = torch.randn(b, pt + dt, hidden, device=DEVICE, dtype=DTYPE)
    prompt, continuation = x[:, :pt], x[:, pt:]
    _, cache = prefill(prompt, weights, hq, hkv, d)
    perf_full = recompute(prompt, continuation, weights, hq, hkv, d)
    perf_cached, _ = cached_decode(continuation, cache, weights, hq, hkv, d)
    torch.testing.assert_close(perf_cached, perf_full, rtol=1e-4, atol=1e-5)
    times = (benchmark(lambda: prefill(prompt, weights, hq, hkv, d)),
             benchmark(lambda: recompute(prompt, continuation, weights, hq, hkv, d)),
             benchmark(lambda: cached_decode(continuation, cache, weights, hq, hkv, d)))

    name = torch.cuda.get_device_name(0) if DEVICE.type == "cuda" else "CPU"
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} name={name} dtype={DTYPE}")
    print(f"correctness cache_KV={tuple(final_cache[0].shape)} max_abs_diff={(cached-full).abs().max():.3e}")
    print(f"timing B={b} prompt={pt} decode={dt} H={hidden} Hq={hq} Hkv={hkv} D={d}")
    print(f"prefill_once_ms={times[0]:.3f}\ndecode_recompute_total_ms={times[1]:.3f}\ndecode_cached_total_ms={times[2]:.3f}")
    if DEVICE.type == "cuda":
        theoretical = 2 * b * (pt + dt) * hkv * d * torch.empty((), dtype=DTYPE).element_size()
        peak_full = peak_memory(lambda: recompute(prompt, continuation, weights, hq, hkv, d))
        peak_path = peak_memory(lambda: cached_decode(continuation, prefill(prompt, weights, hq, hkv, d)[1], weights, hq, hkv, d))
        peak_decode = peak_memory(lambda: cached_decode(continuation, cache, weights, hq, hkv, d))
        print(f"theoretical_final_KV_MiB={theoretical/2**20:.3f}")
        print(f"peak_recompute allocated={peak_full[0]/2**20:.3f} reserved={peak_full[1]/2**20:.3f} MiB")
        print(f"peak_cached_path allocated={peak_path[0]/2**20:.3f} reserved={peak_path[1]/2**20:.3f} MiB")
        print(f"peak_cached_decode_above_prompt allocated={peak_decode[0]/2**20:.3f} reserved={peak_decode[1]/2**20:.3f} MiB")
        print(f"after_run allocated={torch.cuda.memory_allocated()/2**20:.3f} reserved={torch.cuda.memory_reserved()/2**20:.3f} MiB")

if __name__ == "__main__":
    with torch.inference_mode(): main()

在项目根目录可保存为 examples/kv_cache_reference.py 后运行：

source /home/nh/workshop/mylab/bin/activate
python examples/kv_cache_reference.py

2026-08-13 在 Python 3.11.12、PyTorch 2.13.0+cu130、RTX 3080 20 GiB 上重新隔离运行的结果为：小例子最大绝对差 1.788e-07；性能 shape 为 B=2, prompt=256, decode=32, H=512, Hq=8, Hkv=2, D=64, float32，prefill 中位数 0.370 ms，32 步完整重算 8.080 ms，32 步缓存 decode 3.435 ms。九轮计时前有三轮 warmup，每个样本前后都同步 CUDA。

这些数字只描述这个教学 Python 实现：它用 torch.cat 反复复制 cache、用 repeat_interleave 物理展开 GQA KV heads，并使用会物化 score 的 dense attention，不能代表 vLLM、TensorRT-LLM、FlashInfer 或专用 paged/fused kernel。这里的证据是语义一致且在该 shape 上呈现预期趋势，不是“KV Cache 必然获得 2.26× 生产加速”。

prefill 与 decode 必须分开分析

若生成 $G$ 个 token，极简 attention 计数可以帮助看出差异。完整重算在第 $j$ 步再次处理长度 $T_p+j$ 的整个前缀：

增量缓存先付一次 prompt attention，再让每个新 Query 关注当前历史：

这只是 attention 主项的教学模型。真实 Transformer 还包含 QKV/O projection、MLP、norm、采样、kernel launch 和数据搬运；FlashAttention 改变的是中间矩阵的 IO 实现，不改变 causal attention 的数学依赖。缓存也可能让 decode 从重复计算问题转为持续读取不断增长 KV 的带宽问题。

容量公式：先算逻辑 payload，再谈 allocator

标准 MHA/GQA/MQA 每个 token、每层都保存 Key 和 Value，因此逻辑 payload 为：

$L$ 是层数，$B$ 是活跃序列数，$T$ 是每条序列当前缓存长度，$H_{kv}$ 是 KV head 数，$D$ 是 head dimension，$s$ 是每元素字节数；前面的 2 分别对应 K 和 V。非等长 batch 应更精确地写成 $2LH_{kv}Ds\sum_i T_i$，而不是把所有请求都机械按同一 $T$ 计算。

代入 $L=32,B=4,T=4096,H_{kv}=8,D=128,s=2$（FP16/BF16）：

每多生成一个 token，四条序列合计增加：

这只是有效 K/V 元素。生产容量预算还要加入 block 对齐与内部碎片、block table/metadata、allocator reserved-but-unused、workspace、临时 attention buffer、CUDA Graph 私有内存池以及模型权重。实验输出也因此区分：

• theoretical_final_KV：由公式得到的 payload；
• allocated：仍被 live tensor 实际占用的 PyTorch 内存；
• reserved：caching allocator 已向 CUDA 申请、可供 PyTorch 复用的内存；
• peak allocated/reserved：测量窗口内各自的峰值。

示例中理论最终 KV 仅 0.562 MiB，但此次完整 cached path 的 peak allocated 是 13.750 MiB、peak reserved 是 22.000 MiB；纯 decode（prompt cache 已存在）的增量峰值为 allocated 2.414 MiB、reserved 0 MiB。差额来自权重、输入、score、输出、torch.cat 和 allocator 行为；0 MiB reserved 增量只表示复用了 baseline 已保留的 segment，不能把它解释为“decode 不用显存”，也不能把 nvidia-smi 的进程占用直接称为“KV Cache 大小”。

为了避免只用小 tensor 验证公式，examples/kv_capacity_allocation.py 还会在 CUDA 上实际创建并写遍一块合并的 [L,KV,B,H_{kv},T,D] FP16 tensor；KV=2 只是在一个 tensor 中并列表示 K/V，目的是验证 payload，不代表生产框架必须采用该物理 layout。脚本在分配前检查“目标 payload + 4 GiB safety margin”的 device free memory，并在两个 case 之间释放 tensor 与 allocator cache：

source /home/nh/workshop/mylab/bin/activate
python -W error examples/kv_capacity_allocation.py --case both

2026-08-13 在空闲 RTX 3080 上隔离顺序执行的结果为：

两块 tensor 都执行了 zero_()，而不是只依赖 torch.empty 的虚拟地址分配；脚本还读取首尾元素并断言 numel × element_size 与 allocated 增量一致。本次结果确认了下表 2 GiB/8 GiB 的裸 KV payload，但没有模拟权重、prefill activation、paged allocator、CUDA Graph pool 或多请求调度，也不是 admission/OOM 上限测试。

MHA、GQA、MLA 与量化如何改变容量

在 $H_q=32,D=128$ 且其余参数沿用上例时：

GQA/MQA 减少的是 KV head 数，不减少 Query head 数；新 Query 仍要覆盖相同历史 token。量化也不是把公式中的 $s$ 机械替换就结束：per-token/per-head/per-group scale、zero point、packing 对齐、反量化临时量、读写 kernel 支持和误差都会改变真实容量与延迟。若 backend 在使用前先扩展回 FP16，带宽收益可能被削弱。

MLA 不能直接塞进 $2H_{kv}D$。 Multi-head Latent Attention 缓存低维 KV latent，并通常另存解耦的 positional/RoPE key 分量。用一个简化容量写法：

若示意取 $d_c=512,d_{rope}=64$，上例为 $32\times4\times4096\times576\times2=576$ MiB。它看起来小于 GQA 的 2 GiB，但真实缓存字段、矩阵吸收、位置路径与量化格式由具体模型和 kernel 决定；MLA 不是“把 $H_{kv}$ 设成一个更小数字”的同一布局。完整 latent/position 数据路径应由独立 MLA 词条解释。

成本转移、失败场景与常见误解

• 显存换计算：缓存越长、并发越高，占用越大；没有足够容量时会触发更小 batch、抢占、重算或 offload。
• decode 仍随上下文变慢：每个新 Query 读取的历史 K/V 随 $T$ 增长。滑动窗口或稀疏 attention 会改变可见范围，但那是额外的模型/算法约束。
• 固定容量需要改变可见性或转移存储：Attention Sink保留少量开头 tokens 与 recent window，驱逐中段历史；它可让 cache 封顶，但不再等价 full causal attention。
• 短请求不一定显出速度收益：很短上下文、很少生成 token、CPU 调度或小 kernel launch 占主导时，cache 管理开销可能遮住节省的计算。
• cache key 是 token 与模型状态，不是原始字符串：tokenizer、chat template、special tokens、position ids、adapter、模型 revision 等差异都可能使历史状态不可复用。
• RoPE schedule 是缓存语义的一部分：PI、NTK-aware 或 YaRN若在 prefill/decode 间改变 inverse frequencies、attention factor 或 position mapping，已旋转的历史 K 不能继续当成同一坐标系读取。
• Prefix Cache 不是单请求 KV Cache 的同义词：KV Cache 是一次请求跨 decode step 的状态；Prefix Cache 还要识别并让不同请求安全共享相同 token prefix 的 KV blocks。
• PagedAttention 不是更快的 attention 数学公式：它主要解决 KV blocks 的分配、映射、共享和碎片；FlashAttention 主要减少 attention 中间矩阵的显存读写。两者可以与 KV Cache 同时存在。

从教学代码到生产实现还差什么

本页代码只复现单层 causal GQA，没有 embedding、RoPE、residual、norm、MLP、logits、sampling、变长 batch、beam reorder、prefix sharing 或多 rank。尤其有三处不能照搬到生产：

1. torch.cat 每一步重新分配并复制旧 K/V；生产 runtime 通常在预分配连续区或 paged blocks 中就地写入。
2. repeat_interleave 真的复制 GQA heads；生产 kernel 通过 head 映射共享读取。
3. dense PyTorch 算子会产生较大中间张量；生产系统会选择 SDPA/FlashAttention、paged attention 或专用 decode kernel，并把 cache layout 设计成 kernel 可高效读取的形式。

API 名称也不是 backend 证据。框架暴露 past_key_values、use_cache 或 cache object，只说明调用语义；实际是 contiguous tensor、paged blocks、量化 cache，还是哪个 fused kernel 消费它，必须结合 runtime 配置、日志、profiler 和源码版本判断。

与上下游术语的关系

• 上游的 Tokenizer / Chat Template / Special Tokens 决定实际 token id 前缀、长度、停止符和可共享 prefix key；原始字符串相似不代表 cache 状态相同。
• Prefill / Decode 给出 cache 的两个时间阶段，并进一步连接 TTFT、TPOT 与 token/cache 的一步状态错位；本页负责证明缓存语义和容量。
• MHA / MQA / GQA 决定 $H_q$ 与 $H_{kv}$ 的共享关系，直接进入容量公式和 decode kernel 的 head 映射。
• 缩放点积注意力 定义新 Query 读取历史 K/V 后的 score、mask、softmax 与 Value 加权语义；本页负责解释 K/V 如何跨 step 存活。
• PagedAttention 在不改变逻辑 K/V 序列的前提下管理非连续物理 blocks、碎片和共享；本页图中的 slot 是它的上游逻辑视图。
• Sliding Window Attention只有在模型层本身限制可见历史时，才把该层 cache/读取长度封顶为 $W$；它不是对全局 attention 的无损截断。
• FlashAttention / SDPA backend 优化 attention 的计算与 IO；它们可以读取 KV Cache，但不负责跨请求生命周期的 block 分配。
• MLA 与 KV Quantization 改变实际缓存字段、元素数或字节数，需要各自的精度与 kernel 约束，不能只套 MHA 公式。

参考资料

• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need — MQA 与增量解码的 KV 带宽动机。
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — GQA 的 Query/KV head sharing。
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model — MLA 的 latent 与解耦位置路径来源。
• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention — KV block 分配、共享与碎片问题。
• PyTorch — scaled_dot_product_attention — SDPA 与 enable_gqa 的 API 语义。
• TensorRT-LLM — KV Cache Reuse — 生产 prefix/KV block 复用及其限制示例。