如何解析Context Cache在prompt系列54中的应用原理？

Context Cache的使用几乎已经是行业共识，目标是优化大模型首Token的推理延时，在多轮对话，超长System Prompt，超长结构化JSON和Few-shot等应用场景，是不可或缺的。这一章我们主要从原理、一些论文提出的优化项和VLLM开源项目入手，分析下context Cache的实现和适合场景。 重温KV Cache Context Cache的本质其实是KV Cache在多次请求之间的复用，所以我们先重温下KV Cache的原理。 以下是Transformer的基础模型结构由多层layer串联构成，而KV缓存的是每一层用于Self-Attention计算的历史序列的Key和Value取值，所以缓存向量维度是batch_size * seq_len * num_head * head_dim,(以下可视化来自LLM Visualization和transformer-explainer ) 只所以对self-attention的KV进行缓存，因为在Transformer的众多计算单元中只有self-attention是上下文依赖的，也就是在计算第k个token的输出时，需要使用第k个token（Query）和前面K-1个token的（K&V）进行内机计算，使得self-attention的计算复杂度随序列长度平方增长。而如果对历史序列中的KV进行缓存后，每次生成新token只需要计算当前token，这样时间复杂度就可以降为线性（O(n)），显著降低计算量。 而其他Linear, FFN层的计算都是针对dim层，每个token的计算独立，和历史token无关和序列长度无关，因此也没有缓存的必要。 当然KV Cache也不全是优点，虽然能显著降低推理延时，但是会带来较大的显存占用，占用显存和序列长度、模型层数成正比。 那现在众多大模型API厂商都支持的Context Cache能力和KV cache有哪些区别呢？ 首先从cache共享上，因为KV cache只用于单一序列的推理过程中，因此没有任何共享问题，本次推理完成即释放，而context cache会在多次推理请求之间共享，因此对于如何命中Cache，管理Cache, 提升Cache的存储和使用效率，就需要更多的考虑。 其次从使用时机上，KV cache是用于首Token之后的增量预测（auto-regressive Phrase），而Context Cache是用于首Token之前的存量计算（Prompt Phrase）。在context cache出现之前这两个阶段其实有明确的划分，计算prompt的阶段需要对全部序列进行attention计算属于数据计算密集的任务，而解码阶段因为kV Cache的存在更多是存储密集型任务。因此Context Cache是面向首Token延时的优化方案。 KV Cache只是Context Cache的基础使用形式，下面我们会分别就Contxt Cache的几个核心问题包括命中率低，等讨论一些优化方案 并行效率更高: Chunk Attention https://github.com/microsoft/chunk-attention/tree/main ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition 微软的Chunk Attention通过对KV Cache分段存储，实现self-attention的并发计算，提升长序列self-attention计算效率，经评估自注意力内核的速度可以提高3.2-4.8倍。下面我们结合源码简单说下，源码chunk attention是C语言写的，这里简单转换成python pseudo Prefix Aware KV Cache 传统的KV缓存以密集tensor形式存储，大小是batch * head * seq length * dim，当多个序列共享相同的前缀（prompt）时这些kv缓存就是相同的，因此可以进行共享存储，共享存储的大小则是head * chunk length * dim，这里论文选择使用前缀树进行存储。用前缀不用其他存储方式，主要是为了保证位置编码的可用性。这部分就不细说了，就是Trie树，child节点存储的是对应chunk 序列的kv tensor。而这里的分块存储，为后面推理阶段并行self-attention的计算提供了可能性。