KV缓存估算与前缀复用

沿着对象与状态理解本章

这里有三个不同的存储问题：KV缓存跨轮保留历史；FlashAttention安排单次注意力的中间读写；PagedAttention管理物理缓存块。动画分别给出具体数据，避免将三种节约误认为都减少了相同计算。

下方图解的固定输入在“输入与适用条件”中声明。它用于解释本章机制，不与本章其他独立例子的数值暗中接续。

动态图解 · O04

FlashAttention：一边读一边累计，不攒整张表

不把完整注意力表存下来，怎样仍算出同一输出？

想象分批称重汇总：每批到来时更新“最大刻度、总重量、内容加权和”，不用把所有中间单据再写一遍显存。最大刻度变化时，旧累计必须跟着换刻度。这个在线归一化过程解释了省读写如何与同一数学结果兼容。

m截至当前分块的最大分数。

l按当前最大值缩放后的指数总和。

加权累计尚未除以l的Value贡献和。

输入与适用条件：独立单查询例：读取顺序C/A/B，分数[0,2,1]，对应V为[[5,5],[10,0],[0,20]]。为展示最大值变化而设，不接用F09候选顺序。

1 / 5 · 普通写法先想象整张分数表

停留时间

步骤第1步，共5步

窄屏可左右滑动图形，或点击“放大图形”。文字解释紧随图后。

在长序列中，把所有查询对所有Key的分数完整存下来会产生很大的中间表。FlashAttention改变分块与读写组织，目标不是随意丢弃低分位置。

本步状态由图形与说明共同给出。

默认静止。用“下一步”看完解释再继续；播放到最后一步自动停止。键盘聚焦播放器后，可用左右键逐步查看、空格播放/暂停。图中颜色辅助区分，含义同时用文字注明。

为什么需要

大中间矩阵造成显存读写与容量压力。

怎样实现

分块计算并以在线softmax维护归一化统计，减少外部显存访问。

可能的代价与误解

实现复杂；数值顺序变化有舍入差异；不适用的硬件或形状收益可能很小。

如何判断是否做对

与完整softmax加权结果做容差对照，再在目标硬件实测而非把教学计数当加速倍数。

怎样观察与验证

默认例子分别得到1与4 GiB。加入批大小后观察线性项；真实内存还含权重、分配器与其他状态，特殊架构可能不适用。

运行命令

python3 reading-notes/experiments/03_cache_budget.py

完整教学代码

"""标准注意力 KV 缓存估算与前缀匹配；不是厂商计费器。"""
from typing import Sequence, TypeVar
T = TypeVar("T")


def kv_bytes(layers: int, batch: int, length: int, kv_heads: int,
             head_dim: int, bytes_per_element: int) -> int:
    args = [layers, batch, length, kv_heads, head_dim, bytes_per_element]
    if any(not isinstance(x, int) or isinstance(x, bool) or x <= 0 for x in args):
        raise ValueError("所有参数必须是正整数")
    return 2 * layers * batch * length * kv_heads * head_dim * bytes_per_element


def common_prefix(a: Sequence[T], b: Sequence[T]) -> int:
    count = 0
    for x, y in zip(a, b):
        if x != y:
            break
        count += 1
    return count


def self_test() -> None:
    assert kv_bytes(32, 1, 8192, 8, 128, 2) == 1024**3
    assert kv_bytes(32, 1, 8192, 32, 128, 2) == 4 * 1024**3
    assert common_prefix([1,2,3], [1,2,4]) == 2
    assert common_prefix([], [1]) == 0
    assert common_prefix([1,2], [9,1,2]) == 0


def run_demo() -> dict:
    base = ["规则", "工具", "任务A"]
    return {"性质": "简化公式，不含分配器和其他状态，不适用所有架构",
            "8个KV头_GiB": kv_bytes(32,1,8192,8,128,2)/1024**3,
            "32个KV头_GiB": kv_bytes(32,1,8192,32,128,2)/1024**3,
            "仅尾部变化_共同前缀项": common_prefix(base,["规则","工具","任务B"]),
            "首部插入变化_共同前缀项": common_prefix(base,["新时间","规则","工具","任务B"]),
            "限制": "共同前缀项不是实际缓存Token；产品还有阈值和其他条件"}


if __name__ == "__main__":
    import json
    self_test()
    print(json.dumps(run_demo(), ensure_ascii=False, indent=2))

运行结果示例

以下结果由随包脚本产生。它们不是LLM或产品性能数据。

{
  "性质": "简化公式，不含分配器和其他状态，不适用所有架构",
  "8个KV头_GiB": 1.0,
  "32个KV头_GiB": 4.0,
  "仅尾部变化_共同前缀项": 2,
  "首部插入变化_共同前缀项": 0,
  "限制": "共同前缀项不是实际缓存Token；产品还有阈值和其他条件"
}