KV-Cache 空间复杂度怎么算?

空间复杂度

对decoder self-attention，若K/V缓存形状按 [B, n_kv_heads, L, d_head] 计算，总字节数近似为：

2 × B × L × n_layers × n_kv_heads × d_head × bytes_per_element

前面的2代表K和V。因而空间对batch、已缓存序列长度、层数、KV头数和头维都是线性增长。必须使用 n_kv_heads，不能一律代入query头数；MQA/GQA正是通过减少KV头数降低缓存常数。

可核验示例

Llama-2-7B有32层、32个KV头、头维128。batch=1、长度4096、FP16时，缓存约为 2×1×4096×32×32×128×2 bytes = 2 GiB。这并未超过约7B参数的FP16权重，后者约14 GB十进制量级。Llama-2论文中7B/13B使用MHA，70B使用GQA，不能说全系列都是GQA。

优化边界

• MQA/GQA：减少KV头数；质量、带宽和并行效果要按模型验证。
• ：降低每元素字节数，但需评估长上下文误差和反量化开销。