KV Cache 优化的讨论

对于一个 Transformer Block, 我们应该有 token tensor -> projection -> attention -> FFN 这几个核心步骤; 那么使用 kv cache 的话在 decode 阶段可以把 attention 计算的复杂度降低到 O(N)O(N)，并且这个可以同时利好 projection 和后续的 ffn 让其不再需要处理完整的 prompt 而是只有最后一个新的 token; 因此 kv cache 的 存在 可以极大地优化一个 decode 推理流程;

但是我们主要好奇的是 kv cache 的优化究竟占到了多少的性能提升，假设我们有一个弱化版的 kv cache, 对比 vllm 版本究竟弱化化了多少; 然后 FFN 在有 kvcache 的前提下，又能怎么优化; 这是两个不同角度的问题;

而且 KV Cache 的特性是 append only, 这样相当于对于写满的 block 这里的操作是 read-only, 但是经常会涉及 dead -> 回收所有权; KV cache 的并发访问基本上只发生在文章讨论的 Beam Search 和 Parallel Sampling 两个情况下，所以如果发生了，我们要统计 ref_count; 但是应该不会涉及 dirty 的记录，因为我们一般不涉及写入 disk 中的操作;

KV pair

在计算 attention 的时候，我们有 softmax(QKTdk)Vsoftmax(\frac{QK^T}{\sqrt{d_k}})V, 如果利用线性代数知识和对 attention 作用的主观理解，这里的 QK 乘法表示的是 query 和已有的所有 token 的关联关系，然后 V 是逐个对已有的 key 的内容的放大，所以二者要 成对出现，并且我们会发现

[softmax(QK1Tdk)V1softmax(QK2Tdk)V2]=softmax(Q[K1T,K2T])[V1V2]\begin{bmatrix}softmax(\frac{QK_1^T}{\sqrt{d_k}})V_1\\ softmax(\frac{QK_2^T}{\sqrt{d_k}})V_2 \end{bmatrix} = softmax(Q[K_1^T, K_2^T])\begin{bmatrix}V_1 \\ V_2\end{bmatrix}

也就是通过这个公式可以把不同的 KV 拼出来还原出完整的 Q -> K, V 的关联性以及完整的矩阵; 所以我们在存储 K 和 V 的时候最好沿着 sequence id 方向切割，让每一个 token 的所有 hidden dimension 都全部存储进去，并且把对应 token 的 K, V 一块存储;

沿着 token 切割就很容易从绝对位置还原出来每个 token 的 K/V 应该去哪个 id 查找，或者说可以很快看出来这个 token KV 属于哪个 block;