文档摘要

KV 缓存与 Flash Attention:推理优化 本节摘要:训练是并行且受 FLOPs 限制的,推理是串行且受内存带宽限制的——瓶颈不同,技巧也不同。一个朴素的解码器生成 个 token 要做 次注意力:每步都重算整个前缀。一个 4K token 的回复就是 1600 万次注意力操作,绝大多数是冗余的——前缀 token 的隐状态一旦算出来就是确定的,你只需用新 token 的查询对前面所有 token 缓存的键值做一次查询。更糟的是,标准注意力要把 分数矩阵物化到显存, 时注意力先撞内存墙再撞算力墙,经典 kernel 浪费现代 GPU 410 倍。

---