GPT-2 KV 缓存从零实现：理论计算省千倍，实测延迟降至 1/6

📅 2026 年 10 月 18 日 📖 约 2 分钟 Reddit r/MachineLearning 大模型推理性能优化PyTorch

我把 PyTorch 中的 GPT-2 small 和 KV 缓存（Key-Value 缓存）作为一个学习项目实现了一遍，然后将理论上的数学节省与实测的延迟和内存进行了对比。我已经发布了实现代码和详细的文字说明，非常希望能得到关于方法论的反馈。

设置

实验环境：预训练的 GPT-2 small（124M 参数），NVIDIA A40 GPU，fp32 精度且关闭 TF32，batch size 为 1，贪心解码（greedy decoding）。计时由 torch.cuda.synchronize() 包围，先进行预热，每个上下文长度运行五次取中位数。

正确性

与 Hugging Face 实现相比，最大绝对 logit 差异约为 6.1e-5；有缓存和无缓存的贪心生成在 200 个 token 上结果一致。一个有用的失败案例：将缓存拼接顺序反转后，单 token 注意力测试仍然通过，但多个新 token 一起解码时就会失败。

延迟

在上下文长度为 1,000 时，使用缓存后 FLOP（浮点运算量）估计减少约 1,000 倍，但每 token 耗时从 29.69 ms 降至 4.95 ms（约 6 倍）。生成 1,016 个 token 时，无缓存需要 14.66 秒，有缓存为 5.03 秒（约 2.9 倍）。该实现中缓存后的延迟曲线几乎平坦。CPU/内核启动开销目前只是一个假设，尚未用 profiler 进行验证。

内存

根据公式，每个 token 需要 72 KiB 内存。我最初的“申请字节数”测量结果中有一个恒定的 8,519,688 字节偏移：一个惰性分配（首次使用时分配）的 8,519,680 字节 cuBLAS 工作区，加上一个 8 字节的 next-token 张量。在基线之前预热并减去存活的非缓存张量后，所有 16 个检查点上“申请字节数 - 公式值”都变为 0 B。分配器的块大小开销仍然单独可见。